Add files using upload-large-folder tool
Browse files- models/full-128k-bpe-nfc-plus2-balanced-mathcode-scratch/tokenizer.json +0 -0
- models/full-128k-bpe-nfc-plus2-balanced-seed43-mathcode-scratch/README.md +73 -0
- models/full-128k-bpe-nfc-plus2-balanced-seed43-mathcode-scratch/conversion_metadata.json +100 -0
- models/full-128k-bpe-nfc-plus2-balanced-seed43-mathcode-scratch/tokenizer_config.json +7 -0
- models/full-128k-claude-english-bpe-mathcode-scratch/README.md +76 -0
- models/full-128k-claude-english-bpe-mathcode-scratch/config.json +28 -0
- models/full-128k-claude-english-bpe-mathcode-scratch/conversion_metadata.json +250 -0
- models/full-128k-claude-english-bpe-mathcode-scratch/generation_config.json +4 -0
- models/full-128k-claude-english-bpe-mathcode-scratch/nanochat_model.py +361 -0
- models/full-128k-claude-english-bpe-mathcode-scratch/tokenizer.json +0 -0
- models/full-128k-claude-english-bpe-mathcode-scratch/tokenizer_config.json +7 -0
- models/full-128k-gpt4o-codegrad50-bpe-mathcode-scratch/config.json +28 -0
- models/full-128k-gpt4o-codegrad50-bpe-mathcode-scratch/conversion_metadata.json +100 -0
- models/full-128k-gpt4o-codegrad50-bpe-mathcode-scratch/generation_config.json +4 -0
- models/full-128k-gpt4o-codegrad50-bpe-mathcode-scratch/nanochat_model.py +361 -0
- models/full-128k-gpt4o-codegrad50-bpe-mathcode-scratch/tokenizer_config.json +7 -0
- models/full-128k-llama3-mathcode-scratch/config.json +28 -0
- models/full-128k-llama3-mathcode-scratch/generation_config.json +4 -0
- models/full-128k-pabpe-nfc-clean-fineweb2full-plus2-capped-tuned-hybridwindow-consv2-tailcuts-gm90k-v130-sp-eng5g-mathcode-scratch/tokenizer_config.json +7 -0
- models/full-128k-pabpe-nfc-gpt4-balanced-mathcode-scratch/README.md +76 -0
- models/full-128k-punct-balanced-bpe-seed42-mathcode-scratch/README.md +74 -0
- models/full-128k-punct-balanced-bpe-seed42-mathcode-scratch/generation_config.json +4 -0
- models/full-128k-punct-english-bpe-mathcode-scratch/conversion_metadata.json +250 -0
- models/full-128k-punct-english-bpe-mathcode-scratch/nanochat_model.py +361 -0
- models/full-128k-punct-english-bpe-mathcode-scratch/tokenizer_config.json +7 -0
- models/full-128k-rightalign-balanced-bpe-mathcode-scratch/README.md +76 -0
- models/full-128k-rightalign-balanced-bpe-mathcode-scratch/config.json +28 -0
- models/full-128k-rightalign-balanced-bpe-mathcode-scratch/generation_config.json +4 -0
- models/full-128k-rightalign-balanced-bpe-mathcode-scratch/nanochat_model.py +361 -0
- models/full-128k-rightalign-balanced-bpe-mathcode-scratch/tokenizer.json +0 -0
- models/full-128k-rightalign-balanced-bpe-mathcode-scratch/tokenizer_config.json +7 -0
- models/full-128k-scripttok-mingram-scriptenc_cb-mathcode-scratch/conversion_metadata.json +100 -0
- models/full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch/README.md +77 -0
- models/full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch/conversion_metadata.json +100 -0
- models/full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch/generation_config.json +4 -0
- models/full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch/tokenizer_config.json +7 -0
- models/full-128k-superbpe-clean-c2-balanced-mathcode-scratch/conversion_metadata.json +100 -0
- models/full-128k-superbpe-clean-c2-balanced-mathcode-scratch/nanochat_model.py +361 -0
- models/full-128k-superbpe-clean-c2-balanced-mathcode-scratch/tokenizer.json +0 -0
- models/full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch/README.md +77 -0
- models/full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch/config.json +28 -0
- models/full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch/generation_config.json +4 -0
- models/full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch/nanochat_model.py +361 -0
- models/full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch/tokenizer.json +0 -0
- models/full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch/tokenizer_config.json +7 -0
- models/full-128k-whitespace-balanced-bpe-mathcode-scratch/README.md +76 -0
- models/full-128k-whitespace-balanced-bpe-mathcode-scratch/config.json +28 -0
- models/full-128k-whitespace-balanced-bpe-mathcode-scratch/conversion_metadata.json +100 -0
- models/full-128k-whitespace-balanced-bpe-mathcode-scratch/tokenizer.json +0 -0
- models/full-128k-whitespace-balanced-bpe-mathcode-scratch/tokenizer_config.json +7 -0
models/full-128k-bpe-nfc-plus2-balanced-mathcode-scratch/tokenizer.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
models/full-128k-bpe-nfc-plus2-balanced-seed43-mathcode-scratch/README.md
ADDED
|
@@ -0,0 +1,73 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# `full-128k-bpe-nfc-plus2-balanced-seed43-mathcode-scratch`
|
| 2 |
+
|
| 3 |
+
A 1.27B parameter model trained with the `bpe-nfc-plus2-balanced` tokenizer. Regime: mathcode-20B-from-scratch (24 layers, width 1536, 1.27B parameters, 20.0B tokens, math and code mixture, trained from scratch). Reported in: math and code ablations.
|
| 4 |
+
|
| 5 |
+
## Load
|
| 6 |
+
|
| 7 |
+
```python
|
| 8 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
| 9 |
+
|
| 10 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 11 |
+
"cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-bpe-nfc-plus2-balanced-seed43-mathcode-scratch",
|
| 12 |
+
trust_remote_code=True, dtype="float32")
|
| 13 |
+
tokenizer = AutoTokenizer.from_pretrained("cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-bpe-nfc-plus2-balanced-seed43-mathcode-scratch")
|
| 14 |
+
|
| 15 |
+
text = "def add(a, b):\n return a + b\n"
|
| 16 |
+
ids = [0] + tokenizer(text)["input_ids"] # training prepends BOS
|
| 17 |
+
```
|
| 18 |
+
|
| 19 |
+
## Settings
|
| 20 |
+
|
| 21 |
+
| Setting | Value |
|
| 22 |
+
|---|---|
|
| 23 |
+
| layers | 24 |
|
| 24 |
+
| model dimension | 1536 |
|
| 25 |
+
| attention heads | 12 |
|
| 26 |
+
| key and value heads | 12 |
|
| 27 |
+
| context length | 2048 |
|
| 28 |
+
| sliding window pattern | SSSL |
|
| 29 |
+
| value embedding dimension | 128 |
|
| 30 |
+
| logit soft cap | 15 |
|
| 31 |
+
| vocabulary size | 128256 |
|
| 32 |
+
| parameters | 1,272,841,994 |
|
| 33 |
+
| training step | 19073 |
|
| 34 |
+
| token budget | 19,999,490,048 |
|
| 35 |
+
| bytes consumed | 151,753,084,755 |
|
| 36 |
+
| seed | 43 |
|
| 37 |
+
| Muon matrix learning rate | 0.0200 |
|
| 38 |
+
| AdamW embedding learning rate | 0.3000 |
|
| 39 |
+
| AdamW unembedding learning rate | 0.0080 |
|
| 40 |
+
| weight decay | 0.2800 |
|
| 41 |
+
| width learning-rate exponent | -1 |
|
| 42 |
+
| weight dtype | as trained (float32 matrices, bfloat16 embeddings) |
|
| 43 |
+
| BOS token id | 0 |
|
| 44 |
+
|
| 45 |
+
## Results
|
| 46 |
+
|
| 47 |
+
| Metric | Value |
|
| 48 |
+
|---|---|
|
| 49 |
+
| validation BPB | 0.3388 |
|
| 50 |
+
| code BPB, mean over 7 languages | n/a |
|
| 51 |
+
| FLORES-200 BPB, mean | 3.3594 |
|
| 52 |
+
| FLORES-200 languages scored | 215 |
|
| 53 |
+
| BLiMP accuracy | n/a |
|
| 54 |
+
| MBPP pass@1 | 0.2120 |
|
| 55 |
+
| HumanEval pass@1 | 0.1768 |
|
| 56 |
+
| GSM8K exact match, strict | 0.2024 |
|
| 57 |
+
|
| 58 |
+
Generation metrics use specification `{'version': 'v2-2026-07-30', 'bos_context': True, 'eos_stop': True, 'token_healing': True, 'truncation': 'left_to_max_length_minus_max_gen', 'stop_truncation': 'earliest_match_all_terms', 'stop_match_decode': 'specials_kept'}`, token healing on.
|
| 59 |
+
|
| 60 |
+
Missing measurements, with the reason each is absent:
|
| 61 |
+
|
| 62 |
+
- `blimp_acc`: no full-128k-bpe-nfc-plus2-balanced-seed43-mathcode-scratch_blimp_code_bpb.json
|
| 63 |
+
- `code_bpb_mean`: no full-128k-bpe-nfc-plus2-balanced-seed43-mathcode-scratch_blimp_code_bpb.json; no full-128k-bpe-nfc-plus2-balanced-seed43-mathcode-scratch_code_bpb.json
|
| 64 |
+
|
| 65 |
+
BLiMP is comparable within a training scale only. Its scoring sums unnormalized sentence log-likelihoods, and under that scoring a group of quantifier subtasks drops at the 1.27B scale relative to the 300M scale.
|
| 66 |
+
|
| 67 |
+
## Provenance
|
| 68 |
+
|
| 69 |
+
- Source checkpoint: `/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-bpe-nfc-plus2-balanced-seed43-mathcode-scratch`, step 19073
|
| 70 |
+
- Repository commit: `1a9a718`
|
| 71 |
+
- nanochat commit: `f919449` on branch `tokenizer-lm-patches`
|
| 72 |
+
- Tokenizer: `bpe-nfc-plus2-balanced`, released at `tokenizers/bpe-nfc-plus2-balanced`
|
| 73 |
+
|
models/full-128k-bpe-nfc-plus2-balanced-seed43-mathcode-scratch/conversion_metadata.json
ADDED
|
@@ -0,0 +1,100 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"run": "full-128k-bpe-nfc-plus2-balanced-seed43-mathcode-scratch",
|
| 3 |
+
"source_checkpoint": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-bpe-nfc-plus2-balanced-seed43-mathcode-scratch",
|
| 4 |
+
"source_model_file": "model_019073.pt",
|
| 5 |
+
"source_step": 19073,
|
| 6 |
+
"source_meta": {
|
| 7 |
+
"step": 19073,
|
| 8 |
+
"val_bpb": 0.3387550570042015,
|
| 9 |
+
"model_config": {
|
| 10 |
+
"sequence_len": 2048,
|
| 11 |
+
"vocab_size": 128256,
|
| 12 |
+
"n_layer": 24,
|
| 13 |
+
"n_head": 12,
|
| 14 |
+
"n_kv_head": 12,
|
| 15 |
+
"n_embd": 1536,
|
| 16 |
+
"window_pattern": "SSSL",
|
| 17 |
+
"ve_dim": 128,
|
| 18 |
+
"logit_softcap": 15.0
|
| 19 |
+
},
|
| 20 |
+
"tokenizer_path": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/tokenizers/bpe-nfc-plus2-balanced",
|
| 21 |
+
"training_config": {
|
| 22 |
+
"name": "full-128k-mathcode-scratch",
|
| 23 |
+
"model": {
|
| 24 |
+
"n_layer": 24,
|
| 25 |
+
"n_embd": 1536,
|
| 26 |
+
"n_kv_head": 12,
|
| 27 |
+
"sequence_len": 2048,
|
| 28 |
+
"window_pattern": "SSSL",
|
| 29 |
+
"ve_dim": 128
|
| 30 |
+
},
|
| 31 |
+
"training": {
|
| 32 |
+
"target_param_data_ratio": -1,
|
| 33 |
+
"max_tokens": 20000000000,
|
| 34 |
+
"num_iterations": -1,
|
| 35 |
+
"total_batch_size": 1048576,
|
| 36 |
+
"device_batch_size": 16,
|
| 37 |
+
"warmup_steps": 40,
|
| 38 |
+
"warmdown_ratio": 0.65,
|
| 39 |
+
"final_lr_frac": 0.05,
|
| 40 |
+
"eval_tokens": 41943040,
|
| 41 |
+
"compile": true
|
| 42 |
+
},
|
| 43 |
+
"optimizer": {
|
| 44 |
+
"matrix_lr": 0.02,
|
| 45 |
+
"embedding_lr": 0.3,
|
| 46 |
+
"unembedding_lr": 0.008,
|
| 47 |
+
"scalar_lr": 0.5,
|
| 48 |
+
"weight_decay": 0.28,
|
| 49 |
+
"width_lr_exponent": -1.0
|
| 50 |
+
},
|
| 51 |
+
"data": {
|
| 52 |
+
"data_dir": null
|
| 53 |
+
},
|
| 54 |
+
"logging": {
|
| 55 |
+
"eval_every": 500,
|
| 56 |
+
"save_every": 500,
|
| 57 |
+
"max_checkpoints": 0
|
| 58 |
+
},
|
| 59 |
+
"wandb": {
|
| 60 |
+
"enabled": true,
|
| 61 |
+
"project": "tokenizer-lm-experiments",
|
| 62 |
+
"run_name": null
|
| 63 |
+
},
|
| 64 |
+
"output_dir": "/capstor/scratch/cscs/${USER}/tokenizer-lm/checkpoints",
|
| 65 |
+
"seed": 42
|
| 66 |
+
},
|
| 67 |
+
"total_batch_size": 1048576,
|
| 68 |
+
"device_batch_size": 16,
|
| 69 |
+
"max_seq_len": 2048,
|
| 70 |
+
"dataloader_state_dict": {
|
| 71 |
+
"pq_idx": 208,
|
| 72 |
+
"rg_idx": 40,
|
| 73 |
+
"epoch": 2
|
| 74 |
+
},
|
| 75 |
+
"loop_state": {
|
| 76 |
+
"min_val_bpb": 0.3387550570042015,
|
| 77 |
+
"smooth_train_loss": 0.7248042415647279,
|
| 78 |
+
"total_training_time": 93284.78915786743,
|
| 79 |
+
"total_bytes_consumed": 151753084755
|
| 80 |
+
}
|
| 81 |
+
},
|
| 82 |
+
"weight_dtype": "as trained (float32 matrices, bfloat16 embeddings)",
|
| 83 |
+
"weight_dtype_counts": {
|
| 84 |
+
"torch.float32": 174,
|
| 85 |
+
"torch.bfloat16": 13
|
| 86 |
+
},
|
| 87 |
+
"total_parameters": 1272841994,
|
| 88 |
+
"logit_softcap_source": "meta",
|
| 89 |
+
"bos_token": "<s>",
|
| 90 |
+
"bos_token_id": 0,
|
| 91 |
+
"bos_resolved_from": "packaged tokenizer",
|
| 92 |
+
"tokenizer_slug": "bpe-nfc-plus2-balanced",
|
| 93 |
+
"tokenizer_redistributed": true,
|
| 94 |
+
"tokenizer_hub_id": null,
|
| 95 |
+
"weight_key_mapping": "strip _orig_mod. and transformer. prefixes",
|
| 96 |
+
"converted_by": "scripts/hf_release/convert_for_release.py",
|
| 97 |
+
"repo_commit": "1a9a718",
|
| 98 |
+
"nanochat_commit": "f919449",
|
| 99 |
+
"nanochat_branch": "tokenizer-lm-patches"
|
| 100 |
+
}
|
models/full-128k-bpe-nfc-plus2-balanced-seed43-mathcode-scratch/tokenizer_config.json
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"tokenizer_class": "PreTrainedTokenizerFast",
|
| 3 |
+
"model_max_length": 2048,
|
| 4 |
+
"bos_token": "<s>",
|
| 5 |
+
"add_bos_token": false,
|
| 6 |
+
"clean_up_tokenization_spaces": false
|
| 7 |
+
}
|
models/full-128k-claude-english-bpe-mathcode-scratch/README.md
ADDED
|
@@ -0,0 +1,76 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# `full-128k-claude-english-bpe-mathcode-scratch`
|
| 2 |
+
|
| 3 |
+
A 1.27B parameter model trained with the `claude-english-bpe` tokenizer. Regime: mathcode-20B-from-scratch (24 layers, width 1536, 1.27B parameters, 20.0B tokens, math and code mixture, trained from scratch). Reported in: math and code ablations; TokEval.
|
| 4 |
+
|
| 5 |
+
## Load
|
| 6 |
+
|
| 7 |
+
```python
|
| 8 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
| 9 |
+
|
| 10 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 11 |
+
"cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-claude-english-bpe-mathcode-scratch",
|
| 12 |
+
trust_remote_code=True, dtype="float32")
|
| 13 |
+
tokenizer = AutoTokenizer.from_pretrained("cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-claude-english-bpe-mathcode-scratch")
|
| 14 |
+
|
| 15 |
+
text = "def add(a, b):\n return a + b\n"
|
| 16 |
+
ids = [0] + tokenizer(text)["input_ids"] # training prepends BOS
|
| 17 |
+
```
|
| 18 |
+
|
| 19 |
+
## Settings
|
| 20 |
+
|
| 21 |
+
| Setting | Value |
|
| 22 |
+
|---|---|
|
| 23 |
+
| layers | 24 |
|
| 24 |
+
| model dimension | 1536 |
|
| 25 |
+
| attention heads | 12 |
|
| 26 |
+
| key and value heads | 12 |
|
| 27 |
+
| context length | 2048 |
|
| 28 |
+
| sliding window pattern | SSSL |
|
| 29 |
+
| value embedding dimension | 128 |
|
| 30 |
+
| logit soft cap | 15 |
|
| 31 |
+
| vocabulary size | 128256 |
|
| 32 |
+
| parameters | 1,272,841,994 |
|
| 33 |
+
| training step | 19073 |
|
| 34 |
+
| token budget | 19,999,490,048 |
|
| 35 |
+
| bytes consumed | 142,775,328,290 |
|
| 36 |
+
| seed | 42 |
|
| 37 |
+
| Muon matrix learning rate | 0.0200 |
|
| 38 |
+
| AdamW embedding learning rate | 0.3000 |
|
| 39 |
+
| AdamW unembedding learning rate | 0.0080 |
|
| 40 |
+
| weight decay | 0.2800 |
|
| 41 |
+
| width learning-rate exponent | -1 |
|
| 42 |
+
| weight dtype | as trained (float32 matrices, bfloat16 embeddings) |
|
| 43 |
+
| BOS token id | 0 |
|
| 44 |
+
|
| 45 |
+
## Results
|
| 46 |
+
|
| 47 |
+
| Metric | Value |
|
| 48 |
+
|---|---|
|
| 49 |
+
| validation BPB | 0.3405 |
|
| 50 |
+
| code BPB, mean over 7 languages | 0.3652 |
|
| 51 |
+
| FLORES-200 BPB, mean | 3.1993 |
|
| 52 |
+
| FLORES-200 languages scored | 215 |
|
| 53 |
+
| BLiMP accuracy | n/a |
|
| 54 |
+
| MBPP pass@1 | 0.2340 |
|
| 55 |
+
| HumanEval pass@1 | 0.1646 |
|
| 56 |
+
| GSM8K exact match, strict | 0.2214 |
|
| 57 |
+
|
| 58 |
+
Generation metrics use specification `{'version': 'v2-2026-07-30', 'bos_context': True, 'eos_stop': True, 'token_healing': True, 'truncation': 'left_to_max_length_minus_max_gen', 'stop_truncation': 'earliest_match_all_terms', 'stop_match_decode': 'specials_kept'}`, token healing on.
|
| 59 |
+
|
| 60 |
+
Missing measurements, with the reason each is absent:
|
| 61 |
+
|
| 62 |
+
- `blimp_acc`: no full-128k-claude-english-bpe-mathcode-scratch_blimp_code_bpb.json
|
| 63 |
+
|
| 64 |
+
BLiMP is comparable within a training scale only. Its scoring sums unnormalized sentence log-likelihoods, and under that scoring a group of quantifier subtasks drops at the 1.27B scale relative to the 300M scale.
|
| 65 |
+
|
| 66 |
+
## Registry note
|
| 67 |
+
|
| 68 |
+
#51, claude pretok, slug `claude-english-bpe`. Main-grid Claude English-data cell. FINISHED at step 19073, 4-eval chain + code_bpb complete (verified 2026-08-09). In panel.
|
| 69 |
+
|
| 70 |
+
## Provenance
|
| 71 |
+
|
| 72 |
+
- Source checkpoint: `/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-claude-english-bpe-mathcode-scratch`, step 19073
|
| 73 |
+
- Repository commit: `1a9a718`
|
| 74 |
+
- nanochat commit: `f919449` on branch `tokenizer-lm-patches`
|
| 75 |
+
- Tokenizer: `claude-english-bpe`, released at `tokenizers/claude-english-bpe`
|
| 76 |
+
|
models/full-128k-claude-english-bpe-mathcode-scratch/config.json
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"NanoChatForCausalLM"
|
| 4 |
+
],
|
| 5 |
+
"auto_map": {
|
| 6 |
+
"AutoConfig": "nanochat_model.NanoChatConfig",
|
| 7 |
+
"AutoModelForCausalLM": "nanochat_model.NanoChatForCausalLM"
|
| 8 |
+
},
|
| 9 |
+
"bos_token_id": 0,
|
| 10 |
+
"dtype": "float32",
|
| 11 |
+
"eos_token_id": null,
|
| 12 |
+
"head_dim": 128,
|
| 13 |
+
"kv_dim": 1536,
|
| 14 |
+
"logit_softcap": 15.0,
|
| 15 |
+
"model_type": "nanochat",
|
| 16 |
+
"n_embd": 1536,
|
| 17 |
+
"n_head": 12,
|
| 18 |
+
"n_kv_head": 12,
|
| 19 |
+
"n_layer": 24,
|
| 20 |
+
"pad_vocab_size_to": 64,
|
| 21 |
+
"padded_vocab_size": 128256,
|
| 22 |
+
"sequence_len": 2048,
|
| 23 |
+
"tie_word_embeddings": false,
|
| 24 |
+
"transformers_version": "5.3.0",
|
| 25 |
+
"ve_dim": 128,
|
| 26 |
+
"vocab_size": 128256,
|
| 27 |
+
"window_pattern": "SSSL"
|
| 28 |
+
}
|
models/full-128k-claude-english-bpe-mathcode-scratch/conversion_metadata.json
ADDED
|
@@ -0,0 +1,250 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"run": "full-128k-claude-english-bpe-mathcode-scratch",
|
| 3 |
+
"source_checkpoint": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-claude-english-bpe-mathcode-scratch",
|
| 4 |
+
"source_model_file": "model_019073.pt",
|
| 5 |
+
"source_step": 19073,
|
| 6 |
+
"source_meta": {
|
| 7 |
+
"step": 19073,
|
| 8 |
+
"val_bpb": 0.34048861753468623,
|
| 9 |
+
"model_config": {
|
| 10 |
+
"sequence_len": 2048,
|
| 11 |
+
"vocab_size": 128256,
|
| 12 |
+
"n_layer": 24,
|
| 13 |
+
"n_head": 12,
|
| 14 |
+
"n_kv_head": 12,
|
| 15 |
+
"n_embd": 1536,
|
| 16 |
+
"window_pattern": "SSSL",
|
| 17 |
+
"ve_dim": 128,
|
| 18 |
+
"logit_softcap": 15.0
|
| 19 |
+
},
|
| 20 |
+
"tokenizer_path": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/tokenizers/claude-english-bpe",
|
| 21 |
+
"training_config": {
|
| 22 |
+
"name": "full-128k-mathcode-scratch",
|
| 23 |
+
"model": {
|
| 24 |
+
"n_layer": 24,
|
| 25 |
+
"n_embd": 1536,
|
| 26 |
+
"n_kv_head": 12,
|
| 27 |
+
"sequence_len": 2048,
|
| 28 |
+
"window_pattern": "SSSL",
|
| 29 |
+
"ve_dim": 128
|
| 30 |
+
},
|
| 31 |
+
"training": {
|
| 32 |
+
"target_param_data_ratio": -1,
|
| 33 |
+
"max_tokens": 20000000000,
|
| 34 |
+
"num_iterations": -1,
|
| 35 |
+
"total_batch_size": 1048576,
|
| 36 |
+
"device_batch_size": 16,
|
| 37 |
+
"warmup_steps": 40,
|
| 38 |
+
"warmdown_ratio": 0.65,
|
| 39 |
+
"final_lr_frac": 0.05,
|
| 40 |
+
"eval_tokens": 41943040,
|
| 41 |
+
"compile": true
|
| 42 |
+
},
|
| 43 |
+
"optimizer": {
|
| 44 |
+
"matrix_lr": 0.02,
|
| 45 |
+
"embedding_lr": 0.3,
|
| 46 |
+
"unembedding_lr": 0.008,
|
| 47 |
+
"scalar_lr": 0.5,
|
| 48 |
+
"weight_decay": 0.28,
|
| 49 |
+
"width_lr_exponent": -1.0
|
| 50 |
+
},
|
| 51 |
+
"data": {
|
| 52 |
+
"data_dir": null
|
| 53 |
+
},
|
| 54 |
+
"logging": {
|
| 55 |
+
"eval_every": 500,
|
| 56 |
+
"save_every": 500,
|
| 57 |
+
"max_checkpoints": 0
|
| 58 |
+
},
|
| 59 |
+
"wandb": {
|
| 60 |
+
"enabled": true,
|
| 61 |
+
"project": "tokenizer-lm-experiments",
|
| 62 |
+
"run_name": null
|
| 63 |
+
},
|
| 64 |
+
"output_dir": "/capstor/scratch/cscs/${USER}/tokenizer-lm/checkpoints",
|
| 65 |
+
"seed": 42
|
| 66 |
+
},
|
| 67 |
+
"code_provenance": {
|
| 68 |
+
"commit": "6231791054e9c0df5144025df8f8a2117e8b78d2",
|
| 69 |
+
"branch": "main",
|
| 70 |
+
"dirty": true,
|
| 71 |
+
"dirty_files": [
|
| 72 |
+
"M .gitignore",
|
| 73 |
+
" M EXPERIMENTAL_SETUP.md",
|
| 74 |
+
" M RESULT_COMPUTATION_ANALYSIS.md",
|
| 75 |
+
" M eval_runner.py",
|
| 76 |
+
" M mc_evaluation",
|
| 77 |
+
" M scripts/add_mc_math_ci.py",
|
| 78 |
+
" M scripts/bootstrap_mathcode_significance.py",
|
| 79 |
+
" M scripts/correlate_intrinsic_downstream.py",
|
| 80 |
+
" M scripts/correlate_mathcode_intrinsic_downstream.py",
|
| 81 |
+
" M scripts/cross_lingual_analyses.py",
|
| 82 |
+
" M scripts/generate_latex_tables.py",
|
| 83 |
+
" M scripts/launch_mathcode_scratch_chained.sh",
|
| 84 |
+
" M scripts/launch_new6_all_scales.sh",
|
| 85 |
+
" M scripts/reorganize_checkpoints.py",
|
| 86 |
+
" M scripts/slurm_3b_mathcode_resume.sh",
|
| 87 |
+
" M scripts/slurm_eval_mgsm.sh",
|
| 88 |
+
" M scripts/slurm_full.sh",
|
| 89 |
+
" M scripts/slurm_full_1node.sh",
|
| 90 |
+
" M scripts/slurm_full_1node_resume.sh",
|
| 91 |
+
" M scripts/slurm_mathcode_resume_1node.sh",
|
| 92 |
+
" M scripts/slurm_pilot.sh",
|
| 93 |
+
" M scripts/slurm_train_1node.sh",
|
| 94 |
+
" M scripts/slurm_train_tokenizers.sh",
|
| 95 |
+
" M scripts/train_custom_tokenizers.py",
|
| 96 |
+
" M tokenizer_lm/tokenizer.py",
|
| 97 |
+
" M train.py",
|
| 98 |
+
"?? REPRODUCIBILITY_STATUS.md",
|
| 99 |
+
"?? SESSION_STATUS.md",
|
| 100 |
+
"?? TOKEN_HEALING_HANDOFF.md",
|
| 101 |
+
"?? configs/full_128k_bytematched.yaml",
|
| 102 |
+
"?? configs/full_128k_clean_v128k.yaml",
|
| 103 |
+
"?? configs/full_128k_clean_v160k.yaml",
|
| 104 |
+
"?? configs/full_128k_clean_v200k.yaml",
|
| 105 |
+
"?? configs/full_128k_clean_v64k.yaml",
|
| 106 |
+
"?? configs/full_128k_clean_v96k.yaml",
|
| 107 |
+
"?? configs/panels/",
|
| 108 |
+
"?? configs/starcoder_ast_18lang.json",
|
| 109 |
+
"?? correspondence/",
|
| 110 |
+
"?? scripts/_lexer_treesitter_worker.py",
|
| 111 |
+
"?? scripts/add_code_bpb_std.py",
|
| 112 |
+
"?? scripts/analyze_apertus_track_sensitivity.py",
|
| 113 |
+
"?? scripts/analyze_axes2_codegrad.py",
|
| 114 |
+
"?? scripts/analyze_code_boundary_drift.py",
|
| 115 |
+
"?? scripts/analyze_code_gen_failures.py",
|
| 116 |
+
"?? scripts/analyze_code_intrinsic_downstream.py",
|
| 117 |
+
"?? scripts/analyze_cr_byte_drop.py",
|
| 118 |
+
"?? scripts/analyze_cross_scale_ranking.py",
|
| 119 |
+
"?? scripts/analyze_equal_weighting.py",
|
| 120 |
+
"?? scripts/analyze_gen_ablation.py",
|
| 121 |
+
"?? scripts/analyze_seed_variance.py",
|
| 122 |
+
"?? scripts/analyze_seed_variance_mc20b.py",
|
| 123 |
+
"?? scripts/analyze_vocab_size.py",
|
| 124 |
+
"?? scripts/ast_node_discovery.py",
|
| 125 |
+
"?? scripts/ast_scope_extractor.py",
|
| 126 |
+
"?? scripts/audit_bpb_byte_bias.py",
|
| 127 |
+
"?? scripts/audit_corpus_parseability.py",
|
| 128 |
+
"?? scripts/audit_generation_symptoms.py",
|
| 129 |
+
"?? scripts/audit_lexer_boundary_leak.py",
|
| 130 |
+
"?? scripts/audit_logit_softcap.py",
|
| 131 |
+
"?? scripts/audit_maxchars_fragmentation.py",
|
| 132 |
+
"?? scripts/audit_mgsm_prompt_boundary.py",
|
| 133 |
+
"?? scripts/audit_newline_tokens.py",
|
| 134 |
+
"?? scripts/audit_newline_tokens_fullbyte_pair.py",
|
| 135 |
+
"?? scripts/audit_newline_tokens_new11.py",
|
| 136 |
+
"?? scripts/backfill_flores_summary.py",
|
| 137 |
+
"?? scripts/build_axes2_codegrad_tokenizers.sh",
|
| 138 |
+
"?? scripts/build_vocab_ladder_tokenizers.sh",
|
| 139 |
+
"?? scripts/build_whitespace_ladder_tokenizers.sh",
|
| 140 |
+
"?? scripts/characterize_axes2_intrinsic.py",
|
| 141 |
+
"?? scripts/check_hf_dataset_cache.py",
|
| 142 |
+
"?? scripts/check_number_provenance.py",
|
| 143 |
+
"?? scripts/compare_bigram_entropy_variants.py",
|
| 144 |
+
"?? scripts/compute_headline_spreads.py",
|
| 145 |
+
"?? scripts/compute_vocab_ladder_byte_matching.py",
|
| 146 |
+
"?? scripts/correlate_operator_domains.py",
|
| 147 |
+
"?? scripts/d2_candidate_screen.py",
|
| 148 |
+
"?? scripts/diagnose_fullbyte_indent.py",
|
| 149 |
+
"?? scripts/diagnose_operator_isolation.py",
|
| 150 |
+
"?? scripts/dump_axes2_composition_weights.py",
|
| 151 |
+
"?? scripts/estimate_intrinsic_cost.py",
|
| 152 |
+
"?? scripts/fit_quota_response_model.py",
|
| 153 |
+
"?? scripts/gen_v2_spec_audit.py",
|
| 154 |
+
"?? scripts/generate_intrinsic_tables.py",
|
| 155 |
+
"?? scripts/launch_boundary_retrain.sh",
|
| 156 |
+
"?? scripts/launch_mc_pretok_chained.sh",
|
| 157 |
+
"?? scripts/launch_std_pretok_chained.sh",
|
| 158 |
+
"?? scripts/launch_vocab_ladder.sh",
|
| 159 |
+
"?? scripts/launch_vocab_ladder_paramscaled.sh",
|
| 160 |
+
"?? scripts/lexer_pretok_common.py",
|
| 161 |
+
"?? scripts/lexer_pretok_corpus.py",
|
| 162 |
+
"?? scripts/mbpp_mediation.py",
|
| 163 |
+
"?? scripts/merge_intrinsic_components.py",
|
| 164 |
+
"?? scripts/merge_intrinsic_shards.py",
|
| 165 |
+
"?? scripts/operator_isolation_by_corpus_n36.py",
|
| 166 |
+
"?? scripts/optimize_quotas.py",
|
| 167 |
+
"?? scripts/preflight_gen_v2_tokenizers.py",
|
| 168 |
+
"?? scripts/probe_belebele_bos.py",
|
| 169 |
+
"?? scripts/probe_eot_stop_incidence.py",
|
| 170 |
+
"?? scripts/probe_healing_prefill_equivalence.py",
|
| 171 |
+
"?? scripts/recheck_seed_sd_verdicts.py",
|
| 172 |
+
"?? scripts/recompute_val_bpb.py",
|
| 173 |
+
"?? scripts/rejudge_tail_penalty.py",
|
| 174 |
+
"?? scripts/render_code_appendix.py",
|
| 175 |
+
"?? scripts/retry_whitespace_axes2_builds.sh",
|
| 176 |
+
"?? scripts/run_gen_v2_one.sh",
|
| 177 |
+
"?? scripts/run_mc_panels.sh",
|
| 178 |
+
"?? scripts/script_pretok_per_language.py",
|
| 179 |
+
"?? scripts/slurm_corpus_parseability_audit.sh",
|
| 180 |
+
"?? scripts/slurm_d2_candidate_build.sh",
|
| 181 |
+
"?? scripts/slurm_d2_candidate_screen.sh",
|
| 182 |
+
"?? scripts/slurm_eval_code_gen_healed.sh",
|
| 183 |
+
"?? scripts/slurm_eval_gen_ablation.sh",
|
| 184 |
+
"?? scripts/slurm_eval_gen_v2.sh",
|
| 185 |
+
"?? scripts/slurm_eval_gen_v2_packed.sbatch",
|
| 186 |
+
"?? scripts/slurm_full_1node_resume_pretok.sh",
|
| 187 |
+
"?? scripts/slurm_intrinsic_camera_ready.sbatch",
|
| 188 |
+
"?? scripts/slurm_intrinsic_component.sbatch",
|
| 189 |
+
"?? scripts/slurm_probe_mgsm_bos.sbatch",
|
| 190 |
+
"?? scripts/slurm_train_pretok_1node.sh",
|
| 191 |
+
"?? scripts/submit_gen_v2_sweep.sh",
|
| 192 |
+
"?? scripts/submit_lexer_corpus_cache_jobs.sh",
|
| 193 |
+
"?? scripts/summarize_lexer_smoke_run.py",
|
| 194 |
+
"?? scripts/test_ast_scope_extractor.py",
|
| 195 |
+
"?? scripts/test_lexer_pretok_common.py",
|
| 196 |
+
"?? scripts/test_lexer_treesitter_worker.py",
|
| 197 |
+
"?? scripts/test_presegmented_training_path.py",
|
| 198 |
+
"?? scripts/token_healing_failure_modes.py",
|
| 199 |
+
"?? scripts/validate_axes2_codegrad_tokenizers.py",
|
| 200 |
+
"?? scripts/validate_generation_conformance.py",
|
| 201 |
+
"?? scripts/validate_gpt4o_english_fullbyte_bpe.py",
|
| 202 |
+
"?? scripts/validate_token_healing.py",
|
| 203 |
+
"?? scripts/validate_whitespace_ladder_regexes.py",
|
| 204 |
+
"?? scripts/validate_whitespace_ladder_tokenizers.py",
|
| 205 |
+
"?? scripts/verify_gen_v2_smoke.py",
|
| 206 |
+
"?? scripts/verify_iter_source_hardening.py",
|
| 207 |
+
"?? scripts/vet_runtime_regex.py",
|
| 208 |
+
"?? scripts/whitespace_fusion_audit.py",
|
| 209 |
+
"?? std1b_invalid.md",
|
| 210 |
+
"?? std1b_used.md",
|
| 211 |
+
"?? tokenizer_lm/pretok_dataloader.py",
|
| 212 |
+
"?? tokenizer_lm/token_healing.py"
|
| 213 |
+
],
|
| 214 |
+
"nanochat_commit": "f91944934f78d2083c8f5bd6ea8ab4e048786054",
|
| 215 |
+
"torch_version": "2.10.0+cu126"
|
| 216 |
+
},
|
| 217 |
+
"total_batch_size": 1048576,
|
| 218 |
+
"device_batch_size": 16,
|
| 219 |
+
"max_seq_len": 2048,
|
| 220 |
+
"dataloader_state_dict": {
|
| 221 |
+
"pq_idx": 185,
|
| 222 |
+
"rg_idx": 40,
|
| 223 |
+
"epoch": 2
|
| 224 |
+
},
|
| 225 |
+
"loop_state": {
|
| 226 |
+
"min_val_bpb": 0.34048861753468623,
|
| 227 |
+
"smooth_train_loss": 0.7963783171242066,
|
| 228 |
+
"total_training_time": 93372.5398645401,
|
| 229 |
+
"total_bytes_consumed": 142775328290
|
| 230 |
+
}
|
| 231 |
+
},
|
| 232 |
+
"weight_dtype": "as trained (float32 matrices, bfloat16 embeddings)",
|
| 233 |
+
"weight_dtype_counts": {
|
| 234 |
+
"torch.float32": 174,
|
| 235 |
+
"torch.bfloat16": 13
|
| 236 |
+
},
|
| 237 |
+
"total_parameters": 1272841994,
|
| 238 |
+
"logit_softcap_source": "meta",
|
| 239 |
+
"bos_token": "<s>",
|
| 240 |
+
"bos_token_id": 0,
|
| 241 |
+
"bos_resolved_from": "packaged tokenizer",
|
| 242 |
+
"tokenizer_slug": "claude-english-bpe",
|
| 243 |
+
"tokenizer_redistributed": true,
|
| 244 |
+
"tokenizer_hub_id": null,
|
| 245 |
+
"weight_key_mapping": "strip _orig_mod. and transformer. prefixes",
|
| 246 |
+
"converted_by": "scripts/hf_release/convert_for_release.py",
|
| 247 |
+
"repo_commit": "1a9a718",
|
| 248 |
+
"nanochat_commit": "f919449",
|
| 249 |
+
"nanochat_branch": "tokenizer-lm-patches"
|
| 250 |
+
}
|
models/full-128k-claude-english-bpe-mathcode-scratch/generation_config.json
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"bos_token_id": 0,
|
| 3 |
+
"transformers_version": "5.3.0"
|
| 4 |
+
}
|
models/full-128k-claude-english-bpe-mathcode-scratch/nanochat_model.py
ADDED
|
@@ -0,0 +1,361 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
NanoChat model for HuggingFace Transformers — self-contained single file.
|
| 3 |
+
|
| 4 |
+
Includes the config class, model class, and a CLI demo for loading + generation.
|
| 5 |
+
Drop this file into a directory containing config.json, model.safetensors,
|
| 6 |
+
tokenizer.json, and tokenizer_config.json, then run:
|
| 7 |
+
|
| 8 |
+
python nanochat_model.py --model-dir . --prompt "Hello world"
|
| 9 |
+
|
| 10 |
+
Or load programmatically:
|
| 11 |
+
|
| 12 |
+
from nanochat_model import NanoChatConfig, NanoChatForCausalLM
|
| 13 |
+
model = NanoChatForCausalLM.from_pretrained("path/to/model", torch_dtype=torch.bfloat16)
|
| 14 |
+
|
| 15 |
+
Requirements: pip install torch transformers tokenizers
|
| 16 |
+
|
| 17 |
+
Architecture: nanochat GPT (Karpathy, MIT licensed) with value embeddings
|
| 18 |
+
(ResFormer-style), per-layer residual/x0 lambdas, smear, backout, logit
|
| 19 |
+
soft-capping, sliding window attention, QK norm, ReLU^2 MLP, and RoPE.
|
| 20 |
+
|
| 21 |
+
The models were trained with a patched nanochat, not upstream nanochat:
|
| 22 |
+
https://github.com/cimeister/nanochat branch tokenizer-lm-patches, commit f919449.
|
| 23 |
+
The patches touch two files and change three things, two of which affect the forward
|
| 24 |
+
pass and are implemented here:
|
| 25 |
+
1. ve_dim: the value-embedding tables are built at a reduced dimension and projected
|
| 26 |
+
up to kv_dim by a learned ve_proj, instead of being built at full kv_dim. All
|
| 27 |
+
released models use ve_dim=128.
|
| 28 |
+
2. logit_softcap: the soft cap is a config field with 0 meaning disabled, instead of
|
| 29 |
+
a hardcoded 15.
|
| 30 |
+
3. width_lr_exponent: the AdamW learning-rate width scaling exponent is configurable
|
| 31 |
+
for muP. This affects training only and has no inference-time effect.
|
| 32 |
+
|
| 33 |
+
This file is the released model definition. Its agreement with the training-time model
|
| 34 |
+
is checked per released checkpoint by scripts/hf_release/verify_conversion.py, which
|
| 35 |
+
compares logits against the patched nanochat GPT in float32.
|
| 36 |
+
"""
|
| 37 |
+
|
| 38 |
+
from typing import Optional, Tuple, Union
|
| 39 |
+
|
| 40 |
+
import torch
|
| 41 |
+
import torch.nn as nn
|
| 42 |
+
import torch.nn.functional as F
|
| 43 |
+
from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
|
| 44 |
+
from transformers.modeling_outputs import CausalLMOutputWithPast
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
# ---------------------------------------------------------------------------
|
| 48 |
+
# Config
|
| 49 |
+
# ---------------------------------------------------------------------------
|
| 50 |
+
|
| 51 |
+
class NanoChatConfig(PretrainedConfig):
|
| 52 |
+
model_type = "nanochat"
|
| 53 |
+
|
| 54 |
+
def __init__(
|
| 55 |
+
self,
|
| 56 |
+
vocab_size=128256,
|
| 57 |
+
n_embd=1536,
|
| 58 |
+
n_layer=24,
|
| 59 |
+
n_head=12,
|
| 60 |
+
n_kv_head=12,
|
| 61 |
+
sequence_len=2048,
|
| 62 |
+
window_pattern="SSSL",
|
| 63 |
+
ve_dim=0,
|
| 64 |
+
logit_softcap=15.0,
|
| 65 |
+
pad_vocab_size_to=64,
|
| 66 |
+
bos_token_id=0,
|
| 67 |
+
eos_token_id=None,
|
| 68 |
+
tie_word_embeddings=False,
|
| 69 |
+
**kwargs,
|
| 70 |
+
):
|
| 71 |
+
self.vocab_size = vocab_size
|
| 72 |
+
self.n_embd = n_embd
|
| 73 |
+
self.n_layer = n_layer
|
| 74 |
+
self.n_head = n_head
|
| 75 |
+
self.n_kv_head = n_kv_head
|
| 76 |
+
self.sequence_len = sequence_len
|
| 77 |
+
self.window_pattern = window_pattern
|
| 78 |
+
self.ve_dim = ve_dim
|
| 79 |
+
self.logit_softcap = logit_softcap
|
| 80 |
+
self.pad_vocab_size_to = pad_vocab_size_to
|
| 81 |
+
self.head_dim = n_embd // n_head
|
| 82 |
+
self.kv_dim = n_kv_head * self.head_dim
|
| 83 |
+
self.padded_vocab_size = (
|
| 84 |
+
(vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to
|
| 85 |
+
) * pad_vocab_size_to
|
| 86 |
+
super().__init__(
|
| 87 |
+
bos_token_id=bos_token_id,
|
| 88 |
+
eos_token_id=eos_token_id,
|
| 89 |
+
tie_word_embeddings=tie_word_embeddings,
|
| 90 |
+
**kwargs,
|
| 91 |
+
)
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
# ---------------------------------------------------------------------------
|
| 95 |
+
# Model
|
| 96 |
+
# ---------------------------------------------------------------------------
|
| 97 |
+
|
| 98 |
+
def _has_ve(layer_idx: int, n_layer: int) -> bool:
|
| 99 |
+
return layer_idx % 2 == (n_layer - 1) % 2
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def _apply_rotary_emb(x, cos, sin):
|
| 103 |
+
d = x.shape[-1] // 2
|
| 104 |
+
x1, x2 = x[..., :d], x[..., d:]
|
| 105 |
+
y1 = x1 * cos + x2 * sin
|
| 106 |
+
y2 = x1 * (-sin) + x2 * cos
|
| 107 |
+
return torch.cat([y1, y2], dim=-1)
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
def _rms_norm(x):
|
| 111 |
+
return F.rms_norm(x, (x.size(-1),))
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
class NanoChatAttention(nn.Module):
|
| 115 |
+
def __init__(self, config, layer_idx):
|
| 116 |
+
super().__init__()
|
| 117 |
+
self.layer_idx = layer_idx
|
| 118 |
+
self.n_head = config.n_head
|
| 119 |
+
self.n_kv_head = config.n_kv_head
|
| 120 |
+
self.n_embd = config.n_embd
|
| 121 |
+
self.head_dim = config.head_dim
|
| 122 |
+
self.n_kv_groups = self.n_head // self.n_kv_head
|
| 123 |
+
self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
|
| 124 |
+
self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
|
| 125 |
+
self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
|
| 126 |
+
self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
|
| 127 |
+
self.ve_gate_channels = 12
|
| 128 |
+
if _has_ve(layer_idx, config.n_layer):
|
| 129 |
+
self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False)
|
| 130 |
+
kv_dim = self.n_kv_head * self.head_dim
|
| 131 |
+
ve_dim = config.ve_dim if config.ve_dim > 0 else kv_dim
|
| 132 |
+
self.ve_proj = nn.Linear(ve_dim, kv_dim, bias=False) if ve_dim < kv_dim else None
|
| 133 |
+
else:
|
| 134 |
+
self.ve_gate = None
|
| 135 |
+
self.ve_proj = None
|
| 136 |
+
|
| 137 |
+
def forward(self, x, ve, cos, sin, window_size):
|
| 138 |
+
B, T, C = x.size()
|
| 139 |
+
q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
|
| 140 |
+
k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
|
| 141 |
+
v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
|
| 142 |
+
if ve is not None and self.ve_gate is not None:
|
| 143 |
+
if self.ve_proj is not None:
|
| 144 |
+
ve = self.ve_proj(ve)
|
| 145 |
+
ve = ve.view(B, T, self.n_kv_head, self.head_dim)
|
| 146 |
+
gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
|
| 147 |
+
v = v + gate.unsqueeze(-1) * ve
|
| 148 |
+
q = _apply_rotary_emb(q, cos, sin)
|
| 149 |
+
k = _apply_rotary_emb(k, cos, sin)
|
| 150 |
+
q = _rms_norm(q) * 1.2
|
| 151 |
+
k = _rms_norm(k) * 1.2
|
| 152 |
+
q = q.transpose(1, 2)
|
| 153 |
+
k = k.transpose(1, 2)
|
| 154 |
+
v = v.transpose(1, 2)
|
| 155 |
+
if self.n_kv_groups > 1:
|
| 156 |
+
k = k.repeat_interleave(self.n_kv_groups, dim=1)
|
| 157 |
+
v = v.repeat_interleave(self.n_kv_groups, dim=1)
|
| 158 |
+
if window_size < 0 or window_size >= T:
|
| 159 |
+
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
|
| 160 |
+
else:
|
| 161 |
+
row_idx = torch.arange(T, device=x.device).unsqueeze(1)
|
| 162 |
+
col_idx = torch.arange(T, device=x.device).unsqueeze(0)
|
| 163 |
+
mask = (col_idx <= row_idx) & ((row_idx - col_idx) <= window_size)
|
| 164 |
+
y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
|
| 165 |
+
y = y.transpose(1, 2).contiguous().view(B, T, -1)
|
| 166 |
+
y = self.c_proj(y)
|
| 167 |
+
return y
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
class NanoChatMLP(nn.Module):
|
| 171 |
+
def __init__(self, config):
|
| 172 |
+
super().__init__()
|
| 173 |
+
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
|
| 174 |
+
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
|
| 175 |
+
|
| 176 |
+
def forward(self, x):
|
| 177 |
+
x = self.c_fc(x)
|
| 178 |
+
x = F.relu(x).square()
|
| 179 |
+
x = self.c_proj(x)
|
| 180 |
+
return x
|
| 181 |
+
|
| 182 |
+
|
| 183 |
+
class NanoChatBlock(nn.Module):
|
| 184 |
+
def __init__(self, config, layer_idx):
|
| 185 |
+
super().__init__()
|
| 186 |
+
self.attn = NanoChatAttention(config, layer_idx)
|
| 187 |
+
self.mlp = NanoChatMLP(config)
|
| 188 |
+
|
| 189 |
+
def forward(self, x, ve, cos, sin, window_size):
|
| 190 |
+
x = x + self.attn(_rms_norm(x), ve, cos, sin, window_size)
|
| 191 |
+
x = x + self.mlp(_rms_norm(x))
|
| 192 |
+
return x
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
class NanoChatPreTrainedModel(PreTrainedModel):
|
| 196 |
+
config_class = NanoChatConfig
|
| 197 |
+
base_model_prefix = "model"
|
| 198 |
+
supports_gradient_checkpointing = False
|
| 199 |
+
_no_split_modules = ["NanoChatBlock"]
|
| 200 |
+
|
| 201 |
+
def _init_weights(self, module):
|
| 202 |
+
pass
|
| 203 |
+
|
| 204 |
+
|
| 205 |
+
class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin):
|
| 206 |
+
def __init__(self, config):
|
| 207 |
+
super().__init__(config)
|
| 208 |
+
pvs = config.padded_vocab_size
|
| 209 |
+
self.wte = nn.Embedding(pvs, config.n_embd)
|
| 210 |
+
self.h = nn.ModuleList([NanoChatBlock(config, i) for i in range(config.n_layer)])
|
| 211 |
+
self.lm_head = nn.Linear(config.n_embd, pvs, bias=False)
|
| 212 |
+
self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
|
| 213 |
+
self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
|
| 214 |
+
self.smear_gate = nn.Linear(24, 1, bias=False)
|
| 215 |
+
self.smear_lambda = nn.Parameter(torch.zeros(1))
|
| 216 |
+
self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))
|
| 217 |
+
ve_emb_dim = config.ve_dim if config.ve_dim > 0 else config.kv_dim
|
| 218 |
+
self.value_embeds = nn.ModuleDict({
|
| 219 |
+
str(i): nn.Embedding(pvs, ve_emb_dim)
|
| 220 |
+
for i in range(config.n_layer) if _has_ve(i, config.n_layer)
|
| 221 |
+
})
|
| 222 |
+
self.window_sizes = self._compute_window_sizes(config)
|
| 223 |
+
self.rotary_seq_len = config.sequence_len * 10
|
| 224 |
+
self._rotary_initialized = False
|
| 225 |
+
self.register_buffer("cos", torch.zeros(1), persistent=False)
|
| 226 |
+
self.register_buffer("sin", torch.zeros(1), persistent=False)
|
| 227 |
+
self.post_init()
|
| 228 |
+
|
| 229 |
+
def _compute_window_sizes(self, config):
|
| 230 |
+
pattern = config.window_pattern.upper()
|
| 231 |
+
long_window = config.sequence_len
|
| 232 |
+
short_window = -(-long_window // 4 // 128) * 128
|
| 233 |
+
char_to_window = {"L": long_window, "S": short_window}
|
| 234 |
+
sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.n_layer)]
|
| 235 |
+
sizes[-1] = long_window
|
| 236 |
+
return sizes
|
| 237 |
+
|
| 238 |
+
def _precompute_rotary(self, seq_len, head_dim, base=100000.0):
|
| 239 |
+
channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32)
|
| 240 |
+
inv_freq = 1.0 / (base ** (channel_range / head_dim))
|
| 241 |
+
t = torch.arange(seq_len, dtype=torch.float32)
|
| 242 |
+
freqs = torch.outer(t, inv_freq)
|
| 243 |
+
cos = freqs.cos()[None, :, None, :]
|
| 244 |
+
sin = freqs.sin()[None, :, None, :]
|
| 245 |
+
return cos, sin
|
| 246 |
+
|
| 247 |
+
def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs):
|
| 248 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 249 |
+
B, T = input_ids.size()
|
| 250 |
+
if not self._rotary_initialized:
|
| 251 |
+
cos, sin = self._precompute_rotary(self.rotary_seq_len, self.config.head_dim)
|
| 252 |
+
self.cos = cos.to(device=input_ids.device)
|
| 253 |
+
self.sin = sin.to(device=input_ids.device)
|
| 254 |
+
self._rotary_initialized = True
|
| 255 |
+
cos = self.cos[:, :T].to(dtype=self.wte.weight.dtype)
|
| 256 |
+
sin = self.sin[:, :T].to(dtype=self.wte.weight.dtype)
|
| 257 |
+
x = self.wte(input_ids)
|
| 258 |
+
x = _rms_norm(x)
|
| 259 |
+
if T > 1:
|
| 260 |
+
gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
|
| 261 |
+
x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
|
| 262 |
+
x0 = x
|
| 263 |
+
n_layer = self.config.n_layer
|
| 264 |
+
backout_layer = n_layer // 2
|
| 265 |
+
x_backout = None
|
| 266 |
+
for i, block in enumerate(self.h):
|
| 267 |
+
x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
|
| 268 |
+
ve = self.value_embeds[str(i)](input_ids).to(x.dtype) if str(i) in self.value_embeds else None
|
| 269 |
+
x = block(x, ve, cos, sin, self.window_sizes[i])
|
| 270 |
+
if i == backout_layer:
|
| 271 |
+
x_backout = x
|
| 272 |
+
if x_backout is not None:
|
| 273 |
+
x = x - self.backout_lambda.to(x.dtype) * x_backout
|
| 274 |
+
x = _rms_norm(x)
|
| 275 |
+
logits = self.lm_head(x)
|
| 276 |
+
logits = logits[..., :self.config.vocab_size]
|
| 277 |
+
logits = logits.float()
|
| 278 |
+
if self.config.logit_softcap > 0:
|
| 279 |
+
cap = self.config.logit_softcap
|
| 280 |
+
logits = cap * torch.tanh(logits / cap)
|
| 281 |
+
loss = None
|
| 282 |
+
if labels is not None:
|
| 283 |
+
shift_logits = logits[..., :-1, :].contiguous()
|
| 284 |
+
shift_labels = labels[..., 1:].contiguous()
|
| 285 |
+
loss = F.cross_entropy(
|
| 286 |
+
shift_logits.view(-1, shift_logits.size(-1)),
|
| 287 |
+
shift_labels.view(-1),
|
| 288 |
+
ignore_index=-100,
|
| 289 |
+
)
|
| 290 |
+
if not return_dict:
|
| 291 |
+
output = (logits,)
|
| 292 |
+
return (loss,) + output if loss is not None else output
|
| 293 |
+
return CausalLMOutputWithPast(loss=loss, logits=logits)
|
| 294 |
+
|
| 295 |
+
def prepare_inputs_for_generation(self, input_ids, **kwargs):
|
| 296 |
+
return {"input_ids": input_ids}
|
| 297 |
+
|
| 298 |
+
|
| 299 |
+
# ---------------------------------------------------------------------------
|
| 300 |
+
# CLI demo
|
| 301 |
+
# ---------------------------------------------------------------------------
|
| 302 |
+
|
| 303 |
+
if __name__ == "__main__":
|
| 304 |
+
import argparse
|
| 305 |
+
import sys
|
| 306 |
+
import os
|
| 307 |
+
|
| 308 |
+
parser = argparse.ArgumentParser(description="Load and generate with a NanoChat model")
|
| 309 |
+
parser.add_argument("--model-dir", required=True, help="Path to converted model directory")
|
| 310 |
+
parser.add_argument("--prompt", default="The meaning of life is", help="Text prompt")
|
| 311 |
+
parser.add_argument("--max-tokens", type=int, default=100, help="Max tokens to generate")
|
| 312 |
+
parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature (0=greedy)")
|
| 313 |
+
parser.add_argument("--device", default="cpu", help="Device (cpu, cuda, cuda:0, ...)")
|
| 314 |
+
parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
|
| 315 |
+
args = parser.parse_args()
|
| 316 |
+
|
| 317 |
+
dtype_map = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
|
| 318 |
+
|
| 319 |
+
print(f"Loading model from {args.model_dir}...")
|
| 320 |
+
from transformers import AutoTokenizer
|
| 321 |
+
|
| 322 |
+
model = NanoChatForCausalLM.from_pretrained(args.model_dir, torch_dtype=dtype_map[args.dtype])
|
| 323 |
+
model.to(args.device)
|
| 324 |
+
model.eval()
|
| 325 |
+
tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
|
| 326 |
+
|
| 327 |
+
total_params = sum(p.numel() for p in model.parameters())
|
| 328 |
+
print(f"Model loaded: {total_params:,} parameters, dtype={args.dtype}, device={args.device}")
|
| 329 |
+
print(f"Vocab size: {model.config.vocab_size}, ve_dim: {model.config.ve_dim}")
|
| 330 |
+
print()
|
| 331 |
+
print(f"Prompt: {args.prompt}")
|
| 332 |
+
print(f"Generating (max {args.max_tokens} tokens, temp={args.temperature})...")
|
| 333 |
+
print()
|
| 334 |
+
|
| 335 |
+
# Generate
|
| 336 |
+
bos_id = tokenizer.bos_token_id
|
| 337 |
+
input_ids = tokenizer.encode(args.prompt, add_special_tokens=False)
|
| 338 |
+
if bos_id is not None:
|
| 339 |
+
input_ids = [bos_id] + input_ids
|
| 340 |
+
input_ids = torch.tensor([input_ids], dtype=torch.long, device=args.device)
|
| 341 |
+
|
| 342 |
+
generated = []
|
| 343 |
+
with torch.inference_mode():
|
| 344 |
+
for _ in range(args.max_tokens):
|
| 345 |
+
logits = model(input_ids).logits[:, -1, :]
|
| 346 |
+
if args.temperature > 0:
|
| 347 |
+
logits = logits / args.temperature
|
| 348 |
+
v, _ = torch.topk(logits, min(50, logits.size(-1)))
|
| 349 |
+
logits[logits < v[:, [-1]]] = -float("inf")
|
| 350 |
+
probs = torch.softmax(logits, dim=-1)
|
| 351 |
+
next_id = torch.multinomial(probs, num_samples=1)
|
| 352 |
+
else:
|
| 353 |
+
next_id = torch.argmax(logits, dim=-1, keepdim=True)
|
| 354 |
+
input_ids = torch.cat([input_ids, next_id], dim=1)
|
| 355 |
+
token_id = next_id.item()
|
| 356 |
+
generated.append(token_id)
|
| 357 |
+
if tokenizer.eos_token_id is not None and token_id == tokenizer.eos_token_id:
|
| 358 |
+
break
|
| 359 |
+
|
| 360 |
+
output = tokenizer.decode(generated, skip_special_tokens=True)
|
| 361 |
+
print(f"Output: {args.prompt}{output}")
|
models/full-128k-claude-english-bpe-mathcode-scratch/tokenizer.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
models/full-128k-claude-english-bpe-mathcode-scratch/tokenizer_config.json
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"tokenizer_class": "PreTrainedTokenizerFast",
|
| 3 |
+
"model_max_length": 2048,
|
| 4 |
+
"bos_token": "<s>",
|
| 5 |
+
"add_bos_token": false,
|
| 6 |
+
"clean_up_tokenization_spaces": false
|
| 7 |
+
}
|
models/full-128k-gpt4o-codegrad50-bpe-mathcode-scratch/config.json
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"NanoChatForCausalLM"
|
| 4 |
+
],
|
| 5 |
+
"auto_map": {
|
| 6 |
+
"AutoConfig": "nanochat_model.NanoChatConfig",
|
| 7 |
+
"AutoModelForCausalLM": "nanochat_model.NanoChatForCausalLM"
|
| 8 |
+
},
|
| 9 |
+
"bos_token_id": 0,
|
| 10 |
+
"dtype": "float32",
|
| 11 |
+
"eos_token_id": null,
|
| 12 |
+
"head_dim": 128,
|
| 13 |
+
"kv_dim": 1536,
|
| 14 |
+
"logit_softcap": 15.0,
|
| 15 |
+
"model_type": "nanochat",
|
| 16 |
+
"n_embd": 1536,
|
| 17 |
+
"n_head": 12,
|
| 18 |
+
"n_kv_head": 12,
|
| 19 |
+
"n_layer": 24,
|
| 20 |
+
"pad_vocab_size_to": 64,
|
| 21 |
+
"padded_vocab_size": 128256,
|
| 22 |
+
"sequence_len": 2048,
|
| 23 |
+
"tie_word_embeddings": false,
|
| 24 |
+
"transformers_version": "5.3.0",
|
| 25 |
+
"ve_dim": 128,
|
| 26 |
+
"vocab_size": 128256,
|
| 27 |
+
"window_pattern": "SSSL"
|
| 28 |
+
}
|
models/full-128k-gpt4o-codegrad50-bpe-mathcode-scratch/conversion_metadata.json
ADDED
|
@@ -0,0 +1,100 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"run": "full-128k-gpt4o-codegrad50-bpe-mathcode-scratch",
|
| 3 |
+
"source_checkpoint": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-gpt4o-codegrad50-bpe-mathcode-scratch",
|
| 4 |
+
"source_model_file": "model_019073.pt",
|
| 5 |
+
"source_step": 19073,
|
| 6 |
+
"source_meta": {
|
| 7 |
+
"step": 19073,
|
| 8 |
+
"val_bpb": 0.33476615801535653,
|
| 9 |
+
"model_config": {
|
| 10 |
+
"sequence_len": 2048,
|
| 11 |
+
"vocab_size": 128256,
|
| 12 |
+
"n_layer": 24,
|
| 13 |
+
"n_head": 12,
|
| 14 |
+
"n_kv_head": 12,
|
| 15 |
+
"n_embd": 1536,
|
| 16 |
+
"window_pattern": "SSSL",
|
| 17 |
+
"ve_dim": 128,
|
| 18 |
+
"logit_softcap": 15.0
|
| 19 |
+
},
|
| 20 |
+
"tokenizer_path": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/tokenizers/gpt4o-codegrad50-bpe",
|
| 21 |
+
"training_config": {
|
| 22 |
+
"name": "full-128k-mathcode-scratch",
|
| 23 |
+
"model": {
|
| 24 |
+
"n_layer": 24,
|
| 25 |
+
"n_embd": 1536,
|
| 26 |
+
"n_kv_head": 12,
|
| 27 |
+
"sequence_len": 2048,
|
| 28 |
+
"window_pattern": "SSSL",
|
| 29 |
+
"ve_dim": 128
|
| 30 |
+
},
|
| 31 |
+
"training": {
|
| 32 |
+
"target_param_data_ratio": -1,
|
| 33 |
+
"max_tokens": 20000000000,
|
| 34 |
+
"num_iterations": -1,
|
| 35 |
+
"total_batch_size": 1048576,
|
| 36 |
+
"device_batch_size": 16,
|
| 37 |
+
"warmup_steps": 40,
|
| 38 |
+
"warmdown_ratio": 0.65,
|
| 39 |
+
"final_lr_frac": 0.05,
|
| 40 |
+
"eval_tokens": 41943040,
|
| 41 |
+
"compile": true
|
| 42 |
+
},
|
| 43 |
+
"optimizer": {
|
| 44 |
+
"matrix_lr": 0.02,
|
| 45 |
+
"embedding_lr": 0.3,
|
| 46 |
+
"unembedding_lr": 0.008,
|
| 47 |
+
"scalar_lr": 0.5,
|
| 48 |
+
"weight_decay": 0.28,
|
| 49 |
+
"width_lr_exponent": -1.0
|
| 50 |
+
},
|
| 51 |
+
"data": {
|
| 52 |
+
"data_dir": null
|
| 53 |
+
},
|
| 54 |
+
"logging": {
|
| 55 |
+
"eval_every": 500,
|
| 56 |
+
"save_every": 500,
|
| 57 |
+
"max_checkpoints": 0
|
| 58 |
+
},
|
| 59 |
+
"wandb": {
|
| 60 |
+
"enabled": true,
|
| 61 |
+
"project": "tokenizer-lm-experiments",
|
| 62 |
+
"run_name": null
|
| 63 |
+
},
|
| 64 |
+
"output_dir": "/capstor/scratch/cscs/${USER}/tokenizer-lm/checkpoints",
|
| 65 |
+
"seed": 42
|
| 66 |
+
},
|
| 67 |
+
"total_batch_size": 1048576,
|
| 68 |
+
"device_batch_size": 16,
|
| 69 |
+
"max_seq_len": 2048,
|
| 70 |
+
"dataloader_state_dict": {
|
| 71 |
+
"pq_idx": 270,
|
| 72 |
+
"rg_idx": 28,
|
| 73 |
+
"epoch": 2
|
| 74 |
+
},
|
| 75 |
+
"loop_state": {
|
| 76 |
+
"min_val_bpb": 0.33476615801535653,
|
| 77 |
+
"smooth_train_loss": 0.9962024787360929,
|
| 78 |
+
"total_training_time": 93434.8084256649,
|
| 79 |
+
"total_bytes_consumed": 138855769571
|
| 80 |
+
}
|
| 81 |
+
},
|
| 82 |
+
"weight_dtype": "as trained (float32 matrices, bfloat16 embeddings)",
|
| 83 |
+
"weight_dtype_counts": {
|
| 84 |
+
"torch.float32": 174,
|
| 85 |
+
"torch.bfloat16": 13
|
| 86 |
+
},
|
| 87 |
+
"total_parameters": 1272841994,
|
| 88 |
+
"logit_softcap_source": "meta",
|
| 89 |
+
"bos_token": "<s>",
|
| 90 |
+
"bos_token_id": 0,
|
| 91 |
+
"bos_resolved_from": "packaged tokenizer",
|
| 92 |
+
"tokenizer_slug": "gpt4o-codegrad50-bpe",
|
| 93 |
+
"tokenizer_redistributed": true,
|
| 94 |
+
"tokenizer_hub_id": null,
|
| 95 |
+
"weight_key_mapping": "strip _orig_mod. and transformer. prefixes",
|
| 96 |
+
"converted_by": "scripts/hf_release/convert_for_release.py",
|
| 97 |
+
"repo_commit": "1a9a718",
|
| 98 |
+
"nanochat_commit": "f919449",
|
| 99 |
+
"nanochat_branch": "tokenizer-lm-patches"
|
| 100 |
+
}
|
models/full-128k-gpt4o-codegrad50-bpe-mathcode-scratch/generation_config.json
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"bos_token_id": 0,
|
| 3 |
+
"transformers_version": "5.3.0"
|
| 4 |
+
}
|
models/full-128k-gpt4o-codegrad50-bpe-mathcode-scratch/nanochat_model.py
ADDED
|
@@ -0,0 +1,361 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
NanoChat model for HuggingFace Transformers — self-contained single file.
|
| 3 |
+
|
| 4 |
+
Includes the config class, model class, and a CLI demo for loading + generation.
|
| 5 |
+
Drop this file into a directory containing config.json, model.safetensors,
|
| 6 |
+
tokenizer.json, and tokenizer_config.json, then run:
|
| 7 |
+
|
| 8 |
+
python nanochat_model.py --model-dir . --prompt "Hello world"
|
| 9 |
+
|
| 10 |
+
Or load programmatically:
|
| 11 |
+
|
| 12 |
+
from nanochat_model import NanoChatConfig, NanoChatForCausalLM
|
| 13 |
+
model = NanoChatForCausalLM.from_pretrained("path/to/model", torch_dtype=torch.bfloat16)
|
| 14 |
+
|
| 15 |
+
Requirements: pip install torch transformers tokenizers
|
| 16 |
+
|
| 17 |
+
Architecture: nanochat GPT (Karpathy, MIT licensed) with value embeddings
|
| 18 |
+
(ResFormer-style), per-layer residual/x0 lambdas, smear, backout, logit
|
| 19 |
+
soft-capping, sliding window attention, QK norm, ReLU^2 MLP, and RoPE.
|
| 20 |
+
|
| 21 |
+
The models were trained with a patched nanochat, not upstream nanochat:
|
| 22 |
+
https://github.com/cimeister/nanochat branch tokenizer-lm-patches, commit f919449.
|
| 23 |
+
The patches touch two files and change three things, two of which affect the forward
|
| 24 |
+
pass and are implemented here:
|
| 25 |
+
1. ve_dim: the value-embedding tables are built at a reduced dimension and projected
|
| 26 |
+
up to kv_dim by a learned ve_proj, instead of being built at full kv_dim. All
|
| 27 |
+
released models use ve_dim=128.
|
| 28 |
+
2. logit_softcap: the soft cap is a config field with 0 meaning disabled, instead of
|
| 29 |
+
a hardcoded 15.
|
| 30 |
+
3. width_lr_exponent: the AdamW learning-rate width scaling exponent is configurable
|
| 31 |
+
for muP. This affects training only and has no inference-time effect.
|
| 32 |
+
|
| 33 |
+
This file is the released model definition. Its agreement with the training-time model
|
| 34 |
+
is checked per released checkpoint by scripts/hf_release/verify_conversion.py, which
|
| 35 |
+
compares logits against the patched nanochat GPT in float32.
|
| 36 |
+
"""
|
| 37 |
+
|
| 38 |
+
from typing import Optional, Tuple, Union
|
| 39 |
+
|
| 40 |
+
import torch
|
| 41 |
+
import torch.nn as nn
|
| 42 |
+
import torch.nn.functional as F
|
| 43 |
+
from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
|
| 44 |
+
from transformers.modeling_outputs import CausalLMOutputWithPast
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
# ---------------------------------------------------------------------------
|
| 48 |
+
# Config
|
| 49 |
+
# ---------------------------------------------------------------------------
|
| 50 |
+
|
| 51 |
+
class NanoChatConfig(PretrainedConfig):
|
| 52 |
+
model_type = "nanochat"
|
| 53 |
+
|
| 54 |
+
def __init__(
|
| 55 |
+
self,
|
| 56 |
+
vocab_size=128256,
|
| 57 |
+
n_embd=1536,
|
| 58 |
+
n_layer=24,
|
| 59 |
+
n_head=12,
|
| 60 |
+
n_kv_head=12,
|
| 61 |
+
sequence_len=2048,
|
| 62 |
+
window_pattern="SSSL",
|
| 63 |
+
ve_dim=0,
|
| 64 |
+
logit_softcap=15.0,
|
| 65 |
+
pad_vocab_size_to=64,
|
| 66 |
+
bos_token_id=0,
|
| 67 |
+
eos_token_id=None,
|
| 68 |
+
tie_word_embeddings=False,
|
| 69 |
+
**kwargs,
|
| 70 |
+
):
|
| 71 |
+
self.vocab_size = vocab_size
|
| 72 |
+
self.n_embd = n_embd
|
| 73 |
+
self.n_layer = n_layer
|
| 74 |
+
self.n_head = n_head
|
| 75 |
+
self.n_kv_head = n_kv_head
|
| 76 |
+
self.sequence_len = sequence_len
|
| 77 |
+
self.window_pattern = window_pattern
|
| 78 |
+
self.ve_dim = ve_dim
|
| 79 |
+
self.logit_softcap = logit_softcap
|
| 80 |
+
self.pad_vocab_size_to = pad_vocab_size_to
|
| 81 |
+
self.head_dim = n_embd // n_head
|
| 82 |
+
self.kv_dim = n_kv_head * self.head_dim
|
| 83 |
+
self.padded_vocab_size = (
|
| 84 |
+
(vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to
|
| 85 |
+
) * pad_vocab_size_to
|
| 86 |
+
super().__init__(
|
| 87 |
+
bos_token_id=bos_token_id,
|
| 88 |
+
eos_token_id=eos_token_id,
|
| 89 |
+
tie_word_embeddings=tie_word_embeddings,
|
| 90 |
+
**kwargs,
|
| 91 |
+
)
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
# ---------------------------------------------------------------------------
|
| 95 |
+
# Model
|
| 96 |
+
# ---------------------------------------------------------------------------
|
| 97 |
+
|
| 98 |
+
def _has_ve(layer_idx: int, n_layer: int) -> bool:
|
| 99 |
+
return layer_idx % 2 == (n_layer - 1) % 2
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def _apply_rotary_emb(x, cos, sin):
|
| 103 |
+
d = x.shape[-1] // 2
|
| 104 |
+
x1, x2 = x[..., :d], x[..., d:]
|
| 105 |
+
y1 = x1 * cos + x2 * sin
|
| 106 |
+
y2 = x1 * (-sin) + x2 * cos
|
| 107 |
+
return torch.cat([y1, y2], dim=-1)
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
def _rms_norm(x):
|
| 111 |
+
return F.rms_norm(x, (x.size(-1),))
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
class NanoChatAttention(nn.Module):
|
| 115 |
+
def __init__(self, config, layer_idx):
|
| 116 |
+
super().__init__()
|
| 117 |
+
self.layer_idx = layer_idx
|
| 118 |
+
self.n_head = config.n_head
|
| 119 |
+
self.n_kv_head = config.n_kv_head
|
| 120 |
+
self.n_embd = config.n_embd
|
| 121 |
+
self.head_dim = config.head_dim
|
| 122 |
+
self.n_kv_groups = self.n_head // self.n_kv_head
|
| 123 |
+
self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
|
| 124 |
+
self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
|
| 125 |
+
self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
|
| 126 |
+
self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
|
| 127 |
+
self.ve_gate_channels = 12
|
| 128 |
+
if _has_ve(layer_idx, config.n_layer):
|
| 129 |
+
self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False)
|
| 130 |
+
kv_dim = self.n_kv_head * self.head_dim
|
| 131 |
+
ve_dim = config.ve_dim if config.ve_dim > 0 else kv_dim
|
| 132 |
+
self.ve_proj = nn.Linear(ve_dim, kv_dim, bias=False) if ve_dim < kv_dim else None
|
| 133 |
+
else:
|
| 134 |
+
self.ve_gate = None
|
| 135 |
+
self.ve_proj = None
|
| 136 |
+
|
| 137 |
+
def forward(self, x, ve, cos, sin, window_size):
|
| 138 |
+
B, T, C = x.size()
|
| 139 |
+
q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
|
| 140 |
+
k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
|
| 141 |
+
v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
|
| 142 |
+
if ve is not None and self.ve_gate is not None:
|
| 143 |
+
if self.ve_proj is not None:
|
| 144 |
+
ve = self.ve_proj(ve)
|
| 145 |
+
ve = ve.view(B, T, self.n_kv_head, self.head_dim)
|
| 146 |
+
gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
|
| 147 |
+
v = v + gate.unsqueeze(-1) * ve
|
| 148 |
+
q = _apply_rotary_emb(q, cos, sin)
|
| 149 |
+
k = _apply_rotary_emb(k, cos, sin)
|
| 150 |
+
q = _rms_norm(q) * 1.2
|
| 151 |
+
k = _rms_norm(k) * 1.2
|
| 152 |
+
q = q.transpose(1, 2)
|
| 153 |
+
k = k.transpose(1, 2)
|
| 154 |
+
v = v.transpose(1, 2)
|
| 155 |
+
if self.n_kv_groups > 1:
|
| 156 |
+
k = k.repeat_interleave(self.n_kv_groups, dim=1)
|
| 157 |
+
v = v.repeat_interleave(self.n_kv_groups, dim=1)
|
| 158 |
+
if window_size < 0 or window_size >= T:
|
| 159 |
+
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
|
| 160 |
+
else:
|
| 161 |
+
row_idx = torch.arange(T, device=x.device).unsqueeze(1)
|
| 162 |
+
col_idx = torch.arange(T, device=x.device).unsqueeze(0)
|
| 163 |
+
mask = (col_idx <= row_idx) & ((row_idx - col_idx) <= window_size)
|
| 164 |
+
y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
|
| 165 |
+
y = y.transpose(1, 2).contiguous().view(B, T, -1)
|
| 166 |
+
y = self.c_proj(y)
|
| 167 |
+
return y
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
class NanoChatMLP(nn.Module):
|
| 171 |
+
def __init__(self, config):
|
| 172 |
+
super().__init__()
|
| 173 |
+
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
|
| 174 |
+
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
|
| 175 |
+
|
| 176 |
+
def forward(self, x):
|
| 177 |
+
x = self.c_fc(x)
|
| 178 |
+
x = F.relu(x).square()
|
| 179 |
+
x = self.c_proj(x)
|
| 180 |
+
return x
|
| 181 |
+
|
| 182 |
+
|
| 183 |
+
class NanoChatBlock(nn.Module):
|
| 184 |
+
def __init__(self, config, layer_idx):
|
| 185 |
+
super().__init__()
|
| 186 |
+
self.attn = NanoChatAttention(config, layer_idx)
|
| 187 |
+
self.mlp = NanoChatMLP(config)
|
| 188 |
+
|
| 189 |
+
def forward(self, x, ve, cos, sin, window_size):
|
| 190 |
+
x = x + self.attn(_rms_norm(x), ve, cos, sin, window_size)
|
| 191 |
+
x = x + self.mlp(_rms_norm(x))
|
| 192 |
+
return x
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
class NanoChatPreTrainedModel(PreTrainedModel):
|
| 196 |
+
config_class = NanoChatConfig
|
| 197 |
+
base_model_prefix = "model"
|
| 198 |
+
supports_gradient_checkpointing = False
|
| 199 |
+
_no_split_modules = ["NanoChatBlock"]
|
| 200 |
+
|
| 201 |
+
def _init_weights(self, module):
|
| 202 |
+
pass
|
| 203 |
+
|
| 204 |
+
|
| 205 |
+
class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin):
|
| 206 |
+
def __init__(self, config):
|
| 207 |
+
super().__init__(config)
|
| 208 |
+
pvs = config.padded_vocab_size
|
| 209 |
+
self.wte = nn.Embedding(pvs, config.n_embd)
|
| 210 |
+
self.h = nn.ModuleList([NanoChatBlock(config, i) for i in range(config.n_layer)])
|
| 211 |
+
self.lm_head = nn.Linear(config.n_embd, pvs, bias=False)
|
| 212 |
+
self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
|
| 213 |
+
self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
|
| 214 |
+
self.smear_gate = nn.Linear(24, 1, bias=False)
|
| 215 |
+
self.smear_lambda = nn.Parameter(torch.zeros(1))
|
| 216 |
+
self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))
|
| 217 |
+
ve_emb_dim = config.ve_dim if config.ve_dim > 0 else config.kv_dim
|
| 218 |
+
self.value_embeds = nn.ModuleDict({
|
| 219 |
+
str(i): nn.Embedding(pvs, ve_emb_dim)
|
| 220 |
+
for i in range(config.n_layer) if _has_ve(i, config.n_layer)
|
| 221 |
+
})
|
| 222 |
+
self.window_sizes = self._compute_window_sizes(config)
|
| 223 |
+
self.rotary_seq_len = config.sequence_len * 10
|
| 224 |
+
self._rotary_initialized = False
|
| 225 |
+
self.register_buffer("cos", torch.zeros(1), persistent=False)
|
| 226 |
+
self.register_buffer("sin", torch.zeros(1), persistent=False)
|
| 227 |
+
self.post_init()
|
| 228 |
+
|
| 229 |
+
def _compute_window_sizes(self, config):
|
| 230 |
+
pattern = config.window_pattern.upper()
|
| 231 |
+
long_window = config.sequence_len
|
| 232 |
+
short_window = -(-long_window // 4 // 128) * 128
|
| 233 |
+
char_to_window = {"L": long_window, "S": short_window}
|
| 234 |
+
sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.n_layer)]
|
| 235 |
+
sizes[-1] = long_window
|
| 236 |
+
return sizes
|
| 237 |
+
|
| 238 |
+
def _precompute_rotary(self, seq_len, head_dim, base=100000.0):
|
| 239 |
+
channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32)
|
| 240 |
+
inv_freq = 1.0 / (base ** (channel_range / head_dim))
|
| 241 |
+
t = torch.arange(seq_len, dtype=torch.float32)
|
| 242 |
+
freqs = torch.outer(t, inv_freq)
|
| 243 |
+
cos = freqs.cos()[None, :, None, :]
|
| 244 |
+
sin = freqs.sin()[None, :, None, :]
|
| 245 |
+
return cos, sin
|
| 246 |
+
|
| 247 |
+
def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs):
|
| 248 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 249 |
+
B, T = input_ids.size()
|
| 250 |
+
if not self._rotary_initialized:
|
| 251 |
+
cos, sin = self._precompute_rotary(self.rotary_seq_len, self.config.head_dim)
|
| 252 |
+
self.cos = cos.to(device=input_ids.device)
|
| 253 |
+
self.sin = sin.to(device=input_ids.device)
|
| 254 |
+
self._rotary_initialized = True
|
| 255 |
+
cos = self.cos[:, :T].to(dtype=self.wte.weight.dtype)
|
| 256 |
+
sin = self.sin[:, :T].to(dtype=self.wte.weight.dtype)
|
| 257 |
+
x = self.wte(input_ids)
|
| 258 |
+
x = _rms_norm(x)
|
| 259 |
+
if T > 1:
|
| 260 |
+
gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
|
| 261 |
+
x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
|
| 262 |
+
x0 = x
|
| 263 |
+
n_layer = self.config.n_layer
|
| 264 |
+
backout_layer = n_layer // 2
|
| 265 |
+
x_backout = None
|
| 266 |
+
for i, block in enumerate(self.h):
|
| 267 |
+
x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
|
| 268 |
+
ve = self.value_embeds[str(i)](input_ids).to(x.dtype) if str(i) in self.value_embeds else None
|
| 269 |
+
x = block(x, ve, cos, sin, self.window_sizes[i])
|
| 270 |
+
if i == backout_layer:
|
| 271 |
+
x_backout = x
|
| 272 |
+
if x_backout is not None:
|
| 273 |
+
x = x - self.backout_lambda.to(x.dtype) * x_backout
|
| 274 |
+
x = _rms_norm(x)
|
| 275 |
+
logits = self.lm_head(x)
|
| 276 |
+
logits = logits[..., :self.config.vocab_size]
|
| 277 |
+
logits = logits.float()
|
| 278 |
+
if self.config.logit_softcap > 0:
|
| 279 |
+
cap = self.config.logit_softcap
|
| 280 |
+
logits = cap * torch.tanh(logits / cap)
|
| 281 |
+
loss = None
|
| 282 |
+
if labels is not None:
|
| 283 |
+
shift_logits = logits[..., :-1, :].contiguous()
|
| 284 |
+
shift_labels = labels[..., 1:].contiguous()
|
| 285 |
+
loss = F.cross_entropy(
|
| 286 |
+
shift_logits.view(-1, shift_logits.size(-1)),
|
| 287 |
+
shift_labels.view(-1),
|
| 288 |
+
ignore_index=-100,
|
| 289 |
+
)
|
| 290 |
+
if not return_dict:
|
| 291 |
+
output = (logits,)
|
| 292 |
+
return (loss,) + output if loss is not None else output
|
| 293 |
+
return CausalLMOutputWithPast(loss=loss, logits=logits)
|
| 294 |
+
|
| 295 |
+
def prepare_inputs_for_generation(self, input_ids, **kwargs):
|
| 296 |
+
return {"input_ids": input_ids}
|
| 297 |
+
|
| 298 |
+
|
| 299 |
+
# ---------------------------------------------------------------------------
|
| 300 |
+
# CLI demo
|
| 301 |
+
# ---------------------------------------------------------------------------
|
| 302 |
+
|
| 303 |
+
if __name__ == "__main__":
|
| 304 |
+
import argparse
|
| 305 |
+
import sys
|
| 306 |
+
import os
|
| 307 |
+
|
| 308 |
+
parser = argparse.ArgumentParser(description="Load and generate with a NanoChat model")
|
| 309 |
+
parser.add_argument("--model-dir", required=True, help="Path to converted model directory")
|
| 310 |
+
parser.add_argument("--prompt", default="The meaning of life is", help="Text prompt")
|
| 311 |
+
parser.add_argument("--max-tokens", type=int, default=100, help="Max tokens to generate")
|
| 312 |
+
parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature (0=greedy)")
|
| 313 |
+
parser.add_argument("--device", default="cpu", help="Device (cpu, cuda, cuda:0, ...)")
|
| 314 |
+
parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
|
| 315 |
+
args = parser.parse_args()
|
| 316 |
+
|
| 317 |
+
dtype_map = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
|
| 318 |
+
|
| 319 |
+
print(f"Loading model from {args.model_dir}...")
|
| 320 |
+
from transformers import AutoTokenizer
|
| 321 |
+
|
| 322 |
+
model = NanoChatForCausalLM.from_pretrained(args.model_dir, torch_dtype=dtype_map[args.dtype])
|
| 323 |
+
model.to(args.device)
|
| 324 |
+
model.eval()
|
| 325 |
+
tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
|
| 326 |
+
|
| 327 |
+
total_params = sum(p.numel() for p in model.parameters())
|
| 328 |
+
print(f"Model loaded: {total_params:,} parameters, dtype={args.dtype}, device={args.device}")
|
| 329 |
+
print(f"Vocab size: {model.config.vocab_size}, ve_dim: {model.config.ve_dim}")
|
| 330 |
+
print()
|
| 331 |
+
print(f"Prompt: {args.prompt}")
|
| 332 |
+
print(f"Generating (max {args.max_tokens} tokens, temp={args.temperature})...")
|
| 333 |
+
print()
|
| 334 |
+
|
| 335 |
+
# Generate
|
| 336 |
+
bos_id = tokenizer.bos_token_id
|
| 337 |
+
input_ids = tokenizer.encode(args.prompt, add_special_tokens=False)
|
| 338 |
+
if bos_id is not None:
|
| 339 |
+
input_ids = [bos_id] + input_ids
|
| 340 |
+
input_ids = torch.tensor([input_ids], dtype=torch.long, device=args.device)
|
| 341 |
+
|
| 342 |
+
generated = []
|
| 343 |
+
with torch.inference_mode():
|
| 344 |
+
for _ in range(args.max_tokens):
|
| 345 |
+
logits = model(input_ids).logits[:, -1, :]
|
| 346 |
+
if args.temperature > 0:
|
| 347 |
+
logits = logits / args.temperature
|
| 348 |
+
v, _ = torch.topk(logits, min(50, logits.size(-1)))
|
| 349 |
+
logits[logits < v[:, [-1]]] = -float("inf")
|
| 350 |
+
probs = torch.softmax(logits, dim=-1)
|
| 351 |
+
next_id = torch.multinomial(probs, num_samples=1)
|
| 352 |
+
else:
|
| 353 |
+
next_id = torch.argmax(logits, dim=-1, keepdim=True)
|
| 354 |
+
input_ids = torch.cat([input_ids, next_id], dim=1)
|
| 355 |
+
token_id = next_id.item()
|
| 356 |
+
generated.append(token_id)
|
| 357 |
+
if tokenizer.eos_token_id is not None and token_id == tokenizer.eos_token_id:
|
| 358 |
+
break
|
| 359 |
+
|
| 360 |
+
output = tokenizer.decode(generated, skip_special_tokens=True)
|
| 361 |
+
print(f"Output: {args.prompt}{output}")
|
models/full-128k-gpt4o-codegrad50-bpe-mathcode-scratch/tokenizer_config.json
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"tokenizer_class": "PreTrainedTokenizerFast",
|
| 3 |
+
"model_max_length": 2048,
|
| 4 |
+
"bos_token": "<s>",
|
| 5 |
+
"add_bos_token": false,
|
| 6 |
+
"clean_up_tokenization_spaces": false
|
| 7 |
+
}
|
models/full-128k-llama3-mathcode-scratch/config.json
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"NanoChatForCausalLM"
|
| 4 |
+
],
|
| 5 |
+
"auto_map": {
|
| 6 |
+
"AutoConfig": "nanochat_model.NanoChatConfig",
|
| 7 |
+
"AutoModelForCausalLM": "nanochat_model.NanoChatForCausalLM"
|
| 8 |
+
},
|
| 9 |
+
"bos_token_id": 128000,
|
| 10 |
+
"dtype": "float32",
|
| 11 |
+
"eos_token_id": null,
|
| 12 |
+
"head_dim": 128,
|
| 13 |
+
"kv_dim": 1536,
|
| 14 |
+
"logit_softcap": 15.0,
|
| 15 |
+
"model_type": "nanochat",
|
| 16 |
+
"n_embd": 1536,
|
| 17 |
+
"n_head": 12,
|
| 18 |
+
"n_kv_head": 12,
|
| 19 |
+
"n_layer": 24,
|
| 20 |
+
"pad_vocab_size_to": 64,
|
| 21 |
+
"padded_vocab_size": 128256,
|
| 22 |
+
"sequence_len": 2048,
|
| 23 |
+
"tie_word_embeddings": false,
|
| 24 |
+
"transformers_version": "5.3.0",
|
| 25 |
+
"ve_dim": 128,
|
| 26 |
+
"vocab_size": 128256,
|
| 27 |
+
"window_pattern": "SSSL"
|
| 28 |
+
}
|
models/full-128k-llama3-mathcode-scratch/generation_config.json
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"bos_token_id": 128000,
|
| 3 |
+
"transformers_version": "5.3.0"
|
| 4 |
+
}
|
models/full-128k-pabpe-nfc-clean-fineweb2full-plus2-capped-tuned-hybridwindow-consv2-tailcuts-gm90k-v130-sp-eng5g-mathcode-scratch/tokenizer_config.json
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"tokenizer_class": "PreTrainedTokenizerFast",
|
| 3 |
+
"model_max_length": 2048,
|
| 4 |
+
"bos_token": "<s>",
|
| 5 |
+
"add_bos_token": false,
|
| 6 |
+
"clean_up_tokenization_spaces": false
|
| 7 |
+
}
|
models/full-128k-pabpe-nfc-gpt4-balanced-mathcode-scratch/README.md
ADDED
|
@@ -0,0 +1,76 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# `full-128k-pabpe-nfc-gpt4-balanced-mathcode-scratch`
|
| 2 |
+
|
| 3 |
+
A 1.27B parameter model trained with the `pabpe-nfc-gpt4-balanced` tokenizer. Regime: mathcode-20B-from-scratch (24 layers, width 1536, 1.27B parameters, 20.0B tokens, math and code mixture, trained from scratch). Reported in: math and code ablations; TokEval.
|
| 4 |
+
|
| 5 |
+
## Load
|
| 6 |
+
|
| 7 |
+
```python
|
| 8 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
| 9 |
+
|
| 10 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 11 |
+
"cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-pabpe-nfc-gpt4-balanced-mathcode-scratch",
|
| 12 |
+
trust_remote_code=True, dtype="float32")
|
| 13 |
+
tokenizer = AutoTokenizer.from_pretrained("cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-pabpe-nfc-gpt4-balanced-mathcode-scratch")
|
| 14 |
+
|
| 15 |
+
text = "def add(a, b):\n return a + b\n"
|
| 16 |
+
ids = [127822] + tokenizer(text)["input_ids"] # training prepends BOS
|
| 17 |
+
```
|
| 18 |
+
|
| 19 |
+
## Settings
|
| 20 |
+
|
| 21 |
+
| Setting | Value |
|
| 22 |
+
|---|---|
|
| 23 |
+
| layers | 24 |
|
| 24 |
+
| model dimension | 1536 |
|
| 25 |
+
| attention heads | 12 |
|
| 26 |
+
| key and value heads | 12 |
|
| 27 |
+
| context length | 2048 |
|
| 28 |
+
| sliding window pattern | SSSL |
|
| 29 |
+
| value embedding dimension | 128 |
|
| 30 |
+
| logit soft cap | 15 |
|
| 31 |
+
| vocabulary size | 127826 |
|
| 32 |
+
| parameters | 1,271,072,522 |
|
| 33 |
+
| training step | 19073 |
|
| 34 |
+
| token budget | 19,999,490,048 |
|
| 35 |
+
| bytes consumed | 165,451,465,939 |
|
| 36 |
+
| seed | 42 |
|
| 37 |
+
| Muon matrix learning rate | 0.0200 |
|
| 38 |
+
| AdamW embedding learning rate | 0.3000 |
|
| 39 |
+
| AdamW unembedding learning rate | 0.0080 |
|
| 40 |
+
| weight decay | 0.2800 |
|
| 41 |
+
| width learning-rate exponent | -1 |
|
| 42 |
+
| weight dtype | as trained (float32 matrices, bfloat16 embeddings) |
|
| 43 |
+
| BOS token id | 127822 |
|
| 44 |
+
|
| 45 |
+
## Results
|
| 46 |
+
|
| 47 |
+
| Metric | Value |
|
| 48 |
+
|---|---|
|
| 49 |
+
| validation BPB | 0.3469 |
|
| 50 |
+
| code BPB, mean over 7 languages | 0.3878 |
|
| 51 |
+
| FLORES-200 BPB, mean | 3.3580 |
|
| 52 |
+
| FLORES-200 languages scored | 215 |
|
| 53 |
+
| BLiMP accuracy | n/a |
|
| 54 |
+
| MBPP pass@1 | 0.1300 |
|
| 55 |
+
| HumanEval pass@1 | 0.1646 |
|
| 56 |
+
| GSM8K exact match, strict | 0.2032 |
|
| 57 |
+
|
| 58 |
+
Generation metrics use specification `{'version': 'v2-2026-07-30', 'bos_context': True, 'eos_stop': True, 'token_healing': True, 'truncation': 'left_to_max_length_minus_max_gen', 'stop_truncation': 'earliest_match_all_terms', 'stop_match_decode': 'specials_kept'}`, token healing on.
|
| 59 |
+
|
| 60 |
+
Missing measurements, with the reason each is absent:
|
| 61 |
+
|
| 62 |
+
- `blimp_acc`: no full-128k-pabpe-nfc-gpt4-balanced-mathcode-scratch_blimp_code_bpb.json
|
| 63 |
+
|
| 64 |
+
BLiMP is comparable within a training scale only. Its scoring sums unnormalized sentence log-likelihoods, and under that scoring a group of quantifier subtasks drops at the 1.27B scale relative to the 300M scale.
|
| 65 |
+
|
| 66 |
+
## Registry note
|
| 67 |
+
|
| 68 |
+
#56, gpt4 pretok, slug `pabpe-nfc-gpt4-balanced`. Main-grid PA-BPE cell (first PA-BPE run at 20B). FINISHED at step 19073, 4-eval chain + code_bpb complete (verified 2026-08-09). In panel.
|
| 69 |
+
|
| 70 |
+
## Provenance
|
| 71 |
+
|
| 72 |
+
- Source checkpoint: `/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-pabpe-nfc-gpt4-balanced-mathcode-scratch`, step 19073
|
| 73 |
+
- Repository commit: `1a9a718`
|
| 74 |
+
- nanochat commit: `f919449` on branch `tokenizer-lm-patches`
|
| 75 |
+
- Tokenizer: `pabpe-nfc-gpt4-balanced`, released at `tokenizers/pabpe-nfc-gpt4-balanced`
|
| 76 |
+
|
models/full-128k-punct-balanced-bpe-seed42-mathcode-scratch/README.md
ADDED
|
@@ -0,0 +1,74 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# `full-128k-punct-balanced-bpe-seed42-mathcode-scratch`
|
| 2 |
+
|
| 3 |
+
A 1.27B parameter model trained with the `punct-balanced-bpe` tokenizer. Regime: mathcode-20B-from-scratch (24 layers, width 1536, 1.27B parameters, 20.0B tokens, math and code mixture, trained from scratch). Reported in: math and code ablations.
|
| 4 |
+
|
| 5 |
+
## Load
|
| 6 |
+
|
| 7 |
+
```python
|
| 8 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
| 9 |
+
|
| 10 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 11 |
+
"cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-punct-balanced-bpe-seed42-mathcode-scratch",
|
| 12 |
+
trust_remote_code=True, dtype="float32")
|
| 13 |
+
tokenizer = AutoTokenizer.from_pretrained("cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-punct-balanced-bpe-seed42-mathcode-scratch")
|
| 14 |
+
|
| 15 |
+
text = "def add(a, b):\n return a + b\n"
|
| 16 |
+
ids = [0] + tokenizer(text)["input_ids"] # training prepends BOS
|
| 17 |
+
```
|
| 18 |
+
|
| 19 |
+
## Settings
|
| 20 |
+
|
| 21 |
+
| Setting | Value |
|
| 22 |
+
|---|---|
|
| 23 |
+
| layers | 24 |
|
| 24 |
+
| model dimension | 1536 |
|
| 25 |
+
| attention heads | 12 |
|
| 26 |
+
| key and value heads | 12 |
|
| 27 |
+
| context length | 2048 |
|
| 28 |
+
| sliding window pattern | SSSL |
|
| 29 |
+
| value embedding dimension | 128 |
|
| 30 |
+
| logit soft cap | 15 |
|
| 31 |
+
| vocabulary size | 128256 |
|
| 32 |
+
| parameters | 1,272,841,994 |
|
| 33 |
+
| training step | 19073 |
|
| 34 |
+
| token budget | 19,999,490,048 |
|
| 35 |
+
| bytes consumed | 152,301,967,549 |
|
| 36 |
+
| seed | 42 |
|
| 37 |
+
| Muon matrix learning rate | 0.0200 |
|
| 38 |
+
| AdamW embedding learning rate | 0.3000 |
|
| 39 |
+
| AdamW unembedding learning rate | 0.0080 |
|
| 40 |
+
| weight decay | 0.2800 |
|
| 41 |
+
| width learning-rate exponent | -1 |
|
| 42 |
+
| weight dtype | as trained (float32 matrices, bfloat16 embeddings) |
|
| 43 |
+
| BOS token id | 0 |
|
| 44 |
+
|
| 45 |
+
## Results
|
| 46 |
+
|
| 47 |
+
| Metric | Value |
|
| 48 |
+
|---|---|
|
| 49 |
+
| validation BPB | 0.3403 |
|
| 50 |
+
| code BPB, mean over 7 languages | n/a |
|
| 51 |
+
| FLORES-200 BPB, mean | 3.3228 |
|
| 52 |
+
| FLORES-200 languages scored | 215 |
|
| 53 |
+
| BLiMP accuracy | n/a |
|
| 54 |
+
| MBPP pass@1 | 0.1920 |
|
| 55 |
+
| HumanEval pass@1 | 0.1707 |
|
| 56 |
+
| GSM8K exact match, strict | n/a |
|
| 57 |
+
|
| 58 |
+
Generation metrics use specification `{'version': 'v2-2026-07-30', 'bos_context': True, 'eos_stop': True, 'token_healing': True, 'truncation': 'left_to_max_length_minus_max_gen', 'stop_truncation': 'earliest_match_all_terms', 'stop_match_decode': 'specials_kept'}`, token healing on.
|
| 59 |
+
|
| 60 |
+
Missing measurements, with the reason each is absent:
|
| 61 |
+
|
| 62 |
+
- `blimp_acc`: no full-128k-punct-balanced-bpe-seed42-mathcode-scratch_blimp_code_bpb.json
|
| 63 |
+
- `code_bpb_mean`: no full-128k-punct-balanced-bpe-seed42-mathcode-scratch_blimp_code_bpb.json; no full-128k-punct-balanced-bpe-seed42-mathcode-scratch_code_bpb.json
|
| 64 |
+
- `gsm8k_exact_match_strict`: no full-128k-punct-balanced-bpe-seed42-mathcode-scratch_gsm8k_v2.json
|
| 65 |
+
|
| 66 |
+
BLiMP is comparable within a training scale only. Its scoring sums unnormalized sentence log-likelihoods, and under that scoring a group of quantifier subtasks drops at the 1.27B scale relative to the 300M scale.
|
| 67 |
+
|
| 68 |
+
## Provenance
|
| 69 |
+
|
| 70 |
+
- Source checkpoint: `/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-punct-balanced-bpe-seed42-mathcode-scratch`, step 19073
|
| 71 |
+
- Repository commit: `1a9a718`
|
| 72 |
+
- nanochat commit: `f919449` on branch `tokenizer-lm-patches`
|
| 73 |
+
- Tokenizer: `punct-balanced-bpe`, released at `tokenizers/punct-balanced-bpe`
|
| 74 |
+
|
models/full-128k-punct-balanced-bpe-seed42-mathcode-scratch/generation_config.json
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"bos_token_id": 0,
|
| 3 |
+
"transformers_version": "5.3.0"
|
| 4 |
+
}
|
models/full-128k-punct-english-bpe-mathcode-scratch/conversion_metadata.json
ADDED
|
@@ -0,0 +1,250 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"run": "full-128k-punct-english-bpe-mathcode-scratch",
|
| 3 |
+
"source_checkpoint": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-punct-english-bpe-mathcode-scratch",
|
| 4 |
+
"source_model_file": "model_019073.pt",
|
| 5 |
+
"source_step": 19073,
|
| 6 |
+
"source_meta": {
|
| 7 |
+
"step": 19073,
|
| 8 |
+
"val_bpb": 0.3434937319193573,
|
| 9 |
+
"model_config": {
|
| 10 |
+
"sequence_len": 2048,
|
| 11 |
+
"vocab_size": 128256,
|
| 12 |
+
"n_layer": 24,
|
| 13 |
+
"n_head": 12,
|
| 14 |
+
"n_kv_head": 12,
|
| 15 |
+
"n_embd": 1536,
|
| 16 |
+
"window_pattern": "SSSL",
|
| 17 |
+
"ve_dim": 128,
|
| 18 |
+
"logit_softcap": 15.0
|
| 19 |
+
},
|
| 20 |
+
"tokenizer_path": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/tokenizers/punct-english-bpe",
|
| 21 |
+
"training_config": {
|
| 22 |
+
"name": "full-128k-mathcode-scratch",
|
| 23 |
+
"model": {
|
| 24 |
+
"n_layer": 24,
|
| 25 |
+
"n_embd": 1536,
|
| 26 |
+
"n_kv_head": 12,
|
| 27 |
+
"sequence_len": 2048,
|
| 28 |
+
"window_pattern": "SSSL",
|
| 29 |
+
"ve_dim": 128
|
| 30 |
+
},
|
| 31 |
+
"training": {
|
| 32 |
+
"target_param_data_ratio": -1,
|
| 33 |
+
"max_tokens": 20000000000,
|
| 34 |
+
"num_iterations": -1,
|
| 35 |
+
"total_batch_size": 1048576,
|
| 36 |
+
"device_batch_size": 16,
|
| 37 |
+
"warmup_steps": 40,
|
| 38 |
+
"warmdown_ratio": 0.65,
|
| 39 |
+
"final_lr_frac": 0.05,
|
| 40 |
+
"eval_tokens": 41943040,
|
| 41 |
+
"compile": true
|
| 42 |
+
},
|
| 43 |
+
"optimizer": {
|
| 44 |
+
"matrix_lr": 0.02,
|
| 45 |
+
"embedding_lr": 0.3,
|
| 46 |
+
"unembedding_lr": 0.008,
|
| 47 |
+
"scalar_lr": 0.5,
|
| 48 |
+
"weight_decay": 0.28,
|
| 49 |
+
"width_lr_exponent": -1.0
|
| 50 |
+
},
|
| 51 |
+
"data": {
|
| 52 |
+
"data_dir": null
|
| 53 |
+
},
|
| 54 |
+
"logging": {
|
| 55 |
+
"eval_every": 500,
|
| 56 |
+
"save_every": 500,
|
| 57 |
+
"max_checkpoints": 0
|
| 58 |
+
},
|
| 59 |
+
"wandb": {
|
| 60 |
+
"enabled": true,
|
| 61 |
+
"project": "tokenizer-lm-experiments",
|
| 62 |
+
"run_name": null
|
| 63 |
+
},
|
| 64 |
+
"output_dir": "/capstor/scratch/cscs/${USER}/tokenizer-lm/checkpoints",
|
| 65 |
+
"seed": 42
|
| 66 |
+
},
|
| 67 |
+
"code_provenance": {
|
| 68 |
+
"commit": "6231791054e9c0df5144025df8f8a2117e8b78d2",
|
| 69 |
+
"branch": "main",
|
| 70 |
+
"dirty": true,
|
| 71 |
+
"dirty_files": [
|
| 72 |
+
"M .gitignore",
|
| 73 |
+
" M EXPERIMENTAL_SETUP.md",
|
| 74 |
+
" M RESULT_COMPUTATION_ANALYSIS.md",
|
| 75 |
+
" M eval_runner.py",
|
| 76 |
+
" M mc_evaluation",
|
| 77 |
+
" M scripts/add_mc_math_ci.py",
|
| 78 |
+
" M scripts/bootstrap_mathcode_significance.py",
|
| 79 |
+
" M scripts/correlate_intrinsic_downstream.py",
|
| 80 |
+
" M scripts/correlate_mathcode_intrinsic_downstream.py",
|
| 81 |
+
" M scripts/cross_lingual_analyses.py",
|
| 82 |
+
" M scripts/generate_latex_tables.py",
|
| 83 |
+
" M scripts/launch_mathcode_scratch_chained.sh",
|
| 84 |
+
" M scripts/launch_new6_all_scales.sh",
|
| 85 |
+
" M scripts/reorganize_checkpoints.py",
|
| 86 |
+
" M scripts/slurm_3b_mathcode_resume.sh",
|
| 87 |
+
" M scripts/slurm_eval_mgsm.sh",
|
| 88 |
+
" M scripts/slurm_full.sh",
|
| 89 |
+
" M scripts/slurm_full_1node.sh",
|
| 90 |
+
" M scripts/slurm_full_1node_resume.sh",
|
| 91 |
+
" M scripts/slurm_mathcode_resume_1node.sh",
|
| 92 |
+
" M scripts/slurm_pilot.sh",
|
| 93 |
+
" M scripts/slurm_train_1node.sh",
|
| 94 |
+
" M scripts/slurm_train_tokenizers.sh",
|
| 95 |
+
" M scripts/train_custom_tokenizers.py",
|
| 96 |
+
" M tokenizer_lm/tokenizer.py",
|
| 97 |
+
" M train.py",
|
| 98 |
+
"?? REPRODUCIBILITY_STATUS.md",
|
| 99 |
+
"?? SESSION_STATUS.md",
|
| 100 |
+
"?? TOKEN_HEALING_HANDOFF.md",
|
| 101 |
+
"?? configs/full_128k_bytematched.yaml",
|
| 102 |
+
"?? configs/full_128k_clean_v128k.yaml",
|
| 103 |
+
"?? configs/full_128k_clean_v160k.yaml",
|
| 104 |
+
"?? configs/full_128k_clean_v200k.yaml",
|
| 105 |
+
"?? configs/full_128k_clean_v64k.yaml",
|
| 106 |
+
"?? configs/full_128k_clean_v96k.yaml",
|
| 107 |
+
"?? configs/panels/",
|
| 108 |
+
"?? configs/starcoder_ast_18lang.json",
|
| 109 |
+
"?? correspondence/",
|
| 110 |
+
"?? scripts/_lexer_treesitter_worker.py",
|
| 111 |
+
"?? scripts/add_code_bpb_std.py",
|
| 112 |
+
"?? scripts/analyze_apertus_track_sensitivity.py",
|
| 113 |
+
"?? scripts/analyze_axes2_codegrad.py",
|
| 114 |
+
"?? scripts/analyze_code_boundary_drift.py",
|
| 115 |
+
"?? scripts/analyze_code_gen_failures.py",
|
| 116 |
+
"?? scripts/analyze_code_intrinsic_downstream.py",
|
| 117 |
+
"?? scripts/analyze_cr_byte_drop.py",
|
| 118 |
+
"?? scripts/analyze_cross_scale_ranking.py",
|
| 119 |
+
"?? scripts/analyze_equal_weighting.py",
|
| 120 |
+
"?? scripts/analyze_gen_ablation.py",
|
| 121 |
+
"?? scripts/analyze_seed_variance.py",
|
| 122 |
+
"?? scripts/analyze_seed_variance_mc20b.py",
|
| 123 |
+
"?? scripts/analyze_vocab_size.py",
|
| 124 |
+
"?? scripts/ast_node_discovery.py",
|
| 125 |
+
"?? scripts/ast_scope_extractor.py",
|
| 126 |
+
"?? scripts/audit_bpb_byte_bias.py",
|
| 127 |
+
"?? scripts/audit_corpus_parseability.py",
|
| 128 |
+
"?? scripts/audit_generation_symptoms.py",
|
| 129 |
+
"?? scripts/audit_lexer_boundary_leak.py",
|
| 130 |
+
"?? scripts/audit_logit_softcap.py",
|
| 131 |
+
"?? scripts/audit_maxchars_fragmentation.py",
|
| 132 |
+
"?? scripts/audit_mgsm_prompt_boundary.py",
|
| 133 |
+
"?? scripts/audit_newline_tokens.py",
|
| 134 |
+
"?? scripts/audit_newline_tokens_fullbyte_pair.py",
|
| 135 |
+
"?? scripts/audit_newline_tokens_new11.py",
|
| 136 |
+
"?? scripts/backfill_flores_summary.py",
|
| 137 |
+
"?? scripts/build_axes2_codegrad_tokenizers.sh",
|
| 138 |
+
"?? scripts/build_vocab_ladder_tokenizers.sh",
|
| 139 |
+
"?? scripts/build_whitespace_ladder_tokenizers.sh",
|
| 140 |
+
"?? scripts/characterize_axes2_intrinsic.py",
|
| 141 |
+
"?? scripts/check_hf_dataset_cache.py",
|
| 142 |
+
"?? scripts/check_number_provenance.py",
|
| 143 |
+
"?? scripts/compare_bigram_entropy_variants.py",
|
| 144 |
+
"?? scripts/compute_headline_spreads.py",
|
| 145 |
+
"?? scripts/compute_vocab_ladder_byte_matching.py",
|
| 146 |
+
"?? scripts/correlate_operator_domains.py",
|
| 147 |
+
"?? scripts/d2_candidate_screen.py",
|
| 148 |
+
"?? scripts/diagnose_fullbyte_indent.py",
|
| 149 |
+
"?? scripts/diagnose_operator_isolation.py",
|
| 150 |
+
"?? scripts/dump_axes2_composition_weights.py",
|
| 151 |
+
"?? scripts/estimate_intrinsic_cost.py",
|
| 152 |
+
"?? scripts/fit_quota_response_model.py",
|
| 153 |
+
"?? scripts/gen_v2_spec_audit.py",
|
| 154 |
+
"?? scripts/generate_intrinsic_tables.py",
|
| 155 |
+
"?? scripts/launch_boundary_retrain.sh",
|
| 156 |
+
"?? scripts/launch_mc_pretok_chained.sh",
|
| 157 |
+
"?? scripts/launch_std_pretok_chained.sh",
|
| 158 |
+
"?? scripts/launch_vocab_ladder.sh",
|
| 159 |
+
"?? scripts/launch_vocab_ladder_paramscaled.sh",
|
| 160 |
+
"?? scripts/lexer_pretok_common.py",
|
| 161 |
+
"?? scripts/lexer_pretok_corpus.py",
|
| 162 |
+
"?? scripts/mbpp_mediation.py",
|
| 163 |
+
"?? scripts/merge_intrinsic_components.py",
|
| 164 |
+
"?? scripts/merge_intrinsic_shards.py",
|
| 165 |
+
"?? scripts/operator_isolation_by_corpus_n36.py",
|
| 166 |
+
"?? scripts/optimize_quotas.py",
|
| 167 |
+
"?? scripts/preflight_gen_v2_tokenizers.py",
|
| 168 |
+
"?? scripts/probe_belebele_bos.py",
|
| 169 |
+
"?? scripts/probe_eot_stop_incidence.py",
|
| 170 |
+
"?? scripts/probe_healing_prefill_equivalence.py",
|
| 171 |
+
"?? scripts/recheck_seed_sd_verdicts.py",
|
| 172 |
+
"?? scripts/recompute_val_bpb.py",
|
| 173 |
+
"?? scripts/rejudge_tail_penalty.py",
|
| 174 |
+
"?? scripts/render_code_appendix.py",
|
| 175 |
+
"?? scripts/retry_whitespace_axes2_builds.sh",
|
| 176 |
+
"?? scripts/run_gen_v2_one.sh",
|
| 177 |
+
"?? scripts/run_mc_panels.sh",
|
| 178 |
+
"?? scripts/script_pretok_per_language.py",
|
| 179 |
+
"?? scripts/slurm_corpus_parseability_audit.sh",
|
| 180 |
+
"?? scripts/slurm_d2_candidate_build.sh",
|
| 181 |
+
"?? scripts/slurm_d2_candidate_screen.sh",
|
| 182 |
+
"?? scripts/slurm_eval_code_gen_healed.sh",
|
| 183 |
+
"?? scripts/slurm_eval_gen_ablation.sh",
|
| 184 |
+
"?? scripts/slurm_eval_gen_v2.sh",
|
| 185 |
+
"?? scripts/slurm_eval_gen_v2_packed.sbatch",
|
| 186 |
+
"?? scripts/slurm_full_1node_resume_pretok.sh",
|
| 187 |
+
"?? scripts/slurm_intrinsic_camera_ready.sbatch",
|
| 188 |
+
"?? scripts/slurm_intrinsic_component.sbatch",
|
| 189 |
+
"?? scripts/slurm_probe_mgsm_bos.sbatch",
|
| 190 |
+
"?? scripts/slurm_train_pretok_1node.sh",
|
| 191 |
+
"?? scripts/submit_gen_v2_sweep.sh",
|
| 192 |
+
"?? scripts/submit_lexer_corpus_cache_jobs.sh",
|
| 193 |
+
"?? scripts/summarize_lexer_smoke_run.py",
|
| 194 |
+
"?? scripts/test_ast_scope_extractor.py",
|
| 195 |
+
"?? scripts/test_lexer_pretok_common.py",
|
| 196 |
+
"?? scripts/test_lexer_treesitter_worker.py",
|
| 197 |
+
"?? scripts/test_presegmented_training_path.py",
|
| 198 |
+
"?? scripts/token_healing_failure_modes.py",
|
| 199 |
+
"?? scripts/validate_axes2_codegrad_tokenizers.py",
|
| 200 |
+
"?? scripts/validate_generation_conformance.py",
|
| 201 |
+
"?? scripts/validate_gpt4o_english_fullbyte_bpe.py",
|
| 202 |
+
"?? scripts/validate_token_healing.py",
|
| 203 |
+
"?? scripts/validate_whitespace_ladder_regexes.py",
|
| 204 |
+
"?? scripts/validate_whitespace_ladder_tokenizers.py",
|
| 205 |
+
"?? scripts/verify_gen_v2_smoke.py",
|
| 206 |
+
"?? scripts/verify_iter_source_hardening.py",
|
| 207 |
+
"?? scripts/vet_runtime_regex.py",
|
| 208 |
+
"?? scripts/whitespace_fusion_audit.py",
|
| 209 |
+
"?? std1b_invalid.md",
|
| 210 |
+
"?? std1b_used.md",
|
| 211 |
+
"?? tokenizer_lm/pretok_dataloader.py",
|
| 212 |
+
"?? tokenizer_lm/token_healing.py"
|
| 213 |
+
],
|
| 214 |
+
"nanochat_commit": "f91944934f78d2083c8f5bd6ea8ab4e048786054",
|
| 215 |
+
"torch_version": "2.10.0+cu126"
|
| 216 |
+
},
|
| 217 |
+
"total_batch_size": 1048576,
|
| 218 |
+
"device_batch_size": 16,
|
| 219 |
+
"max_seq_len": 2048,
|
| 220 |
+
"dataloader_state_dict": {
|
| 221 |
+
"pq_idx": 152,
|
| 222 |
+
"rg_idx": 40,
|
| 223 |
+
"epoch": 2
|
| 224 |
+
},
|
| 225 |
+
"loop_state": {
|
| 226 |
+
"min_val_bpb": 0.3434937319193573,
|
| 227 |
+
"smooth_train_loss": 0.6776276294672776,
|
| 228 |
+
"total_training_time": 94782.62604355812,
|
| 229 |
+
"total_bytes_consumed": 141498848692
|
| 230 |
+
}
|
| 231 |
+
},
|
| 232 |
+
"weight_dtype": "as trained (float32 matrices, bfloat16 embeddings)",
|
| 233 |
+
"weight_dtype_counts": {
|
| 234 |
+
"torch.float32": 174,
|
| 235 |
+
"torch.bfloat16": 13
|
| 236 |
+
},
|
| 237 |
+
"total_parameters": 1272841994,
|
| 238 |
+
"logit_softcap_source": "meta",
|
| 239 |
+
"bos_token": "<s>",
|
| 240 |
+
"bos_token_id": 0,
|
| 241 |
+
"bos_resolved_from": "packaged tokenizer",
|
| 242 |
+
"tokenizer_slug": "punct-english-bpe",
|
| 243 |
+
"tokenizer_redistributed": true,
|
| 244 |
+
"tokenizer_hub_id": null,
|
| 245 |
+
"weight_key_mapping": "strip _orig_mod. and transformer. prefixes",
|
| 246 |
+
"converted_by": "scripts/hf_release/convert_for_release.py",
|
| 247 |
+
"repo_commit": "1a9a718",
|
| 248 |
+
"nanochat_commit": "f919449",
|
| 249 |
+
"nanochat_branch": "tokenizer-lm-patches"
|
| 250 |
+
}
|
models/full-128k-punct-english-bpe-mathcode-scratch/nanochat_model.py
ADDED
|
@@ -0,0 +1,361 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
NanoChat model for HuggingFace Transformers — self-contained single file.
|
| 3 |
+
|
| 4 |
+
Includes the config class, model class, and a CLI demo for loading + generation.
|
| 5 |
+
Drop this file into a directory containing config.json, model.safetensors,
|
| 6 |
+
tokenizer.json, and tokenizer_config.json, then run:
|
| 7 |
+
|
| 8 |
+
python nanochat_model.py --model-dir . --prompt "Hello world"
|
| 9 |
+
|
| 10 |
+
Or load programmatically:
|
| 11 |
+
|
| 12 |
+
from nanochat_model import NanoChatConfig, NanoChatForCausalLM
|
| 13 |
+
model = NanoChatForCausalLM.from_pretrained("path/to/model", torch_dtype=torch.bfloat16)
|
| 14 |
+
|
| 15 |
+
Requirements: pip install torch transformers tokenizers
|
| 16 |
+
|
| 17 |
+
Architecture: nanochat GPT (Karpathy, MIT licensed) with value embeddings
|
| 18 |
+
(ResFormer-style), per-layer residual/x0 lambdas, smear, backout, logit
|
| 19 |
+
soft-capping, sliding window attention, QK norm, ReLU^2 MLP, and RoPE.
|
| 20 |
+
|
| 21 |
+
The models were trained with a patched nanochat, not upstream nanochat:
|
| 22 |
+
https://github.com/cimeister/nanochat branch tokenizer-lm-patches, commit f919449.
|
| 23 |
+
The patches touch two files and change three things, two of which affect the forward
|
| 24 |
+
pass and are implemented here:
|
| 25 |
+
1. ve_dim: the value-embedding tables are built at a reduced dimension and projected
|
| 26 |
+
up to kv_dim by a learned ve_proj, instead of being built at full kv_dim. All
|
| 27 |
+
released models use ve_dim=128.
|
| 28 |
+
2. logit_softcap: the soft cap is a config field with 0 meaning disabled, instead of
|
| 29 |
+
a hardcoded 15.
|
| 30 |
+
3. width_lr_exponent: the AdamW learning-rate width scaling exponent is configurable
|
| 31 |
+
for muP. This affects training only and has no inference-time effect.
|
| 32 |
+
|
| 33 |
+
This file is the released model definition. Its agreement with the training-time model
|
| 34 |
+
is checked per released checkpoint by scripts/hf_release/verify_conversion.py, which
|
| 35 |
+
compares logits against the patched nanochat GPT in float32.
|
| 36 |
+
"""
|
| 37 |
+
|
| 38 |
+
from typing import Optional, Tuple, Union
|
| 39 |
+
|
| 40 |
+
import torch
|
| 41 |
+
import torch.nn as nn
|
| 42 |
+
import torch.nn.functional as F
|
| 43 |
+
from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
|
| 44 |
+
from transformers.modeling_outputs import CausalLMOutputWithPast
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
# ---------------------------------------------------------------------------
|
| 48 |
+
# Config
|
| 49 |
+
# ---------------------------------------------------------------------------
|
| 50 |
+
|
| 51 |
+
class NanoChatConfig(PretrainedConfig):
|
| 52 |
+
model_type = "nanochat"
|
| 53 |
+
|
| 54 |
+
def __init__(
|
| 55 |
+
self,
|
| 56 |
+
vocab_size=128256,
|
| 57 |
+
n_embd=1536,
|
| 58 |
+
n_layer=24,
|
| 59 |
+
n_head=12,
|
| 60 |
+
n_kv_head=12,
|
| 61 |
+
sequence_len=2048,
|
| 62 |
+
window_pattern="SSSL",
|
| 63 |
+
ve_dim=0,
|
| 64 |
+
logit_softcap=15.0,
|
| 65 |
+
pad_vocab_size_to=64,
|
| 66 |
+
bos_token_id=0,
|
| 67 |
+
eos_token_id=None,
|
| 68 |
+
tie_word_embeddings=False,
|
| 69 |
+
**kwargs,
|
| 70 |
+
):
|
| 71 |
+
self.vocab_size = vocab_size
|
| 72 |
+
self.n_embd = n_embd
|
| 73 |
+
self.n_layer = n_layer
|
| 74 |
+
self.n_head = n_head
|
| 75 |
+
self.n_kv_head = n_kv_head
|
| 76 |
+
self.sequence_len = sequence_len
|
| 77 |
+
self.window_pattern = window_pattern
|
| 78 |
+
self.ve_dim = ve_dim
|
| 79 |
+
self.logit_softcap = logit_softcap
|
| 80 |
+
self.pad_vocab_size_to = pad_vocab_size_to
|
| 81 |
+
self.head_dim = n_embd // n_head
|
| 82 |
+
self.kv_dim = n_kv_head * self.head_dim
|
| 83 |
+
self.padded_vocab_size = (
|
| 84 |
+
(vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to
|
| 85 |
+
) * pad_vocab_size_to
|
| 86 |
+
super().__init__(
|
| 87 |
+
bos_token_id=bos_token_id,
|
| 88 |
+
eos_token_id=eos_token_id,
|
| 89 |
+
tie_word_embeddings=tie_word_embeddings,
|
| 90 |
+
**kwargs,
|
| 91 |
+
)
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
# ---------------------------------------------------------------------------
|
| 95 |
+
# Model
|
| 96 |
+
# ---------------------------------------------------------------------------
|
| 97 |
+
|
| 98 |
+
def _has_ve(layer_idx: int, n_layer: int) -> bool:
|
| 99 |
+
return layer_idx % 2 == (n_layer - 1) % 2
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def _apply_rotary_emb(x, cos, sin):
|
| 103 |
+
d = x.shape[-1] // 2
|
| 104 |
+
x1, x2 = x[..., :d], x[..., d:]
|
| 105 |
+
y1 = x1 * cos + x2 * sin
|
| 106 |
+
y2 = x1 * (-sin) + x2 * cos
|
| 107 |
+
return torch.cat([y1, y2], dim=-1)
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
def _rms_norm(x):
|
| 111 |
+
return F.rms_norm(x, (x.size(-1),))
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
class NanoChatAttention(nn.Module):
|
| 115 |
+
def __init__(self, config, layer_idx):
|
| 116 |
+
super().__init__()
|
| 117 |
+
self.layer_idx = layer_idx
|
| 118 |
+
self.n_head = config.n_head
|
| 119 |
+
self.n_kv_head = config.n_kv_head
|
| 120 |
+
self.n_embd = config.n_embd
|
| 121 |
+
self.head_dim = config.head_dim
|
| 122 |
+
self.n_kv_groups = self.n_head // self.n_kv_head
|
| 123 |
+
self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
|
| 124 |
+
self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
|
| 125 |
+
self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
|
| 126 |
+
self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
|
| 127 |
+
self.ve_gate_channels = 12
|
| 128 |
+
if _has_ve(layer_idx, config.n_layer):
|
| 129 |
+
self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False)
|
| 130 |
+
kv_dim = self.n_kv_head * self.head_dim
|
| 131 |
+
ve_dim = config.ve_dim if config.ve_dim > 0 else kv_dim
|
| 132 |
+
self.ve_proj = nn.Linear(ve_dim, kv_dim, bias=False) if ve_dim < kv_dim else None
|
| 133 |
+
else:
|
| 134 |
+
self.ve_gate = None
|
| 135 |
+
self.ve_proj = None
|
| 136 |
+
|
| 137 |
+
def forward(self, x, ve, cos, sin, window_size):
|
| 138 |
+
B, T, C = x.size()
|
| 139 |
+
q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
|
| 140 |
+
k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
|
| 141 |
+
v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
|
| 142 |
+
if ve is not None and self.ve_gate is not None:
|
| 143 |
+
if self.ve_proj is not None:
|
| 144 |
+
ve = self.ve_proj(ve)
|
| 145 |
+
ve = ve.view(B, T, self.n_kv_head, self.head_dim)
|
| 146 |
+
gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
|
| 147 |
+
v = v + gate.unsqueeze(-1) * ve
|
| 148 |
+
q = _apply_rotary_emb(q, cos, sin)
|
| 149 |
+
k = _apply_rotary_emb(k, cos, sin)
|
| 150 |
+
q = _rms_norm(q) * 1.2
|
| 151 |
+
k = _rms_norm(k) * 1.2
|
| 152 |
+
q = q.transpose(1, 2)
|
| 153 |
+
k = k.transpose(1, 2)
|
| 154 |
+
v = v.transpose(1, 2)
|
| 155 |
+
if self.n_kv_groups > 1:
|
| 156 |
+
k = k.repeat_interleave(self.n_kv_groups, dim=1)
|
| 157 |
+
v = v.repeat_interleave(self.n_kv_groups, dim=1)
|
| 158 |
+
if window_size < 0 or window_size >= T:
|
| 159 |
+
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
|
| 160 |
+
else:
|
| 161 |
+
row_idx = torch.arange(T, device=x.device).unsqueeze(1)
|
| 162 |
+
col_idx = torch.arange(T, device=x.device).unsqueeze(0)
|
| 163 |
+
mask = (col_idx <= row_idx) & ((row_idx - col_idx) <= window_size)
|
| 164 |
+
y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
|
| 165 |
+
y = y.transpose(1, 2).contiguous().view(B, T, -1)
|
| 166 |
+
y = self.c_proj(y)
|
| 167 |
+
return y
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
class NanoChatMLP(nn.Module):
|
| 171 |
+
def __init__(self, config):
|
| 172 |
+
super().__init__()
|
| 173 |
+
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
|
| 174 |
+
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
|
| 175 |
+
|
| 176 |
+
def forward(self, x):
|
| 177 |
+
x = self.c_fc(x)
|
| 178 |
+
x = F.relu(x).square()
|
| 179 |
+
x = self.c_proj(x)
|
| 180 |
+
return x
|
| 181 |
+
|
| 182 |
+
|
| 183 |
+
class NanoChatBlock(nn.Module):
|
| 184 |
+
def __init__(self, config, layer_idx):
|
| 185 |
+
super().__init__()
|
| 186 |
+
self.attn = NanoChatAttention(config, layer_idx)
|
| 187 |
+
self.mlp = NanoChatMLP(config)
|
| 188 |
+
|
| 189 |
+
def forward(self, x, ve, cos, sin, window_size):
|
| 190 |
+
x = x + self.attn(_rms_norm(x), ve, cos, sin, window_size)
|
| 191 |
+
x = x + self.mlp(_rms_norm(x))
|
| 192 |
+
return x
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
class NanoChatPreTrainedModel(PreTrainedModel):
|
| 196 |
+
config_class = NanoChatConfig
|
| 197 |
+
base_model_prefix = "model"
|
| 198 |
+
supports_gradient_checkpointing = False
|
| 199 |
+
_no_split_modules = ["NanoChatBlock"]
|
| 200 |
+
|
| 201 |
+
def _init_weights(self, module):
|
| 202 |
+
pass
|
| 203 |
+
|
| 204 |
+
|
| 205 |
+
class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin):
|
| 206 |
+
def __init__(self, config):
|
| 207 |
+
super().__init__(config)
|
| 208 |
+
pvs = config.padded_vocab_size
|
| 209 |
+
self.wte = nn.Embedding(pvs, config.n_embd)
|
| 210 |
+
self.h = nn.ModuleList([NanoChatBlock(config, i) for i in range(config.n_layer)])
|
| 211 |
+
self.lm_head = nn.Linear(config.n_embd, pvs, bias=False)
|
| 212 |
+
self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
|
| 213 |
+
self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
|
| 214 |
+
self.smear_gate = nn.Linear(24, 1, bias=False)
|
| 215 |
+
self.smear_lambda = nn.Parameter(torch.zeros(1))
|
| 216 |
+
self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))
|
| 217 |
+
ve_emb_dim = config.ve_dim if config.ve_dim > 0 else config.kv_dim
|
| 218 |
+
self.value_embeds = nn.ModuleDict({
|
| 219 |
+
str(i): nn.Embedding(pvs, ve_emb_dim)
|
| 220 |
+
for i in range(config.n_layer) if _has_ve(i, config.n_layer)
|
| 221 |
+
})
|
| 222 |
+
self.window_sizes = self._compute_window_sizes(config)
|
| 223 |
+
self.rotary_seq_len = config.sequence_len * 10
|
| 224 |
+
self._rotary_initialized = False
|
| 225 |
+
self.register_buffer("cos", torch.zeros(1), persistent=False)
|
| 226 |
+
self.register_buffer("sin", torch.zeros(1), persistent=False)
|
| 227 |
+
self.post_init()
|
| 228 |
+
|
| 229 |
+
def _compute_window_sizes(self, config):
|
| 230 |
+
pattern = config.window_pattern.upper()
|
| 231 |
+
long_window = config.sequence_len
|
| 232 |
+
short_window = -(-long_window // 4 // 128) * 128
|
| 233 |
+
char_to_window = {"L": long_window, "S": short_window}
|
| 234 |
+
sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.n_layer)]
|
| 235 |
+
sizes[-1] = long_window
|
| 236 |
+
return sizes
|
| 237 |
+
|
| 238 |
+
def _precompute_rotary(self, seq_len, head_dim, base=100000.0):
|
| 239 |
+
channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32)
|
| 240 |
+
inv_freq = 1.0 / (base ** (channel_range / head_dim))
|
| 241 |
+
t = torch.arange(seq_len, dtype=torch.float32)
|
| 242 |
+
freqs = torch.outer(t, inv_freq)
|
| 243 |
+
cos = freqs.cos()[None, :, None, :]
|
| 244 |
+
sin = freqs.sin()[None, :, None, :]
|
| 245 |
+
return cos, sin
|
| 246 |
+
|
| 247 |
+
def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs):
|
| 248 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 249 |
+
B, T = input_ids.size()
|
| 250 |
+
if not self._rotary_initialized:
|
| 251 |
+
cos, sin = self._precompute_rotary(self.rotary_seq_len, self.config.head_dim)
|
| 252 |
+
self.cos = cos.to(device=input_ids.device)
|
| 253 |
+
self.sin = sin.to(device=input_ids.device)
|
| 254 |
+
self._rotary_initialized = True
|
| 255 |
+
cos = self.cos[:, :T].to(dtype=self.wte.weight.dtype)
|
| 256 |
+
sin = self.sin[:, :T].to(dtype=self.wte.weight.dtype)
|
| 257 |
+
x = self.wte(input_ids)
|
| 258 |
+
x = _rms_norm(x)
|
| 259 |
+
if T > 1:
|
| 260 |
+
gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
|
| 261 |
+
x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
|
| 262 |
+
x0 = x
|
| 263 |
+
n_layer = self.config.n_layer
|
| 264 |
+
backout_layer = n_layer // 2
|
| 265 |
+
x_backout = None
|
| 266 |
+
for i, block in enumerate(self.h):
|
| 267 |
+
x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
|
| 268 |
+
ve = self.value_embeds[str(i)](input_ids).to(x.dtype) if str(i) in self.value_embeds else None
|
| 269 |
+
x = block(x, ve, cos, sin, self.window_sizes[i])
|
| 270 |
+
if i == backout_layer:
|
| 271 |
+
x_backout = x
|
| 272 |
+
if x_backout is not None:
|
| 273 |
+
x = x - self.backout_lambda.to(x.dtype) * x_backout
|
| 274 |
+
x = _rms_norm(x)
|
| 275 |
+
logits = self.lm_head(x)
|
| 276 |
+
logits = logits[..., :self.config.vocab_size]
|
| 277 |
+
logits = logits.float()
|
| 278 |
+
if self.config.logit_softcap > 0:
|
| 279 |
+
cap = self.config.logit_softcap
|
| 280 |
+
logits = cap * torch.tanh(logits / cap)
|
| 281 |
+
loss = None
|
| 282 |
+
if labels is not None:
|
| 283 |
+
shift_logits = logits[..., :-1, :].contiguous()
|
| 284 |
+
shift_labels = labels[..., 1:].contiguous()
|
| 285 |
+
loss = F.cross_entropy(
|
| 286 |
+
shift_logits.view(-1, shift_logits.size(-1)),
|
| 287 |
+
shift_labels.view(-1),
|
| 288 |
+
ignore_index=-100,
|
| 289 |
+
)
|
| 290 |
+
if not return_dict:
|
| 291 |
+
output = (logits,)
|
| 292 |
+
return (loss,) + output if loss is not None else output
|
| 293 |
+
return CausalLMOutputWithPast(loss=loss, logits=logits)
|
| 294 |
+
|
| 295 |
+
def prepare_inputs_for_generation(self, input_ids, **kwargs):
|
| 296 |
+
return {"input_ids": input_ids}
|
| 297 |
+
|
| 298 |
+
|
| 299 |
+
# ---------------------------------------------------------------------------
|
| 300 |
+
# CLI demo
|
| 301 |
+
# ---------------------------------------------------------------------------
|
| 302 |
+
|
| 303 |
+
if __name__ == "__main__":
|
| 304 |
+
import argparse
|
| 305 |
+
import sys
|
| 306 |
+
import os
|
| 307 |
+
|
| 308 |
+
parser = argparse.ArgumentParser(description="Load and generate with a NanoChat model")
|
| 309 |
+
parser.add_argument("--model-dir", required=True, help="Path to converted model directory")
|
| 310 |
+
parser.add_argument("--prompt", default="The meaning of life is", help="Text prompt")
|
| 311 |
+
parser.add_argument("--max-tokens", type=int, default=100, help="Max tokens to generate")
|
| 312 |
+
parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature (0=greedy)")
|
| 313 |
+
parser.add_argument("--device", default="cpu", help="Device (cpu, cuda, cuda:0, ...)")
|
| 314 |
+
parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
|
| 315 |
+
args = parser.parse_args()
|
| 316 |
+
|
| 317 |
+
dtype_map = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
|
| 318 |
+
|
| 319 |
+
print(f"Loading model from {args.model_dir}...")
|
| 320 |
+
from transformers import AutoTokenizer
|
| 321 |
+
|
| 322 |
+
model = NanoChatForCausalLM.from_pretrained(args.model_dir, torch_dtype=dtype_map[args.dtype])
|
| 323 |
+
model.to(args.device)
|
| 324 |
+
model.eval()
|
| 325 |
+
tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
|
| 326 |
+
|
| 327 |
+
total_params = sum(p.numel() for p in model.parameters())
|
| 328 |
+
print(f"Model loaded: {total_params:,} parameters, dtype={args.dtype}, device={args.device}")
|
| 329 |
+
print(f"Vocab size: {model.config.vocab_size}, ve_dim: {model.config.ve_dim}")
|
| 330 |
+
print()
|
| 331 |
+
print(f"Prompt: {args.prompt}")
|
| 332 |
+
print(f"Generating (max {args.max_tokens} tokens, temp={args.temperature})...")
|
| 333 |
+
print()
|
| 334 |
+
|
| 335 |
+
# Generate
|
| 336 |
+
bos_id = tokenizer.bos_token_id
|
| 337 |
+
input_ids = tokenizer.encode(args.prompt, add_special_tokens=False)
|
| 338 |
+
if bos_id is not None:
|
| 339 |
+
input_ids = [bos_id] + input_ids
|
| 340 |
+
input_ids = torch.tensor([input_ids], dtype=torch.long, device=args.device)
|
| 341 |
+
|
| 342 |
+
generated = []
|
| 343 |
+
with torch.inference_mode():
|
| 344 |
+
for _ in range(args.max_tokens):
|
| 345 |
+
logits = model(input_ids).logits[:, -1, :]
|
| 346 |
+
if args.temperature > 0:
|
| 347 |
+
logits = logits / args.temperature
|
| 348 |
+
v, _ = torch.topk(logits, min(50, logits.size(-1)))
|
| 349 |
+
logits[logits < v[:, [-1]]] = -float("inf")
|
| 350 |
+
probs = torch.softmax(logits, dim=-1)
|
| 351 |
+
next_id = torch.multinomial(probs, num_samples=1)
|
| 352 |
+
else:
|
| 353 |
+
next_id = torch.argmax(logits, dim=-1, keepdim=True)
|
| 354 |
+
input_ids = torch.cat([input_ids, next_id], dim=1)
|
| 355 |
+
token_id = next_id.item()
|
| 356 |
+
generated.append(token_id)
|
| 357 |
+
if tokenizer.eos_token_id is not None and token_id == tokenizer.eos_token_id:
|
| 358 |
+
break
|
| 359 |
+
|
| 360 |
+
output = tokenizer.decode(generated, skip_special_tokens=True)
|
| 361 |
+
print(f"Output: {args.prompt}{output}")
|
models/full-128k-punct-english-bpe-mathcode-scratch/tokenizer_config.json
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"tokenizer_class": "PreTrainedTokenizerFast",
|
| 3 |
+
"model_max_length": 2048,
|
| 4 |
+
"bos_token": "<s>",
|
| 5 |
+
"add_bos_token": false,
|
| 6 |
+
"clean_up_tokenization_spaces": false
|
| 7 |
+
}
|
models/full-128k-rightalign-balanced-bpe-mathcode-scratch/README.md
ADDED
|
@@ -0,0 +1,76 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# `full-128k-rightalign-balanced-bpe-mathcode-scratch`
|
| 2 |
+
|
| 3 |
+
A 1.27B parameter model trained with the `rightalign-balanced-bpe` tokenizer. Regime: mathcode-20B-from-scratch (24 layers, width 1536, 1.27B parameters, 20.0B tokens, math and code mixture, trained from scratch). Reported in: math and code ablations; TokEval.
|
| 4 |
+
|
| 5 |
+
## Load
|
| 6 |
+
|
| 7 |
+
```python
|
| 8 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
| 9 |
+
|
| 10 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 11 |
+
"cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-rightalign-balanced-bpe-mathcode-scratch",
|
| 12 |
+
trust_remote_code=True, dtype="float32")
|
| 13 |
+
tokenizer = AutoTokenizer.from_pretrained("cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-rightalign-balanced-bpe-mathcode-scratch")
|
| 14 |
+
|
| 15 |
+
text = "def add(a, b):\n return a + b\n"
|
| 16 |
+
ids = [0] + tokenizer(text)["input_ids"] # training prepends BOS
|
| 17 |
+
```
|
| 18 |
+
|
| 19 |
+
## Settings
|
| 20 |
+
|
| 21 |
+
| Setting | Value |
|
| 22 |
+
|---|---|
|
| 23 |
+
| layers | 24 |
|
| 24 |
+
| model dimension | 1536 |
|
| 25 |
+
| attention heads | 12 |
|
| 26 |
+
| key and value heads | 12 |
|
| 27 |
+
| context length | 2048 |
|
| 28 |
+
| sliding window pattern | SSSL |
|
| 29 |
+
| value embedding dimension | 128 |
|
| 30 |
+
| logit soft cap | 15 |
|
| 31 |
+
| vocabulary size | 128256 |
|
| 32 |
+
| parameters | 1,272,841,994 |
|
| 33 |
+
| training step | 19073 |
|
| 34 |
+
| token budget | 19,999,490,048 |
|
| 35 |
+
| bytes consumed | 147,637,512,815 |
|
| 36 |
+
| seed | 42 |
|
| 37 |
+
| Muon matrix learning rate | 0.0200 |
|
| 38 |
+
| AdamW embedding learning rate | 0.3000 |
|
| 39 |
+
| AdamW unembedding learning rate | 0.0080 |
|
| 40 |
+
| weight decay | 0.2800 |
|
| 41 |
+
| width learning-rate exponent | -1 |
|
| 42 |
+
| weight dtype | as trained (float32 matrices, bfloat16 embeddings) |
|
| 43 |
+
| BOS token id | 0 |
|
| 44 |
+
|
| 45 |
+
## Results
|
| 46 |
+
|
| 47 |
+
| Metric | Value |
|
| 48 |
+
|---|---|
|
| 49 |
+
| validation BPB | 0.3353 |
|
| 50 |
+
| code BPB, mean over 7 languages | 0.3798 |
|
| 51 |
+
| FLORES-200 BPB, mean | 3.3832 |
|
| 52 |
+
| FLORES-200 languages scored | 215 |
|
| 53 |
+
| BLiMP accuracy | n/a |
|
| 54 |
+
| MBPP pass@1 | 0.0720 |
|
| 55 |
+
| HumanEval pass@1 | 0.1524 |
|
| 56 |
+
| GSM8K exact match, strict | 0.2335 |
|
| 57 |
+
|
| 58 |
+
Generation metrics use specification `{'version': 'v2-2026-07-30', 'bos_context': True, 'eos_stop': True, 'token_healing': True, 'truncation': 'left_to_max_length_minus_max_gen', 'stop_truncation': 'earliest_match_all_terms', 'stop_match_decode': 'specials_kept'}`, token healing on.
|
| 59 |
+
|
| 60 |
+
Missing measurements, with the reason each is absent:
|
| 61 |
+
|
| 62 |
+
- `blimp_acc`: no full-128k-rightalign-balanced-bpe-mathcode-scratch_blimp_code_bpb.json
|
| 63 |
+
|
| 64 |
+
BLiMP is comparable within a training scale only. Its scoring sums unnormalized sentence log-likelihoods, and under that scoring a group of quantifier subtasks drops at the 1.27B scale relative to the 300M scale.
|
| 65 |
+
|
| 66 |
+
## Registry note
|
| 67 |
+
|
| 68 |
+
#52, rightalign pretok, slug `rightalign-balanced-bpe`. Main-grid RightAlign balanced BPE cell. FINISHED at step 19073, 4-eval chain + code_bpb complete (verified 2026-08-09). In panel.
|
| 69 |
+
|
| 70 |
+
## Provenance
|
| 71 |
+
|
| 72 |
+
- Source checkpoint: `/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-rightalign-balanced-bpe-mathcode-scratch`, step 19073
|
| 73 |
+
- Repository commit: `1a9a718`
|
| 74 |
+
- nanochat commit: `f919449` on branch `tokenizer-lm-patches`
|
| 75 |
+
- Tokenizer: `rightalign-balanced-bpe`, released at `tokenizers/rightalign-balanced-bpe`
|
| 76 |
+
|
models/full-128k-rightalign-balanced-bpe-mathcode-scratch/config.json
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"NanoChatForCausalLM"
|
| 4 |
+
],
|
| 5 |
+
"auto_map": {
|
| 6 |
+
"AutoConfig": "nanochat_model.NanoChatConfig",
|
| 7 |
+
"AutoModelForCausalLM": "nanochat_model.NanoChatForCausalLM"
|
| 8 |
+
},
|
| 9 |
+
"bos_token_id": 0,
|
| 10 |
+
"dtype": "float32",
|
| 11 |
+
"eos_token_id": null,
|
| 12 |
+
"head_dim": 128,
|
| 13 |
+
"kv_dim": 1536,
|
| 14 |
+
"logit_softcap": 15.0,
|
| 15 |
+
"model_type": "nanochat",
|
| 16 |
+
"n_embd": 1536,
|
| 17 |
+
"n_head": 12,
|
| 18 |
+
"n_kv_head": 12,
|
| 19 |
+
"n_layer": 24,
|
| 20 |
+
"pad_vocab_size_to": 64,
|
| 21 |
+
"padded_vocab_size": 128256,
|
| 22 |
+
"sequence_len": 2048,
|
| 23 |
+
"tie_word_embeddings": false,
|
| 24 |
+
"transformers_version": "5.3.0",
|
| 25 |
+
"ve_dim": 128,
|
| 26 |
+
"vocab_size": 128256,
|
| 27 |
+
"window_pattern": "SSSL"
|
| 28 |
+
}
|
models/full-128k-rightalign-balanced-bpe-mathcode-scratch/generation_config.json
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"bos_token_id": 0,
|
| 3 |
+
"transformers_version": "5.3.0"
|
| 4 |
+
}
|
models/full-128k-rightalign-balanced-bpe-mathcode-scratch/nanochat_model.py
ADDED
|
@@ -0,0 +1,361 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
NanoChat model for HuggingFace Transformers — self-contained single file.
|
| 3 |
+
|
| 4 |
+
Includes the config class, model class, and a CLI demo for loading + generation.
|
| 5 |
+
Drop this file into a directory containing config.json, model.safetensors,
|
| 6 |
+
tokenizer.json, and tokenizer_config.json, then run:
|
| 7 |
+
|
| 8 |
+
python nanochat_model.py --model-dir . --prompt "Hello world"
|
| 9 |
+
|
| 10 |
+
Or load programmatically:
|
| 11 |
+
|
| 12 |
+
from nanochat_model import NanoChatConfig, NanoChatForCausalLM
|
| 13 |
+
model = NanoChatForCausalLM.from_pretrained("path/to/model", torch_dtype=torch.bfloat16)
|
| 14 |
+
|
| 15 |
+
Requirements: pip install torch transformers tokenizers
|
| 16 |
+
|
| 17 |
+
Architecture: nanochat GPT (Karpathy, MIT licensed) with value embeddings
|
| 18 |
+
(ResFormer-style), per-layer residual/x0 lambdas, smear, backout, logit
|
| 19 |
+
soft-capping, sliding window attention, QK norm, ReLU^2 MLP, and RoPE.
|
| 20 |
+
|
| 21 |
+
The models were trained with a patched nanochat, not upstream nanochat:
|
| 22 |
+
https://github.com/cimeister/nanochat branch tokenizer-lm-patches, commit f919449.
|
| 23 |
+
The patches touch two files and change three things, two of which affect the forward
|
| 24 |
+
pass and are implemented here:
|
| 25 |
+
1. ve_dim: the value-embedding tables are built at a reduced dimension and projected
|
| 26 |
+
up to kv_dim by a learned ve_proj, instead of being built at full kv_dim. All
|
| 27 |
+
released models use ve_dim=128.
|
| 28 |
+
2. logit_softcap: the soft cap is a config field with 0 meaning disabled, instead of
|
| 29 |
+
a hardcoded 15.
|
| 30 |
+
3. width_lr_exponent: the AdamW learning-rate width scaling exponent is configurable
|
| 31 |
+
for muP. This affects training only and has no inference-time effect.
|
| 32 |
+
|
| 33 |
+
This file is the released model definition. Its agreement with the training-time model
|
| 34 |
+
is checked per released checkpoint by scripts/hf_release/verify_conversion.py, which
|
| 35 |
+
compares logits against the patched nanochat GPT in float32.
|
| 36 |
+
"""
|
| 37 |
+
|
| 38 |
+
from typing import Optional, Tuple, Union
|
| 39 |
+
|
| 40 |
+
import torch
|
| 41 |
+
import torch.nn as nn
|
| 42 |
+
import torch.nn.functional as F
|
| 43 |
+
from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
|
| 44 |
+
from transformers.modeling_outputs import CausalLMOutputWithPast
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
# ---------------------------------------------------------------------------
|
| 48 |
+
# Config
|
| 49 |
+
# ---------------------------------------------------------------------------
|
| 50 |
+
|
| 51 |
+
class NanoChatConfig(PretrainedConfig):
|
| 52 |
+
model_type = "nanochat"
|
| 53 |
+
|
| 54 |
+
def __init__(
|
| 55 |
+
self,
|
| 56 |
+
vocab_size=128256,
|
| 57 |
+
n_embd=1536,
|
| 58 |
+
n_layer=24,
|
| 59 |
+
n_head=12,
|
| 60 |
+
n_kv_head=12,
|
| 61 |
+
sequence_len=2048,
|
| 62 |
+
window_pattern="SSSL",
|
| 63 |
+
ve_dim=0,
|
| 64 |
+
logit_softcap=15.0,
|
| 65 |
+
pad_vocab_size_to=64,
|
| 66 |
+
bos_token_id=0,
|
| 67 |
+
eos_token_id=None,
|
| 68 |
+
tie_word_embeddings=False,
|
| 69 |
+
**kwargs,
|
| 70 |
+
):
|
| 71 |
+
self.vocab_size = vocab_size
|
| 72 |
+
self.n_embd = n_embd
|
| 73 |
+
self.n_layer = n_layer
|
| 74 |
+
self.n_head = n_head
|
| 75 |
+
self.n_kv_head = n_kv_head
|
| 76 |
+
self.sequence_len = sequence_len
|
| 77 |
+
self.window_pattern = window_pattern
|
| 78 |
+
self.ve_dim = ve_dim
|
| 79 |
+
self.logit_softcap = logit_softcap
|
| 80 |
+
self.pad_vocab_size_to = pad_vocab_size_to
|
| 81 |
+
self.head_dim = n_embd // n_head
|
| 82 |
+
self.kv_dim = n_kv_head * self.head_dim
|
| 83 |
+
self.padded_vocab_size = (
|
| 84 |
+
(vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to
|
| 85 |
+
) * pad_vocab_size_to
|
| 86 |
+
super().__init__(
|
| 87 |
+
bos_token_id=bos_token_id,
|
| 88 |
+
eos_token_id=eos_token_id,
|
| 89 |
+
tie_word_embeddings=tie_word_embeddings,
|
| 90 |
+
**kwargs,
|
| 91 |
+
)
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
# ---------------------------------------------------------------------------
|
| 95 |
+
# Model
|
| 96 |
+
# ---------------------------------------------------------------------------
|
| 97 |
+
|
| 98 |
+
def _has_ve(layer_idx: int, n_layer: int) -> bool:
|
| 99 |
+
return layer_idx % 2 == (n_layer - 1) % 2
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def _apply_rotary_emb(x, cos, sin):
|
| 103 |
+
d = x.shape[-1] // 2
|
| 104 |
+
x1, x2 = x[..., :d], x[..., d:]
|
| 105 |
+
y1 = x1 * cos + x2 * sin
|
| 106 |
+
y2 = x1 * (-sin) + x2 * cos
|
| 107 |
+
return torch.cat([y1, y2], dim=-1)
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
def _rms_norm(x):
|
| 111 |
+
return F.rms_norm(x, (x.size(-1),))
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
class NanoChatAttention(nn.Module):
|
| 115 |
+
def __init__(self, config, layer_idx):
|
| 116 |
+
super().__init__()
|
| 117 |
+
self.layer_idx = layer_idx
|
| 118 |
+
self.n_head = config.n_head
|
| 119 |
+
self.n_kv_head = config.n_kv_head
|
| 120 |
+
self.n_embd = config.n_embd
|
| 121 |
+
self.head_dim = config.head_dim
|
| 122 |
+
self.n_kv_groups = self.n_head // self.n_kv_head
|
| 123 |
+
self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
|
| 124 |
+
self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
|
| 125 |
+
self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
|
| 126 |
+
self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
|
| 127 |
+
self.ve_gate_channels = 12
|
| 128 |
+
if _has_ve(layer_idx, config.n_layer):
|
| 129 |
+
self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False)
|
| 130 |
+
kv_dim = self.n_kv_head * self.head_dim
|
| 131 |
+
ve_dim = config.ve_dim if config.ve_dim > 0 else kv_dim
|
| 132 |
+
self.ve_proj = nn.Linear(ve_dim, kv_dim, bias=False) if ve_dim < kv_dim else None
|
| 133 |
+
else:
|
| 134 |
+
self.ve_gate = None
|
| 135 |
+
self.ve_proj = None
|
| 136 |
+
|
| 137 |
+
def forward(self, x, ve, cos, sin, window_size):
|
| 138 |
+
B, T, C = x.size()
|
| 139 |
+
q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
|
| 140 |
+
k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
|
| 141 |
+
v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
|
| 142 |
+
if ve is not None and self.ve_gate is not None:
|
| 143 |
+
if self.ve_proj is not None:
|
| 144 |
+
ve = self.ve_proj(ve)
|
| 145 |
+
ve = ve.view(B, T, self.n_kv_head, self.head_dim)
|
| 146 |
+
gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
|
| 147 |
+
v = v + gate.unsqueeze(-1) * ve
|
| 148 |
+
q = _apply_rotary_emb(q, cos, sin)
|
| 149 |
+
k = _apply_rotary_emb(k, cos, sin)
|
| 150 |
+
q = _rms_norm(q) * 1.2
|
| 151 |
+
k = _rms_norm(k) * 1.2
|
| 152 |
+
q = q.transpose(1, 2)
|
| 153 |
+
k = k.transpose(1, 2)
|
| 154 |
+
v = v.transpose(1, 2)
|
| 155 |
+
if self.n_kv_groups > 1:
|
| 156 |
+
k = k.repeat_interleave(self.n_kv_groups, dim=1)
|
| 157 |
+
v = v.repeat_interleave(self.n_kv_groups, dim=1)
|
| 158 |
+
if window_size < 0 or window_size >= T:
|
| 159 |
+
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
|
| 160 |
+
else:
|
| 161 |
+
row_idx = torch.arange(T, device=x.device).unsqueeze(1)
|
| 162 |
+
col_idx = torch.arange(T, device=x.device).unsqueeze(0)
|
| 163 |
+
mask = (col_idx <= row_idx) & ((row_idx - col_idx) <= window_size)
|
| 164 |
+
y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
|
| 165 |
+
y = y.transpose(1, 2).contiguous().view(B, T, -1)
|
| 166 |
+
y = self.c_proj(y)
|
| 167 |
+
return y
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
class NanoChatMLP(nn.Module):
|
| 171 |
+
def __init__(self, config):
|
| 172 |
+
super().__init__()
|
| 173 |
+
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
|
| 174 |
+
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
|
| 175 |
+
|
| 176 |
+
def forward(self, x):
|
| 177 |
+
x = self.c_fc(x)
|
| 178 |
+
x = F.relu(x).square()
|
| 179 |
+
x = self.c_proj(x)
|
| 180 |
+
return x
|
| 181 |
+
|
| 182 |
+
|
| 183 |
+
class NanoChatBlock(nn.Module):
|
| 184 |
+
def __init__(self, config, layer_idx):
|
| 185 |
+
super().__init__()
|
| 186 |
+
self.attn = NanoChatAttention(config, layer_idx)
|
| 187 |
+
self.mlp = NanoChatMLP(config)
|
| 188 |
+
|
| 189 |
+
def forward(self, x, ve, cos, sin, window_size):
|
| 190 |
+
x = x + self.attn(_rms_norm(x), ve, cos, sin, window_size)
|
| 191 |
+
x = x + self.mlp(_rms_norm(x))
|
| 192 |
+
return x
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
class NanoChatPreTrainedModel(PreTrainedModel):
|
| 196 |
+
config_class = NanoChatConfig
|
| 197 |
+
base_model_prefix = "model"
|
| 198 |
+
supports_gradient_checkpointing = False
|
| 199 |
+
_no_split_modules = ["NanoChatBlock"]
|
| 200 |
+
|
| 201 |
+
def _init_weights(self, module):
|
| 202 |
+
pass
|
| 203 |
+
|
| 204 |
+
|
| 205 |
+
class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin):
|
| 206 |
+
def __init__(self, config):
|
| 207 |
+
super().__init__(config)
|
| 208 |
+
pvs = config.padded_vocab_size
|
| 209 |
+
self.wte = nn.Embedding(pvs, config.n_embd)
|
| 210 |
+
self.h = nn.ModuleList([NanoChatBlock(config, i) for i in range(config.n_layer)])
|
| 211 |
+
self.lm_head = nn.Linear(config.n_embd, pvs, bias=False)
|
| 212 |
+
self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
|
| 213 |
+
self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
|
| 214 |
+
self.smear_gate = nn.Linear(24, 1, bias=False)
|
| 215 |
+
self.smear_lambda = nn.Parameter(torch.zeros(1))
|
| 216 |
+
self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))
|
| 217 |
+
ve_emb_dim = config.ve_dim if config.ve_dim > 0 else config.kv_dim
|
| 218 |
+
self.value_embeds = nn.ModuleDict({
|
| 219 |
+
str(i): nn.Embedding(pvs, ve_emb_dim)
|
| 220 |
+
for i in range(config.n_layer) if _has_ve(i, config.n_layer)
|
| 221 |
+
})
|
| 222 |
+
self.window_sizes = self._compute_window_sizes(config)
|
| 223 |
+
self.rotary_seq_len = config.sequence_len * 10
|
| 224 |
+
self._rotary_initialized = False
|
| 225 |
+
self.register_buffer("cos", torch.zeros(1), persistent=False)
|
| 226 |
+
self.register_buffer("sin", torch.zeros(1), persistent=False)
|
| 227 |
+
self.post_init()
|
| 228 |
+
|
| 229 |
+
def _compute_window_sizes(self, config):
|
| 230 |
+
pattern = config.window_pattern.upper()
|
| 231 |
+
long_window = config.sequence_len
|
| 232 |
+
short_window = -(-long_window // 4 // 128) * 128
|
| 233 |
+
char_to_window = {"L": long_window, "S": short_window}
|
| 234 |
+
sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.n_layer)]
|
| 235 |
+
sizes[-1] = long_window
|
| 236 |
+
return sizes
|
| 237 |
+
|
| 238 |
+
def _precompute_rotary(self, seq_len, head_dim, base=100000.0):
|
| 239 |
+
channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32)
|
| 240 |
+
inv_freq = 1.0 / (base ** (channel_range / head_dim))
|
| 241 |
+
t = torch.arange(seq_len, dtype=torch.float32)
|
| 242 |
+
freqs = torch.outer(t, inv_freq)
|
| 243 |
+
cos = freqs.cos()[None, :, None, :]
|
| 244 |
+
sin = freqs.sin()[None, :, None, :]
|
| 245 |
+
return cos, sin
|
| 246 |
+
|
| 247 |
+
def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs):
|
| 248 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 249 |
+
B, T = input_ids.size()
|
| 250 |
+
if not self._rotary_initialized:
|
| 251 |
+
cos, sin = self._precompute_rotary(self.rotary_seq_len, self.config.head_dim)
|
| 252 |
+
self.cos = cos.to(device=input_ids.device)
|
| 253 |
+
self.sin = sin.to(device=input_ids.device)
|
| 254 |
+
self._rotary_initialized = True
|
| 255 |
+
cos = self.cos[:, :T].to(dtype=self.wte.weight.dtype)
|
| 256 |
+
sin = self.sin[:, :T].to(dtype=self.wte.weight.dtype)
|
| 257 |
+
x = self.wte(input_ids)
|
| 258 |
+
x = _rms_norm(x)
|
| 259 |
+
if T > 1:
|
| 260 |
+
gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
|
| 261 |
+
x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
|
| 262 |
+
x0 = x
|
| 263 |
+
n_layer = self.config.n_layer
|
| 264 |
+
backout_layer = n_layer // 2
|
| 265 |
+
x_backout = None
|
| 266 |
+
for i, block in enumerate(self.h):
|
| 267 |
+
x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
|
| 268 |
+
ve = self.value_embeds[str(i)](input_ids).to(x.dtype) if str(i) in self.value_embeds else None
|
| 269 |
+
x = block(x, ve, cos, sin, self.window_sizes[i])
|
| 270 |
+
if i == backout_layer:
|
| 271 |
+
x_backout = x
|
| 272 |
+
if x_backout is not None:
|
| 273 |
+
x = x - self.backout_lambda.to(x.dtype) * x_backout
|
| 274 |
+
x = _rms_norm(x)
|
| 275 |
+
logits = self.lm_head(x)
|
| 276 |
+
logits = logits[..., :self.config.vocab_size]
|
| 277 |
+
logits = logits.float()
|
| 278 |
+
if self.config.logit_softcap > 0:
|
| 279 |
+
cap = self.config.logit_softcap
|
| 280 |
+
logits = cap * torch.tanh(logits / cap)
|
| 281 |
+
loss = None
|
| 282 |
+
if labels is not None:
|
| 283 |
+
shift_logits = logits[..., :-1, :].contiguous()
|
| 284 |
+
shift_labels = labels[..., 1:].contiguous()
|
| 285 |
+
loss = F.cross_entropy(
|
| 286 |
+
shift_logits.view(-1, shift_logits.size(-1)),
|
| 287 |
+
shift_labels.view(-1),
|
| 288 |
+
ignore_index=-100,
|
| 289 |
+
)
|
| 290 |
+
if not return_dict:
|
| 291 |
+
output = (logits,)
|
| 292 |
+
return (loss,) + output if loss is not None else output
|
| 293 |
+
return CausalLMOutputWithPast(loss=loss, logits=logits)
|
| 294 |
+
|
| 295 |
+
def prepare_inputs_for_generation(self, input_ids, **kwargs):
|
| 296 |
+
return {"input_ids": input_ids}
|
| 297 |
+
|
| 298 |
+
|
| 299 |
+
# ---------------------------------------------------------------------------
|
| 300 |
+
# CLI demo
|
| 301 |
+
# ---------------------------------------------------------------------------
|
| 302 |
+
|
| 303 |
+
if __name__ == "__main__":
|
| 304 |
+
import argparse
|
| 305 |
+
import sys
|
| 306 |
+
import os
|
| 307 |
+
|
| 308 |
+
parser = argparse.ArgumentParser(description="Load and generate with a NanoChat model")
|
| 309 |
+
parser.add_argument("--model-dir", required=True, help="Path to converted model directory")
|
| 310 |
+
parser.add_argument("--prompt", default="The meaning of life is", help="Text prompt")
|
| 311 |
+
parser.add_argument("--max-tokens", type=int, default=100, help="Max tokens to generate")
|
| 312 |
+
parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature (0=greedy)")
|
| 313 |
+
parser.add_argument("--device", default="cpu", help="Device (cpu, cuda, cuda:0, ...)")
|
| 314 |
+
parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
|
| 315 |
+
args = parser.parse_args()
|
| 316 |
+
|
| 317 |
+
dtype_map = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
|
| 318 |
+
|
| 319 |
+
print(f"Loading model from {args.model_dir}...")
|
| 320 |
+
from transformers import AutoTokenizer
|
| 321 |
+
|
| 322 |
+
model = NanoChatForCausalLM.from_pretrained(args.model_dir, torch_dtype=dtype_map[args.dtype])
|
| 323 |
+
model.to(args.device)
|
| 324 |
+
model.eval()
|
| 325 |
+
tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
|
| 326 |
+
|
| 327 |
+
total_params = sum(p.numel() for p in model.parameters())
|
| 328 |
+
print(f"Model loaded: {total_params:,} parameters, dtype={args.dtype}, device={args.device}")
|
| 329 |
+
print(f"Vocab size: {model.config.vocab_size}, ve_dim: {model.config.ve_dim}")
|
| 330 |
+
print()
|
| 331 |
+
print(f"Prompt: {args.prompt}")
|
| 332 |
+
print(f"Generating (max {args.max_tokens} tokens, temp={args.temperature})...")
|
| 333 |
+
print()
|
| 334 |
+
|
| 335 |
+
# Generate
|
| 336 |
+
bos_id = tokenizer.bos_token_id
|
| 337 |
+
input_ids = tokenizer.encode(args.prompt, add_special_tokens=False)
|
| 338 |
+
if bos_id is not None:
|
| 339 |
+
input_ids = [bos_id] + input_ids
|
| 340 |
+
input_ids = torch.tensor([input_ids], dtype=torch.long, device=args.device)
|
| 341 |
+
|
| 342 |
+
generated = []
|
| 343 |
+
with torch.inference_mode():
|
| 344 |
+
for _ in range(args.max_tokens):
|
| 345 |
+
logits = model(input_ids).logits[:, -1, :]
|
| 346 |
+
if args.temperature > 0:
|
| 347 |
+
logits = logits / args.temperature
|
| 348 |
+
v, _ = torch.topk(logits, min(50, logits.size(-1)))
|
| 349 |
+
logits[logits < v[:, [-1]]] = -float("inf")
|
| 350 |
+
probs = torch.softmax(logits, dim=-1)
|
| 351 |
+
next_id = torch.multinomial(probs, num_samples=1)
|
| 352 |
+
else:
|
| 353 |
+
next_id = torch.argmax(logits, dim=-1, keepdim=True)
|
| 354 |
+
input_ids = torch.cat([input_ids, next_id], dim=1)
|
| 355 |
+
token_id = next_id.item()
|
| 356 |
+
generated.append(token_id)
|
| 357 |
+
if tokenizer.eos_token_id is not None and token_id == tokenizer.eos_token_id:
|
| 358 |
+
break
|
| 359 |
+
|
| 360 |
+
output = tokenizer.decode(generated, skip_special_tokens=True)
|
| 361 |
+
print(f"Output: {args.prompt}{output}")
|
models/full-128k-rightalign-balanced-bpe-mathcode-scratch/tokenizer.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
models/full-128k-rightalign-balanced-bpe-mathcode-scratch/tokenizer_config.json
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"tokenizer_class": "PreTrainedTokenizerFast",
|
| 3 |
+
"model_max_length": 2048,
|
| 4 |
+
"bos_token": "<s>",
|
| 5 |
+
"add_bos_token": false,
|
| 6 |
+
"clean_up_tokenization_spaces": false
|
| 7 |
+
}
|
models/full-128k-scripttok-mingram-scriptenc_cb-mathcode-scratch/conversion_metadata.json
ADDED
|
@@ -0,0 +1,100 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"run": "full-128k-scripttok-mingram-scriptenc_cb-mathcode-scratch",
|
| 3 |
+
"source_checkpoint": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-scripttok-mingram-scriptenc_cb-mathcode-scratch",
|
| 4 |
+
"source_model_file": "model_019073.pt",
|
| 5 |
+
"source_step": 19073,
|
| 6 |
+
"source_meta": {
|
| 7 |
+
"step": 19073,
|
| 8 |
+
"val_bpb": 0.33485716807277754,
|
| 9 |
+
"model_config": {
|
| 10 |
+
"sequence_len": 2048,
|
| 11 |
+
"vocab_size": 128256,
|
| 12 |
+
"n_layer": 24,
|
| 13 |
+
"n_head": 12,
|
| 14 |
+
"n_kv_head": 12,
|
| 15 |
+
"n_embd": 1536,
|
| 16 |
+
"window_pattern": "SSSL",
|
| 17 |
+
"ve_dim": 128,
|
| 18 |
+
"logit_softcap": 15.0
|
| 19 |
+
},
|
| 20 |
+
"tokenizer_path": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/tokenizers/scripttok-mingram-scriptenc_cb",
|
| 21 |
+
"training_config": {
|
| 22 |
+
"name": "full-128k-mathcode-scratch",
|
| 23 |
+
"model": {
|
| 24 |
+
"n_layer": 24,
|
| 25 |
+
"n_embd": 1536,
|
| 26 |
+
"n_kv_head": 12,
|
| 27 |
+
"sequence_len": 2048,
|
| 28 |
+
"window_pattern": "SSSL",
|
| 29 |
+
"ve_dim": 128
|
| 30 |
+
},
|
| 31 |
+
"training": {
|
| 32 |
+
"target_param_data_ratio": -1,
|
| 33 |
+
"max_tokens": 20000000000,
|
| 34 |
+
"num_iterations": -1,
|
| 35 |
+
"total_batch_size": 1048576,
|
| 36 |
+
"device_batch_size": 16,
|
| 37 |
+
"warmup_steps": 40,
|
| 38 |
+
"warmdown_ratio": 0.65,
|
| 39 |
+
"final_lr_frac": 0.05,
|
| 40 |
+
"eval_tokens": 41943040,
|
| 41 |
+
"compile": true
|
| 42 |
+
},
|
| 43 |
+
"optimizer": {
|
| 44 |
+
"matrix_lr": 0.02,
|
| 45 |
+
"embedding_lr": 0.3,
|
| 46 |
+
"unembedding_lr": 0.008,
|
| 47 |
+
"scalar_lr": 0.5,
|
| 48 |
+
"weight_decay": 0.28,
|
| 49 |
+
"width_lr_exponent": -1.0
|
| 50 |
+
},
|
| 51 |
+
"data": {
|
| 52 |
+
"data_dir": null
|
| 53 |
+
},
|
| 54 |
+
"logging": {
|
| 55 |
+
"eval_every": 500,
|
| 56 |
+
"save_every": 500,
|
| 57 |
+
"max_checkpoints": 0
|
| 58 |
+
},
|
| 59 |
+
"wandb": {
|
| 60 |
+
"enabled": true,
|
| 61 |
+
"project": "tokenizer-lm-experiments",
|
| 62 |
+
"run_name": null
|
| 63 |
+
},
|
| 64 |
+
"output_dir": "/capstor/scratch/cscs/${USER}/tokenizer-lm/checkpoints",
|
| 65 |
+
"seed": 42
|
| 66 |
+
},
|
| 67 |
+
"total_batch_size": 1048576,
|
| 68 |
+
"device_batch_size": 16,
|
| 69 |
+
"max_seq_len": 2048,
|
| 70 |
+
"dataloader_state_dict": {
|
| 71 |
+
"pq_idx": 238,
|
| 72 |
+
"rg_idx": 48,
|
| 73 |
+
"epoch": 2
|
| 74 |
+
},
|
| 75 |
+
"loop_state": {
|
| 76 |
+
"min_val_bpb": 0.33485716807277754,
|
| 77 |
+
"smooth_train_loss": 0.9416237912354406,
|
| 78 |
+
"total_training_time": 93482.01090359688,
|
| 79 |
+
"total_bytes_consumed": 153083502797
|
| 80 |
+
}
|
| 81 |
+
},
|
| 82 |
+
"weight_dtype": "as trained (float32 matrices, bfloat16 embeddings)",
|
| 83 |
+
"weight_dtype_counts": {
|
| 84 |
+
"torch.float32": 174,
|
| 85 |
+
"torch.bfloat16": 13
|
| 86 |
+
},
|
| 87 |
+
"total_parameters": 1272841994,
|
| 88 |
+
"logit_softcap_source": "meta",
|
| 89 |
+
"bos_token": "<|bos|>",
|
| 90 |
+
"bos_token_id": 128255,
|
| 91 |
+
"bos_resolved_from": "script_tok loader at /users/cmeister747/script_tok",
|
| 92 |
+
"tokenizer_slug": "scripttok-mingram-scriptenc_cb",
|
| 93 |
+
"tokenizer_redistributed": true,
|
| 94 |
+
"tokenizer_hub_id": null,
|
| 95 |
+
"weight_key_mapping": "strip _orig_mod. and transformer. prefixes",
|
| 96 |
+
"converted_by": "scripts/hf_release/convert_for_release.py",
|
| 97 |
+
"repo_commit": "1a9a718",
|
| 98 |
+
"nanochat_commit": "f919449",
|
| 99 |
+
"nanochat_branch": "tokenizer-lm-patches"
|
| 100 |
+
}
|
models/full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch/README.md
ADDED
|
@@ -0,0 +1,77 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# `full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch`
|
| 2 |
+
|
| 3 |
+
A 1.27B parameter model trained with the `superbpe-apertus-fineweb2full-capped-hybridwindow` tokenizer. Regime: mathcode-20B-from-scratch (24 layers, width 1536, 1.27B parameters, 20.0B tokens, math and code mixture, trained from scratch). Reported in: math and code ablations.
|
| 4 |
+
|
| 5 |
+
## Load
|
| 6 |
+
|
| 7 |
+
```python
|
| 8 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
| 9 |
+
|
| 10 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 11 |
+
"cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch",
|
| 12 |
+
trust_remote_code=True, dtype="float32")
|
| 13 |
+
tokenizer = AutoTokenizer.from_pretrained("cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch")
|
| 14 |
+
|
| 15 |
+
text = "def add(a, b):\n return a + b\n"
|
| 16 |
+
ids = [128000] + tokenizer(text)["input_ids"] # training prepends BOS
|
| 17 |
+
```
|
| 18 |
+
|
| 19 |
+
## Settings
|
| 20 |
+
|
| 21 |
+
| Setting | Value |
|
| 22 |
+
|---|---|
|
| 23 |
+
| layers | 24 |
|
| 24 |
+
| model dimension | 1536 |
|
| 25 |
+
| attention heads | 12 |
|
| 26 |
+
| key and value heads | 12 |
|
| 27 |
+
| context length | 2048 |
|
| 28 |
+
| sliding window pattern | SSSL |
|
| 29 |
+
| value embedding dimension | 128 |
|
| 30 |
+
| logit soft cap | 15 |
|
| 31 |
+
| vocabulary size | 128004 |
|
| 32 |
+
| parameters | 1,271,957,258 |
|
| 33 |
+
| training step | 19073 |
|
| 34 |
+
| token budget | 19,999,490,048 |
|
| 35 |
+
| bytes consumed | 126,018,305,461 |
|
| 36 |
+
| seed | 42 |
|
| 37 |
+
| Muon matrix learning rate | 0.0200 |
|
| 38 |
+
| AdamW embedding learning rate | 0.3000 |
|
| 39 |
+
| AdamW unembedding learning rate | 0.0080 |
|
| 40 |
+
| weight decay | 0.2800 |
|
| 41 |
+
| width learning-rate exponent | -1 |
|
| 42 |
+
| weight dtype | as trained (float32 matrices, bfloat16 embeddings) |
|
| 43 |
+
| BOS token id | 128000 |
|
| 44 |
+
|
| 45 |
+
## Results
|
| 46 |
+
|
| 47 |
+
| Metric | Value |
|
| 48 |
+
|---|---|
|
| 49 |
+
| validation BPB | 0.3426 |
|
| 50 |
+
| code BPB, mean over 7 languages | n/a |
|
| 51 |
+
| FLORES-200 BPB, mean | 3.6638 |
|
| 52 |
+
| FLORES-200 languages scored | 214 |
|
| 53 |
+
| BLiMP accuracy | n/a |
|
| 54 |
+
| MBPP pass@1 | 0.0120 |
|
| 55 |
+
| HumanEval pass@1 | 0.1585 |
|
| 56 |
+
| GSM8K exact match, strict | 0.1926 |
|
| 57 |
+
|
| 58 |
+
Generation metrics use specification `{'version': 'v2-2026-07-30', 'bos_context': True, 'eos_stop': True, 'token_healing': True, 'truncation': 'left_to_max_length_minus_max_gen', 'stop_truncation': 'earliest_match_all_terms', 'stop_match_decode': 'specials_kept'}`, token healing on.
|
| 59 |
+
|
| 60 |
+
Missing measurements, with the reason each is absent:
|
| 61 |
+
|
| 62 |
+
- `blimp_acc`: no full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch_blimp_code_bpb.json
|
| 63 |
+
- `code_bpb_mean`: no full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch_blimp_code_bpb.json; no full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch_code_bpb.json
|
| 64 |
+
|
| 65 |
+
BLiMP is comparable within a training scale only. Its scoring sums unnormalized sentence log-likelihoods, and under that scoring a group of quantifier subtasks drops at the 1.27B scale relative to the 300M scale.
|
| 66 |
+
|
| 67 |
+
## Registry note
|
| 68 |
+
|
| 69 |
+
#11, apertus pretok, slug `superbpe-apertus-fineweb2full-capped-hybridwindow`. SuperBPE candidate on apertus pretok.
|
| 70 |
+
|
| 71 |
+
## Provenance
|
| 72 |
+
|
| 73 |
+
- Source checkpoint: `/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch`, step 19073
|
| 74 |
+
- Repository commit: `1a9a718`
|
| 75 |
+
- nanochat commit: `f919449` on branch `tokenizer-lm-patches`
|
| 76 |
+
- Tokenizer: `superbpe-apertus-fineweb2full-capped-hybridwindow`, released at `tokenizers/superbpe-apertus-fineweb2full-capped-hybridwindow`
|
| 77 |
+
|
models/full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch/conversion_metadata.json
ADDED
|
@@ -0,0 +1,100 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"run": "full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch",
|
| 3 |
+
"source_checkpoint": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch",
|
| 4 |
+
"source_model_file": "model_019073.pt",
|
| 5 |
+
"source_step": 19073,
|
| 6 |
+
"source_meta": {
|
| 7 |
+
"step": 19073,
|
| 8 |
+
"val_bpb": 0.3426277336443733,
|
| 9 |
+
"model_config": {
|
| 10 |
+
"sequence_len": 2048,
|
| 11 |
+
"vocab_size": 128004,
|
| 12 |
+
"n_layer": 24,
|
| 13 |
+
"n_head": 12,
|
| 14 |
+
"n_kv_head": 12,
|
| 15 |
+
"n_embd": 1536,
|
| 16 |
+
"window_pattern": "SSSL",
|
| 17 |
+
"ve_dim": 128,
|
| 18 |
+
"logit_softcap": 15.0
|
| 19 |
+
},
|
| 20 |
+
"tokenizer_path": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/tokenizers/superbpe-apertus-fineweb2full-capped-hybridwindow",
|
| 21 |
+
"training_config": {
|
| 22 |
+
"name": "full-128k-mathcode-scratch",
|
| 23 |
+
"model": {
|
| 24 |
+
"n_layer": 24,
|
| 25 |
+
"n_embd": 1536,
|
| 26 |
+
"n_kv_head": 12,
|
| 27 |
+
"sequence_len": 2048,
|
| 28 |
+
"window_pattern": "SSSL",
|
| 29 |
+
"ve_dim": 128
|
| 30 |
+
},
|
| 31 |
+
"training": {
|
| 32 |
+
"target_param_data_ratio": -1,
|
| 33 |
+
"max_tokens": 20000000000,
|
| 34 |
+
"num_iterations": -1,
|
| 35 |
+
"total_batch_size": 1048576,
|
| 36 |
+
"device_batch_size": 16,
|
| 37 |
+
"warmup_steps": 40,
|
| 38 |
+
"warmdown_ratio": 0.65,
|
| 39 |
+
"final_lr_frac": 0.05,
|
| 40 |
+
"eval_tokens": 41943040,
|
| 41 |
+
"compile": true
|
| 42 |
+
},
|
| 43 |
+
"optimizer": {
|
| 44 |
+
"matrix_lr": 0.02,
|
| 45 |
+
"embedding_lr": 0.3,
|
| 46 |
+
"unembedding_lr": 0.008,
|
| 47 |
+
"scalar_lr": 0.5,
|
| 48 |
+
"weight_decay": 0.28,
|
| 49 |
+
"width_lr_exponent": -1.0
|
| 50 |
+
},
|
| 51 |
+
"data": {
|
| 52 |
+
"data_dir": null
|
| 53 |
+
},
|
| 54 |
+
"logging": {
|
| 55 |
+
"eval_every": 500,
|
| 56 |
+
"save_every": 500,
|
| 57 |
+
"max_checkpoints": 0
|
| 58 |
+
},
|
| 59 |
+
"wandb": {
|
| 60 |
+
"enabled": true,
|
| 61 |
+
"project": "tokenizer-lm-experiments",
|
| 62 |
+
"run_name": null
|
| 63 |
+
},
|
| 64 |
+
"output_dir": "/capstor/scratch/cscs/${USER}/tokenizer-lm/checkpoints",
|
| 65 |
+
"seed": 42
|
| 66 |
+
},
|
| 67 |
+
"total_batch_size": 1048576,
|
| 68 |
+
"device_batch_size": 16,
|
| 69 |
+
"max_seq_len": 2048,
|
| 70 |
+
"dataloader_state_dict": {
|
| 71 |
+
"pq_idx": 284,
|
| 72 |
+
"rg_idx": 44,
|
| 73 |
+
"epoch": 2
|
| 74 |
+
},
|
| 75 |
+
"loop_state": {
|
| 76 |
+
"min_val_bpb": 0.34260583858556143,
|
| 77 |
+
"smooth_train_loss": 1.0121825854128466,
|
| 78 |
+
"total_training_time": 94313.82984352112,
|
| 79 |
+
"total_bytes_consumed": 126018305461
|
| 80 |
+
}
|
| 81 |
+
},
|
| 82 |
+
"weight_dtype": "as trained (float32 matrices, bfloat16 embeddings)",
|
| 83 |
+
"weight_dtype_counts": {
|
| 84 |
+
"torch.float32": 174,
|
| 85 |
+
"torch.bfloat16": 13
|
| 86 |
+
},
|
| 87 |
+
"total_parameters": 1271957258,
|
| 88 |
+
"logit_softcap_source": "meta",
|
| 89 |
+
"bos_token": "<s>",
|
| 90 |
+
"bos_token_id": 128000,
|
| 91 |
+
"bos_resolved_from": "packaged tokenizer",
|
| 92 |
+
"tokenizer_slug": "superbpe-apertus-fineweb2full-capped-hybridwindow",
|
| 93 |
+
"tokenizer_redistributed": true,
|
| 94 |
+
"tokenizer_hub_id": null,
|
| 95 |
+
"weight_key_mapping": "strip _orig_mod. and transformer. prefixes",
|
| 96 |
+
"converted_by": "scripts/hf_release/convert_for_release.py",
|
| 97 |
+
"repo_commit": "1a9a718",
|
| 98 |
+
"nanochat_commit": "f919449",
|
| 99 |
+
"nanochat_branch": "tokenizer-lm-patches"
|
| 100 |
+
}
|
models/full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch/generation_config.json
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"bos_token_id": 128000,
|
| 3 |
+
"transformers_version": "5.3.0"
|
| 4 |
+
}
|
models/full-128k-superbpe-apertus-fineweb2full-capped-hybridwindow-mathcode-scratch/tokenizer_config.json
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"tokenizer_class": "PreTrainedTokenizerFast",
|
| 3 |
+
"model_max_length": 2048,
|
| 4 |
+
"bos_token": "<s>",
|
| 5 |
+
"add_bos_token": false,
|
| 6 |
+
"clean_up_tokenization_spaces": false
|
| 7 |
+
}
|
models/full-128k-superbpe-clean-c2-balanced-mathcode-scratch/conversion_metadata.json
ADDED
|
@@ -0,0 +1,100 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"run": "full-128k-superbpe-clean-c2-balanced-mathcode-scratch",
|
| 3 |
+
"source_checkpoint": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-superbpe-clean-c2-balanced-mathcode-scratch",
|
| 4 |
+
"source_model_file": "model_019073.pt",
|
| 5 |
+
"source_step": 19073,
|
| 6 |
+
"source_meta": {
|
| 7 |
+
"step": 19073,
|
| 8 |
+
"val_bpb": 0.34319379203223876,
|
| 9 |
+
"model_config": {
|
| 10 |
+
"sequence_len": 2048,
|
| 11 |
+
"vocab_size": 128004,
|
| 12 |
+
"n_layer": 24,
|
| 13 |
+
"n_head": 12,
|
| 14 |
+
"n_kv_head": 12,
|
| 15 |
+
"n_embd": 1536,
|
| 16 |
+
"window_pattern": "SSSL",
|
| 17 |
+
"ve_dim": 128,
|
| 18 |
+
"logit_softcap": 15.0
|
| 19 |
+
},
|
| 20 |
+
"tokenizer_path": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/tokenizers/superbpe-clean-c2-balanced",
|
| 21 |
+
"training_config": {
|
| 22 |
+
"name": "full-128k-mathcode-scratch",
|
| 23 |
+
"model": {
|
| 24 |
+
"n_layer": 24,
|
| 25 |
+
"n_embd": 1536,
|
| 26 |
+
"n_kv_head": 12,
|
| 27 |
+
"sequence_len": 2048,
|
| 28 |
+
"window_pattern": "SSSL",
|
| 29 |
+
"ve_dim": 128
|
| 30 |
+
},
|
| 31 |
+
"training": {
|
| 32 |
+
"target_param_data_ratio": -1,
|
| 33 |
+
"max_tokens": 20000000000,
|
| 34 |
+
"num_iterations": -1,
|
| 35 |
+
"total_batch_size": 1048576,
|
| 36 |
+
"device_batch_size": 16,
|
| 37 |
+
"warmup_steps": 40,
|
| 38 |
+
"warmdown_ratio": 0.65,
|
| 39 |
+
"final_lr_frac": 0.05,
|
| 40 |
+
"eval_tokens": 41943040,
|
| 41 |
+
"compile": true
|
| 42 |
+
},
|
| 43 |
+
"optimizer": {
|
| 44 |
+
"matrix_lr": 0.02,
|
| 45 |
+
"embedding_lr": 0.3,
|
| 46 |
+
"unembedding_lr": 0.008,
|
| 47 |
+
"scalar_lr": 0.5,
|
| 48 |
+
"weight_decay": 0.28,
|
| 49 |
+
"width_lr_exponent": -1.0
|
| 50 |
+
},
|
| 51 |
+
"data": {
|
| 52 |
+
"data_dir": null
|
| 53 |
+
},
|
| 54 |
+
"logging": {
|
| 55 |
+
"eval_every": 500,
|
| 56 |
+
"save_every": 500,
|
| 57 |
+
"max_checkpoints": 0
|
| 58 |
+
},
|
| 59 |
+
"wandb": {
|
| 60 |
+
"enabled": true,
|
| 61 |
+
"project": "tokenizer-lm-experiments",
|
| 62 |
+
"run_name": null
|
| 63 |
+
},
|
| 64 |
+
"output_dir": "/capstor/scratch/cscs/${USER}/tokenizer-lm/checkpoints",
|
| 65 |
+
"seed": 42
|
| 66 |
+
},
|
| 67 |
+
"total_batch_size": 1048576,
|
| 68 |
+
"device_batch_size": 16,
|
| 69 |
+
"max_seq_len": 2048,
|
| 70 |
+
"dataloader_state_dict": {
|
| 71 |
+
"pq_idx": 226,
|
| 72 |
+
"rg_idx": 0,
|
| 73 |
+
"epoch": 2
|
| 74 |
+
},
|
| 75 |
+
"loop_state": {
|
| 76 |
+
"min_val_bpb": 0.34319379203223876,
|
| 77 |
+
"smooth_train_loss": 0.9382318900362904,
|
| 78 |
+
"total_training_time": 94756.40242409706,
|
| 79 |
+
"total_bytes_consumed": 157305292274
|
| 80 |
+
}
|
| 81 |
+
},
|
| 82 |
+
"weight_dtype": "as trained (float32 matrices, bfloat16 embeddings)",
|
| 83 |
+
"weight_dtype_counts": {
|
| 84 |
+
"torch.float32": 174,
|
| 85 |
+
"torch.bfloat16": 13
|
| 86 |
+
},
|
| 87 |
+
"total_parameters": 1271957258,
|
| 88 |
+
"logit_softcap_source": "meta",
|
| 89 |
+
"bos_token": "<s>",
|
| 90 |
+
"bos_token_id": 128000,
|
| 91 |
+
"bos_resolved_from": "packaged tokenizer",
|
| 92 |
+
"tokenizer_slug": "superbpe-clean-c2-balanced",
|
| 93 |
+
"tokenizer_redistributed": true,
|
| 94 |
+
"tokenizer_hub_id": null,
|
| 95 |
+
"weight_key_mapping": "strip _orig_mod. and transformer. prefixes",
|
| 96 |
+
"converted_by": "scripts/hf_release/convert_for_release.py",
|
| 97 |
+
"repo_commit": "1a9a718",
|
| 98 |
+
"nanochat_commit": "f919449",
|
| 99 |
+
"nanochat_branch": "tokenizer-lm-patches"
|
| 100 |
+
}
|
models/full-128k-superbpe-clean-c2-balanced-mathcode-scratch/nanochat_model.py
ADDED
|
@@ -0,0 +1,361 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
NanoChat model for HuggingFace Transformers — self-contained single file.
|
| 3 |
+
|
| 4 |
+
Includes the config class, model class, and a CLI demo for loading + generation.
|
| 5 |
+
Drop this file into a directory containing config.json, model.safetensors,
|
| 6 |
+
tokenizer.json, and tokenizer_config.json, then run:
|
| 7 |
+
|
| 8 |
+
python nanochat_model.py --model-dir . --prompt "Hello world"
|
| 9 |
+
|
| 10 |
+
Or load programmatically:
|
| 11 |
+
|
| 12 |
+
from nanochat_model import NanoChatConfig, NanoChatForCausalLM
|
| 13 |
+
model = NanoChatForCausalLM.from_pretrained("path/to/model", torch_dtype=torch.bfloat16)
|
| 14 |
+
|
| 15 |
+
Requirements: pip install torch transformers tokenizers
|
| 16 |
+
|
| 17 |
+
Architecture: nanochat GPT (Karpathy, MIT licensed) with value embeddings
|
| 18 |
+
(ResFormer-style), per-layer residual/x0 lambdas, smear, backout, logit
|
| 19 |
+
soft-capping, sliding window attention, QK norm, ReLU^2 MLP, and RoPE.
|
| 20 |
+
|
| 21 |
+
The models were trained with a patched nanochat, not upstream nanochat:
|
| 22 |
+
https://github.com/cimeister/nanochat branch tokenizer-lm-patches, commit f919449.
|
| 23 |
+
The patches touch two files and change three things, two of which affect the forward
|
| 24 |
+
pass and are implemented here:
|
| 25 |
+
1. ve_dim: the value-embedding tables are built at a reduced dimension and projected
|
| 26 |
+
up to kv_dim by a learned ve_proj, instead of being built at full kv_dim. All
|
| 27 |
+
released models use ve_dim=128.
|
| 28 |
+
2. logit_softcap: the soft cap is a config field with 0 meaning disabled, instead of
|
| 29 |
+
a hardcoded 15.
|
| 30 |
+
3. width_lr_exponent: the AdamW learning-rate width scaling exponent is configurable
|
| 31 |
+
for muP. This affects training only and has no inference-time effect.
|
| 32 |
+
|
| 33 |
+
This file is the released model definition. Its agreement with the training-time model
|
| 34 |
+
is checked per released checkpoint by scripts/hf_release/verify_conversion.py, which
|
| 35 |
+
compares logits against the patched nanochat GPT in float32.
|
| 36 |
+
"""
|
| 37 |
+
|
| 38 |
+
from typing import Optional, Tuple, Union
|
| 39 |
+
|
| 40 |
+
import torch
|
| 41 |
+
import torch.nn as nn
|
| 42 |
+
import torch.nn.functional as F
|
| 43 |
+
from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
|
| 44 |
+
from transformers.modeling_outputs import CausalLMOutputWithPast
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
# ---------------------------------------------------------------------------
|
| 48 |
+
# Config
|
| 49 |
+
# ---------------------------------------------------------------------------
|
| 50 |
+
|
| 51 |
+
class NanoChatConfig(PretrainedConfig):
|
| 52 |
+
model_type = "nanochat"
|
| 53 |
+
|
| 54 |
+
def __init__(
|
| 55 |
+
self,
|
| 56 |
+
vocab_size=128256,
|
| 57 |
+
n_embd=1536,
|
| 58 |
+
n_layer=24,
|
| 59 |
+
n_head=12,
|
| 60 |
+
n_kv_head=12,
|
| 61 |
+
sequence_len=2048,
|
| 62 |
+
window_pattern="SSSL",
|
| 63 |
+
ve_dim=0,
|
| 64 |
+
logit_softcap=15.0,
|
| 65 |
+
pad_vocab_size_to=64,
|
| 66 |
+
bos_token_id=0,
|
| 67 |
+
eos_token_id=None,
|
| 68 |
+
tie_word_embeddings=False,
|
| 69 |
+
**kwargs,
|
| 70 |
+
):
|
| 71 |
+
self.vocab_size = vocab_size
|
| 72 |
+
self.n_embd = n_embd
|
| 73 |
+
self.n_layer = n_layer
|
| 74 |
+
self.n_head = n_head
|
| 75 |
+
self.n_kv_head = n_kv_head
|
| 76 |
+
self.sequence_len = sequence_len
|
| 77 |
+
self.window_pattern = window_pattern
|
| 78 |
+
self.ve_dim = ve_dim
|
| 79 |
+
self.logit_softcap = logit_softcap
|
| 80 |
+
self.pad_vocab_size_to = pad_vocab_size_to
|
| 81 |
+
self.head_dim = n_embd // n_head
|
| 82 |
+
self.kv_dim = n_kv_head * self.head_dim
|
| 83 |
+
self.padded_vocab_size = (
|
| 84 |
+
(vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to
|
| 85 |
+
) * pad_vocab_size_to
|
| 86 |
+
super().__init__(
|
| 87 |
+
bos_token_id=bos_token_id,
|
| 88 |
+
eos_token_id=eos_token_id,
|
| 89 |
+
tie_word_embeddings=tie_word_embeddings,
|
| 90 |
+
**kwargs,
|
| 91 |
+
)
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
# ---------------------------------------------------------------------------
|
| 95 |
+
# Model
|
| 96 |
+
# ---------------------------------------------------------------------------
|
| 97 |
+
|
| 98 |
+
def _has_ve(layer_idx: int, n_layer: int) -> bool:
|
| 99 |
+
return layer_idx % 2 == (n_layer - 1) % 2
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def _apply_rotary_emb(x, cos, sin):
|
| 103 |
+
d = x.shape[-1] // 2
|
| 104 |
+
x1, x2 = x[..., :d], x[..., d:]
|
| 105 |
+
y1 = x1 * cos + x2 * sin
|
| 106 |
+
y2 = x1 * (-sin) + x2 * cos
|
| 107 |
+
return torch.cat([y1, y2], dim=-1)
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
def _rms_norm(x):
|
| 111 |
+
return F.rms_norm(x, (x.size(-1),))
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
class NanoChatAttention(nn.Module):
|
| 115 |
+
def __init__(self, config, layer_idx):
|
| 116 |
+
super().__init__()
|
| 117 |
+
self.layer_idx = layer_idx
|
| 118 |
+
self.n_head = config.n_head
|
| 119 |
+
self.n_kv_head = config.n_kv_head
|
| 120 |
+
self.n_embd = config.n_embd
|
| 121 |
+
self.head_dim = config.head_dim
|
| 122 |
+
self.n_kv_groups = self.n_head // self.n_kv_head
|
| 123 |
+
self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
|
| 124 |
+
self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
|
| 125 |
+
self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
|
| 126 |
+
self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
|
| 127 |
+
self.ve_gate_channels = 12
|
| 128 |
+
if _has_ve(layer_idx, config.n_layer):
|
| 129 |
+
self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False)
|
| 130 |
+
kv_dim = self.n_kv_head * self.head_dim
|
| 131 |
+
ve_dim = config.ve_dim if config.ve_dim > 0 else kv_dim
|
| 132 |
+
self.ve_proj = nn.Linear(ve_dim, kv_dim, bias=False) if ve_dim < kv_dim else None
|
| 133 |
+
else:
|
| 134 |
+
self.ve_gate = None
|
| 135 |
+
self.ve_proj = None
|
| 136 |
+
|
| 137 |
+
def forward(self, x, ve, cos, sin, window_size):
|
| 138 |
+
B, T, C = x.size()
|
| 139 |
+
q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
|
| 140 |
+
k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
|
| 141 |
+
v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
|
| 142 |
+
if ve is not None and self.ve_gate is not None:
|
| 143 |
+
if self.ve_proj is not None:
|
| 144 |
+
ve = self.ve_proj(ve)
|
| 145 |
+
ve = ve.view(B, T, self.n_kv_head, self.head_dim)
|
| 146 |
+
gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
|
| 147 |
+
v = v + gate.unsqueeze(-1) * ve
|
| 148 |
+
q = _apply_rotary_emb(q, cos, sin)
|
| 149 |
+
k = _apply_rotary_emb(k, cos, sin)
|
| 150 |
+
q = _rms_norm(q) * 1.2
|
| 151 |
+
k = _rms_norm(k) * 1.2
|
| 152 |
+
q = q.transpose(1, 2)
|
| 153 |
+
k = k.transpose(1, 2)
|
| 154 |
+
v = v.transpose(1, 2)
|
| 155 |
+
if self.n_kv_groups > 1:
|
| 156 |
+
k = k.repeat_interleave(self.n_kv_groups, dim=1)
|
| 157 |
+
v = v.repeat_interleave(self.n_kv_groups, dim=1)
|
| 158 |
+
if window_size < 0 or window_size >= T:
|
| 159 |
+
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
|
| 160 |
+
else:
|
| 161 |
+
row_idx = torch.arange(T, device=x.device).unsqueeze(1)
|
| 162 |
+
col_idx = torch.arange(T, device=x.device).unsqueeze(0)
|
| 163 |
+
mask = (col_idx <= row_idx) & ((row_idx - col_idx) <= window_size)
|
| 164 |
+
y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
|
| 165 |
+
y = y.transpose(1, 2).contiguous().view(B, T, -1)
|
| 166 |
+
y = self.c_proj(y)
|
| 167 |
+
return y
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
class NanoChatMLP(nn.Module):
|
| 171 |
+
def __init__(self, config):
|
| 172 |
+
super().__init__()
|
| 173 |
+
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
|
| 174 |
+
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
|
| 175 |
+
|
| 176 |
+
def forward(self, x):
|
| 177 |
+
x = self.c_fc(x)
|
| 178 |
+
x = F.relu(x).square()
|
| 179 |
+
x = self.c_proj(x)
|
| 180 |
+
return x
|
| 181 |
+
|
| 182 |
+
|
| 183 |
+
class NanoChatBlock(nn.Module):
|
| 184 |
+
def __init__(self, config, layer_idx):
|
| 185 |
+
super().__init__()
|
| 186 |
+
self.attn = NanoChatAttention(config, layer_idx)
|
| 187 |
+
self.mlp = NanoChatMLP(config)
|
| 188 |
+
|
| 189 |
+
def forward(self, x, ve, cos, sin, window_size):
|
| 190 |
+
x = x + self.attn(_rms_norm(x), ve, cos, sin, window_size)
|
| 191 |
+
x = x + self.mlp(_rms_norm(x))
|
| 192 |
+
return x
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
class NanoChatPreTrainedModel(PreTrainedModel):
|
| 196 |
+
config_class = NanoChatConfig
|
| 197 |
+
base_model_prefix = "model"
|
| 198 |
+
supports_gradient_checkpointing = False
|
| 199 |
+
_no_split_modules = ["NanoChatBlock"]
|
| 200 |
+
|
| 201 |
+
def _init_weights(self, module):
|
| 202 |
+
pass
|
| 203 |
+
|
| 204 |
+
|
| 205 |
+
class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin):
|
| 206 |
+
def __init__(self, config):
|
| 207 |
+
super().__init__(config)
|
| 208 |
+
pvs = config.padded_vocab_size
|
| 209 |
+
self.wte = nn.Embedding(pvs, config.n_embd)
|
| 210 |
+
self.h = nn.ModuleList([NanoChatBlock(config, i) for i in range(config.n_layer)])
|
| 211 |
+
self.lm_head = nn.Linear(config.n_embd, pvs, bias=False)
|
| 212 |
+
self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
|
| 213 |
+
self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
|
| 214 |
+
self.smear_gate = nn.Linear(24, 1, bias=False)
|
| 215 |
+
self.smear_lambda = nn.Parameter(torch.zeros(1))
|
| 216 |
+
self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))
|
| 217 |
+
ve_emb_dim = config.ve_dim if config.ve_dim > 0 else config.kv_dim
|
| 218 |
+
self.value_embeds = nn.ModuleDict({
|
| 219 |
+
str(i): nn.Embedding(pvs, ve_emb_dim)
|
| 220 |
+
for i in range(config.n_layer) if _has_ve(i, config.n_layer)
|
| 221 |
+
})
|
| 222 |
+
self.window_sizes = self._compute_window_sizes(config)
|
| 223 |
+
self.rotary_seq_len = config.sequence_len * 10
|
| 224 |
+
self._rotary_initialized = False
|
| 225 |
+
self.register_buffer("cos", torch.zeros(1), persistent=False)
|
| 226 |
+
self.register_buffer("sin", torch.zeros(1), persistent=False)
|
| 227 |
+
self.post_init()
|
| 228 |
+
|
| 229 |
+
def _compute_window_sizes(self, config):
|
| 230 |
+
pattern = config.window_pattern.upper()
|
| 231 |
+
long_window = config.sequence_len
|
| 232 |
+
short_window = -(-long_window // 4 // 128) * 128
|
| 233 |
+
char_to_window = {"L": long_window, "S": short_window}
|
| 234 |
+
sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.n_layer)]
|
| 235 |
+
sizes[-1] = long_window
|
| 236 |
+
return sizes
|
| 237 |
+
|
| 238 |
+
def _precompute_rotary(self, seq_len, head_dim, base=100000.0):
|
| 239 |
+
channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32)
|
| 240 |
+
inv_freq = 1.0 / (base ** (channel_range / head_dim))
|
| 241 |
+
t = torch.arange(seq_len, dtype=torch.float32)
|
| 242 |
+
freqs = torch.outer(t, inv_freq)
|
| 243 |
+
cos = freqs.cos()[None, :, None, :]
|
| 244 |
+
sin = freqs.sin()[None, :, None, :]
|
| 245 |
+
return cos, sin
|
| 246 |
+
|
| 247 |
+
def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs):
|
| 248 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 249 |
+
B, T = input_ids.size()
|
| 250 |
+
if not self._rotary_initialized:
|
| 251 |
+
cos, sin = self._precompute_rotary(self.rotary_seq_len, self.config.head_dim)
|
| 252 |
+
self.cos = cos.to(device=input_ids.device)
|
| 253 |
+
self.sin = sin.to(device=input_ids.device)
|
| 254 |
+
self._rotary_initialized = True
|
| 255 |
+
cos = self.cos[:, :T].to(dtype=self.wte.weight.dtype)
|
| 256 |
+
sin = self.sin[:, :T].to(dtype=self.wte.weight.dtype)
|
| 257 |
+
x = self.wte(input_ids)
|
| 258 |
+
x = _rms_norm(x)
|
| 259 |
+
if T > 1:
|
| 260 |
+
gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
|
| 261 |
+
x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
|
| 262 |
+
x0 = x
|
| 263 |
+
n_layer = self.config.n_layer
|
| 264 |
+
backout_layer = n_layer // 2
|
| 265 |
+
x_backout = None
|
| 266 |
+
for i, block in enumerate(self.h):
|
| 267 |
+
x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
|
| 268 |
+
ve = self.value_embeds[str(i)](input_ids).to(x.dtype) if str(i) in self.value_embeds else None
|
| 269 |
+
x = block(x, ve, cos, sin, self.window_sizes[i])
|
| 270 |
+
if i == backout_layer:
|
| 271 |
+
x_backout = x
|
| 272 |
+
if x_backout is not None:
|
| 273 |
+
x = x - self.backout_lambda.to(x.dtype) * x_backout
|
| 274 |
+
x = _rms_norm(x)
|
| 275 |
+
logits = self.lm_head(x)
|
| 276 |
+
logits = logits[..., :self.config.vocab_size]
|
| 277 |
+
logits = logits.float()
|
| 278 |
+
if self.config.logit_softcap > 0:
|
| 279 |
+
cap = self.config.logit_softcap
|
| 280 |
+
logits = cap * torch.tanh(logits / cap)
|
| 281 |
+
loss = None
|
| 282 |
+
if labels is not None:
|
| 283 |
+
shift_logits = logits[..., :-1, :].contiguous()
|
| 284 |
+
shift_labels = labels[..., 1:].contiguous()
|
| 285 |
+
loss = F.cross_entropy(
|
| 286 |
+
shift_logits.view(-1, shift_logits.size(-1)),
|
| 287 |
+
shift_labels.view(-1),
|
| 288 |
+
ignore_index=-100,
|
| 289 |
+
)
|
| 290 |
+
if not return_dict:
|
| 291 |
+
output = (logits,)
|
| 292 |
+
return (loss,) + output if loss is not None else output
|
| 293 |
+
return CausalLMOutputWithPast(loss=loss, logits=logits)
|
| 294 |
+
|
| 295 |
+
def prepare_inputs_for_generation(self, input_ids, **kwargs):
|
| 296 |
+
return {"input_ids": input_ids}
|
| 297 |
+
|
| 298 |
+
|
| 299 |
+
# ---------------------------------------------------------------------------
|
| 300 |
+
# CLI demo
|
| 301 |
+
# ---------------------------------------------------------------------------
|
| 302 |
+
|
| 303 |
+
if __name__ == "__main__":
|
| 304 |
+
import argparse
|
| 305 |
+
import sys
|
| 306 |
+
import os
|
| 307 |
+
|
| 308 |
+
parser = argparse.ArgumentParser(description="Load and generate with a NanoChat model")
|
| 309 |
+
parser.add_argument("--model-dir", required=True, help="Path to converted model directory")
|
| 310 |
+
parser.add_argument("--prompt", default="The meaning of life is", help="Text prompt")
|
| 311 |
+
parser.add_argument("--max-tokens", type=int, default=100, help="Max tokens to generate")
|
| 312 |
+
parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature (0=greedy)")
|
| 313 |
+
parser.add_argument("--device", default="cpu", help="Device (cpu, cuda, cuda:0, ...)")
|
| 314 |
+
parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
|
| 315 |
+
args = parser.parse_args()
|
| 316 |
+
|
| 317 |
+
dtype_map = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
|
| 318 |
+
|
| 319 |
+
print(f"Loading model from {args.model_dir}...")
|
| 320 |
+
from transformers import AutoTokenizer
|
| 321 |
+
|
| 322 |
+
model = NanoChatForCausalLM.from_pretrained(args.model_dir, torch_dtype=dtype_map[args.dtype])
|
| 323 |
+
model.to(args.device)
|
| 324 |
+
model.eval()
|
| 325 |
+
tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
|
| 326 |
+
|
| 327 |
+
total_params = sum(p.numel() for p in model.parameters())
|
| 328 |
+
print(f"Model loaded: {total_params:,} parameters, dtype={args.dtype}, device={args.device}")
|
| 329 |
+
print(f"Vocab size: {model.config.vocab_size}, ve_dim: {model.config.ve_dim}")
|
| 330 |
+
print()
|
| 331 |
+
print(f"Prompt: {args.prompt}")
|
| 332 |
+
print(f"Generating (max {args.max_tokens} tokens, temp={args.temperature})...")
|
| 333 |
+
print()
|
| 334 |
+
|
| 335 |
+
# Generate
|
| 336 |
+
bos_id = tokenizer.bos_token_id
|
| 337 |
+
input_ids = tokenizer.encode(args.prompt, add_special_tokens=False)
|
| 338 |
+
if bos_id is not None:
|
| 339 |
+
input_ids = [bos_id] + input_ids
|
| 340 |
+
input_ids = torch.tensor([input_ids], dtype=torch.long, device=args.device)
|
| 341 |
+
|
| 342 |
+
generated = []
|
| 343 |
+
with torch.inference_mode():
|
| 344 |
+
for _ in range(args.max_tokens):
|
| 345 |
+
logits = model(input_ids).logits[:, -1, :]
|
| 346 |
+
if args.temperature > 0:
|
| 347 |
+
logits = logits / args.temperature
|
| 348 |
+
v, _ = torch.topk(logits, min(50, logits.size(-1)))
|
| 349 |
+
logits[logits < v[:, [-1]]] = -float("inf")
|
| 350 |
+
probs = torch.softmax(logits, dim=-1)
|
| 351 |
+
next_id = torch.multinomial(probs, num_samples=1)
|
| 352 |
+
else:
|
| 353 |
+
next_id = torch.argmax(logits, dim=-1, keepdim=True)
|
| 354 |
+
input_ids = torch.cat([input_ids, next_id], dim=1)
|
| 355 |
+
token_id = next_id.item()
|
| 356 |
+
generated.append(token_id)
|
| 357 |
+
if tokenizer.eos_token_id is not None and token_id == tokenizer.eos_token_id:
|
| 358 |
+
break
|
| 359 |
+
|
| 360 |
+
output = tokenizer.decode(generated, skip_special_tokens=True)
|
| 361 |
+
print(f"Output: {args.prompt}{output}")
|
models/full-128k-superbpe-clean-c2-balanced-mathcode-scratch/tokenizer.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
models/full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch/README.md
ADDED
|
@@ -0,0 +1,77 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# `full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch`
|
| 2 |
+
|
| 3 |
+
A 1.28B parameter model trained with the `superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2` tokenizer. Regime: mathcode-20B-from-scratch (24 layers, width 1536, 1.27B parameters, 20.0B tokens, math and code mixture, trained from scratch). Reported in: math and code ablations.
|
| 4 |
+
|
| 5 |
+
## Load
|
| 6 |
+
|
| 7 |
+
```python
|
| 8 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
| 9 |
+
|
| 10 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 11 |
+
"cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch",
|
| 12 |
+
trust_remote_code=True, dtype="float32")
|
| 13 |
+
tokenizer = AutoTokenizer.from_pretrained("cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch")
|
| 14 |
+
|
| 15 |
+
text = "def add(a, b):\n return a + b\n"
|
| 16 |
+
ids = [130000] + tokenizer(text)["input_ids"] # training prepends BOS
|
| 17 |
+
```
|
| 18 |
+
|
| 19 |
+
## Settings
|
| 20 |
+
|
| 21 |
+
| Setting | Value |
|
| 22 |
+
|---|---|
|
| 23 |
+
| layers | 24 |
|
| 24 |
+
| model dimension | 1536 |
|
| 25 |
+
| attention heads | 12 |
|
| 26 |
+
| key and value heads | 12 |
|
| 27 |
+
| context length | 2048 |
|
| 28 |
+
| sliding window pattern | SSSL |
|
| 29 |
+
| value embedding dimension | 128 |
|
| 30 |
+
| logit soft cap | 15 |
|
| 31 |
+
| vocabulary size | 130004 |
|
| 32 |
+
| parameters | 1,281,099,530 |
|
| 33 |
+
| training step | 19073 |
|
| 34 |
+
| token budget | 19,999,490,048 |
|
| 35 |
+
| bytes consumed | 132,980,980,403 |
|
| 36 |
+
| seed | 42 |
|
| 37 |
+
| Muon matrix learning rate | 0.0200 |
|
| 38 |
+
| AdamW embedding learning rate | 0.3000 |
|
| 39 |
+
| AdamW unembedding learning rate | 0.0080 |
|
| 40 |
+
| weight decay | 0.2800 |
|
| 41 |
+
| width learning-rate exponent | -1 |
|
| 42 |
+
| weight dtype | as trained (float32 matrices, bfloat16 embeddings) |
|
| 43 |
+
| BOS token id | 130000 |
|
| 44 |
+
|
| 45 |
+
## Results
|
| 46 |
+
|
| 47 |
+
| Metric | Value |
|
| 48 |
+
|---|---|
|
| 49 |
+
| validation BPB | 0.3433 |
|
| 50 |
+
| code BPB, mean over 7 languages | n/a |
|
| 51 |
+
| FLORES-200 BPB, mean | 3.5595 |
|
| 52 |
+
| FLORES-200 languages scored | 214 |
|
| 53 |
+
| BLiMP accuracy | n/a |
|
| 54 |
+
| MBPP pass@1 | 0.2460 |
|
| 55 |
+
| HumanEval pass@1 | 0.1463 |
|
| 56 |
+
| GSM8K exact match, strict | 0.1660 |
|
| 57 |
+
|
| 58 |
+
Generation metrics use specification `{'version': 'v2-2026-07-30', 'bos_context': True, 'eos_stop': True, 'token_healing': True, 'truncation': 'left_to_max_length_minus_max_gen', 'stop_truncation': 'earliest_match_all_terms', 'stop_match_decode': 'specials_kept'}`, token healing on.
|
| 59 |
+
|
| 60 |
+
Missing measurements, with the reason each is absent:
|
| 61 |
+
|
| 62 |
+
- `blimp_acc`: no full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch_blimp_code_bpb.json
|
| 63 |
+
- `code_bpb_mean`: no full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch_blimp_code_bpb.json; no full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch_code_bpb.json
|
| 64 |
+
|
| 65 |
+
BLiMP is comparable within a training scale only. Its scoring sums unnormalized sentence log-likelihoods, and under that scoring a group of quantifier subtasks drops at the 1.27B scale relative to the 300M scale.
|
| 66 |
+
|
| 67 |
+
## Registry note
|
| 68 |
+
|
| 69 |
+
#16, clean pretok, slug `superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2`. SuperBPE on clean-multi + `plus2` + `cap_v2` + hw + `consv2` data, stage-2 transition 110k → 130k. Specials added post-hoc (vocab 130000 → 130004, bos_id=130000).
|
| 70 |
+
|
| 71 |
+
## Provenance
|
| 72 |
+
|
| 73 |
+
- Source checkpoint: `/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch`, step 19073
|
| 74 |
+
- Repository commit: `1a9a718`
|
| 75 |
+
- nanochat commit: `f919449` on branch `tokenizer-lm-patches`
|
| 76 |
+
- Tokenizer: `superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2`, released at `tokenizers/superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2`
|
| 77 |
+
|
models/full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch/config.json
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"NanoChatForCausalLM"
|
| 4 |
+
],
|
| 5 |
+
"auto_map": {
|
| 6 |
+
"AutoConfig": "nanochat_model.NanoChatConfig",
|
| 7 |
+
"AutoModelForCausalLM": "nanochat_model.NanoChatForCausalLM"
|
| 8 |
+
},
|
| 9 |
+
"bos_token_id": 130000,
|
| 10 |
+
"dtype": "float32",
|
| 11 |
+
"eos_token_id": null,
|
| 12 |
+
"head_dim": 128,
|
| 13 |
+
"kv_dim": 1536,
|
| 14 |
+
"logit_softcap": 15.0,
|
| 15 |
+
"model_type": "nanochat",
|
| 16 |
+
"n_embd": 1536,
|
| 17 |
+
"n_head": 12,
|
| 18 |
+
"n_kv_head": 12,
|
| 19 |
+
"n_layer": 24,
|
| 20 |
+
"pad_vocab_size_to": 64,
|
| 21 |
+
"padded_vocab_size": 130048,
|
| 22 |
+
"sequence_len": 2048,
|
| 23 |
+
"tie_word_embeddings": false,
|
| 24 |
+
"transformers_version": "5.3.0",
|
| 25 |
+
"ve_dim": 128,
|
| 26 |
+
"vocab_size": 130004,
|
| 27 |
+
"window_pattern": "SSSL"
|
| 28 |
+
}
|
models/full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch/generation_config.json
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"bos_token_id": 130000,
|
| 3 |
+
"transformers_version": "5.3.0"
|
| 4 |
+
}
|
models/full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch/nanochat_model.py
ADDED
|
@@ -0,0 +1,361 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
NanoChat model for HuggingFace Transformers — self-contained single file.
|
| 3 |
+
|
| 4 |
+
Includes the config class, model class, and a CLI demo for loading + generation.
|
| 5 |
+
Drop this file into a directory containing config.json, model.safetensors,
|
| 6 |
+
tokenizer.json, and tokenizer_config.json, then run:
|
| 7 |
+
|
| 8 |
+
python nanochat_model.py --model-dir . --prompt "Hello world"
|
| 9 |
+
|
| 10 |
+
Or load programmatically:
|
| 11 |
+
|
| 12 |
+
from nanochat_model import NanoChatConfig, NanoChatForCausalLM
|
| 13 |
+
model = NanoChatForCausalLM.from_pretrained("path/to/model", torch_dtype=torch.bfloat16)
|
| 14 |
+
|
| 15 |
+
Requirements: pip install torch transformers tokenizers
|
| 16 |
+
|
| 17 |
+
Architecture: nanochat GPT (Karpathy, MIT licensed) with value embeddings
|
| 18 |
+
(ResFormer-style), per-layer residual/x0 lambdas, smear, backout, logit
|
| 19 |
+
soft-capping, sliding window attention, QK norm, ReLU^2 MLP, and RoPE.
|
| 20 |
+
|
| 21 |
+
The models were trained with a patched nanochat, not upstream nanochat:
|
| 22 |
+
https://github.com/cimeister/nanochat branch tokenizer-lm-patches, commit f919449.
|
| 23 |
+
The patches touch two files and change three things, two of which affect the forward
|
| 24 |
+
pass and are implemented here:
|
| 25 |
+
1. ve_dim: the value-embedding tables are built at a reduced dimension and projected
|
| 26 |
+
up to kv_dim by a learned ve_proj, instead of being built at full kv_dim. All
|
| 27 |
+
released models use ve_dim=128.
|
| 28 |
+
2. logit_softcap: the soft cap is a config field with 0 meaning disabled, instead of
|
| 29 |
+
a hardcoded 15.
|
| 30 |
+
3. width_lr_exponent: the AdamW learning-rate width scaling exponent is configurable
|
| 31 |
+
for muP. This affects training only and has no inference-time effect.
|
| 32 |
+
|
| 33 |
+
This file is the released model definition. Its agreement with the training-time model
|
| 34 |
+
is checked per released checkpoint by scripts/hf_release/verify_conversion.py, which
|
| 35 |
+
compares logits against the patched nanochat GPT in float32.
|
| 36 |
+
"""
|
| 37 |
+
|
| 38 |
+
from typing import Optional, Tuple, Union
|
| 39 |
+
|
| 40 |
+
import torch
|
| 41 |
+
import torch.nn as nn
|
| 42 |
+
import torch.nn.functional as F
|
| 43 |
+
from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
|
| 44 |
+
from transformers.modeling_outputs import CausalLMOutputWithPast
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
# ---------------------------------------------------------------------------
|
| 48 |
+
# Config
|
| 49 |
+
# ---------------------------------------------------------------------------
|
| 50 |
+
|
| 51 |
+
class NanoChatConfig(PretrainedConfig):
|
| 52 |
+
model_type = "nanochat"
|
| 53 |
+
|
| 54 |
+
def __init__(
|
| 55 |
+
self,
|
| 56 |
+
vocab_size=128256,
|
| 57 |
+
n_embd=1536,
|
| 58 |
+
n_layer=24,
|
| 59 |
+
n_head=12,
|
| 60 |
+
n_kv_head=12,
|
| 61 |
+
sequence_len=2048,
|
| 62 |
+
window_pattern="SSSL",
|
| 63 |
+
ve_dim=0,
|
| 64 |
+
logit_softcap=15.0,
|
| 65 |
+
pad_vocab_size_to=64,
|
| 66 |
+
bos_token_id=0,
|
| 67 |
+
eos_token_id=None,
|
| 68 |
+
tie_word_embeddings=False,
|
| 69 |
+
**kwargs,
|
| 70 |
+
):
|
| 71 |
+
self.vocab_size = vocab_size
|
| 72 |
+
self.n_embd = n_embd
|
| 73 |
+
self.n_layer = n_layer
|
| 74 |
+
self.n_head = n_head
|
| 75 |
+
self.n_kv_head = n_kv_head
|
| 76 |
+
self.sequence_len = sequence_len
|
| 77 |
+
self.window_pattern = window_pattern
|
| 78 |
+
self.ve_dim = ve_dim
|
| 79 |
+
self.logit_softcap = logit_softcap
|
| 80 |
+
self.pad_vocab_size_to = pad_vocab_size_to
|
| 81 |
+
self.head_dim = n_embd // n_head
|
| 82 |
+
self.kv_dim = n_kv_head * self.head_dim
|
| 83 |
+
self.padded_vocab_size = (
|
| 84 |
+
(vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to
|
| 85 |
+
) * pad_vocab_size_to
|
| 86 |
+
super().__init__(
|
| 87 |
+
bos_token_id=bos_token_id,
|
| 88 |
+
eos_token_id=eos_token_id,
|
| 89 |
+
tie_word_embeddings=tie_word_embeddings,
|
| 90 |
+
**kwargs,
|
| 91 |
+
)
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
# ---------------------------------------------------------------------------
|
| 95 |
+
# Model
|
| 96 |
+
# ---------------------------------------------------------------------------
|
| 97 |
+
|
| 98 |
+
def _has_ve(layer_idx: int, n_layer: int) -> bool:
|
| 99 |
+
return layer_idx % 2 == (n_layer - 1) % 2
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def _apply_rotary_emb(x, cos, sin):
|
| 103 |
+
d = x.shape[-1] // 2
|
| 104 |
+
x1, x2 = x[..., :d], x[..., d:]
|
| 105 |
+
y1 = x1 * cos + x2 * sin
|
| 106 |
+
y2 = x1 * (-sin) + x2 * cos
|
| 107 |
+
return torch.cat([y1, y2], dim=-1)
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
def _rms_norm(x):
|
| 111 |
+
return F.rms_norm(x, (x.size(-1),))
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
class NanoChatAttention(nn.Module):
|
| 115 |
+
def __init__(self, config, layer_idx):
|
| 116 |
+
super().__init__()
|
| 117 |
+
self.layer_idx = layer_idx
|
| 118 |
+
self.n_head = config.n_head
|
| 119 |
+
self.n_kv_head = config.n_kv_head
|
| 120 |
+
self.n_embd = config.n_embd
|
| 121 |
+
self.head_dim = config.head_dim
|
| 122 |
+
self.n_kv_groups = self.n_head // self.n_kv_head
|
| 123 |
+
self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
|
| 124 |
+
self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
|
| 125 |
+
self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
|
| 126 |
+
self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
|
| 127 |
+
self.ve_gate_channels = 12
|
| 128 |
+
if _has_ve(layer_idx, config.n_layer):
|
| 129 |
+
self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False)
|
| 130 |
+
kv_dim = self.n_kv_head * self.head_dim
|
| 131 |
+
ve_dim = config.ve_dim if config.ve_dim > 0 else kv_dim
|
| 132 |
+
self.ve_proj = nn.Linear(ve_dim, kv_dim, bias=False) if ve_dim < kv_dim else None
|
| 133 |
+
else:
|
| 134 |
+
self.ve_gate = None
|
| 135 |
+
self.ve_proj = None
|
| 136 |
+
|
| 137 |
+
def forward(self, x, ve, cos, sin, window_size):
|
| 138 |
+
B, T, C = x.size()
|
| 139 |
+
q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
|
| 140 |
+
k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
|
| 141 |
+
v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
|
| 142 |
+
if ve is not None and self.ve_gate is not None:
|
| 143 |
+
if self.ve_proj is not None:
|
| 144 |
+
ve = self.ve_proj(ve)
|
| 145 |
+
ve = ve.view(B, T, self.n_kv_head, self.head_dim)
|
| 146 |
+
gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
|
| 147 |
+
v = v + gate.unsqueeze(-1) * ve
|
| 148 |
+
q = _apply_rotary_emb(q, cos, sin)
|
| 149 |
+
k = _apply_rotary_emb(k, cos, sin)
|
| 150 |
+
q = _rms_norm(q) * 1.2
|
| 151 |
+
k = _rms_norm(k) * 1.2
|
| 152 |
+
q = q.transpose(1, 2)
|
| 153 |
+
k = k.transpose(1, 2)
|
| 154 |
+
v = v.transpose(1, 2)
|
| 155 |
+
if self.n_kv_groups > 1:
|
| 156 |
+
k = k.repeat_interleave(self.n_kv_groups, dim=1)
|
| 157 |
+
v = v.repeat_interleave(self.n_kv_groups, dim=1)
|
| 158 |
+
if window_size < 0 or window_size >= T:
|
| 159 |
+
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
|
| 160 |
+
else:
|
| 161 |
+
row_idx = torch.arange(T, device=x.device).unsqueeze(1)
|
| 162 |
+
col_idx = torch.arange(T, device=x.device).unsqueeze(0)
|
| 163 |
+
mask = (col_idx <= row_idx) & ((row_idx - col_idx) <= window_size)
|
| 164 |
+
y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
|
| 165 |
+
y = y.transpose(1, 2).contiguous().view(B, T, -1)
|
| 166 |
+
y = self.c_proj(y)
|
| 167 |
+
return y
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
class NanoChatMLP(nn.Module):
|
| 171 |
+
def __init__(self, config):
|
| 172 |
+
super().__init__()
|
| 173 |
+
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
|
| 174 |
+
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
|
| 175 |
+
|
| 176 |
+
def forward(self, x):
|
| 177 |
+
x = self.c_fc(x)
|
| 178 |
+
x = F.relu(x).square()
|
| 179 |
+
x = self.c_proj(x)
|
| 180 |
+
return x
|
| 181 |
+
|
| 182 |
+
|
| 183 |
+
class NanoChatBlock(nn.Module):
|
| 184 |
+
def __init__(self, config, layer_idx):
|
| 185 |
+
super().__init__()
|
| 186 |
+
self.attn = NanoChatAttention(config, layer_idx)
|
| 187 |
+
self.mlp = NanoChatMLP(config)
|
| 188 |
+
|
| 189 |
+
def forward(self, x, ve, cos, sin, window_size):
|
| 190 |
+
x = x + self.attn(_rms_norm(x), ve, cos, sin, window_size)
|
| 191 |
+
x = x + self.mlp(_rms_norm(x))
|
| 192 |
+
return x
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
class NanoChatPreTrainedModel(PreTrainedModel):
|
| 196 |
+
config_class = NanoChatConfig
|
| 197 |
+
base_model_prefix = "model"
|
| 198 |
+
supports_gradient_checkpointing = False
|
| 199 |
+
_no_split_modules = ["NanoChatBlock"]
|
| 200 |
+
|
| 201 |
+
def _init_weights(self, module):
|
| 202 |
+
pass
|
| 203 |
+
|
| 204 |
+
|
| 205 |
+
class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin):
|
| 206 |
+
def __init__(self, config):
|
| 207 |
+
super().__init__(config)
|
| 208 |
+
pvs = config.padded_vocab_size
|
| 209 |
+
self.wte = nn.Embedding(pvs, config.n_embd)
|
| 210 |
+
self.h = nn.ModuleList([NanoChatBlock(config, i) for i in range(config.n_layer)])
|
| 211 |
+
self.lm_head = nn.Linear(config.n_embd, pvs, bias=False)
|
| 212 |
+
self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
|
| 213 |
+
self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
|
| 214 |
+
self.smear_gate = nn.Linear(24, 1, bias=False)
|
| 215 |
+
self.smear_lambda = nn.Parameter(torch.zeros(1))
|
| 216 |
+
self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))
|
| 217 |
+
ve_emb_dim = config.ve_dim if config.ve_dim > 0 else config.kv_dim
|
| 218 |
+
self.value_embeds = nn.ModuleDict({
|
| 219 |
+
str(i): nn.Embedding(pvs, ve_emb_dim)
|
| 220 |
+
for i in range(config.n_layer) if _has_ve(i, config.n_layer)
|
| 221 |
+
})
|
| 222 |
+
self.window_sizes = self._compute_window_sizes(config)
|
| 223 |
+
self.rotary_seq_len = config.sequence_len * 10
|
| 224 |
+
self._rotary_initialized = False
|
| 225 |
+
self.register_buffer("cos", torch.zeros(1), persistent=False)
|
| 226 |
+
self.register_buffer("sin", torch.zeros(1), persistent=False)
|
| 227 |
+
self.post_init()
|
| 228 |
+
|
| 229 |
+
def _compute_window_sizes(self, config):
|
| 230 |
+
pattern = config.window_pattern.upper()
|
| 231 |
+
long_window = config.sequence_len
|
| 232 |
+
short_window = -(-long_window // 4 // 128) * 128
|
| 233 |
+
char_to_window = {"L": long_window, "S": short_window}
|
| 234 |
+
sizes = [char_to_window[pattern[i % len(pattern)]] for i in range(config.n_layer)]
|
| 235 |
+
sizes[-1] = long_window
|
| 236 |
+
return sizes
|
| 237 |
+
|
| 238 |
+
def _precompute_rotary(self, seq_len, head_dim, base=100000.0):
|
| 239 |
+
channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32)
|
| 240 |
+
inv_freq = 1.0 / (base ** (channel_range / head_dim))
|
| 241 |
+
t = torch.arange(seq_len, dtype=torch.float32)
|
| 242 |
+
freqs = torch.outer(t, inv_freq)
|
| 243 |
+
cos = freqs.cos()[None, :, None, :]
|
| 244 |
+
sin = freqs.sin()[None, :, None, :]
|
| 245 |
+
return cos, sin
|
| 246 |
+
|
| 247 |
+
def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs):
|
| 248 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 249 |
+
B, T = input_ids.size()
|
| 250 |
+
if not self._rotary_initialized:
|
| 251 |
+
cos, sin = self._precompute_rotary(self.rotary_seq_len, self.config.head_dim)
|
| 252 |
+
self.cos = cos.to(device=input_ids.device)
|
| 253 |
+
self.sin = sin.to(device=input_ids.device)
|
| 254 |
+
self._rotary_initialized = True
|
| 255 |
+
cos = self.cos[:, :T].to(dtype=self.wte.weight.dtype)
|
| 256 |
+
sin = self.sin[:, :T].to(dtype=self.wte.weight.dtype)
|
| 257 |
+
x = self.wte(input_ids)
|
| 258 |
+
x = _rms_norm(x)
|
| 259 |
+
if T > 1:
|
| 260 |
+
gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
|
| 261 |
+
x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
|
| 262 |
+
x0 = x
|
| 263 |
+
n_layer = self.config.n_layer
|
| 264 |
+
backout_layer = n_layer // 2
|
| 265 |
+
x_backout = None
|
| 266 |
+
for i, block in enumerate(self.h):
|
| 267 |
+
x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
|
| 268 |
+
ve = self.value_embeds[str(i)](input_ids).to(x.dtype) if str(i) in self.value_embeds else None
|
| 269 |
+
x = block(x, ve, cos, sin, self.window_sizes[i])
|
| 270 |
+
if i == backout_layer:
|
| 271 |
+
x_backout = x
|
| 272 |
+
if x_backout is not None:
|
| 273 |
+
x = x - self.backout_lambda.to(x.dtype) * x_backout
|
| 274 |
+
x = _rms_norm(x)
|
| 275 |
+
logits = self.lm_head(x)
|
| 276 |
+
logits = logits[..., :self.config.vocab_size]
|
| 277 |
+
logits = logits.float()
|
| 278 |
+
if self.config.logit_softcap > 0:
|
| 279 |
+
cap = self.config.logit_softcap
|
| 280 |
+
logits = cap * torch.tanh(logits / cap)
|
| 281 |
+
loss = None
|
| 282 |
+
if labels is not None:
|
| 283 |
+
shift_logits = logits[..., :-1, :].contiguous()
|
| 284 |
+
shift_labels = labels[..., 1:].contiguous()
|
| 285 |
+
loss = F.cross_entropy(
|
| 286 |
+
shift_logits.view(-1, shift_logits.size(-1)),
|
| 287 |
+
shift_labels.view(-1),
|
| 288 |
+
ignore_index=-100,
|
| 289 |
+
)
|
| 290 |
+
if not return_dict:
|
| 291 |
+
output = (logits,)
|
| 292 |
+
return (loss,) + output if loss is not None else output
|
| 293 |
+
return CausalLMOutputWithPast(loss=loss, logits=logits)
|
| 294 |
+
|
| 295 |
+
def prepare_inputs_for_generation(self, input_ids, **kwargs):
|
| 296 |
+
return {"input_ids": input_ids}
|
| 297 |
+
|
| 298 |
+
|
| 299 |
+
# ---------------------------------------------------------------------------
|
| 300 |
+
# CLI demo
|
| 301 |
+
# ---------------------------------------------------------------------------
|
| 302 |
+
|
| 303 |
+
if __name__ == "__main__":
|
| 304 |
+
import argparse
|
| 305 |
+
import sys
|
| 306 |
+
import os
|
| 307 |
+
|
| 308 |
+
parser = argparse.ArgumentParser(description="Load and generate with a NanoChat model")
|
| 309 |
+
parser.add_argument("--model-dir", required=True, help="Path to converted model directory")
|
| 310 |
+
parser.add_argument("--prompt", default="The meaning of life is", help="Text prompt")
|
| 311 |
+
parser.add_argument("--max-tokens", type=int, default=100, help="Max tokens to generate")
|
| 312 |
+
parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature (0=greedy)")
|
| 313 |
+
parser.add_argument("--device", default="cpu", help="Device (cpu, cuda, cuda:0, ...)")
|
| 314 |
+
parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
|
| 315 |
+
args = parser.parse_args()
|
| 316 |
+
|
| 317 |
+
dtype_map = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
|
| 318 |
+
|
| 319 |
+
print(f"Loading model from {args.model_dir}...")
|
| 320 |
+
from transformers import AutoTokenizer
|
| 321 |
+
|
| 322 |
+
model = NanoChatForCausalLM.from_pretrained(args.model_dir, torch_dtype=dtype_map[args.dtype])
|
| 323 |
+
model.to(args.device)
|
| 324 |
+
model.eval()
|
| 325 |
+
tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
|
| 326 |
+
|
| 327 |
+
total_params = sum(p.numel() for p in model.parameters())
|
| 328 |
+
print(f"Model loaded: {total_params:,} parameters, dtype={args.dtype}, device={args.device}")
|
| 329 |
+
print(f"Vocab size: {model.config.vocab_size}, ve_dim: {model.config.ve_dim}")
|
| 330 |
+
print()
|
| 331 |
+
print(f"Prompt: {args.prompt}")
|
| 332 |
+
print(f"Generating (max {args.max_tokens} tokens, temp={args.temperature})...")
|
| 333 |
+
print()
|
| 334 |
+
|
| 335 |
+
# Generate
|
| 336 |
+
bos_id = tokenizer.bos_token_id
|
| 337 |
+
input_ids = tokenizer.encode(args.prompt, add_special_tokens=False)
|
| 338 |
+
if bos_id is not None:
|
| 339 |
+
input_ids = [bos_id] + input_ids
|
| 340 |
+
input_ids = torch.tensor([input_ids], dtype=torch.long, device=args.device)
|
| 341 |
+
|
| 342 |
+
generated = []
|
| 343 |
+
with torch.inference_mode():
|
| 344 |
+
for _ in range(args.max_tokens):
|
| 345 |
+
logits = model(input_ids).logits[:, -1, :]
|
| 346 |
+
if args.temperature > 0:
|
| 347 |
+
logits = logits / args.temperature
|
| 348 |
+
v, _ = torch.topk(logits, min(50, logits.size(-1)))
|
| 349 |
+
logits[logits < v[:, [-1]]] = -float("inf")
|
| 350 |
+
probs = torch.softmax(logits, dim=-1)
|
| 351 |
+
next_id = torch.multinomial(probs, num_samples=1)
|
| 352 |
+
else:
|
| 353 |
+
next_id = torch.argmax(logits, dim=-1, keepdim=True)
|
| 354 |
+
input_ids = torch.cat([input_ids, next_id], dim=1)
|
| 355 |
+
token_id = next_id.item()
|
| 356 |
+
generated.append(token_id)
|
| 357 |
+
if tokenizer.eos_token_id is not None and token_id == tokenizer.eos_token_id:
|
| 358 |
+
break
|
| 359 |
+
|
| 360 |
+
output = tokenizer.decode(generated, skip_special_tokens=True)
|
| 361 |
+
print(f"Output: {args.prompt}{output}")
|
models/full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch/tokenizer.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
models/full-128k-superbpe-clean-fineweb2full-plus2-cappedv2-hybridwindow-t110k-v130k-consv2-mathcode-scratch/tokenizer_config.json
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"tokenizer_class": "PreTrainedTokenizerFast",
|
| 3 |
+
"model_max_length": 2048,
|
| 4 |
+
"bos_token": "<s>",
|
| 5 |
+
"add_bos_token": false,
|
| 6 |
+
"clean_up_tokenization_spaces": false
|
| 7 |
+
}
|
models/full-128k-whitespace-balanced-bpe-mathcode-scratch/README.md
ADDED
|
@@ -0,0 +1,76 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# `full-128k-whitespace-balanced-bpe-mathcode-scratch`
|
| 2 |
+
|
| 3 |
+
A 1.27B parameter model trained with the `whitespace-balanced-bpe` tokenizer. Regime: mathcode-20B-from-scratch (24 layers, width 1536, 1.27B parameters, 20.0B tokens, math and code mixture, trained from scratch). Reported in: math and code ablations.
|
| 4 |
+
|
| 5 |
+
## Load
|
| 6 |
+
|
| 7 |
+
```python
|
| 8 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
| 9 |
+
|
| 10 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 11 |
+
"cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-whitespace-balanced-bpe-mathcode-scratch",
|
| 12 |
+
trust_remote_code=True, dtype="float32")
|
| 13 |
+
tokenizer = AutoTokenizer.from_pretrained("cmeister/tokenizer-lm-ablations", subfolder="models/full-128k-whitespace-balanced-bpe-mathcode-scratch")
|
| 14 |
+
|
| 15 |
+
text = "def add(a, b):\n return a + b\n"
|
| 16 |
+
ids = [0] + tokenizer(text)["input_ids"] # training prepends BOS
|
| 17 |
+
```
|
| 18 |
+
|
| 19 |
+
## Settings
|
| 20 |
+
|
| 21 |
+
| Setting | Value |
|
| 22 |
+
|---|---|
|
| 23 |
+
| layers | 24 |
|
| 24 |
+
| model dimension | 1536 |
|
| 25 |
+
| attention heads | 12 |
|
| 26 |
+
| key and value heads | 12 |
|
| 27 |
+
| context length | 2048 |
|
| 28 |
+
| sliding window pattern | SSSL |
|
| 29 |
+
| value embedding dimension | 128 |
|
| 30 |
+
| logit soft cap | 15 |
|
| 31 |
+
| vocabulary size | 128256 |
|
| 32 |
+
| parameters | 1,272,841,994 |
|
| 33 |
+
| training step | 19073 |
|
| 34 |
+
| token budget | 19,999,490,048 |
|
| 35 |
+
| bytes consumed | 143,259,038,643 |
|
| 36 |
+
| seed | 42 |
|
| 37 |
+
| Muon matrix learning rate | 0.0200 |
|
| 38 |
+
| AdamW embedding learning rate | 0.3000 |
|
| 39 |
+
| AdamW unembedding learning rate | 0.0080 |
|
| 40 |
+
| weight decay | 0.2800 |
|
| 41 |
+
| width learning-rate exponent | -1 |
|
| 42 |
+
| weight dtype | as trained (float32 matrices, bfloat16 embeddings) |
|
| 43 |
+
| BOS token id | 0 |
|
| 44 |
+
|
| 45 |
+
## Results
|
| 46 |
+
|
| 47 |
+
| Metric | Value |
|
| 48 |
+
|---|---|
|
| 49 |
+
| validation BPB | 0.3395 |
|
| 50 |
+
| code BPB, mean over 7 languages | 0.3742 |
|
| 51 |
+
| FLORES-200 BPB, mean | 3.3220 |
|
| 52 |
+
| FLORES-200 languages scored | 214 |
|
| 53 |
+
| BLiMP accuracy | n/a |
|
| 54 |
+
| MBPP pass@1 | 0.1020 |
|
| 55 |
+
| HumanEval pass@1 | 0.1768 |
|
| 56 |
+
| GSM8K exact match, strict | 0.1774 |
|
| 57 |
+
|
| 58 |
+
Generation metrics use specification `{'version': 'v2-2026-07-30', 'bos_context': True, 'eos_stop': True, 'token_healing': True, 'truncation': 'left_to_max_length_minus_max_gen', 'stop_truncation': 'earliest_match_all_terms', 'stop_match_decode': 'specials_kept'}`, token healing on.
|
| 59 |
+
|
| 60 |
+
Missing measurements, with the reason each is absent:
|
| 61 |
+
|
| 62 |
+
- `blimp_acc`: no full-128k-whitespace-balanced-bpe-mathcode-scratch_blimp_code_bpb.json
|
| 63 |
+
|
| 64 |
+
BLiMP is comparable within a training scale only. Its scoring sums unnormalized sentence log-likelihoods, and under that scoring a group of quantifier subtasks drops at the 1.27B scale relative to the 300M scale.
|
| 65 |
+
|
| 66 |
+
## Registry note
|
| 67 |
+
|
| 68 |
+
#36, whitespace pretok, slug `whitespace-balanced-bpe`. BPE with minimal WhitespaceSplit + ByteLevel pretok (lower-end pretok anchor), `balanced`. FINISHED (step 19073). MBPP 0.006 [0.000, 0.014].
|
| 69 |
+
|
| 70 |
+
## Provenance
|
| 71 |
+
|
| 72 |
+
- Source checkpoint: `/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-whitespace-balanced-bpe-mathcode-scratch`, step 19073
|
| 73 |
+
- Repository commit: `1a9a718`
|
| 74 |
+
- nanochat commit: `f919449` on branch `tokenizer-lm-patches`
|
| 75 |
+
- Tokenizer: `whitespace-balanced-bpe`, released at `tokenizers/whitespace-balanced-bpe`
|
| 76 |
+
|
models/full-128k-whitespace-balanced-bpe-mathcode-scratch/config.json
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"NanoChatForCausalLM"
|
| 4 |
+
],
|
| 5 |
+
"auto_map": {
|
| 6 |
+
"AutoConfig": "nanochat_model.NanoChatConfig",
|
| 7 |
+
"AutoModelForCausalLM": "nanochat_model.NanoChatForCausalLM"
|
| 8 |
+
},
|
| 9 |
+
"bos_token_id": 0,
|
| 10 |
+
"dtype": "float32",
|
| 11 |
+
"eos_token_id": null,
|
| 12 |
+
"head_dim": 128,
|
| 13 |
+
"kv_dim": 1536,
|
| 14 |
+
"logit_softcap": 15.0,
|
| 15 |
+
"model_type": "nanochat",
|
| 16 |
+
"n_embd": 1536,
|
| 17 |
+
"n_head": 12,
|
| 18 |
+
"n_kv_head": 12,
|
| 19 |
+
"n_layer": 24,
|
| 20 |
+
"pad_vocab_size_to": 64,
|
| 21 |
+
"padded_vocab_size": 128256,
|
| 22 |
+
"sequence_len": 2048,
|
| 23 |
+
"tie_word_embeddings": false,
|
| 24 |
+
"transformers_version": "5.3.0",
|
| 25 |
+
"ve_dim": 128,
|
| 26 |
+
"vocab_size": 128256,
|
| 27 |
+
"window_pattern": "SSSL"
|
| 28 |
+
}
|
models/full-128k-whitespace-balanced-bpe-mathcode-scratch/conversion_metadata.json
ADDED
|
@@ -0,0 +1,100 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"run": "full-128k-whitespace-balanced-bpe-mathcode-scratch",
|
| 3 |
+
"source_checkpoint": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/checkpoints/full-128k-whitespace-balanced-bpe-mathcode-scratch",
|
| 4 |
+
"source_model_file": "model_019073.pt",
|
| 5 |
+
"source_step": 19073,
|
| 6 |
+
"source_meta": {
|
| 7 |
+
"step": 19073,
|
| 8 |
+
"val_bpb": 0.3395392358308889,
|
| 9 |
+
"model_config": {
|
| 10 |
+
"sequence_len": 2048,
|
| 11 |
+
"vocab_size": 128256,
|
| 12 |
+
"n_layer": 24,
|
| 13 |
+
"n_head": 12,
|
| 14 |
+
"n_kv_head": 12,
|
| 15 |
+
"n_embd": 1536,
|
| 16 |
+
"window_pattern": "SSSL",
|
| 17 |
+
"ve_dim": 128,
|
| 18 |
+
"logit_softcap": 15.0
|
| 19 |
+
},
|
| 20 |
+
"tokenizer_path": "/capstor/store/cscs/swissai/a139/cmeister/tokenizer-lm/tokenizers/whitespace-balanced-bpe",
|
| 21 |
+
"training_config": {
|
| 22 |
+
"name": "full-128k-mathcode-scratch",
|
| 23 |
+
"model": {
|
| 24 |
+
"n_layer": 24,
|
| 25 |
+
"n_embd": 1536,
|
| 26 |
+
"n_kv_head": 12,
|
| 27 |
+
"sequence_len": 2048,
|
| 28 |
+
"window_pattern": "SSSL",
|
| 29 |
+
"ve_dim": 128
|
| 30 |
+
},
|
| 31 |
+
"training": {
|
| 32 |
+
"target_param_data_ratio": -1,
|
| 33 |
+
"max_tokens": 20000000000,
|
| 34 |
+
"num_iterations": -1,
|
| 35 |
+
"total_batch_size": 1048576,
|
| 36 |
+
"device_batch_size": 16,
|
| 37 |
+
"warmup_steps": 40,
|
| 38 |
+
"warmdown_ratio": 0.65,
|
| 39 |
+
"final_lr_frac": 0.05,
|
| 40 |
+
"eval_tokens": 41943040,
|
| 41 |
+
"compile": true
|
| 42 |
+
},
|
| 43 |
+
"optimizer": {
|
| 44 |
+
"matrix_lr": 0.02,
|
| 45 |
+
"embedding_lr": 0.3,
|
| 46 |
+
"unembedding_lr": 0.008,
|
| 47 |
+
"scalar_lr": 0.5,
|
| 48 |
+
"weight_decay": 0.28,
|
| 49 |
+
"width_lr_exponent": -1.0
|
| 50 |
+
},
|
| 51 |
+
"data": {
|
| 52 |
+
"data_dir": null
|
| 53 |
+
},
|
| 54 |
+
"logging": {
|
| 55 |
+
"eval_every": 500,
|
| 56 |
+
"save_every": 500,
|
| 57 |
+
"max_checkpoints": 0
|
| 58 |
+
},
|
| 59 |
+
"wandb": {
|
| 60 |
+
"enabled": true,
|
| 61 |
+
"project": "tokenizer-lm-experiments",
|
| 62 |
+
"run_name": null
|
| 63 |
+
},
|
| 64 |
+
"output_dir": "/capstor/scratch/cscs/${USER}/tokenizer-lm/checkpoints",
|
| 65 |
+
"seed": 42
|
| 66 |
+
},
|
| 67 |
+
"total_batch_size": 1048576,
|
| 68 |
+
"device_batch_size": 16,
|
| 69 |
+
"max_seq_len": 2048,
|
| 70 |
+
"dataloader_state_dict": {
|
| 71 |
+
"pq_idx": 344,
|
| 72 |
+
"rg_idx": 12,
|
| 73 |
+
"epoch": 2
|
| 74 |
+
},
|
| 75 |
+
"loop_state": {
|
| 76 |
+
"min_val_bpb": 0.3395392358308889,
|
| 77 |
+
"smooth_train_loss": 1.1001896170367296,
|
| 78 |
+
"total_training_time": 94029.81803798676,
|
| 79 |
+
"total_bytes_consumed": 143259038643
|
| 80 |
+
}
|
| 81 |
+
},
|
| 82 |
+
"weight_dtype": "as trained (float32 matrices, bfloat16 embeddings)",
|
| 83 |
+
"weight_dtype_counts": {
|
| 84 |
+
"torch.float32": 174,
|
| 85 |
+
"torch.bfloat16": 13
|
| 86 |
+
},
|
| 87 |
+
"total_parameters": 1272841994,
|
| 88 |
+
"logit_softcap_source": "meta",
|
| 89 |
+
"bos_token": "<s>",
|
| 90 |
+
"bos_token_id": 0,
|
| 91 |
+
"bos_resolved_from": "packaged tokenizer",
|
| 92 |
+
"tokenizer_slug": "whitespace-balanced-bpe",
|
| 93 |
+
"tokenizer_redistributed": true,
|
| 94 |
+
"tokenizer_hub_id": null,
|
| 95 |
+
"weight_key_mapping": "strip _orig_mod. and transformer. prefixes",
|
| 96 |
+
"converted_by": "scripts/hf_release/convert_for_release.py",
|
| 97 |
+
"repo_commit": "1a9a718",
|
| 98 |
+
"nanochat_commit": "f919449",
|
| 99 |
+
"nanochat_branch": "tokenizer-lm-patches"
|
| 100 |
+
}
|
models/full-128k-whitespace-balanced-bpe-mathcode-scratch/tokenizer.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
models/full-128k-whitespace-balanced-bpe-mathcode-scratch/tokenizer_config.json
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"tokenizer_class": "PreTrainedTokenizerFast",
|
| 3 |
+
"model_max_length": 2048,
|
| 4 |
+
"bos_token": "<s>",
|
| 5 |
+
"add_bos_token": false,
|
| 6 |
+
"clean_up_tokenization_spaces": false
|
| 7 |
+
}
|