diff --git a/.gitignore b/.gitignore index 3423c416a7..a040f4db3b 100644 --- a/.gitignore +++ b/.gitignore @@ -8,4 +8,7 @@ data/manifest.json data/docs_selected.jsonl .mypy_cache/ .venv -logs/ \ No newline at end of file +logs/ +recent_no_cheating_literature.md +top_pr_recipe_notes.md +recent_advances.md \ No newline at end of file diff --git a/pr_report_canon_acd_k3.md b/pr_report_canon_acd_k3.md new file mode 100644 index 0000000000..105dbcd256 --- /dev/null +++ b/pr_report_canon_acd_k3.md @@ -0,0 +1,67 @@ +Submission: Int6 + Canon ACD (K=3) + Muon WD 0.04 + SWA + Sliding Eval (val_bpb=1.1668) + +This PR reports a standalone run using Canon ACD (`CANON_SET=ACD`) with `CANON_KERNEL=3` and mixed int6 quantization for `mlp,attn`. + +Approach summary: +- Architecture: 9-layer decoder-only Transformer, `model_dim=512`, `num_heads=8`, `num_kv_heads=4`, `MLP_MULT=3.0`. +- MLP nonlinearity: ReLU-squared style MLP as used in this repo. +- Context modules: Bigram hash embedding (`bigram_vocab_size=2048`, `bigram_dim=128`) and SmearGate. +- Quantization: mixed post-training quantization where `mlp/attn` are int6 and remaining large tensors stay int8; +- Optimizer: mixed Muon + Adam setup. Muon handles matrix-like parameters; Adam handles token/scalar/head groups. +- Schedule: momentum warmup (`0.92 -> 0.99`), warmdown tail (`WARMDOWN_ITERS=3000`), and SWA averaging near the end. +- Eval: report both final roundtrip and sliding-window eval (`EVAL_STRIDE=64`), where sliding bpb is the key comparison metric. + +Canon details: +- Canon layer is a depthwise causal 1D conv with residual connection. +- `A`: Canon before attention. +- `B`: Canon on concatenated QKV stream (expensive because width is larger). +- `C`: Canon before MLP. +- `D`: Canon in the widened MLP hidden stream. +- This run uses `ACD` to keep most Canon effect while avoiding `B` compute cost. + +Configuration highlights: +- 8x GPU (`torchrun --nproc_per_node=8`) +- `TRAIN_BATCH_TOKENS=524288`, `TRAIN_SEQ_LEN=2048` +- `EVAL_SEQ_LEN=2048`, `EVAL_STRIDE=64`, `EVAL_BATCH_SEQS=32` +- `INT6_CATEGORIES=mlp,attn` +- `CANON_SET=ACD`, `CANON_KERNEL=3`, `CANON_RESIDUAL=1`, `CANON_ACTIVATION=0`, `CANON_BIAS=0` +- `MATRIX_LR=0.025`, `SCALAR_LR=0.025`, `TIED_EMBED_LR=0.035` +- `MUON_MOMENTUM=0.99`, `MUON_WEIGHT_DECAY=0.04`, `ADAM_WEIGHT_DECAY=0.04` +- `SWA_ENABLED=1`, `SWA_EVERY=200`, `SWA_START_LRMUL=0.5` +- `ITERATIONS=7200`, `WARMUP_STEPS=20`, `WARMDOWN_ITERS=3000`, `MAX_WALLCLOCK_SECONDS=600` +- `VOCAB_SIZE=1024`, `SEED=1337` + +Run output: +- `final_int6_sliding_window val_bpb` (stride=64): **1.16682362** +- `final_model.int6.ptz`: **13,196,032 bytes** +- Code size (`train_gpt.py`): **71,315 bytes** +- Total submission size: **13,267,347 bytes** (under 16,000,000 limit) + +Notes: +- Main score to compare is sliding-window bpb: **1.16682362**. +- SWA applied with 8 checkpoints. +- Data loader overhead is low (`data_loading_step_avg=0.64ms`). +- End-of-run metrics: + - `data_loading_total:4637ms` + - `final_int6_sliding_window eval_time:246064ms` + +Repro command: +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ +env \ + RUN_ID=frontier_canon_acd_k3_8gpu \ + DATA_PATH=./data/datasets/fineweb10B_sp1024/ \ + TOKENIZER_PATH=./data/tokenizers/fineweb_1024_bpe.model \ + VOCAB_SIZE=1024 SEED=1337 \ + TRAIN_BATCH_TOKENS=524288 TRAIN_SEQ_LEN=2048 \ + EVAL_SEQ_LEN=2048 EVAL_STRIDE=64 EVAL_BATCH_SEQS=32 \ + ITERATIONS=7200 WARMUP_STEPS=20 WARMDOWN_ITERS=3000 MAX_WALLCLOCK_SECONDS=600 \ + MATRIX_LR=0.025 SCALAR_LR=0.025 TIED_EMBED_LR=0.035 \ + MUON_MOMENTUM=0.99 MUON_MOMENTUM_WARMUP_START=0.92 MUON_MOMENTUM_WARMUP_STEPS=1500 \ + MUON_WEIGHT_DECAY=0.04 ADAM_WEIGHT_DECAY=0.04 \ + SWA_ENABLED=1 SWA_EVERY=200 SWA_START_LRMUL=0.5 \ + INT6_CATEGORIES=mlp,attn \ + CANON_SET=ACD CANON_KERNEL=3 CANON_RESIDUAL=1 CANON_ACTIVATION=0 CANON_BIAS=0 \ + TORCHINDUCTOR_CACHE_DIR=/tmp/torchinductor \ + torchrun --standalone --nproc_per_node=8 train_gpt.py +``` diff --git a/pr_report_canon_deltagate_11layer.md b/pr_report_canon_deltagate_11layer.md new file mode 100644 index 0000000000..ed296229a8 --- /dev/null +++ b/pr_report_canon_deltagate_11layer.md @@ -0,0 +1,139 @@ +# Record: 11L CANON-AC(last5)+DeltaGate Report (Humble Record Attempt) + +## Summary + +This run builds on the **current leaderboard-aligned stack** (official + pending-validated direction) and focuses on a scoped CANON placement with CANON delta gate. + +Best observed result in this sweep: + +- `final_int6_sliding_window_exact val_bpb: 1.12961770` + +Compared to my previous PR [#312](https://github.com/openai/parameter-golf/pull/312): + +- `1.16682362 -> 1.12961770` (large improvement) + +## Quick Comparison (vs #312) + + +| Run | Setup | Steps Before Wallclock Stop | Final sliding-window val_bpb | Submission size (int6+zstd) | +| ------------------------------------------------------------------ | ------------------------ | ---------------------------- | ---------------------------- | --------------------------- | +| Previous [#312](https://github.com/openai/parameter-golf/pull/312) | ACD (all) + SWA | 7210 (batch size is default) | `1.16682362` | `13,267,347` bytes | +| This work (seed 1337) | AC(last5)+delta+tightSWA | `6278` | `1.12961770` | `15,581,348` bytes | +| This work (seed 1336) | AC(last5)+delta+tightSWA | `6243` | `1.1303` | `15,505,544` bytes | +| This work (seed 1335) | AC(last5)+delta+tightSWA | `6252` | `1.12970337` | `15,579,865` bytes | + + +## What Was Reused From Current Leaderboard (not unofficial-only additions) + +This run intentionally reuses patterns already common in official/pending leaderboard entries, to check the possibility of Canon layers.: + +- 11L / 512-dim / GQA (8 heads, 4 KV heads), MLP 3x +- BigramHash + SmearGate +- XSA on last 4 layers (`XSA_LAST_N=4`) +- Partial RoPE (`ROPE_DIMS=16`) + LN Scale +- Late QAT +- WD 0.04, Tight SWA schedule +- Sliding-window eval (`stride=64`) + +## Main Configuration (this report) + +- `CANON_SET=AC` +- `CANON_LAST_N=5` +- `CANON_DELTA_GATE=1` +- `SWA_ENABLED=1`, `TIGHT_SWA=1`, `TIGHT_SWA_EVERY=50`, `TIGHT_SWA_START_LRMUL=0.2`, `TIGHT_SWA_MAX_CHECKPOINTS=12` +- `TRAIN_BATCH_TOKENS=786432`, wallclock-capped run (`MAX_WALLCLOCK_SECONDS=600`) + +## Definitions (for this report) + +- `Delta` (in `AC(last5)+delta`) means **CANON delta gate**: + - `CANON_DELTA_GATE=1` + - each CANON branch output is scaled by a learnable sigmoid gate before residual add. +- `Last 4` means **XSA is enabled only on the last 4 transformer blocks**: + - `XSA_LAST_N=4` +- `XSA learnable gate` means an extra learnable scalar that mixes normal attention output and XSA output: + - `y <- y + sigmoid(g) * (y_xsa - y)` + - controlled by `XSA_LEARNABLE_GATE` and `XSA_GATE_INIT` + +## Final Run Command (renamed RUN_ID) + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ +env \ + RUN_ID=frontier_canon_ac_k3_8gpu_final_report_seed1336 \ + DATA_PATH=./data/datasets/fineweb10B_sp1024/ \ + TOKENIZER_PATH=./data/tokenizers/fineweb_1024_bpe.model \ + VOCAB_SIZE=1024 SEED=1336 \ + NUM_LAYERS=11 MODEL_DIM=512 NUM_HEADS=8 NUM_KV_HEADS=4 MLP_MULT=3.0 \ + TRAIN_SEQ_LEN=2048 \ + EVAL_SEQ_LEN=2048 EVAL_STRIDE=64 EVAL_BATCH_SEQS=32 \ + ITERATIONS=7000 WARMUP_STEPS=20 WARMDOWN_ITERS=3000 MAX_WALLCLOCK_SECONDS=600 \ + MATRIX_LR=0.025 SCALAR_LR=0.025 TIED_EMBED_LR=0.035 \ + MUON_MOMENTUM=0.99 MUON_MOMENTUM_WARMUP_START=0.92 MUON_MOMENTUM_WARMUP_STEPS=1500 \ + MUON_WEIGHT_DECAY=0.04 ADAM_WEIGHT_DECAY=0.04 \ + EMA_ENABLED=0 \ + SWA_ENABLED=1 TIGHT_SWA=1 TIGHT_SWA_EVERY=50 TIGHT_SWA_START_LRMUL=0.2 TIGHT_SWA_MAX_CHECKPOINTS=12 \ + XSA_LAST_N=4 ROPE_DIMS=16 LN_SCALE=1 \ + LATE_QAT=1 QAT_THRESHOLD=0.1 \ + INT6_CATEGORIES=mlp,attn TRAIN_BATCH_TOKENS=786432 GRAD_CLIP_NORM=0.3 \ + CANON_SET=AC CANON_KERNEL=3 CANON_RESIDUAL=1 CANON_ACTIVATION=0 CANON_BIAS=0 \ + CANON_FIRST_N=0 CANON_LAST_N=5 CANON_DELTA_GATE=1 CANON_DELTA_GATE_INIT=-4.0 \ + TORCHINDUCTOR_CACHE_DIR=/tmp/torchinductor \ + torchrun --standalone --nproc_per_node=8 train_gpt.py +``` + +## Results + +### Seed-level excerpts + +- Seed `1337`: + - `step:6278/7000 val_loss:1.9339 val_bpb:1.1454` + - `final_int6_sliding_window_exact val_loss:1.90730712 val_bpb:1.12961770` + - `Total submission size int6+zstd: 15581348 bytes` +- Seed `1335`: + - `step:6252/7000 val_loss:1.9349 val_bpb:1.1460` + - `final_int6_sliding_window_exact val_loss:1.90745178 val_bpb:1.12970337` + - `Total submission size int6+zstd: 15579865 bytes` +- Seed `1336`: + - `step:6243/7000 val_loss:1.9365 val_bpb:1.1469` + - `final_int6_sliding_window_exact val_bpb: 1.1303` + - `Total submission size int6+zstd: 15505544 bytes` + +### Wallclock / speed notes + +- AC(last5)+delta runs stopped around ~6250-6280 steps due to 600s wallclock cap. +- No-canon run reached `6930` steps under the same cap (faster, but lower quality). + +## Ablations (sliding-window val_bpb) + +- Full CANON `ACD`: `1.14083538` +- CANON `AC` (broad): `1.13218808` +- CANON `AC` (first 4 layers): `1.1314` +- No CANON: `1.13587538` -- it was faster, but it doesn't have a better bpb. +- CANON `AC(last5)+delta`: best observed `1.1296` +- XSA learnable gate (`XSA_LEARNABLE_GATE=1`): not helpful here (`~1.131`) + +## Comparison vs Previous PR + +Previous: [#312](https://github.com/openai/parameter-golf/pull/312) + +- `final_int6_sliding_window_exact val_bpb: 1.16682362` + +Current best in this report: + +- `final_int6_sliding_window_exact val_bpb: 1.12961770` + +Approx improvement: + +- `Δ bpb = -0.03720592` +- `Δ nats ≈ 0.0258` (using `bpb * ln(2)` conversion) + +## Significance Note + +Against official SOTA context (`1.1428 BPB`), this run clears the `>=0.005 nat` improvement margin by a comfortable amount in point estimate. +For formal `p < 0.01` reporting, include the completed 3-seed list (1335/1336/1337) and test output in PR comments. + +## Humble Notes + +- This is an incremental engineering result built on existing leaderboard-proven ideas plus scoped CANON placement and gating. +- The strongest gain seems to come from the interaction of AC(last5), CANON delta gate, and tight SWA under the same compute budget. + diff --git a/records/track_10min_16mb/2026-03-17_NaiveBaseline/README.md b/records/track_10min_16mb/2026-03-17_NaiveBaseline/README.md deleted file mode 100644 index 1eb352a8a3..0000000000 --- a/records/track_10min_16mb/2026-03-17_NaiveBaseline/README.md +++ /dev/null @@ -1,46 +0,0 @@ -This record captures the `Simple Baseline`. - -Trainer changes in this snapshot: -- current repository `train_gpt.py` snapshot copied into the record folder -- published `fineweb10B_sp1024` dataset and tokenizer loaded from the new Hugging Face export -- 10-minute wallclock cap on `8xH100` -- periodic validation every `200` steps on the full `fineweb_val_*` split - -Configuration: -- Layout: `VOCAB_SIZE=1024 NUM_LAYERS=9 MODEL_DIM=512 NUM_HEADS=8 NUM_KV_HEADS=4 MLP_MULT=2` -- Tied output/input embeddings: `TIE_EMBEDDINGS=1` -- Tied embedding LR: `TIED_EMBED_LR=0.05` -- Batching: `TRAIN_BATCH_TOKENS=524288 TRAIN_SEQ_LEN=1024` - -Command (track-relevant params): -```bash -NCCL_IB_DISABLE=1 \ -RUN_ID=hf_verify_sp1024_8gpu \ -DATA_PATH=/root/code/parameter-golf/data/datasets/fineweb10B_sp1024 \ -TOKENIZER_PATH=/root/code/parameter-golf/data/tokenizers/fineweb_1024_bpe.model \ -VOCAB_SIZE=1024 \ -MAX_WALLCLOCK_SECONDS=600 \ -TRAIN_LOG_EVERY=50 \ -VAL_LOSS_EVERY=200 \ -torchrun --standalone --nproc_per_node=8 /root/code/parameter-golf/train_gpt.py -``` - -Key metrics (from `train.log`): -- Timed training stopped at `13780/20000` steps due to the wallclock cap. -- Pre-quant eval at stop: `val_loss:2.0606`, `val_bpb:1.2172` -- Post-quant roundtrip eval: `val_loss:2.0727`, `val_bpb:1.2244` -- Exact printed metric: `final_int8_zlib_roundtrip_exact val_bpb:1.22436570` -- Train time: `600038ms` (`step_avg:43.54ms`) -- Peak memory: `10184 MiB allocated`, `10200 MiB reserved` -- Serialized model int8+zlib: `15815847 bytes` -- Code size: `47642 bytes` -- Total submission size int8+zlib: `15863489 bytes` - -Training volume: -- Global batch: `524288` tokens/step -- Total train tokens seen: `7224688640` - -Included files: -- `train_gpt.py` (code snapshot used for the run) -- `train.log` (exact remote training log) -- `submission.json` (leaderboard metadata) diff --git a/records/track_10min_16mb/2026-03-17_NaiveBaseline/submission.json b/records/track_10min_16mb/2026-03-17_NaiveBaseline/submission.json deleted file mode 100644 index faffa83ea8..0000000000 --- a/records/track_10min_16mb/2026-03-17_NaiveBaseline/submission.json +++ /dev/null @@ -1,11 +0,0 @@ -{ - "author": "Baseline", - "github_id": "openai", - "name": "Naive Baseline", - "blurb": "SP-1024 9x512 KV4 run on pgut1 using the published Hugging Face fineweb10B_sp1024 export and the current train_gpt.py; score is the default final int8+zlib roundtrip metric under the 16,000,000-byte cap.", - "date": "2026-03-18T14:56:29Z", - "val_loss": 2.07269931, - "val_bpb": 1.2243657, - "bytes_total": 15863489, - "bytes_code": 47642 -} diff --git a/records/track_10min_16mb/2026-03-17_NaiveBaseline/train.log b/records/track_10min_16mb/2026-03-17_NaiveBaseline/train.log deleted file mode 100644 index 69b17b6c7f..0000000000 --- a/records/track_10min_16mb/2026-03-17_NaiveBaseline/train.log +++ /dev/null @@ -1,448 +0,0 @@ -W0318 14:37:59.159000 871689 site-packages/torch/distributed/run.py:852] -W0318 14:37:59.159000 871689 site-packages/torch/distributed/run.py:852] ***************************************** -W0318 14:37:59.159000 871689 site-packages/torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. -W0318 14:37:59.159000 871689 site-packages/torch/distributed/run.py:852] ***************************************** -[W318 14:38:11.514156940 Utils.hpp:137] Warning: Environment variable NCCL_ASYNC_ERROR_HANDLING is deprecated; use TORCH_NCCL_ASYNC_ERROR_HANDLING instead (function operator()) -[W318 14:38:11.543417305 Utils.hpp:137] Warning: Environment variable NCCL_ASYNC_ERROR_HANDLING is deprecated; use TORCH_NCCL_ASYNC_ERROR_HANDLING instead (function operator()) -[W318 14:38:11.552597211 Utils.hpp:137] Warning: Environment variable NCCL_ASYNC_ERROR_HANDLING is deprecated; use TORCH_NCCL_ASYNC_ERROR_HANDLING instead (function operator()) -NCCL version 2.27.5+cuda12.9 -[W318 14:38:11.832390267 Utils.hpp:137] Warning: Environment variable NCCL_ASYNC_ERROR_HANDLING is deprecated; use TORCH_NCCL_ASYNC_ERROR_HANDLING instead (function operator()) -[W318 14:38:11.842257581 Utils.hpp:137] Warning: Environment variable NCCL_ASYNC_ERROR_HANDLING is deprecated; use TORCH_NCCL_ASYNC_ERROR_HANDLING instead (function operator()) -[W318 14:38:11.842253680 Utils.hpp:137] Warning: Environment variable NCCL_ASYNC_ERROR_HANDLING is deprecated; use TORCH_NCCL_ASYNC_ERROR_HANDLING instead (function operator()) -[W318 14:38:11.899166383 Utils.hpp:137] Warning: Environment variable NCCL_ASYNC_ERROR_HANDLING is deprecated; use TORCH_NCCL_ASYNC_ERROR_HANDLING instead (function operator()) -[W318 14:38:11.901800020 Utils.hpp:137] Warning: Environment variable NCCL_ASYNC_ERROR_HANDLING is deprecated; use TORCH_NCCL_ASYNC_ERROR_HANDLING instead (function operator()) - -[2026-03-18 14:38:12] pgut1-0:871784:871848 [5] ibvwrap.c:94 NCCL WARN Call to ibv_open_device failed - -[2026-03-18 14:38:12] pgut1-0:871784:871848 [5] p2p_plugin.c:565 NCCL WARN NET/IB : Unable to open device mlx5_an0 - -[2026-03-18 14:38:12] pgut1-0:871786:871849 [7] ibvwrap.c:94 NCCL WARN Call to ibv_open_device failed - -[2026-03-18 14:38:12] pgut1-0:871786:871849 [7] p2p_plugin.c:565 NCCL WARN NET/IB : Unable to open device mlx5_an0 - -[2026-03-18 14:38:12] pgut1-0:871779:871850 [0] ibvwrap.c:94 NCCL WARN Call to ibv_open_device failed - -[2026-03-18 14:38:12] pgut1-0:871779:871850 [0] p2p_plugin.c:565 NCCL WARN NET/IB : Unable to open device mlx5_an0 - -[2026-03-18 14:38:12] pgut1-0:871780:871857 [1] ibvwrap.c:94 NCCL WARN Call to ibv_open_device failed - -[2026-03-18 14:38:12] pgut1-0:871780:871857 [1] p2p_plugin.c:565 NCCL WARN NET/IB : Unable to open device mlx5_an0 - -[2026-03-18 14:38:12] pgut1-0:871781:871858 [2] ibvwrap.c:94 NCCL WARN Call to ibv_open_device failed - -[2026-03-18 14:38:12] pgut1-0:871781:871858 [2] p2p_plugin.c:565 NCCL WARN NET/IB : Unable to open device mlx5_an0 - -[2026-03-18 14:38:12] pgut1-0:871783:871859 [4] ibvwrap.c:94 NCCL WARN Call to ibv_open_device failed - -[2026-03-18 14:38:12] pgut1-0:871783:871859 [4] p2p_plugin.c:565 NCCL WARN NET/IB : Unable to open device mlx5_an0 - -[2026-03-18 14:38:12] pgut1-0:871782:871864 [3] ibvwrap.c:94 NCCL WARN Call to ibv_open_device failed - -[2026-03-18 14:38:12] pgut1-0:871782:871864 [3] p2p_plugin.c:565 NCCL WARN NET/IB : Unable to open device mlx5_an0 - -[2026-03-18 14:38:12] pgut1-0:871785:871865 [6] ibvwrap.c:94 NCCL WARN Call to ibv_open_device failed - -[2026-03-18 14:38:12] pgut1-0:871785:871865 [6] p2p_plugin.c:565 NCCL WARN NET/IB : Unable to open device mlx5_an0 -logs/hf_verify_sp1024_8gpu.txt -val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path=/root/code/parameter-golf/data/tokenizers/fineweb_1024_bpe.model -train_loader:dataset:fineweb10B_sp1024 train_shards:25 -val_loader:shards pattern=/root/code/parameter-golf/data/datasets/fineweb10B_sp1024/fineweb_val_*.bin tokens:63779840 -[rank0]:[W318 14:38:18.833454927 Utils.hpp:112] Warning: Environment variable NCCL_ASYNC_ERROR_HANDLING is deprecated; use TORCH_NCCL_ASYNC_ERROR_HANDLING instead (function operator()) -model_params:17059912 -world_size:8 grad_accum_steps:1 -sdp_backends:cudnn=False flash=True mem_efficient=False math=False -attention_mode:gqa num_heads:8 num_kv_heads:4 -tie_embeddings:True embed_lr:0.05 head_lr:0.0 matrix_lr:0.04 scalar_lr:0.04 -train_batch_tokens:524288 train_seq_len:1024 iterations:20000 warmup_steps:20 max_wallclock_seconds:600.000 -seed:1337 -[rank3]:[W318 14:38:18.835915381 Utils.hpp:112] Warning: Environment variable NCCL_ASYNC_ERROR_HANDLING is deprecated; use TORCH_NCCL_ASYNC_ERROR_HANDLING instead (function operator()) -[rank7]:[W318 14:38:18.835951425 Utils.hpp:112] Warning: Environment variable NCCL_ASYNC_ERROR_HANDLING is deprecated; use TORCH_NCCL_ASYNC_ERROR_HANDLING instead (function operator()) -[rank6]:[W318 14:38:18.835967008 Utils.hpp:112] Warning: Environment variable NCCL_ASYNC_ERROR_HANDLING is deprecated; use TORCH_NCCL_ASYNC_ERROR_HANDLING instead (function operator()) -[rank2]:[W318 14:38:18.836023454 Utils.hpp:112] Warning: Environment variable NCCL_ASYNC_ERROR_HANDLING is deprecated; use TORCH_NCCL_ASYNC_ERROR_HANDLING instead (function operator()) -[rank5]:[W318 14:38:18.836119632 Utils.hpp:112] Warning: Environment variable NCCL_ASYNC_ERROR_HANDLING is deprecated; use TORCH_NCCL_ASYNC_ERROR_HANDLING instead (function operator()) -[rank4]:[W318 14:38:18.836127772 Utils.hpp:112] Warning: Environment variable NCCL_ASYNC_ERROR_HANDLING is deprecated; use TORCH_NCCL_ASYNC_ERROR_HANDLING instead (function operator()) -[rank1]:[W318 14:38:18.836354967 Utils.hpp:112] Warning: Environment variable NCCL_ASYNC_ERROR_HANDLING is deprecated; use TORCH_NCCL_ASYNC_ERROR_HANDLING instead (function operator()) -warmup_step:1/20 -warmup_step:2/20 -warmup_step:3/20 -warmup_step:4/20 -warmup_step:5/20 -warmup_step:6/20 -warmup_step:7/20 -warmup_step:8/20 -warmup_step:9/20 -warmup_step:10/20 -warmup_step:11/20 -warmup_step:12/20 -warmup_step:13/20 -warmup_step:14/20 -warmup_step:15/20 -warmup_step:16/20 -warmup_step:17/20 -warmup_step:18/20 -warmup_step:19/20 -warmup_step:20/20 -step:0/20000 val_loss:6.9370 val_bpb:4.0978 train_time:0ms step_avg:0.01ms -step:1/20000 train_loss:6.9408 train_time:24ms step_avg:23.99ms -step:2/20000 train_loss:16.8763 train_time:67ms step_avg:33.39ms -step:3/20000 train_loss:9.0044 train_time:110ms step_avg:36.62ms -step:4/20000 train_loss:6.5686 train_time:152ms step_avg:37.99ms -step:5/20000 train_loss:6.6665 train_time:195ms step_avg:38.97ms -step:6/20000 train_loss:6.5027 train_time:239ms step_avg:39.81ms -step:7/20000 train_loss:6.2808 train_time:280ms step_avg:40.05ms -step:8/20000 train_loss:5.9951 train_time:324ms step_avg:40.52ms -step:9/20000 train_loss:6.0187 train_time:367ms step_avg:40.77ms -step:10/20000 train_loss:5.9718 train_time:409ms step_avg:40.93ms -step:50/20000 train_loss:3.9508 train_time:2126ms step_avg:42.52ms -step:100/20000 train_loss:3.3373 train_time:4267ms step_avg:42.67ms -step:150/20000 train_loss:2.9651 train_time:6414ms step_avg:42.76ms -step:200/20000 train_loss:2.8041 train_time:8677ms step_avg:43.38ms -step:200/20000 val_loss:2.8397 val_bpb:1.6774 train_time:8699ms step_avg:43.49ms -step:250/20000 train_loss:2.7379 train_time:10816ms step_avg:43.27ms -step:300/20000 train_loss:2.6613 train_time:12958ms step_avg:43.19ms -step:350/20000 train_loss:2.6434 train_time:15097ms step_avg:43.13ms -step:400/20000 train_loss:2.7684 train_time:17357ms step_avg:43.39ms -step:400/20000 val_loss:2.5687 val_bpb:1.5174 train_time:17382ms step_avg:43.45ms -step:450/20000 train_loss:2.6035 train_time:19502ms step_avg:43.34ms -step:500/20000 train_loss:2.5265 train_time:21643ms step_avg:43.29ms -step:550/20000 train_loss:2.4803 train_time:23782ms step_avg:43.24ms -step:600/20000 train_loss:2.4731 train_time:26034ms step_avg:43.39ms -step:600/20000 val_loss:2.4456 val_bpb:1.4447 train_time:26059ms step_avg:43.43ms -step:650/20000 train_loss:2.3204 train_time:28175ms step_avg:43.35ms -step:700/20000 train_loss:2.5926 train_time:30315ms step_avg:43.31ms -step:750/20000 train_loss:2.4301 train_time:32457ms step_avg:43.28ms -step:800/20000 train_loss:2.4775 train_time:34707ms step_avg:43.38ms -step:800/20000 val_loss:2.3868 val_bpb:1.4099 train_time:34732ms step_avg:43.42ms -step:850/20000 train_loss:2.3941 train_time:36851ms step_avg:43.35ms -step:900/20000 train_loss:2.3716 train_time:38990ms step_avg:43.32ms -step:950/20000 train_loss:2.3216 train_time:41131ms step_avg:43.30ms -step:1000/20000 train_loss:2.3030 train_time:43390ms step_avg:43.39ms -step:1000/20000 val_loss:2.3370 val_bpb:1.3805 train_time:43415ms step_avg:43.42ms -step:1050/20000 train_loss:2.3893 train_time:45532ms step_avg:43.36ms -step:1100/20000 train_loss:2.4145 train_time:47675ms step_avg:43.34ms -step:1150/20000 train_loss:2.2261 train_time:49933ms step_avg:43.42ms -step:1200/20000 train_loss:2.2607 train_time:52072ms step_avg:43.39ms -step:1200/20000 val_loss:2.3026 val_bpb:1.3602 train_time:52097ms step_avg:43.41ms -step:1250/20000 train_loss:2.3312 train_time:54219ms step_avg:43.38ms -step:1300/20000 train_loss:2.3575 train_time:56363ms step_avg:43.36ms -step:1350/20000 train_loss:2.2774 train_time:58628ms step_avg:43.43ms -step:1400/20000 train_loss:2.2436 train_time:60772ms step_avg:43.41ms -step:1400/20000 val_loss:2.2812 val_bpb:1.3475 train_time:60797ms step_avg:43.43ms -step:1450/20000 train_loss:2.3006 train_time:62917ms step_avg:43.39ms -step:1500/20000 train_loss:2.2831 train_time:65060ms step_avg:43.37ms -step:1550/20000 train_loss:2.2957 train_time:67324ms step_avg:43.43ms -step:1600/20000 train_loss:2.2187 train_time:69467ms step_avg:43.42ms -step:1600/20000 val_loss:2.2631 val_bpb:1.3368 train_time:69491ms step_avg:43.43ms -step:1650/20000 train_loss:2.2629 train_time:71614ms step_avg:43.40ms -step:1700/20000 train_loss:2.2619 train_time:73759ms step_avg:43.39ms -step:1750/20000 train_loss:2.1068 train_time:76028ms step_avg:43.44ms -step:1800/20000 train_loss:2.3312 train_time:78171ms step_avg:43.43ms -step:1800/20000 val_loss:2.2479 val_bpb:1.3279 train_time:78197ms step_avg:43.44ms -step:1850/20000 train_loss:2.2211 train_time:80317ms step_avg:43.41ms -step:1900/20000 train_loss:2.2477 train_time:82462ms step_avg:43.40ms -step:1950/20000 train_loss:2.2707 train_time:84723ms step_avg:43.45ms -step:2000/20000 train_loss:2.2346 train_time:86867ms step_avg:43.43ms -step:2000/20000 val_loss:2.2368 val_bpb:1.3213 train_time:86892ms step_avg:43.45ms -step:2050/20000 train_loss:2.0689 train_time:89013ms step_avg:43.42ms -step:2100/20000 train_loss:2.3382 train_time:91276ms step_avg:43.46ms -step:2150/20000 train_loss:2.1161 train_time:93418ms step_avg:43.45ms -step:2200/20000 train_loss:2.2380 train_time:95565ms step_avg:43.44ms -step:2200/20000 val_loss:2.2251 val_bpb:1.3144 train_time:95590ms step_avg:43.45ms -step:2250/20000 train_loss:2.2362 train_time:97711ms step_avg:43.43ms -step:2300/20000 train_loss:2.2390 train_time:99973ms step_avg:43.47ms -step:2350/20000 train_loss:2.1494 train_time:102118ms step_avg:43.45ms -step:2400/20000 train_loss:2.1004 train_time:104264ms step_avg:43.44ms -step:2400/20000 val_loss:2.2158 val_bpb:1.3089 train_time:104288ms step_avg:43.45ms -step:2450/20000 train_loss:2.2078 train_time:106409ms step_avg:43.43ms -step:2500/20000 train_loss:2.2990 train_time:108679ms step_avg:43.47ms -step:2550/20000 train_loss:2.3510 train_time:110825ms step_avg:43.46ms -step:2600/20000 train_loss:2.1989 train_time:112969ms step_avg:43.45ms -step:2600/20000 val_loss:2.2097 val_bpb:1.3053 train_time:112994ms step_avg:43.46ms -step:2650/20000 train_loss:2.0953 train_time:115115ms step_avg:43.44ms -step:2700/20000 train_loss:2.2119 train_time:117382ms step_avg:43.47ms -step:2750/20000 train_loss:2.2833 train_time:119524ms step_avg:43.46ms -step:2800/20000 train_loss:2.2056 train_time:121673ms step_avg:43.45ms -step:2800/20000 val_loss:2.2011 val_bpb:1.3002 train_time:121697ms step_avg:43.46ms -step:2850/20000 train_loss:2.1613 train_time:123815ms step_avg:43.44ms -step:2900/20000 train_loss:2.2400 train_time:126078ms step_avg:43.48ms -step:2950/20000 train_loss:2.2531 train_time:128222ms step_avg:43.47ms -step:3000/20000 train_loss:2.1098 train_time:130368ms step_avg:43.46ms -step:3000/20000 val_loss:2.1953 val_bpb:1.2968 train_time:130392ms step_avg:43.46ms -step:3050/20000 train_loss:2.4246 train_time:132514ms step_avg:43.45ms -step:3100/20000 train_loss:2.1884 train_time:134780ms step_avg:43.48ms -step:3150/20000 train_loss:2.2749 train_time:136926ms step_avg:43.47ms -step:3200/20000 train_loss:2.1492 train_time:139071ms step_avg:43.46ms -step:3200/20000 val_loss:2.1881 val_bpb:1.2925 train_time:139096ms step_avg:43.47ms -step:3250/20000 train_loss:2.1286 train_time:141341ms step_avg:43.49ms -step:3300/20000 train_loss:2.1058 train_time:143485ms step_avg:43.48ms -step:3350/20000 train_loss:2.2214 train_time:145628ms step_avg:43.47ms -step:3400/20000 train_loss:2.2454 train_time:147773ms step_avg:43.46ms -step:3400/20000 val_loss:2.1854 val_bpb:1.2909 train_time:147798ms step_avg:43.47ms -step:3450/20000 train_loss:2.2601 train_time:150039ms step_avg:43.49ms -step:3500/20000 train_loss:2.1183 train_time:152184ms step_avg:43.48ms -step:3550/20000 train_loss:2.0846 train_time:154329ms step_avg:43.47ms -step:3600/20000 train_loss:2.2507 train_time:156472ms step_avg:43.46ms -step:3600/20000 val_loss:2.1784 val_bpb:1.2868 train_time:156496ms step_avg:43.47ms -step:3650/20000 train_loss:2.1383 train_time:158738ms step_avg:43.49ms -step:3700/20000 train_loss:2.2848 train_time:160882ms step_avg:43.48ms -step:3750/20000 train_loss:2.1982 train_time:163029ms step_avg:43.47ms -step:3800/20000 train_loss:2.1399 train_time:165176ms step_avg:43.47ms -step:3800/20000 val_loss:2.1767 val_bpb:1.2858 train_time:165200ms step_avg:43.47ms -step:3850/20000 train_loss:2.3361 train_time:167438ms step_avg:43.49ms -step:3900/20000 train_loss:2.2756 train_time:169582ms step_avg:43.48ms -step:3950/20000 train_loss:2.1261 train_time:171729ms step_avg:43.48ms -step:4000/20000 train_loss:2.1437 train_time:173878ms step_avg:43.47ms -step:4000/20000 val_loss:2.1718 val_bpb:1.2829 train_time:173903ms step_avg:43.48ms -step:4050/20000 train_loss:2.1718 train_time:176147ms step_avg:43.49ms -step:4100/20000 train_loss:2.1899 train_time:178291ms step_avg:43.49ms -step:4150/20000 train_loss:2.1285 train_time:180438ms step_avg:43.48ms -step:4200/20000 train_loss:2.0498 train_time:182707ms step_avg:43.50ms -step:4200/20000 val_loss:2.1666 val_bpb:1.2798 train_time:182731ms step_avg:43.51ms -step:4250/20000 train_loss:2.2487 train_time:184852ms step_avg:43.49ms -step:4300/20000 train_loss:2.1979 train_time:186996ms step_avg:43.49ms -step:4350/20000 train_loss:2.1314 train_time:189141ms step_avg:43.48ms -step:4400/20000 train_loss:2.1727 train_time:191402ms step_avg:43.50ms -step:4400/20000 val_loss:2.1625 val_bpb:1.2774 train_time:191427ms step_avg:43.51ms -step:4450/20000 train_loss:2.1882 train_time:193549ms step_avg:43.49ms -step:4500/20000 train_loss:2.0735 train_time:195696ms step_avg:43.49ms -step:4550/20000 train_loss:2.1347 train_time:197840ms step_avg:43.48ms -step:4600/20000 train_loss:2.1710 train_time:200091ms step_avg:43.50ms -step:4600/20000 val_loss:2.1597 val_bpb:1.2757 train_time:200114ms step_avg:43.50ms -step:4650/20000 train_loss:2.2563 train_time:202236ms step_avg:43.49ms -step:4700/20000 train_loss:2.2077 train_time:204381ms step_avg:43.49ms -step:4750/20000 train_loss:2.1328 train_time:206643ms step_avg:43.50ms -step:4800/20000 train_loss:2.1473 train_time:208788ms step_avg:43.50ms -step:4800/20000 val_loss:2.1579 val_bpb:1.2747 train_time:208812ms step_avg:43.50ms -step:4850/20000 train_loss:2.2067 train_time:210933ms step_avg:43.49ms -step:4900/20000 train_loss:2.1119 train_time:213078ms step_avg:43.49ms -step:4950/20000 train_loss:2.0031 train_time:215339ms step_avg:43.50ms -step:5000/20000 train_loss:2.1104 train_time:217483ms step_avg:43.50ms -step:5000/20000 val_loss:2.1532 val_bpb:1.2719 train_time:217508ms step_avg:43.50ms -step:5050/20000 train_loss:2.0232 train_time:219627ms step_avg:43.49ms -step:5100/20000 train_loss:2.1995 train_time:221774ms step_avg:43.49ms -step:5150/20000 train_loss:2.0709 train_time:224038ms step_avg:43.50ms -step:5200/20000 train_loss:2.0972 train_time:226182ms step_avg:43.50ms -step:5200/20000 val_loss:2.1501 val_bpb:1.2701 train_time:226207ms step_avg:43.50ms -step:5250/20000 train_loss:2.1395 train_time:228330ms step_avg:43.49ms -step:5300/20000 train_loss:2.0947 train_time:230476ms step_avg:43.49ms -step:5350/20000 train_loss:2.0819 train_time:232740ms step_avg:43.50ms -step:5400/20000 train_loss:2.2099 train_time:234884ms step_avg:43.50ms -step:5400/20000 val_loss:2.1475 val_bpb:1.2685 train_time:234909ms step_avg:43.50ms -step:5450/20000 train_loss:2.1314 train_time:237031ms step_avg:43.49ms -step:5500/20000 train_loss:2.2057 train_time:239295ms step_avg:43.51ms -step:5550/20000 train_loss:2.0856 train_time:241437ms step_avg:43.50ms -step:5600/20000 train_loss:2.1448 train_time:243583ms step_avg:43.50ms -step:5600/20000 val_loss:2.1455 val_bpb:1.2674 train_time:243608ms step_avg:43.50ms -step:5650/20000 train_loss:2.0312 train_time:245730ms step_avg:43.49ms -step:5700/20000 train_loss:2.1392 train_time:247996ms step_avg:43.51ms -step:5750/20000 train_loss:2.0206 train_time:250140ms step_avg:43.50ms -step:5800/20000 train_loss:2.2107 train_time:252283ms step_avg:43.50ms -step:5800/20000 val_loss:2.1439 val_bpb:1.2664 train_time:252308ms step_avg:43.50ms -step:5850/20000 train_loss:2.0973 train_time:254429ms step_avg:43.49ms -step:5900/20000 train_loss:2.1270 train_time:256697ms step_avg:43.51ms -step:5950/20000 train_loss:2.0899 train_time:258840ms step_avg:43.50ms -step:6000/20000 train_loss:2.2182 train_time:260985ms step_avg:43.50ms -step:6000/20000 val_loss:2.1445 val_bpb:1.2668 train_time:261009ms step_avg:43.50ms -step:6050/20000 train_loss:2.1230 train_time:263130ms step_avg:43.49ms -step:6100/20000 train_loss:2.1640 train_time:265401ms step_avg:43.51ms -step:6150/20000 train_loss:2.1960 train_time:267547ms step_avg:43.50ms -step:6200/20000 train_loss:2.1217 train_time:269692ms step_avg:43.50ms -step:6200/20000 val_loss:2.1416 val_bpb:1.2651 train_time:269717ms step_avg:43.50ms -step:6250/20000 train_loss:2.1106 train_time:271837ms step_avg:43.49ms -step:6300/20000 train_loss:2.1989 train_time:274105ms step_avg:43.51ms -step:6350/20000 train_loss:2.1738 train_time:276249ms step_avg:43.50ms -step:6400/20000 train_loss:2.1333 train_time:278396ms step_avg:43.50ms -step:6400/20000 val_loss:2.1377 val_bpb:1.2628 train_time:278421ms step_avg:43.50ms -step:6450/20000 train_loss:1.9696 train_time:280544ms step_avg:43.50ms -step:6500/20000 train_loss:2.1279 train_time:282815ms step_avg:43.51ms -step:6550/20000 train_loss:2.2768 train_time:284958ms step_avg:43.51ms -step:6600/20000 train_loss:2.1060 train_time:287102ms step_avg:43.50ms -step:6600/20000 val_loss:2.1354 val_bpb:1.2614 train_time:287126ms step_avg:43.50ms -step:6650/20000 train_loss:2.1036 train_time:289368ms step_avg:43.51ms -step:6700/20000 train_loss:2.1438 train_time:291511ms step_avg:43.51ms -step:6750/20000 train_loss:1.8938 train_time:293654ms step_avg:43.50ms -step:6800/20000 train_loss:2.1809 train_time:295799ms step_avg:43.50ms -step:6800/20000 val_loss:2.1342 val_bpb:1.2607 train_time:295824ms step_avg:43.50ms -step:6850/20000 train_loss:2.0978 train_time:298068ms step_avg:43.51ms -step:6900/20000 train_loss:2.1146 train_time:300210ms step_avg:43.51ms -step:6950/20000 train_loss:2.1328 train_time:302354ms step_avg:43.50ms -step:7000/20000 train_loss:2.1537 train_time:304499ms step_avg:43.50ms -step:7000/20000 val_loss:2.1326 val_bpb:1.2598 train_time:304523ms step_avg:43.50ms -step:7050/20000 train_loss:2.1382 train_time:306765ms step_avg:43.51ms -step:7100/20000 train_loss:2.1078 train_time:308911ms step_avg:43.51ms -step:7150/20000 train_loss:2.1952 train_time:311056ms step_avg:43.50ms -step:7200/20000 train_loss:2.1143 train_time:313204ms step_avg:43.50ms -step:7200/20000 val_loss:2.1299 val_bpb:1.2582 train_time:313228ms step_avg:43.50ms -step:7250/20000 train_loss:2.1009 train_time:315469ms step_avg:43.51ms -step:7300/20000 train_loss:2.1529 train_time:317612ms step_avg:43.51ms -step:7350/20000 train_loss:2.1532 train_time:319759ms step_avg:43.50ms -step:7400/20000 train_loss:2.1137 train_time:321901ms step_avg:43.50ms -step:7400/20000 val_loss:2.1282 val_bpb:1.2572 train_time:321927ms step_avg:43.50ms -step:7450/20000 train_loss:2.4067 train_time:324167ms step_avg:43.51ms -step:7500/20000 train_loss:2.0751 train_time:326311ms step_avg:43.51ms -step:7550/20000 train_loss:2.1258 train_time:328457ms step_avg:43.50ms -step:7600/20000 train_loss:2.1723 train_time:330730ms step_avg:43.52ms -step:7600/20000 val_loss:2.1289 val_bpb:1.2576 train_time:330754ms step_avg:43.52ms -step:7650/20000 train_loss:2.2193 train_time:332878ms step_avg:43.51ms -step:7700/20000 train_loss:2.1329 train_time:335023ms step_avg:43.51ms -step:7750/20000 train_loss:2.0562 train_time:337169ms step_avg:43.51ms -step:7800/20000 train_loss:2.1669 train_time:339436ms step_avg:43.52ms -step:7800/20000 val_loss:2.1252 val_bpb:1.2554 train_time:339460ms step_avg:43.52ms -step:7850/20000 train_loss:2.0994 train_time:341583ms step_avg:43.51ms -step:7900/20000 train_loss:2.1585 train_time:343729ms step_avg:43.51ms -step:7950/20000 train_loss:2.1319 train_time:345873ms step_avg:43.51ms -step:8000/20000 train_loss:2.2613 train_time:348141ms step_avg:43.52ms -step:8000/20000 val_loss:2.1232 val_bpb:1.2542 train_time:348165ms step_avg:43.52ms -step:8050/20000 train_loss:2.1775 train_time:350287ms step_avg:43.51ms -step:8100/20000 train_loss:1.9587 train_time:352431ms step_avg:43.51ms -step:8150/20000 train_loss:2.0401 train_time:354575ms step_avg:43.51ms -step:8200/20000 train_loss:2.1076 train_time:356845ms step_avg:43.52ms -step:8200/20000 val_loss:2.1228 val_bpb:1.2540 train_time:356869ms step_avg:43.52ms -step:8250/20000 train_loss:2.0951 train_time:358988ms step_avg:43.51ms -step:8300/20000 train_loss:2.2244 train_time:361133ms step_avg:43.51ms -step:8350/20000 train_loss:2.0681 train_time:363279ms step_avg:43.51ms -step:8400/20000 train_loss:2.1494 train_time:365552ms step_avg:43.52ms -step:8400/20000 val_loss:2.1201 val_bpb:1.2524 train_time:365577ms step_avg:43.52ms -step:8450/20000 train_loss:2.1278 train_time:367698ms step_avg:43.51ms -step:8500/20000 train_loss:2.0289 train_time:369845ms step_avg:43.51ms -step:8550/20000 train_loss:2.0465 train_time:372114ms step_avg:43.52ms -step:8600/20000 train_loss:2.0682 train_time:374259ms step_avg:43.52ms -step:8600/20000 val_loss:2.1206 val_bpb:1.2526 train_time:374282ms step_avg:43.52ms -step:8650/20000 train_loss:2.2717 train_time:376403ms step_avg:43.51ms -step:8700/20000 train_loss:2.1795 train_time:378549ms step_avg:43.51ms -step:8750/20000 train_loss:2.0492 train_time:380817ms step_avg:43.52ms -step:8800/20000 train_loss:2.1100 train_time:382964ms step_avg:43.52ms -step:8800/20000 val_loss:2.1192 val_bpb:1.2518 train_time:382989ms step_avg:43.52ms -step:8850/20000 train_loss:2.4323 train_time:385110ms step_avg:43.52ms -step:8900/20000 train_loss:2.1016 train_time:387258ms step_avg:43.51ms -step:8950/20000 train_loss:2.0290 train_time:389530ms step_avg:43.52ms -step:9000/20000 train_loss:2.1119 train_time:391675ms step_avg:43.52ms -step:9000/20000 val_loss:2.1204 val_bpb:1.2525 train_time:391698ms step_avg:43.52ms -step:9050/20000 train_loss:2.0826 train_time:393819ms step_avg:43.52ms -step:9100/20000 train_loss:2.0427 train_time:395963ms step_avg:43.51ms -step:9150/20000 train_loss:2.1201 train_time:398238ms step_avg:43.52ms -step:9200/20000 train_loss:2.1490 train_time:400385ms step_avg:43.52ms -step:9200/20000 val_loss:2.1170 val_bpb:1.2505 train_time:400409ms step_avg:43.52ms -step:9250/20000 train_loss:2.1221 train_time:402534ms step_avg:43.52ms -step:9300/20000 train_loss:2.4550 train_time:404680ms step_avg:43.51ms -step:9350/20000 train_loss:2.0384 train_time:406932ms step_avg:43.52ms -step:9400/20000 train_loss:2.0736 train_time:409077ms step_avg:43.52ms -step:9400/20000 val_loss:2.1139 val_bpb:1.2487 train_time:409102ms step_avg:43.52ms -step:9450/20000 train_loss:2.1096 train_time:411223ms step_avg:43.52ms -step:9500/20000 train_loss:2.1070 train_time:413493ms step_avg:43.53ms -step:9550/20000 train_loss:2.0249 train_time:415641ms step_avg:43.52ms -step:9600/20000 train_loss:2.1141 train_time:417785ms step_avg:43.52ms -step:9600/20000 val_loss:2.1138 val_bpb:1.2486 train_time:417809ms step_avg:43.52ms -step:9650/20000 train_loss:2.0183 train_time:419932ms step_avg:43.52ms -step:9700/20000 train_loss:2.1482 train_time:422212ms step_avg:43.53ms -step:9750/20000 train_loss:2.1811 train_time:424359ms step_avg:43.52ms -step:9800/20000 train_loss:2.1011 train_time:426503ms step_avg:43.52ms -step:9800/20000 val_loss:2.1143 val_bpb:1.2489 train_time:426528ms step_avg:43.52ms -step:9850/20000 train_loss:2.1134 train_time:428771ms step_avg:43.53ms -step:9900/20000 train_loss:2.0497 train_time:430915ms step_avg:43.53ms -step:9950/20000 train_loss:2.1989 train_time:433061ms step_avg:43.52ms -step:10000/20000 train_loss:2.1982 train_time:435207ms step_avg:43.52ms -step:10000/20000 val_loss:2.1122 val_bpb:1.2477 train_time:435232ms step_avg:43.52ms -step:10050/20000 train_loss:2.0940 train_time:437485ms step_avg:43.53ms -step:10100/20000 train_loss:2.1277 train_time:439630ms step_avg:43.53ms -step:10150/20000 train_loss:2.0896 train_time:441773ms step_avg:43.52ms -step:10200/20000 train_loss:2.0642 train_time:443918ms step_avg:43.52ms -step:10200/20000 val_loss:2.1112 val_bpb:1.2471 train_time:443941ms step_avg:43.52ms -step:10250/20000 train_loss:2.0627 train_time:446192ms step_avg:43.53ms -step:10300/20000 train_loss:2.2191 train_time:448339ms step_avg:43.53ms -step:10350/20000 train_loss:2.1354 train_time:450485ms step_avg:43.53ms -step:10400/20000 train_loss:2.0705 train_time:452630ms step_avg:43.52ms -step:10400/20000 val_loss:2.1098 val_bpb:1.2463 train_time:452654ms step_avg:43.52ms -step:10450/20000 train_loss:2.0663 train_time:454900ms step_avg:43.53ms -step:10500/20000 train_loss:2.1334 train_time:457046ms step_avg:43.53ms -step:10550/20000 train_loss:2.1931 train_time:459192ms step_avg:43.53ms -step:10600/20000 train_loss:2.0978 train_time:461337ms step_avg:43.52ms -step:10600/20000 val_loss:2.1081 val_bpb:1.2453 train_time:461361ms step_avg:43.52ms -step:10650/20000 train_loss:2.0676 train_time:463610ms step_avg:43.53ms -step:10700/20000 train_loss:2.2333 train_time:465754ms step_avg:43.53ms -step:10750/20000 train_loss:2.1661 train_time:467899ms step_avg:43.53ms -step:10800/20000 train_loss:2.0966 train_time:470044ms step_avg:43.52ms -step:10800/20000 val_loss:2.1081 val_bpb:1.2453 train_time:470069ms step_avg:43.52ms -step:10850/20000 train_loss:2.0708 train_time:472323ms step_avg:43.53ms -step:10900/20000 train_loss:2.1666 train_time:474468ms step_avg:43.53ms -step:10950/20000 train_loss:2.1079 train_time:476615ms step_avg:43.53ms -step:11000/20000 train_loss:2.0774 train_time:478893ms step_avg:43.54ms -step:11000/20000 val_loss:2.1069 val_bpb:1.2446 train_time:478917ms step_avg:43.54ms -step:11050/20000 train_loss:2.1288 train_time:481038ms step_avg:43.53ms -step:11100/20000 train_loss:2.0801 train_time:483185ms step_avg:43.53ms -step:11150/20000 train_loss:1.8743 train_time:485331ms step_avg:43.53ms -step:11200/20000 train_loss:2.1471 train_time:487603ms step_avg:43.54ms -step:11200/20000 val_loss:2.1080 val_bpb:1.2452 train_time:487627ms step_avg:43.54ms -step:11250/20000 train_loss:2.2046 train_time:489748ms step_avg:43.53ms -step:11300/20000 train_loss:2.0957 train_time:491892ms step_avg:43.53ms -step:11350/20000 train_loss:2.0963 train_time:494038ms step_avg:43.53ms -step:11400/20000 train_loss:2.3223 train_time:496318ms step_avg:43.54ms -step:11400/20000 val_loss:2.1051 val_bpb:1.2435 train_time:496342ms step_avg:43.54ms -step:11450/20000 train_loss:2.0724 train_time:498464ms step_avg:43.53ms -step:11500/20000 train_loss:2.1197 train_time:500609ms step_avg:43.53ms -step:11550/20000 train_loss:2.0975 train_time:502754ms step_avg:43.53ms -step:11600/20000 train_loss:2.1091 train_time:505029ms step_avg:43.54ms -step:11600/20000 val_loss:2.1054 val_bpb:1.2437 train_time:505053ms step_avg:43.54ms -step:11650/20000 train_loss:2.1235 train_time:507175ms step_avg:43.53ms -step:11700/20000 train_loss:2.0795 train_time:509324ms step_avg:43.53ms -step:11750/20000 train_loss:2.0662 train_time:511469ms step_avg:43.53ms -step:11800/20000 train_loss:2.0765 train_time:513742ms step_avg:43.54ms -step:11800/20000 val_loss:2.1048 val_bpb:1.2433 train_time:513766ms step_avg:43.54ms -step:11850/20000 train_loss:2.1202 train_time:515888ms step_avg:43.53ms -step:11900/20000 train_loss:2.1029 train_time:518033ms step_avg:43.53ms -step:11950/20000 train_loss:2.1512 train_time:520308ms step_avg:43.54ms -step:12000/20000 train_loss:2.1814 train_time:522453ms step_avg:43.54ms -step:12000/20000 val_loss:2.1029 val_bpb:1.2422 train_time:522477ms step_avg:43.54ms -step:12050/20000 train_loss:2.1085 train_time:524601ms step_avg:43.54ms -step:12100/20000 train_loss:2.0347 train_time:526747ms step_avg:43.53ms -step:12150/20000 train_loss:2.0601 train_time:529018ms step_avg:43.54ms -step:12200/20000 train_loss:2.0387 train_time:531162ms step_avg:43.54ms -step:12200/20000 val_loss:2.1021 val_bpb:1.2418 train_time:531186ms step_avg:43.54ms -step:12250/20000 train_loss:2.0381 train_time:533312ms step_avg:43.54ms -step:12300/20000 train_loss:2.1302 train_time:535458ms step_avg:43.53ms -step:12350/20000 train_loss:2.1272 train_time:537727ms step_avg:43.54ms -step:12400/20000 train_loss:2.1828 train_time:539873ms step_avg:43.54ms -step:12400/20000 val_loss:2.1001 val_bpb:1.2406 train_time:539897ms step_avg:43.54ms -step:12450/20000 train_loss:2.1003 train_time:542019ms step_avg:43.54ms -step:12500/20000 train_loss:2.0696 train_time:544164ms step_avg:43.53ms -step:12550/20000 train_loss:2.1302 train_time:546436ms step_avg:43.54ms -step:12600/20000 train_loss:2.0527 train_time:548582ms step_avg:43.54ms -step:12600/20000 val_loss:2.0998 val_bpb:1.2404 train_time:548606ms step_avg:43.54ms -step:12650/20000 train_loss:2.1438 train_time:550728ms step_avg:43.54ms -step:12700/20000 train_loss:2.2689 train_time:552877ms step_avg:43.53ms -step:12750/20000 train_loss:2.1438 train_time:555147ms step_avg:43.54ms -step:12800/20000 train_loss:2.0105 train_time:557293ms step_avg:43.54ms -step:12800/20000 val_loss:2.0930 val_bpb:1.2364 train_time:557317ms step_avg:43.54ms -step:12850/20000 train_loss:2.0413 train_time:559440ms step_avg:43.54ms -step:12900/20000 train_loss:2.0630 train_time:561586ms step_avg:43.53ms -step:12950/20000 train_loss:2.1627 train_time:563863ms step_avg:43.54ms -step:13000/20000 train_loss:1.9579 train_time:566009ms step_avg:43.54ms -step:13000/20000 val_loss:2.0859 val_bpb:1.2322 train_time:566032ms step_avg:43.54ms -step:13050/20000 train_loss:2.0206 train_time:568155ms step_avg:43.54ms -step:13100/20000 train_loss:1.9294 train_time:570432ms step_avg:43.54ms -step:13150/20000 train_loss:2.0689 train_time:572576ms step_avg:43.54ms -step:13200/20000 train_loss:2.0074 train_time:574722ms step_avg:43.54ms -step:13200/20000 val_loss:2.0790 val_bpb:1.2281 train_time:574747ms step_avg:43.54ms -step:13250/20000 train_loss:2.0596 train_time:576871ms step_avg:43.54ms -step:13300/20000 train_loss:1.9474 train_time:579143ms step_avg:43.54ms -step:13350/20000 train_loss:2.0459 train_time:581289ms step_avg:43.54ms -step:13400/20000 train_loss:2.0441 train_time:583434ms step_avg:43.54ms -step:13400/20000 val_loss:2.0718 val_bpb:1.2239 train_time:583458ms step_avg:43.54ms -step:13450/20000 train_loss:2.1638 train_time:585582ms step_avg:43.54ms -step:13500/20000 train_loss:2.1216 train_time:587857ms step_avg:43.54ms -step:13550/20000 train_loss:2.1855 train_time:590003ms step_avg:43.54ms -step:13600/20000 train_loss:2.0234 train_time:592147ms step_avg:43.54ms -step:13600/20000 val_loss:2.0649 val_bpb:1.2197 train_time:592172ms step_avg:43.54ms -step:13650/20000 train_loss:2.0316 train_time:594295ms step_avg:43.54ms -step:13700/20000 train_loss:2.0323 train_time:596577ms step_avg:43.55ms -step:13750/20000 train_loss:1.9910 train_time:598726ms step_avg:43.54ms -step:13780/20000 val_loss:2.0606 val_bpb:1.2172 train_time:600038ms step_avg:43.54ms -stopping_early: wallclock_cap train_time:600038ms step:13780/20000 -peak memory allocated: 10184 MiB reserved: 10200 MiB -Serialized model: 67224983 bytes -Code size: 47642 bytes -Total submission size: 67272625 bytes -Serialized model int8+zlib: 15815847 bytes (payload:17178912 raw_torch:17224025 payload_ratio:3.91x) -Total submission size int8+zlib: 15863489 bytes -final_int8_zlib_roundtrip val_loss:2.0727 val_bpb:1.2244 eval_time:1401ms -final_int8_zlib_roundtrip_exact val_loss:2.07269931 val_bpb:1.22436570 diff --git a/records/track_10min_16mb/2026-03-17_NaiveBaseline/train_gpt.py b/records/track_10min_16mb/2026-03-17_NaiveBaseline/train_gpt.py deleted file mode 100644 index 0deb0565f5..0000000000 --- a/records/track_10min_16mb/2026-03-17_NaiveBaseline/train_gpt.py +++ /dev/null @@ -1,1126 +0,0 @@ -""" -The `train_gpt.py` and `train_gpt_mlx.py` scripts are intended as good launching-off points for new participants, not SOTA configs. We'll accept PRs that tune, improve, or simplify these scripts without significantly increasing complexity, but competitive submissions should stay in the `/records` folder. - -Hard stop: `train_gpt.py` and `train_gpt_mlx.py` must never be longer than 1500 lines. -""" - -from __future__ import annotations - -import copy -import glob -import io -import math -import os -import random -import subprocess -import sys -import time -import uuid -import zlib -from pathlib import Path - -import numpy as np -import sentencepiece as spm -import torch -import torch.distributed as dist -import torch.nn.functional as F -from torch import Tensor, nn -from torch.nn.parallel import DistributedDataParallel as DDP - -# ----------------------------- -# HYPERPARAMETERS -# ----------------------------- -# Default Simple Baseline run: -# - 9 transformer blocks at width 512 -# - 8 attention heads with 4 KV heads (GQA) and 2x MLP expansion -# - vocab size 1024, sequence length 1024, tied embeddings -# - 524,288 train tokens per step for 20,000 iterations with a ~10 minute cap - -class Hyperparameters: - # Data paths are shard globs produced by the existing preprocessing pipeline. - data_path = os.environ.get("DATA_PATH", "./data/datasets/fineweb10B_sp1024") - train_files = os.path.join(data_path, "fineweb_train_*.bin") - val_files = os.path.join(data_path, "fineweb_val_*.bin") - tokenizer_path = os.environ.get("TOKENIZER_PATH", "./data/tokenizers/fineweb_1024_bpe.model") - run_id = os.environ.get("RUN_ID", str(uuid.uuid4())) - seed = int(os.environ.get("SEED", 1337)) - - # Validation cadence and batch size. Validation always uses the full fineweb_val split. - val_batch_size = int(os.environ.get("VAL_BATCH_SIZE", 524_288)) - val_loss_every = int(os.environ.get("VAL_LOSS_EVERY", 1000)) - train_log_every = int(os.environ.get("TRAIN_LOG_EVERY", 200)) - - # Training length. - iterations = int(os.environ.get("ITERATIONS", 20000)) - warmdown_iters = int(os.environ.get("WARMDOWN_ITERS", 1200)) - warmup_steps = int(os.environ.get("WARMUP_STEPS", 20)) - train_batch_tokens = int(os.environ.get("TRAIN_BATCH_TOKENS", 524_288)) - train_seq_len = int(os.environ.get("TRAIN_SEQ_LEN", 1024)) - max_wallclock_seconds = float(os.environ.get("MAX_WALLCLOCK_SECONDS", 600.0)) - qk_gain_init = float(os.environ.get("QK_GAIN_INIT", 1.5)) - - # Model shape. - vocab_size = int(os.environ.get("VOCAB_SIZE", 1024)) - num_layers = int(os.environ.get("NUM_LAYERS", 9)) - num_kv_heads = int(os.environ.get("NUM_KV_HEADS", 4)) - model_dim = int(os.environ.get("MODEL_DIM", 512)) - num_heads = int(os.environ.get("NUM_HEADS", 8)) - mlp_mult = int(os.environ.get("MLP_MULT", 2)) - tie_embeddings = bool(int(os.environ.get("TIE_EMBEDDINGS", "1"))) - rope_base = float(os.environ.get("ROPE_BASE", 10000.0)) - logit_softcap = float(os.environ.get("LOGIT_SOFTCAP", 30.0)) - - # Optimizer hyperparameters. - embed_lr = float(os.environ.get("EMBED_LR", 0.6)) - head_lr = float(os.environ.get("HEAD_LR", 0.008)) - tied_embed_lr = float(os.environ.get("TIED_EMBED_LR", 0.05)) - tied_embed_init_std = float(os.environ.get("TIED_EMBED_INIT_STD", 0.005)) - matrix_lr = float(os.environ.get("MATRIX_LR", 0.04)) - scalar_lr = float(os.environ.get("SCALAR_LR", 0.04)) - muon_momentum = float(os.environ.get("MUON_MOMENTUM", 0.95)) - muon_backend_steps = int(os.environ.get("MUON_BACKEND_STEPS", 5)) - muon_momentum_warmup_start = float(os.environ.get("MUON_MOMENTUM_WARMUP_START", 0.85)) - muon_momentum_warmup_steps = int(os.environ.get("MUON_MOMENTUM_WARMUP_STEPS", 500)) - beta1 = float(os.environ.get("BETA1", 0.9)) - beta2 = float(os.environ.get("BETA2", 0.95)) - adam_eps = float(os.environ.get("ADAM_EPS", 1e-8)) - grad_clip_norm = float(os.environ.get("GRAD_CLIP_NORM", 0.0)) - -# ----------------------------- -# MUON OPTIMIZER -# ----------------------------- -# -# As borrowed from modded-nanogpt -# Background on Muon: https://kellerjordan.github.io/posts/muon/ - -def zeropower_via_newtonschulz5(G: Tensor, steps: int = 10, eps: float = 1e-7) -> Tensor: - # Orthogonalize a 2D update matrix with a fast Newton-Schulz iteration. - # Muon uses this to normalize matrix-shaped gradients before applying them. - a, b, c = (3.4445, -4.7750, 2.0315) - X = G.bfloat16() - X /= X.norm() + eps - transposed = G.size(0) > G.size(1) - if transposed: - X = X.T - for _ in range(steps): - A = X @ X.T - B = b * A + c * A @ A - X = a * X + B @ X - return X.T if transposed else X - - -class Muon(torch.optim.Optimizer): - def __init__(self, params, lr: float, momentum: float, backend_steps: int, nesterov: bool = True): - super().__init__( - params, - dict(lr=lr, momentum=momentum, backend_steps=backend_steps, nesterov=nesterov), - ) - - @torch.no_grad() - def step(self, closure=None): - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - distributed = dist.is_available() and dist.is_initialized() - world_size = dist.get_world_size() if distributed else 1 - rank = dist.get_rank() if distributed else 0 - - for group in self.param_groups: - params = group["params"] - if not params: - continue - lr = group["lr"] - momentum = group["momentum"] - backend_steps = group["backend_steps"] - nesterov = group["nesterov"] - - total_params = sum(int(p.numel()) for p in params) - updates_flat = torch.zeros(total_params, device=params[0].device, dtype=torch.bfloat16) - - curr = 0 - for i, p in enumerate(params): - if i % world_size == rank and p.grad is not None: - g = p.grad - state = self.state[p] - if "momentum_buffer" not in state: - state["momentum_buffer"] = torch.zeros_like(g) - buf = state["momentum_buffer"] - buf.mul_(momentum).add_(g) - if nesterov: - g = g.add(buf, alpha=momentum) - g = zeropower_via_newtonschulz5(g, steps=backend_steps) - # Scale correction from Muon reference implementations. - g *= max(1, g.size(0) / g.size(1)) ** 0.5 - updates_flat[curr : curr + p.numel()] = g.reshape(-1) - curr += p.numel() - - if distributed: - dist.all_reduce(updates_flat, op=dist.ReduceOp.SUM) - - curr = 0 - for p in params: - g = updates_flat[curr : curr + p.numel()].view_as(p).to(dtype=p.dtype) - p.add_(g, alpha=-lr) - curr += p.numel() - - return loss - - -# ----------------------------- -# TOKENIZER-AGNOSTIC EVALUATION SETUP -# ----------------------------- -# -# It's common for small models have a large fraction of their parameters be embeddings, since the 2 * d_model * d_vocab vectors can be gigantic. -# Instead of locking the tokenizer, we let you bring your own and calculate our validation metrics on the average compression of the validation set. -# We calculate BPB (bits-per-byte) instead of validation loss, so we need methods to count the number of bits per token in the tokenizer. -# Note: Submissions that edit the tokenizer will be examined more carefully, since screwing this up might unjustly improve your score. - -def build_sentencepiece_luts( - sp: spm.SentencePieceProcessor, vocab_size: int, device: torch.device -) -> tuple[Tensor, Tensor, Tensor]: - sp_vocab_size = int(sp.vocab_size()) - table_size = max(sp_vocab_size, vocab_size) - base_bytes_np = np.zeros((table_size,), dtype=np.int16) - has_leading_space_np = np.zeros((table_size,), dtype=np.bool_) - is_boundary_token_np = np.ones((table_size,), dtype=np.bool_) - for token_id in range(sp_vocab_size): - if sp.is_control(token_id) or sp.is_unknown(token_id) or sp.is_unused(token_id): - continue - is_boundary_token_np[token_id] = False - if sp.is_byte(token_id): - base_bytes_np[token_id] = 1 - continue - piece = sp.id_to_piece(token_id) - if piece.startswith("▁"): - has_leading_space_np[token_id] = True - piece = piece[1:] - base_bytes_np[token_id] = len(piece.encode("utf-8")) - return ( - torch.tensor(base_bytes_np, dtype=torch.int16, device=device), - torch.tensor(has_leading_space_np, dtype=torch.bool, device=device), - torch.tensor(is_boundary_token_np, dtype=torch.bool, device=device), - ) - - -def load_validation_tokens(pattern: str, seq_len: int) -> Tensor: - files = [Path(p) for p in sorted(glob.glob(pattern))] - if not files: - raise FileNotFoundError(f"No files found for pattern: {pattern}") - # The export pipeline writes the fixed first-50k-doc validation set to fineweb_val_*. - tokens = torch.cat([load_data_shard(file) for file in files]).contiguous() - usable = ((tokens.numel() - 1) // seq_len) * seq_len - if usable <= 0: - raise ValueError(f"Validation split is too short for TRAIN_SEQ_LEN={seq_len}") - return tokens[: usable + 1] - - -def eval_val( - args: Hyperparameters, - model: nn.Module, - rank: int, - world_size: int, - device: torch.device, - grad_accum_steps: int, - val_tokens: Tensor, - base_bytes_lut: Tensor, - has_leading_space_lut: Tensor, - is_boundary_token_lut: Tensor, -) -> tuple[float, float]: - # Validation computes two metrics: - # - val_loss: token cross-entropy (natural log) - # - val_bpb: tokenizer-agnostic compression metric used by the challenge - local_batch_tokens = args.val_batch_size // (world_size * grad_accum_steps) - if local_batch_tokens < args.train_seq_len: - raise ValueError( - "VAL_BATCH_SIZE must provide at least one sequence per rank; " - f"got VAL_BATCH_SIZE={args.val_batch_size}, WORLD_SIZE={world_size}, " - f"GRAD_ACCUM_STEPS={grad_accum_steps}, TRAIN_SEQ_LEN={args.train_seq_len}" - ) - local_batch_seqs = local_batch_tokens // args.train_seq_len - total_seqs = (val_tokens.numel() - 1) // args.train_seq_len - seq_start = (total_seqs * rank) // world_size - seq_end = (total_seqs * (rank + 1)) // world_size - val_loss_sum = torch.zeros((), device=device, dtype=torch.float64) - val_token_count = torch.zeros((), device=device, dtype=torch.float64) - val_byte_count = torch.zeros((), device=device, dtype=torch.float64) - - model.eval() - with torch.inference_mode(): - for batch_seq_start in range(seq_start, seq_end, local_batch_seqs): - batch_seq_end = min(batch_seq_start + local_batch_seqs, seq_end) - raw_start = batch_seq_start * args.train_seq_len - raw_end = batch_seq_end * args.train_seq_len + 1 - local = val_tokens[raw_start:raw_end].to(device=device, dtype=torch.int64, non_blocking=True) - x = local[:-1].reshape(-1, args.train_seq_len) - y = local[1:].reshape(-1, args.train_seq_len) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - batch_loss = model(x, y).detach() - batch_token_count = float(y.numel()) - val_loss_sum += batch_loss.to(torch.float64) * batch_token_count - val_token_count += batch_token_count - prev_ids = x.reshape(-1) - tgt_ids = y.reshape(-1) - token_bytes = base_bytes_lut[tgt_ids].to(dtype=torch.int16) - token_bytes += (has_leading_space_lut[tgt_ids] & ~is_boundary_token_lut[prev_ids]).to(dtype=torch.int16) - val_byte_count += token_bytes.to(torch.float64).sum() - - if dist.is_available() and dist.is_initialized(): - dist.all_reduce(val_loss_sum, op=dist.ReduceOp.SUM) - dist.all_reduce(val_token_count, op=dist.ReduceOp.SUM) - dist.all_reduce(val_byte_count, op=dist.ReduceOp.SUM) - - val_loss = val_loss_sum / val_token_count - bits_per_token = val_loss.item() / math.log(2.0) - tokens_per_byte = val_token_count.item() / val_byte_count.item() - model.train() - return float(val_loss.item()), float(bits_per_token * tokens_per_byte) - -# ----------------------------- -# POST-TRAINING QUANTIZATION -# ----------------------------- -# -# It's silly to export our model, which is trained in bf16 and fp32, at that same precision. -# Instead, we get approximately the same model (with a small hit) by quantizing the model to int8 & zlib compressing. -# We can then decompress the model and run in higher precision for evaluation, after closing in under the size limit. - -CONTROL_TENSOR_NAME_PATTERNS = tuple( - pattern - for pattern in os.environ.get( - "CONTROL_TENSOR_NAME_PATTERNS", - "attn_scale,attn_scales,mlp_scale,mlp_scales,resid_mix,resid_mixes,q_gain,skip_weight,skip_weights", - ).split(",") - if pattern -) -INT8_KEEP_FLOAT_FP32_NAME_PATTERNS = tuple( - pattern - for pattern in os.environ.get( - "INT8_KEEP_FLOAT_FP32_NAME_PATTERNS", - ",".join(CONTROL_TENSOR_NAME_PATTERNS), - ).split(",") - if pattern -) -INT8_KEEP_FLOAT_MAX_NUMEL = 65_536 -INT8_KEEP_FLOAT_STORE_DTYPE = torch.float16 -INT8_PER_ROW_SCALE_DTYPE = torch.float16 -INT8_CLIP_PERCENTILE = 99.99984 -INT8_CLIP_Q = INT8_CLIP_PERCENTILE / 100.0 - -def tensor_nbytes(t: Tensor) -> int: - return int(t.numel()) * int(t.element_size()) - -def keep_float_tensor(name: str, t: Tensor, passthrough_orig_dtypes: dict[str, str]) -> Tensor: - if any(pattern in name for pattern in INT8_KEEP_FLOAT_FP32_NAME_PATTERNS): - return t.float().contiguous() - if t.dtype in {torch.float32, torch.bfloat16}: - passthrough_orig_dtypes[name] = str(t.dtype).removeprefix("torch.") - return t.to(dtype=INT8_KEEP_FLOAT_STORE_DTYPE).contiguous() - return t - -def quantize_float_tensor(t: Tensor) -> tuple[Tensor, Tensor]: - t32 = t.float() - if t32.ndim == 2: - # Matrices get one scale per row, which usually tracks output-channel - # ranges much better than a single tensor-wide scale. - clip_abs = ( - torch.quantile(t32.abs(), INT8_CLIP_Q, dim=1) - if t32.numel() - else torch.empty((t32.shape[0],), dtype=torch.float32) - ) - clipped = torch.maximum(torch.minimum(t32, clip_abs[:, None]), -clip_abs[:, None]) - scale = (clip_abs / 127.0).clamp_min(1.0 / 127.0) - q = torch.clamp(torch.round(clipped / scale[:, None]), -127, 127).to(torch.int8).contiguous() - return q, scale.to(dtype=INT8_PER_ROW_SCALE_DTYPE).contiguous() - - # Vectors / scalars use a simpler per-tensor scale. - clip_abs = float(torch.quantile(t32.abs().flatten(), INT8_CLIP_Q).item()) if t32.numel() else 0.0 - scale = torch.tensor(clip_abs / 127.0 if clip_abs > 0 else 1.0, dtype=torch.float32) - q = torch.clamp(torch.round(torch.clamp(t32, -clip_abs, clip_abs) / scale), -127, 127).to(torch.int8).contiguous() - return q, scale - -def quantize_state_dict_int8(state_dict: dict[str, Tensor]): - # Single supported clean-script export format: - # - per-row int8 for 2D float tensors - # - per-tensor int8 for other float tensors - # - exact passthrough for non-floats - # - passthrough for small float tensors, stored as fp16 to save bytes - quantized: dict[str, Tensor] = {} - scales: dict[str, Tensor] = {} - dtypes: dict[str, str] = {} - passthrough: dict[str, Tensor] = {} - passthrough_orig_dtypes: dict[str, str] = {} - qmeta: dict[str, dict[str, object]] = {} - stats = dict.fromkeys( - ("param_count", "num_tensors", "num_float_tensors", "num_nonfloat_tensors", "baseline_tensor_bytes", "int8_payload_bytes"), - 0, - ) - - for name, tensor in state_dict.items(): - t = tensor.detach().to("cpu").contiguous() - stats["param_count"] += int(t.numel()) - stats["num_tensors"] += 1 - stats["baseline_tensor_bytes"] += tensor_nbytes(t) - - if not t.is_floating_point(): - stats["num_nonfloat_tensors"] += 1 - passthrough[name] = t - stats["int8_payload_bytes"] += tensor_nbytes(t) - continue - - # Small float tensors are cheap enough to keep directly. We still downcast - # fp32/bf16 passthrough tensors to fp16 so metadata does not dominate size. - if t.numel() <= INT8_KEEP_FLOAT_MAX_NUMEL: - kept = keep_float_tensor(name, t, passthrough_orig_dtypes) - passthrough[name] = kept - stats["int8_payload_bytes"] += tensor_nbytes(kept) - continue - - stats["num_float_tensors"] += 1 - q, s = quantize_float_tensor(t) - if s.ndim > 0: - qmeta[name] = {"scheme": "per_row", "axis": 0} - quantized[name] = q - scales[name] = s - dtypes[name] = str(t.dtype).removeprefix("torch.") - stats["int8_payload_bytes"] += tensor_nbytes(q) + tensor_nbytes(s) - - obj: dict[str, object] = { - "__quant_format__": "int8_clean_per_row_v1", - "quantized": quantized, - "scales": scales, - "dtypes": dtypes, - "passthrough": passthrough, - } - if qmeta: - obj["qmeta"] = qmeta - if passthrough_orig_dtypes: - obj["passthrough_orig_dtypes"] = passthrough_orig_dtypes - return obj, stats - -def dequantize_state_dict_int8(obj: dict[str, object]) -> dict[str, Tensor]: - out: dict[str, Tensor] = {} - qmeta = obj.get("qmeta", {}) - passthrough_orig_dtypes = obj.get("passthrough_orig_dtypes", {}) - for name, q in obj["quantized"].items(): - dtype = getattr(torch, obj["dtypes"][name]) - s = obj["scales"][name] - if qmeta.get(name, {}).get("scheme") == "per_row" or s.ndim > 0: - s = s.to(dtype=torch.float32) - # Broadcast the saved row scale back across trailing dimensions. - out[name] = (q.float() * s.view(q.shape[0], *([1] * (q.ndim - 1)))).to(dtype=dtype).contiguous() - else: - scale = float(s.item()) - out[name] = (q.float() * scale).to(dtype=dtype).contiguous() - for name, t in obj["passthrough"].items(): - # Restore small tensors, undoing the temporary fp16 storage cast if needed. - out_t = t.detach().to("cpu").contiguous() - orig_dtype = passthrough_orig_dtypes.get(name) - if isinstance(orig_dtype, str): - out_t = out_t.to(dtype=getattr(torch, orig_dtype)).contiguous() - out[name] = out_t - return out - - -# ----------------------------- -# DATA LOADING -# ----------------------------- - -def load_data_shard(file: Path) -> Tensor: - header_bytes = 256 * np.dtype(" None: - self.file_idx = (self.file_idx + 1) % len(self.files) - self.tokens = load_data_shard(self.files[self.file_idx]) - self.pos = 0 - - def take(self, n: int) -> Tensor: - chunks: list[Tensor] = [] - remaining = n - while remaining > 0: - avail = self.tokens.numel() - self.pos - if avail <= 0: - self._advance_file() - continue - k = min(remaining, avail) - chunks.append(self.tokens[self.pos : self.pos + k]) - self.pos += k - remaining -= k - return chunks[0] if len(chunks) == 1 else torch.cat(chunks) - - -class DistributedTokenLoader: - # Each call consumes a contiguous chunk from the shared token stream, then slices out - # one disjoint span per rank. The extra "+1" token lets us build (x, y) by shifting. - def __init__(self, pattern: str, rank: int, world_size: int, device: torch.device): - self.rank = rank - self.world_size = world_size - self.device = device - self.stream = TokenStream(pattern) - - def next_batch(self, global_tokens: int, seq_len: int, grad_accum_steps: int) -> tuple[Tensor, Tensor]: - local_tokens = global_tokens // (self.world_size * grad_accum_steps) - per_rank_span = local_tokens + 1 - chunk = self.stream.take(per_rank_span * self.world_size) - start = self.rank * per_rank_span - local = chunk[start : start + per_rank_span].to(dtype=torch.int64) - x = local[:-1].reshape(-1, seq_len) - y = local[1:].reshape(-1, seq_len) - return x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True) - -# ----------------------------- -# TRANSFORMER MODULES -# ----------------------------- - -class RMSNorm(nn.Module): - def __init__(self, eps: float | None = None): - super().__init__() - self.eps = eps - - def forward(self, x: Tensor) -> Tensor: - return F.rms_norm(x, (x.size(-1),), eps=self.eps) - - -class CastedLinear(nn.Linear): - # Keep weights in fp32 for optimizer/state quality, cast at matmul time for bf16 compute. - def forward(self, x: Tensor) -> Tensor: - bias = self.bias.to(x.dtype) if self.bias is not None else None - return F.linear(x, self.weight.to(x.dtype), bias) - - -def restore_low_dim_params_to_fp32(module: nn.Module) -> None: - # Keep small/control parameters in fp32 even when the model body runs in bf16. - with torch.no_grad(): - for name, param in module.named_parameters(): - if (param.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS)) and param.dtype != torch.float32: - param.data = param.data.float() - - -class Rotary(nn.Module): - # Caches cos/sin tables per sequence length on the current device. - def __init__(self, dim: int, base: float = 10000.0): - super().__init__() - inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self._seq_len_cached = 0 - self._cos_cached: Tensor | None = None - self._sin_cached: Tensor | None = None - - def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> tuple[Tensor, Tensor]: - if ( - self._cos_cached is None - or self._sin_cached is None - or self._seq_len_cached != seq_len - or self._cos_cached.device != device - ): - t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) - freqs = torch.outer(t, self.inv_freq.to(device)) - self._cos_cached = freqs.cos()[None, None, :, :] - self._sin_cached = freqs.sin()[None, None, :, :] - self._seq_len_cached = seq_len - return self._cos_cached.to(dtype=dtype), self._sin_cached.to(dtype=dtype) - - -def apply_rotary_emb(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: - half = x.size(-1) // 2 - x1, x2 = x[..., :half], x[..., half:] - return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1) - - -class CausalSelfAttention(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - num_kv_heads: int, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - if dim % num_heads != 0: - raise ValueError("model_dim must be divisible by num_heads") - if num_heads % num_kv_heads != 0: - raise ValueError("num_heads must be divisible by num_kv_heads") - self.num_heads = num_heads - self.num_kv_heads = num_kv_heads - self.head_dim = dim // num_heads - if self.head_dim % 2 != 0: - raise ValueError("head_dim must be even for RoPE") - kv_dim = self.num_kv_heads * self.head_dim - self.c_q = CastedLinear(dim, dim, bias=False) - self.c_k = CastedLinear(dim, kv_dim, bias=False) - self.c_v = CastedLinear(dim, kv_dim, bias=False) - self.proj = CastedLinear(dim, dim, bias=False) - self.proj._zero_init = True - self.q_gain = nn.Parameter(torch.full((num_heads,), qk_gain_init, dtype=torch.float32)) - self.rotary = Rotary(self.head_dim, base=rope_base) - - def forward(self, x: Tensor) -> Tensor: - bsz, seqlen, dim = x.shape - q = self.c_q(x).reshape(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2) - k = self.c_k(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - v = self.c_v(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - q = F.rms_norm(q, (q.size(-1),)) - k = F.rms_norm(k, (k.size(-1),)) - cos, sin = self.rotary(seqlen, x.device, q.dtype) - q = apply_rotary_emb(q, cos, sin) - k = apply_rotary_emb(k, cos, sin) - q = q * self.q_gain.to(dtype=q.dtype)[None, :, None, None] - y = F.scaled_dot_product_attention( - q, - k, - v, - attn_mask=None, - is_causal=True, - enable_gqa=(self.num_kv_heads != self.num_heads), - ) - y = y.transpose(1, 2).contiguous().reshape(bsz, seqlen, dim) - return self.proj(y) - - -class MLP(nn.Module): - # relu^2 MLP from the original modded-nanogpt setup - def __init__(self, dim: int, mlp_mult: int): - super().__init__() - hidden = mlp_mult * dim - self.fc = CastedLinear(dim, hidden, bias=False) - self.proj = CastedLinear(hidden, dim, bias=False) - self.proj._zero_init = True - - def forward(self, x: Tensor) -> Tensor: - x = torch.relu(self.fc(x)) - return self.proj(x.square()) - - -class Block(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - num_kv_heads: int, - mlp_mult: int, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - self.attn_norm = RMSNorm() - self.mlp_norm = RMSNorm() - self.attn = CausalSelfAttention(dim, num_heads, num_kv_heads, rope_base, qk_gain_init) - self.mlp = MLP(dim, mlp_mult) - self.attn_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) - self.mlp_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) - self.resid_mix = nn.Parameter(torch.stack((torch.ones(dim), torch.zeros(dim))).float()) - - def forward(self, x: Tensor, x0: Tensor) -> Tensor: - mix = self.resid_mix.to(dtype=x.dtype) - x = mix[0][None, None, :] * x + mix[1][None, None, :] * x0 - attn_out = self.attn(self.attn_norm(x)) - x = x + self.attn_scale.to(dtype=x.dtype)[None, None, :] * attn_out - x = x + self.mlp_scale.to(dtype=x.dtype)[None, None, :] * self.mlp(self.mlp_norm(x)) - return x - - -class GPT(nn.Module): - def __init__( - self, - vocab_size: int, - num_layers: int, - model_dim: int, - num_heads: int, - num_kv_heads: int, - mlp_mult: int, - tie_embeddings: bool, - tied_embed_init_std: float, - logit_softcap: float, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - if logit_softcap <= 0.0: - raise ValueError(f"logit_softcap must be positive, got {logit_softcap}") - self.tie_embeddings = tie_embeddings - self.tied_embed_init_std = tied_embed_init_std - self.logit_softcap = logit_softcap - self.tok_emb = nn.Embedding(vocab_size, model_dim) - self.num_encoder_layers = num_layers // 2 - self.num_decoder_layers = num_layers - self.num_encoder_layers - self.num_skip_weights = min(self.num_encoder_layers, self.num_decoder_layers) - self.skip_weights = nn.Parameter(torch.ones(self.num_skip_weights, model_dim, dtype=torch.float32)) - self.blocks = nn.ModuleList( - [ - Block( - model_dim, - num_heads, - num_kv_heads, - mlp_mult, - rope_base, - qk_gain_init, - ) - for i in range(num_layers) - ] - ) - self.final_norm = RMSNorm() - self.lm_head = None if tie_embeddings else CastedLinear(model_dim, vocab_size, bias=False) - if self.lm_head is not None: - self.lm_head._zero_init = True - self._init_weights() - - def _init_weights(self) -> None: - if self.tie_embeddings: - nn.init.normal_(self.tok_emb.weight, mean=0.0, std=self.tied_embed_init_std) - for module in self.modules(): - if isinstance(module, nn.Linear) and getattr(module, "_zero_init", False): - nn.init.zeros_(module.weight) - - def forward(self, input_ids: Tensor, target_ids: Tensor) -> Tensor: - x = self.tok_emb(input_ids) - x = F.rms_norm(x, (x.size(-1),)) - x0 = x - skips: list[Tensor] = [] - - # First half stores skips; second half reuses them in reverse order. - for i in range(self.num_encoder_layers): - x = self.blocks[i](x, x0) - skips.append(x) - for i in range(self.num_decoder_layers): - if skips: - x = x + self.skip_weights[i].to(dtype=x.dtype)[None, None, :] * skips.pop() - x = self.blocks[self.num_encoder_layers + i](x, x0) - - x = self.final_norm(x).reshape(-1, x.size(-1)) - targets = target_ids.reshape(-1) - if self.tie_embeddings: - logits_proj = F.linear(x, self.tok_emb.weight) - else: - if self.lm_head is None: - raise RuntimeError("lm_head is required when tie_embeddings=False") - logits_proj = self.lm_head(x) - logits = self.logit_softcap * torch.tanh(logits_proj / self.logit_softcap) - return F.cross_entropy(logits.float(), targets, reduction="mean") - - -# ----------------------------- -# TRAINING -# ----------------------------- - -def main() -> None: - global zeropower_via_newtonschulz5 - - code = Path(__file__).read_text(encoding="utf-8") - args = Hyperparameters() - zeropower_via_newtonschulz5 = torch.compile(zeropower_via_newtonschulz5) - - # ----------------------------- - # DISTRIBUTED + CUDA SETUP - # ----------------------------- - - distributed = "RANK" in os.environ and "WORLD_SIZE" in os.environ - rank = int(os.environ.get("RANK", "0")) - world_size = int(os.environ.get("WORLD_SIZE", "1")) - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - if world_size <= 0: - raise ValueError(f"WORLD_SIZE must be positive, got {world_size}") - if 8 % world_size != 0: - raise ValueError(f"WORLD_SIZE={world_size} must divide 8 so grad_accum_steps stays integral") - grad_accum_steps = 8 // world_size - grad_scale = 1.0 / grad_accum_steps - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required") - device = torch.device("cuda", local_rank) - torch.cuda.set_device(device) - if distributed: - dist.init_process_group(backend="nccl", device_id=device) - dist.barrier() - master_process = rank == 0 - - # Fast math knobs - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True - from torch.backends.cuda import enable_cudnn_sdp, enable_flash_sdp, enable_math_sdp, enable_mem_efficient_sdp - - enable_cudnn_sdp(False) - enable_flash_sdp(True) - enable_mem_efficient_sdp(False) - enable_math_sdp(False) - - logfile = None - if master_process: - os.makedirs("logs", exist_ok=True) - logfile = f"logs/{args.run_id}.txt" - print(logfile) - - def log0(msg: str, console: bool = True) -> None: - if not master_process: - return - if console: - print(msg) - if logfile is not None: - with open(logfile, "a", encoding="utf-8") as f: - print(msg, file=f) - - log0(code, console=False) - log0("=" * 100, console=False) - log0(f"Running Python {sys.version}", console=False) - log0(f"Running PyTorch {torch.__version__}", console=False) - log0( - subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False).stdout, - console=False, - ) - log0("=" * 100, console=False) - - # ----------------------------- - # TOKENIZER + VALIDATION METRIC SETUP - # ----------------------------- - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed_all(args.seed) - - if not args.tokenizer_path.endswith(".model"): - raise ValueError(f"Script only setup for SentencePiece .model file: {args.tokenizer_path}") - sp = spm.SentencePieceProcessor(model_file=args.tokenizer_path) - if int(sp.vocab_size()) != args.vocab_size: - raise ValueError( - f"VOCAB_SIZE={args.vocab_size} does not match tokenizer vocab_size={int(sp.vocab_size())}" - ) - dataset_dir = Path(args.data_path).resolve() - actual_train_files = len(list(dataset_dir.glob("fineweb_train_*.bin"))) - val_tokens = load_validation_tokens(args.val_files, args.train_seq_len) - base_bytes_lut, has_leading_space_lut, is_boundary_token_lut = build_sentencepiece_luts( - sp, args.vocab_size, device - ) - log0(f"val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path={args.tokenizer_path}") - log0(f"train_loader:dataset:{dataset_dir.name} train_shards:{actual_train_files}") - log0(f"val_loader:shards pattern={args.val_files} tokens:{val_tokens.numel() - 1}") - - # ----------------------------- - # MODEL + OPTIMIZER SETUP - # ----------------------------- - - base_model = GPT( - vocab_size=args.vocab_size, - num_layers=args.num_layers, - model_dim=args.model_dim, - num_heads=args.num_heads, - num_kv_heads=args.num_kv_heads, - mlp_mult=args.mlp_mult, - tie_embeddings=args.tie_embeddings, - tied_embed_init_std=args.tied_embed_init_std, - logit_softcap=args.logit_softcap, - rope_base=args.rope_base, - qk_gain_init=args.qk_gain_init, - ).to(device).bfloat16() - for module in base_model.modules(): - if isinstance(module, CastedLinear): - module.float() - restore_low_dim_params_to_fp32(base_model) - compiled_model = torch.compile(base_model, dynamic=False, fullgraph=True) - model: nn.Module = DDP(compiled_model, device_ids=[local_rank], broadcast_buffers=False) if distributed else compiled_model - - # Optimizer split: - # - token embedding (Adam) uses EMBED_LR - # - untied lm_head (Adam) uses HEAD_LR - # - matrix params in transformer blocks use MATRIX_LR via Muon - # - vectors/scalars use SCALAR_LR via Adam - block_named_params = list(base_model.blocks.named_parameters()) - matrix_params = [ - p - for name, p in block_named_params - if p.ndim == 2 and not any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - scalar_params = [ - p - for name, p in block_named_params - if p.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - if base_model.skip_weights.numel() > 0: - scalar_params.append(base_model.skip_weights) - token_lr = args.tied_embed_lr if args.tie_embeddings else args.embed_lr - optimizer_tok = torch.optim.Adam( - [{"params": [base_model.tok_emb.weight], "lr": token_lr, "base_lr": token_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizer_muon = Muon( - matrix_params, - lr=args.matrix_lr, - momentum=args.muon_momentum, - backend_steps=args.muon_backend_steps, - ) - for group in optimizer_muon.param_groups: - group["base_lr"] = args.matrix_lr - optimizer_scalar = torch.optim.Adam( - [{"params": scalar_params, "lr": args.scalar_lr, "base_lr": args.scalar_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers: list[torch.optim.Optimizer] = [optimizer_tok, optimizer_muon, optimizer_scalar] - if base_model.lm_head is not None: - optimizer_head = torch.optim.Adam( - [{"params": [base_model.lm_head.weight], "lr": args.head_lr, "base_lr": args.head_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers.insert(1, optimizer_head) - - n_params = sum(p.numel() for p in base_model.parameters()) - log0(f"model_params:{n_params}") - log0(f"world_size:{world_size} grad_accum_steps:{grad_accum_steps}") - log0("sdp_backends:cudnn=False flash=True mem_efficient=False math=False") - log0(f"attention_mode:gqa num_heads:{args.num_heads} num_kv_heads:{args.num_kv_heads}") - log0( - f"tie_embeddings:{args.tie_embeddings} embed_lr:{token_lr} " - f"head_lr:{args.head_lr if base_model.lm_head is not None else 0.0} " - f"matrix_lr:{args.matrix_lr} scalar_lr:{args.scalar_lr}" - ) - log0( - f"train_batch_tokens:{args.train_batch_tokens} train_seq_len:{args.train_seq_len} " - f"iterations:{args.iterations} warmup_steps:{args.warmup_steps} " - f"max_wallclock_seconds:{args.max_wallclock_seconds:.3f}" - ) - log0(f"seed:{args.seed}") - - # ----------------------------- - # DATA LOADER & MODEL WARMUP - # ----------------------------- - - train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device) - - def zero_grad_all() -> None: - for opt in optimizers: - opt.zero_grad(set_to_none=True) - - max_wallclock_ms = 1000.0 * args.max_wallclock_seconds if args.max_wallclock_seconds > 0 else None - - def lr_mul(step: int, elapsed_ms: float) -> float: - if args.warmdown_iters <= 0: - return 1.0 - if max_wallclock_ms is None: - warmdown_start = max(args.iterations - args.warmdown_iters, 0) - return max((args.iterations - step) / max(args.warmdown_iters, 1), 0.0) if warmdown_start <= step < args.iterations else 1.0 - step_ms = elapsed_ms / max(step, 1) - warmdown_ms = args.warmdown_iters * step_ms - remaining_ms = max(max_wallclock_ms - elapsed_ms, 0.0) - return remaining_ms / max(warmdown_ms, 1e-9) if remaining_ms <= warmdown_ms else 1.0 - - # Warmup primes the compiled forward/backward/optimizer paths, then we restore the - # initial weights/optimizer state so measured training starts from the true init. - if args.warmup_steps > 0: - initial_model_state = {name: tensor.detach().cpu().clone() for name, tensor in base_model.state_dict().items()} - initial_optimizer_states = [copy.deepcopy(opt.state_dict()) for opt in optimizers] - model.train() - for warmup_step in range(args.warmup_steps): - zero_grad_all() - for micro_step in range(grad_accum_steps): - if distributed: - model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 - x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - warmup_loss = model(x, y) - (warmup_loss * grad_scale).backward() - for opt in optimizers: - opt.step() - zero_grad_all() - if args.warmup_steps <= 20 or (warmup_step + 1) % 10 == 0 or warmup_step + 1 == args.warmup_steps: - log0(f"warmup_step:{warmup_step + 1}/{args.warmup_steps}") - base_model.load_state_dict(initial_model_state, strict=True) - for opt, state in zip(optimizers, initial_optimizer_states, strict=True): - opt.load_state_dict(state) - zero_grad_all() - if distributed: - model.require_backward_grad_sync = True - train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device) - - # ----------------------------- - # MAIN TRAINING LOOP - # ----------------------------- - - training_time_ms = 0.0 - stop_after_step: int | None = None - torch.cuda.synchronize() - t0 = time.perf_counter() - - step = 0 - while True: - last_step = step == args.iterations or (stop_after_step is not None and step >= stop_after_step) - - should_validate = last_step or (args.val_loss_every > 0 and step % args.val_loss_every == 0) - if should_validate: - torch.cuda.synchronize() - training_time_ms += 1000.0 * (time.perf_counter() - t0) - val_loss, val_bpb = eval_val( - args, - model, - rank, - world_size, - device, - grad_accum_steps, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - ) - log0( - f"step:{step}/{args.iterations} val_loss:{val_loss:.4f} val_bpb:{val_bpb:.4f} " - f"train_time:{training_time_ms:.0f}ms step_avg:{training_time_ms / max(step, 1):.2f}ms" - ) - torch.cuda.synchronize() - t0 = time.perf_counter() - - if last_step: - if stop_after_step is not None and step < args.iterations: - log0( - f"stopping_early: wallclock_cap train_time:{training_time_ms:.0f}ms " - f"step:{step}/{args.iterations}" - ) - break - - elapsed_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) - scale = lr_mul(step, elapsed_ms) - zero_grad_all() - train_loss = torch.zeros((), device=device) - for micro_step in range(grad_accum_steps): - if distributed: - model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 - x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - loss = model(x, y) - train_loss += loss.detach() - (loss * grad_scale).backward() - train_loss /= grad_accum_steps - - frac = min(step / args.muon_momentum_warmup_steps, 1.0) if args.muon_momentum_warmup_steps > 0 else 1.0 - muon_momentum = (1 - frac) * args.muon_momentum_warmup_start + frac * args.muon_momentum - for group in optimizer_muon.param_groups: - group["momentum"] = muon_momentum - - for opt in optimizers: - for group in opt.param_groups: - group["lr"] = group["base_lr"] * scale - - if args.grad_clip_norm > 0: - torch.nn.utils.clip_grad_norm_(base_model.parameters(), args.grad_clip_norm) - for opt in optimizers: - opt.step() - zero_grad_all() - - step += 1 - approx_training_time_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) - should_log_train = ( - args.train_log_every > 0 - and (step <= 10 or step % args.train_log_every == 0 or stop_after_step is not None) - ) - if should_log_train: - log0( - f"step:{step}/{args.iterations} train_loss:{train_loss.item():.4f} " - f"train_time:{approx_training_time_ms:.0f}ms step_avg:{approx_training_time_ms / step:.2f}ms" - ) - - # Needed to sync whether we've reached the wallclock cap. - reached_cap = max_wallclock_ms is not None and approx_training_time_ms >= max_wallclock_ms - if distributed and max_wallclock_ms is not None: - reached_cap_tensor = torch.tensor(int(reached_cap), device=device) - dist.all_reduce(reached_cap_tensor, op=dist.ReduceOp.MAX) - reached_cap = bool(reached_cap_tensor.item()) - if stop_after_step is None and reached_cap: - stop_after_step = step - - log0( - f"peak memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB " - f"reserved: {torch.cuda.max_memory_reserved() // 1024 // 1024} MiB" - ) - - # ----------------------------- - # SERIALIZATION + ROUNDTRIP VALIDATION - # ----------------------------- - # Save the raw state (useful for debugging/loading in PyTorch directly), then always produce - # the compressed int8+zlib artifact and validate the round-tripped weights. - - if master_process: - torch.save(base_model.state_dict(), "final_model.pt") - model_bytes = os.path.getsize("final_model.pt") - code_bytes = len(code.encode("utf-8")) - log0(f"Serialized model: {model_bytes} bytes") - log0(f"Code size: {code_bytes} bytes") - log0(f"Total submission size: {model_bytes + code_bytes} bytes") - - quant_obj, quant_stats = quantize_state_dict_int8(base_model.state_dict()) - quant_buf = io.BytesIO() - torch.save(quant_obj, quant_buf) - quant_raw = quant_buf.getvalue() - quant_blob = zlib.compress(quant_raw, level=9) - quant_raw_bytes = len(quant_raw) - if master_process: - with open("final_model.int8.ptz", "wb") as f: - f.write(quant_blob) - quant_file_bytes = os.path.getsize("final_model.int8.ptz") - code_bytes = len(code.encode("utf-8")) - ratio = quant_stats["baseline_tensor_bytes"] / max(quant_stats["int8_payload_bytes"], 1) - log0( - f"Serialized model int8+zlib: {quant_file_bytes} bytes " - f"(payload:{quant_stats['int8_payload_bytes']} raw_torch:{quant_raw_bytes} payload_ratio:{ratio:.2f}x)" - ) - log0(f"Total submission size int8+zlib: {quant_file_bytes + code_bytes} bytes") - - if distributed: - dist.barrier() - with open("final_model.int8.ptz", "rb") as f: - quant_blob_disk = f.read() - quant_state = torch.load(io.BytesIO(zlib.decompress(quant_blob_disk)), map_location="cpu") - base_model.load_state_dict(dequantize_state_dict_int8(quant_state), strict=True) - torch.cuda.synchronize() - t_qeval = time.perf_counter() - q_val_loss, q_val_bpb = eval_val( - args, - model, - rank, - world_size, - device, - grad_accum_steps, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - ) - torch.cuda.synchronize() - log0( - f"final_int8_zlib_roundtrip val_loss:{q_val_loss:.4f} val_bpb:{q_val_bpb:.4f} " - f"eval_time:{1000.0 * (time.perf_counter() - t_qeval):.0f}ms" - ) - log0(f"final_int8_zlib_roundtrip_exact val_loss:{q_val_loss:.8f} val_bpb:{q_val_bpb:.8f}") - - if distributed: - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/records/track_10min_16mb/2026-03-18_FP16Embed_WD3600/README.md b/records/track_10min_16mb/2026-03-18_FP16Embed_WD3600/README.md deleted file mode 100644 index 0e5004df77..0000000000 --- a/records/track_10min_16mb/2026-03-18_FP16Embed_WD3600/README.md +++ /dev/null @@ -1,59 +0,0 @@ -Kept the tied embedding in fp16 instead of quantizing it to int8, and tuned the LR schedule. Turns out the embedding is by far the most sensitive tensor to quantize — it's pulling double duty as the output head, so every bit of precision matters. - -## what changed - -**fp16 embedding passthrough**: one-line change in the quantization function. Instead of int8-quantizing `tok_emb.weight`, I pass it through as fp16. This drops the post-quant BPB degradation from ~0.007 to basically nothing (~0.0005). The tradeoff is ~500KB extra in the artifact, so I shrank the MLP hidden from 1024 to 992 to stay under 16MB. - -**warmdown + LR**: bumped `WARMDOWN_ITERS` from 1200 to 3600 and `MATRIX_LR` from 0.04 to 0.06. The default schedule assumes way more steps than you actually get in 10 minutes, so a longer warmdown and higher LR help the model converge properly. - -## config - -``` -VOCAB_SIZE=1024 NUM_LAYERS=9 MODEL_DIM=512 NUM_HEADS=8 NUM_KV_HEADS=4 -MLP_HIDDEN=992 TIE_EMBEDDINGS=1 WARMDOWN_ITERS=3600 MATRIX_LR=0.06 -``` - -## run command - -```bash -RUN_ID=fp16embed \ -DATA_PATH=./data/datasets/fineweb10B_sp1024/ \ -TOKENIZER_PATH=./data/tokenizers/fineweb_1024_bpe.model \ -VOCAB_SIZE=1024 \ -MLP_HIDDEN=992 \ -WARMDOWN_ITERS=3600 \ -MATRIX_LR=0.06 \ -torchrun --standalone --nproc_per_node=8 train_gpt.py -``` - -Note: don't set `NCCL_IB_DISABLE=1` — it tanks step throughput on pods with IB/NVLink (~60ms vs ~44ms per step). - -## results - -8xH100 SXM (RunPod secure cloud): - -| seed | steps | val_loss | val_bpb | artifact size | -|------|-------|----------|---------|---------------| -| 1337 | 13,692 | 2.0595 | 1.2197 | 15.90MB | -| 42 | 13,722 | 2.0600 | 1.2201 | 15.90MB | - -Pre-quant vs post-quant gap: ~0.0005 BPB (baseline gap is ~0.007). - -Improvement over baseline: ~0.013 nats. - -Also ran 3 seeds on 8xH200 SXM (all consistent, 1.2163-1.2179 BPB). - -## things I tried that didn't work - -- **SwiGLU**: better per-step quality but 45% slower on 8-GPU, so fewer total steps. Net negative. -- **depth recurrence** (looping layers): promising idea but needs way more steps than 10 min allows. -- **QAT**: tried both full-training and late-stage. The overhead per step wasn't worth the small quant gap reduction. -- **lzma compression**: actually compresses worse than zlib for int8 weight data. -- **higher embed LR** (0.08 vs 0.05): hurt convergence. - -## files - -- `train_gpt.py` — modified training script -- `train.log` — 8xH100 log (seed 1337) -- `train_seed42.log` — 8xH100 log (seed 42) -- `submission.json` diff --git a/records/track_10min_16mb/2026-03-18_FP16Embed_WD3600/submission.json b/records/track_10min_16mb/2026-03-18_FP16Embed_WD3600/submission.json deleted file mode 100644 index 6ff07bab72..0000000000 --- a/records/track_10min_16mb/2026-03-18_FP16Embed_WD3600/submission.json +++ /dev/null @@ -1,11 +0,0 @@ -{ - "author": "Renier Velazco", - "github_id": "chonchiog", - "name": "FP16 Tied Embedding + LR/Warmdown Tuning", - "blurb": "Keep tok_emb.weight in fp16 during int8 quantization to eliminate the output-head quantization gap (0.007 -> 0.0005 BPB). Slightly reduce MLP hidden (992 vs 1024) to fit within 16MB. Tune warmdown (3600 vs 1200) and matrix LR (0.06 vs 0.04) for better convergence under the 10-min wallclock cap.", - "date": "2026-03-18", - "val_loss": 2.05945460, - "val_bpb": 1.21972502, - "bytes_total": 15896222, - "bytes_code": 48125 -} diff --git a/records/track_10min_16mb/2026-03-18_FP16Embed_WD3600/train.log b/records/track_10min_16mb/2026-03-18_FP16Embed_WD3600/train.log deleted file mode 100644 index b4cf10e39f..0000000000 --- a/records/track_10min_16mb/2026-03-18_FP16Embed_WD3600/train.log +++ /dev/null @@ -1,1263 +0,0 @@ -""" -The `train_gpt.py` and `train_gpt_mlx.py` scripts are intended as good launching-off points for new participants, not SOTA configs. We'll accept PRs that tune, improve, or simplify these scripts without significantly increasing complexity, but competitive submissions should stay in the `/records` folder. - -Hard stop: `train_gpt.py` and `train_gpt_mlx.py` must never be longer than 1500 lines. -""" - -from __future__ import annotations - -import copy -import glob -import io -import math -import os -import random -import subprocess -import sys -import time -import uuid -import zlib -from pathlib import Path - -import numpy as np -import sentencepiece as spm -import torch -import torch.distributed as dist -import torch.nn.functional as F -from torch import Tensor, nn -from torch.nn.parallel import DistributedDataParallel as DDP - -# ----------------------------- -# HYPERPARAMETERS -# ----------------------------- -# Default Simple Baseline run: -# - 9 transformer blocks at width 512 -# - 8 attention heads with 4 KV heads (GQA) and 2x MLP expansion -# - vocab size 1024, sequence length 1024, tied embeddings -# - 524,288 train tokens per step for 20,000 iterations with a ~10 minute cap - -class Hyperparameters: - # Data paths are shard globs produced by the existing preprocessing pipeline. - data_path = os.environ.get("DATA_PATH", "./data/datasets/fineweb10B_sp1024") - train_files = os.path.join(data_path, "fineweb_train_*.bin") - val_files = os.path.join(data_path, "fineweb_val_*.bin") - tokenizer_path = os.environ.get("TOKENIZER_PATH", "./data/tokenizers/fineweb_1024_bpe.model") - run_id = os.environ.get("RUN_ID", str(uuid.uuid4())) - seed = int(os.environ.get("SEED", 1337)) - - # Validation cadence and batch size. Validation always uses the full fineweb_val split. - val_batch_size = int(os.environ.get("VAL_BATCH_SIZE", 524_288)) - val_loss_every = int(os.environ.get("VAL_LOSS_EVERY", 1000)) - train_log_every = int(os.environ.get("TRAIN_LOG_EVERY", 200)) - - # Training length. - iterations = int(os.environ.get("ITERATIONS", 20000)) - warmdown_iters = int(os.environ.get("WARMDOWN_ITERS", 1200)) - warmup_steps = int(os.environ.get("WARMUP_STEPS", 20)) - train_batch_tokens = int(os.environ.get("TRAIN_BATCH_TOKENS", 524_288)) - train_seq_len = int(os.environ.get("TRAIN_SEQ_LEN", 1024)) - max_wallclock_seconds = float(os.environ.get("MAX_WALLCLOCK_SECONDS", 600.0)) - qk_gain_init = float(os.environ.get("QK_GAIN_INIT", 1.5)) - - # Model shape. - vocab_size = int(os.environ.get("VOCAB_SIZE", 1024)) - num_layers = int(os.environ.get("NUM_LAYERS", 9)) - num_kv_heads = int(os.environ.get("NUM_KV_HEADS", 4)) - model_dim = int(os.environ.get("MODEL_DIM", 512)) - num_heads = int(os.environ.get("NUM_HEADS", 8)) - mlp_mult = int(os.environ.get("MLP_MULT", 2)) - mlp_hidden = int(os.environ.get("MLP_HIDDEN", 0)) # override mlp_mult if > 0 - tie_embeddings = bool(int(os.environ.get("TIE_EMBEDDINGS", "1"))) - rope_base = float(os.environ.get("ROPE_BASE", 10000.0)) - logit_softcap = float(os.environ.get("LOGIT_SOFTCAP", 30.0)) - - # Optimizer hyperparameters. - embed_lr = float(os.environ.get("EMBED_LR", 0.6)) - head_lr = float(os.environ.get("HEAD_LR", 0.008)) - tied_embed_lr = float(os.environ.get("TIED_EMBED_LR", 0.05)) - tied_embed_init_std = float(os.environ.get("TIED_EMBED_INIT_STD", 0.005)) - matrix_lr = float(os.environ.get("MATRIX_LR", 0.04)) - scalar_lr = float(os.environ.get("SCALAR_LR", 0.04)) - muon_momentum = float(os.environ.get("MUON_MOMENTUM", 0.95)) - muon_backend_steps = int(os.environ.get("MUON_BACKEND_STEPS", 5)) - muon_momentum_warmup_start = float(os.environ.get("MUON_MOMENTUM_WARMUP_START", 0.85)) - muon_momentum_warmup_steps = int(os.environ.get("MUON_MOMENTUM_WARMUP_STEPS", 500)) - beta1 = float(os.environ.get("BETA1", 0.9)) - beta2 = float(os.environ.get("BETA2", 0.95)) - adam_eps = float(os.environ.get("ADAM_EPS", 1e-8)) - grad_clip_norm = float(os.environ.get("GRAD_CLIP_NORM", 0.0)) - -# ----------------------------- -# MUON OPTIMIZER -# ----------------------------- -# -# As borrowed from modded-nanogpt -# Background on Muon: https://kellerjordan.github.io/posts/muon/ - -def zeropower_via_newtonschulz5(G: Tensor, steps: int = 10, eps: float = 1e-7) -> Tensor: - # Orthogonalize a 2D update matrix with a fast Newton-Schulz iteration. - # Muon uses this to normalize matrix-shaped gradients before applying them. - a, b, c = (3.4445, -4.7750, 2.0315) - X = G.bfloat16() - X /= X.norm() + eps - transposed = G.size(0) > G.size(1) - if transposed: - X = X.T - for _ in range(steps): - A = X @ X.T - B = b * A + c * A @ A - X = a * X + B @ X - return X.T if transposed else X - - -class Muon(torch.optim.Optimizer): - def __init__(self, params, lr: float, momentum: float, backend_steps: int, nesterov: bool = True): - super().__init__( - params, - dict(lr=lr, momentum=momentum, backend_steps=backend_steps, nesterov=nesterov), - ) - - @torch.no_grad() - def step(self, closure=None): - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - distributed = dist.is_available() and dist.is_initialized() - world_size = dist.get_world_size() if distributed else 1 - rank = dist.get_rank() if distributed else 0 - - for group in self.param_groups: - params = group["params"] - if not params: - continue - lr = group["lr"] - momentum = group["momentum"] - backend_steps = group["backend_steps"] - nesterov = group["nesterov"] - - total_params = sum(int(p.numel()) for p in params) - updates_flat = torch.zeros(total_params, device=params[0].device, dtype=torch.bfloat16) - - curr = 0 - for i, p in enumerate(params): - if i % world_size == rank and p.grad is not None: - g = p.grad - state = self.state[p] - if "momentum_buffer" not in state: - state["momentum_buffer"] = torch.zeros_like(g) - buf = state["momentum_buffer"] - buf.mul_(momentum).add_(g) - if nesterov: - g = g.add(buf, alpha=momentum) - g = zeropower_via_newtonschulz5(g, steps=backend_steps) - # Scale correction from Muon reference implementations. - g *= max(1, g.size(0) / g.size(1)) ** 0.5 - updates_flat[curr : curr + p.numel()] = g.reshape(-1) - curr += p.numel() - - if distributed: - dist.all_reduce(updates_flat, op=dist.ReduceOp.SUM) - - curr = 0 - for p in params: - g = updates_flat[curr : curr + p.numel()].view_as(p).to(dtype=p.dtype) - p.add_(g, alpha=-lr) - curr += p.numel() - - return loss - - -# ----------------------------- -# TOKENIZER-AGNOSTIC EVALUATION SETUP -# ----------------------------- -# -# It's common for small models have a large fraction of their parameters be embeddings, since the 2 * d_model * d_vocab vectors can be gigantic. -# Instead of locking the tokenizer, we let you bring your own and calculate our validation metrics on the average compression of the validation set. -# We calculate BPB (bits-per-byte) instead of validation loss, so we need methods to count the number of bits per token in the tokenizer. -# Note: Submissions that edit the tokenizer will be examined more carefully, since screwing this up might unjustly improve your score. - -def build_sentencepiece_luts( - sp: spm.SentencePieceProcessor, vocab_size: int, device: torch.device -) -> tuple[Tensor, Tensor, Tensor]: - sp_vocab_size = int(sp.vocab_size()) - table_size = max(sp_vocab_size, vocab_size) - base_bytes_np = np.zeros((table_size,), dtype=np.int16) - has_leading_space_np = np.zeros((table_size,), dtype=np.bool_) - is_boundary_token_np = np.ones((table_size,), dtype=np.bool_) - for token_id in range(sp_vocab_size): - if sp.is_control(token_id) or sp.is_unknown(token_id) or sp.is_unused(token_id): - continue - is_boundary_token_np[token_id] = False - if sp.is_byte(token_id): - base_bytes_np[token_id] = 1 - continue - piece = sp.id_to_piece(token_id) - if piece.startswith("▁"): - has_leading_space_np[token_id] = True - piece = piece[1:] - base_bytes_np[token_id] = len(piece.encode("utf-8")) - return ( - torch.tensor(base_bytes_np, dtype=torch.int16, device=device), - torch.tensor(has_leading_space_np, dtype=torch.bool, device=device), - torch.tensor(is_boundary_token_np, dtype=torch.bool, device=device), - ) - - -def load_validation_tokens(pattern: str, seq_len: int) -> Tensor: - files = [Path(p) for p in sorted(glob.glob(pattern))] - if not files: - raise FileNotFoundError(f"No files found for pattern: {pattern}") - # The export pipeline writes the fixed first-50k-doc validation set to fineweb_val_*. - tokens = torch.cat([load_data_shard(file) for file in files]).contiguous() - usable = ((tokens.numel() - 1) // seq_len) * seq_len - if usable <= 0: - raise ValueError(f"Validation split is too short for TRAIN_SEQ_LEN={seq_len}") - return tokens[: usable + 1] - - -def eval_val( - args: Hyperparameters, - model: nn.Module, - rank: int, - world_size: int, - device: torch.device, - grad_accum_steps: int, - val_tokens: Tensor, - base_bytes_lut: Tensor, - has_leading_space_lut: Tensor, - is_boundary_token_lut: Tensor, -) -> tuple[float, float]: - # Validation computes two metrics: - # - val_loss: token cross-entropy (natural log) - # - val_bpb: tokenizer-agnostic compression metric used by the challenge - local_batch_tokens = args.val_batch_size // (world_size * grad_accum_steps) - if local_batch_tokens < args.train_seq_len: - raise ValueError( - "VAL_BATCH_SIZE must provide at least one sequence per rank; " - f"got VAL_BATCH_SIZE={args.val_batch_size}, WORLD_SIZE={world_size}, " - f"GRAD_ACCUM_STEPS={grad_accum_steps}, TRAIN_SEQ_LEN={args.train_seq_len}" - ) - local_batch_seqs = local_batch_tokens // args.train_seq_len - total_seqs = (val_tokens.numel() - 1) // args.train_seq_len - seq_start = (total_seqs * rank) // world_size - seq_end = (total_seqs * (rank + 1)) // world_size - val_loss_sum = torch.zeros((), device=device, dtype=torch.float64) - val_token_count = torch.zeros((), device=device, dtype=torch.float64) - val_byte_count = torch.zeros((), device=device, dtype=torch.float64) - - model.eval() - with torch.inference_mode(): - for batch_seq_start in range(seq_start, seq_end, local_batch_seqs): - batch_seq_end = min(batch_seq_start + local_batch_seqs, seq_end) - raw_start = batch_seq_start * args.train_seq_len - raw_end = batch_seq_end * args.train_seq_len + 1 - local = val_tokens[raw_start:raw_end].to(device=device, dtype=torch.int64, non_blocking=True) - x = local[:-1].reshape(-1, args.train_seq_len) - y = local[1:].reshape(-1, args.train_seq_len) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - batch_loss = model(x, y).detach() - batch_token_count = float(y.numel()) - val_loss_sum += batch_loss.to(torch.float64) * batch_token_count - val_token_count += batch_token_count - prev_ids = x.reshape(-1) - tgt_ids = y.reshape(-1) - token_bytes = base_bytes_lut[tgt_ids].to(dtype=torch.int16) - token_bytes += (has_leading_space_lut[tgt_ids] & ~is_boundary_token_lut[prev_ids]).to(dtype=torch.int16) - val_byte_count += token_bytes.to(torch.float64).sum() - - if dist.is_available() and dist.is_initialized(): - dist.all_reduce(val_loss_sum, op=dist.ReduceOp.SUM) - dist.all_reduce(val_token_count, op=dist.ReduceOp.SUM) - dist.all_reduce(val_byte_count, op=dist.ReduceOp.SUM) - - val_loss = val_loss_sum / val_token_count - bits_per_token = val_loss.item() / math.log(2.0) - tokens_per_byte = val_token_count.item() / val_byte_count.item() - model.train() - return float(val_loss.item()), float(bits_per_token * tokens_per_byte) - -# ----------------------------- -# POST-TRAINING QUANTIZATION -# ----------------------------- -# -# It's silly to export our model, which is trained in bf16 and fp32, at that same precision. -# Instead, we get approximately the same model (with a small hit) by quantizing the model to int8 & zlib compressing. -# We can then decompress the model and run in higher precision for evaluation, after closing in under the size limit. - -CONTROL_TENSOR_NAME_PATTERNS = tuple( - pattern - for pattern in os.environ.get( - "CONTROL_TENSOR_NAME_PATTERNS", - "attn_scale,attn_scales,mlp_scale,mlp_scales,resid_mix,resid_mixes,q_gain,skip_weight,skip_weights", - ).split(",") - if pattern -) -INT8_KEEP_FLOAT_FP32_NAME_PATTERNS = tuple( - pattern - for pattern in os.environ.get( - "INT8_KEEP_FLOAT_FP32_NAME_PATTERNS", - ",".join(CONTROL_TENSOR_NAME_PATTERNS), - ).split(",") - if pattern -) -INT8_KEEP_FLOAT_MAX_NUMEL = 65_536 -INT8_KEEP_FLOAT_STORE_DTYPE = torch.float16 -INT8_PER_ROW_SCALE_DTYPE = torch.float16 -INT8_CLIP_PERCENTILE = 99.99984 -INT8_CLIP_Q = INT8_CLIP_PERCENTILE / 100.0 - -def tensor_nbytes(t: Tensor) -> int: - return int(t.numel()) * int(t.element_size()) - -def keep_float_tensor(name: str, t: Tensor, passthrough_orig_dtypes: dict[str, str]) -> Tensor: - if any(pattern in name for pattern in INT8_KEEP_FLOAT_FP32_NAME_PATTERNS): - return t.float().contiguous() - if t.dtype in {torch.float32, torch.bfloat16}: - passthrough_orig_dtypes[name] = str(t.dtype).removeprefix("torch.") - return t.to(dtype=INT8_KEEP_FLOAT_STORE_DTYPE).contiguous() - return t - -def quantize_float_tensor(t: Tensor) -> tuple[Tensor, Tensor]: - t32 = t.float() - if t32.ndim == 2: - # Matrices get one scale per row, which usually tracks output-channel - # ranges much better than a single tensor-wide scale. - clip_abs = ( - torch.quantile(t32.abs(), INT8_CLIP_Q, dim=1) - if t32.numel() - else torch.empty((t32.shape[0],), dtype=torch.float32) - ) - clipped = torch.maximum(torch.minimum(t32, clip_abs[:, None]), -clip_abs[:, None]) - scale = (clip_abs / 127.0).clamp_min(1.0 / 127.0) - q = torch.clamp(torch.round(clipped / scale[:, None]), -127, 127).to(torch.int8).contiguous() - return q, scale.to(dtype=INT8_PER_ROW_SCALE_DTYPE).contiguous() - - # Vectors / scalars use a simpler per-tensor scale. - clip_abs = float(torch.quantile(t32.abs().flatten(), INT8_CLIP_Q).item()) if t32.numel() else 0.0 - scale = torch.tensor(clip_abs / 127.0 if clip_abs > 0 else 1.0, dtype=torch.float32) - q = torch.clamp(torch.round(torch.clamp(t32, -clip_abs, clip_abs) / scale), -127, 127).to(torch.int8).contiguous() - return q, scale - -def quantize_state_dict_int8(state_dict: dict[str, Tensor]): - # Single supported clean-script export format: - # - per-row int8 for 2D float tensors - # - per-tensor int8 for other float tensors - # - exact passthrough for non-floats - # - passthrough for small float tensors, stored as fp16 to save bytes - quantized: dict[str, Tensor] = {} - scales: dict[str, Tensor] = {} - dtypes: dict[str, str] = {} - passthrough: dict[str, Tensor] = {} - passthrough_orig_dtypes: dict[str, str] = {} - qmeta: dict[str, dict[str, object]] = {} - stats = dict.fromkeys( - ("param_count", "num_tensors", "num_float_tensors", "num_nonfloat_tensors", "baseline_tensor_bytes", "int8_payload_bytes"), - 0, - ) - - for name, tensor in state_dict.items(): - t = tensor.detach().to("cpu").contiguous() - stats["param_count"] += int(t.numel()) - stats["num_tensors"] += 1 - stats["baseline_tensor_bytes"] += tensor_nbytes(t) - - if not t.is_floating_point(): - stats["num_nonfloat_tensors"] += 1 - passthrough[name] = t - stats["int8_payload_bytes"] += tensor_nbytes(t) - continue - - # Small float tensors are cheap enough to keep directly. We still downcast - # fp32/bf16 passthrough tensors to fp16 so metadata does not dominate size. - # Also keep the tied embedding in fp16 to minimize quantization degradation - # on the output head, which is the most sensitive tensor for BPB. - if t.numel() <= INT8_KEEP_FLOAT_MAX_NUMEL or name == "tok_emb.weight": - kept = keep_float_tensor(name, t, passthrough_orig_dtypes) - passthrough[name] = kept - stats["int8_payload_bytes"] += tensor_nbytes(kept) - continue - - stats["num_float_tensors"] += 1 - q, s = quantize_float_tensor(t) - if s.ndim > 0: - qmeta[name] = {"scheme": "per_row", "axis": 0} - quantized[name] = q - scales[name] = s - dtypes[name] = str(t.dtype).removeprefix("torch.") - stats["int8_payload_bytes"] += tensor_nbytes(q) + tensor_nbytes(s) - - obj: dict[str, object] = { - "__quant_format__": "int8_clean_per_row_v1", - "quantized": quantized, - "scales": scales, - "dtypes": dtypes, - "passthrough": passthrough, - } - if qmeta: - obj["qmeta"] = qmeta - if passthrough_orig_dtypes: - obj["passthrough_orig_dtypes"] = passthrough_orig_dtypes - return obj, stats - -def dequantize_state_dict_int8(obj: dict[str, object]) -> dict[str, Tensor]: - out: dict[str, Tensor] = {} - qmeta = obj.get("qmeta", {}) - passthrough_orig_dtypes = obj.get("passthrough_orig_dtypes", {}) - for name, q in obj["quantized"].items(): - dtype = getattr(torch, obj["dtypes"][name]) - s = obj["scales"][name] - if qmeta.get(name, {}).get("scheme") == "per_row" or s.ndim > 0: - s = s.to(dtype=torch.float32) - # Broadcast the saved row scale back across trailing dimensions. - out[name] = (q.float() * s.view(q.shape[0], *([1] * (q.ndim - 1)))).to(dtype=dtype).contiguous() - else: - scale = float(s.item()) - out[name] = (q.float() * scale).to(dtype=dtype).contiguous() - for name, t in obj["passthrough"].items(): - # Restore small tensors, undoing the temporary fp16 storage cast if needed. - out_t = t.detach().to("cpu").contiguous() - orig_dtype = passthrough_orig_dtypes.get(name) - if isinstance(orig_dtype, str): - out_t = out_t.to(dtype=getattr(torch, orig_dtype)).contiguous() - out[name] = out_t - return out - - -# ----------------------------- -# DATA LOADING -# ----------------------------- - -def load_data_shard(file: Path) -> Tensor: - header_bytes = 256 * np.dtype(" None: - self.file_idx = (self.file_idx + 1) % len(self.files) - self.tokens = load_data_shard(self.files[self.file_idx]) - self.pos = 0 - - def take(self, n: int) -> Tensor: - chunks: list[Tensor] = [] - remaining = n - while remaining > 0: - avail = self.tokens.numel() - self.pos - if avail <= 0: - self._advance_file() - continue - k = min(remaining, avail) - chunks.append(self.tokens[self.pos : self.pos + k]) - self.pos += k - remaining -= k - return chunks[0] if len(chunks) == 1 else torch.cat(chunks) - - -class DistributedTokenLoader: - # Each call consumes a contiguous chunk from the shared token stream, then slices out - # one disjoint span per rank. The extra "+1" token lets us build (x, y) by shifting. - def __init__(self, pattern: str, rank: int, world_size: int, device: torch.device): - self.rank = rank - self.world_size = world_size - self.device = device - self.stream = TokenStream(pattern) - - def next_batch(self, global_tokens: int, seq_len: int, grad_accum_steps: int) -> tuple[Tensor, Tensor]: - local_tokens = global_tokens // (self.world_size * grad_accum_steps) - per_rank_span = local_tokens + 1 - chunk = self.stream.take(per_rank_span * self.world_size) - start = self.rank * per_rank_span - local = chunk[start : start + per_rank_span].to(dtype=torch.int64) - x = local[:-1].reshape(-1, seq_len) - y = local[1:].reshape(-1, seq_len) - return x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True) - -# ----------------------------- -# TRANSFORMER MODULES -# ----------------------------- - -class RMSNorm(nn.Module): - def __init__(self, eps: float | None = None): - super().__init__() - self.eps = eps - - def forward(self, x: Tensor) -> Tensor: - return F.rms_norm(x, (x.size(-1),), eps=self.eps) - - -class CastedLinear(nn.Linear): - # Keep weights in fp32 for optimizer/state quality, cast at matmul time for bf16 compute. - def forward(self, x: Tensor) -> Tensor: - bias = self.bias.to(x.dtype) if self.bias is not None else None - return F.linear(x, self.weight.to(x.dtype), bias) - - -def restore_low_dim_params_to_fp32(module: nn.Module) -> None: - # Keep small/control parameters in fp32 even when the model body runs in bf16. - with torch.no_grad(): - for name, param in module.named_parameters(): - if (param.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS)) and param.dtype != torch.float32: - param.data = param.data.float() - - -class Rotary(nn.Module): - # Caches cos/sin tables per sequence length on the current device. - def __init__(self, dim: int, base: float = 10000.0): - super().__init__() - inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self._seq_len_cached = 0 - self._cos_cached: Tensor | None = None - self._sin_cached: Tensor | None = None - - def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> tuple[Tensor, Tensor]: - if ( - self._cos_cached is None - or self._sin_cached is None - or self._seq_len_cached != seq_len - or self._cos_cached.device != device - ): - t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) - freqs = torch.outer(t, self.inv_freq.to(device)) - self._cos_cached = freqs.cos()[None, None, :, :] - self._sin_cached = freqs.sin()[None, None, :, :] - self._seq_len_cached = seq_len - return self._cos_cached.to(dtype=dtype), self._sin_cached.to(dtype=dtype) - - -def apply_rotary_emb(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: - half = x.size(-1) // 2 - x1, x2 = x[..., :half], x[..., half:] - return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1) - - -class CausalSelfAttention(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - num_kv_heads: int, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - if dim % num_heads != 0: - raise ValueError("model_dim must be divisible by num_heads") - if num_heads % num_kv_heads != 0: - raise ValueError("num_heads must be divisible by num_kv_heads") - self.num_heads = num_heads - self.num_kv_heads = num_kv_heads - self.head_dim = dim // num_heads - if self.head_dim % 2 != 0: - raise ValueError("head_dim must be even for RoPE") - kv_dim = self.num_kv_heads * self.head_dim - self.c_q = CastedLinear(dim, dim, bias=False) - self.c_k = CastedLinear(dim, kv_dim, bias=False) - self.c_v = CastedLinear(dim, kv_dim, bias=False) - self.proj = CastedLinear(dim, dim, bias=False) - self.proj._zero_init = True - self.q_gain = nn.Parameter(torch.full((num_heads,), qk_gain_init, dtype=torch.float32)) - self.rotary = Rotary(self.head_dim, base=rope_base) - - def forward(self, x: Tensor) -> Tensor: - bsz, seqlen, dim = x.shape - q = self.c_q(x).reshape(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2) - k = self.c_k(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - v = self.c_v(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - q = F.rms_norm(q, (q.size(-1),)) - k = F.rms_norm(k, (k.size(-1),)) - cos, sin = self.rotary(seqlen, x.device, q.dtype) - q = apply_rotary_emb(q, cos, sin) - k = apply_rotary_emb(k, cos, sin) - q = q * self.q_gain.to(dtype=q.dtype)[None, :, None, None] - y = F.scaled_dot_product_attention( - q, - k, - v, - attn_mask=None, - is_causal=True, - enable_gqa=(self.num_kv_heads != self.num_heads), - ) - y = y.transpose(1, 2).contiguous().reshape(bsz, seqlen, dim) - return self.proj(y) - - -class MLP(nn.Module): - # relu^2 MLP from the original modded-nanogpt setup - def __init__(self, dim: int, mlp_mult: int, mlp_hidden: int = 0): - super().__init__() - hidden = mlp_hidden if mlp_hidden > 0 else mlp_mult * dim - self.fc = CastedLinear(dim, hidden, bias=False) - self.proj = CastedLinear(hidden, dim, bias=False) - self.proj._zero_init = True - - def forward(self, x: Tensor) -> Tensor: - x = torch.relu(self.fc(x)) - return self.proj(x.square()) - - -class Block(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - num_kv_heads: int, - mlp_mult: int, - rope_base: float, - qk_gain_init: float, - mlp_hidden: int = 0, - ): - super().__init__() - self.attn_norm = RMSNorm() - self.mlp_norm = RMSNorm() - self.attn = CausalSelfAttention(dim, num_heads, num_kv_heads, rope_base, qk_gain_init) - self.mlp = MLP(dim, mlp_mult, mlp_hidden=mlp_hidden) - self.attn_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) - self.mlp_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) - self.resid_mix = nn.Parameter(torch.stack((torch.ones(dim), torch.zeros(dim))).float()) - - def forward(self, x: Tensor, x0: Tensor) -> Tensor: - mix = self.resid_mix.to(dtype=x.dtype) - x = mix[0][None, None, :] * x + mix[1][None, None, :] * x0 - attn_out = self.attn(self.attn_norm(x)) - x = x + self.attn_scale.to(dtype=x.dtype)[None, None, :] * attn_out - x = x + self.mlp_scale.to(dtype=x.dtype)[None, None, :] * self.mlp(self.mlp_norm(x)) - return x - - -class GPT(nn.Module): - def __init__( - self, - vocab_size: int, - num_layers: int, - model_dim: int, - num_heads: int, - num_kv_heads: int, - mlp_mult: int, - tie_embeddings: bool, - tied_embed_init_std: float, - logit_softcap: float, - rope_base: float, - qk_gain_init: float, - mlp_hidden: int = 0, - ): - super().__init__() - if logit_softcap <= 0.0: - raise ValueError(f"logit_softcap must be positive, got {logit_softcap}") - self.tie_embeddings = tie_embeddings - self.tied_embed_init_std = tied_embed_init_std - self.logit_softcap = logit_softcap - self.tok_emb = nn.Embedding(vocab_size, model_dim) - self.num_encoder_layers = num_layers // 2 - self.num_decoder_layers = num_layers - self.num_encoder_layers - self.num_skip_weights = min(self.num_encoder_layers, self.num_decoder_layers) - self.skip_weights = nn.Parameter(torch.ones(self.num_skip_weights, model_dim, dtype=torch.float32)) - self.blocks = nn.ModuleList( - [ - Block( - model_dim, - num_heads, - num_kv_heads, - mlp_mult, - rope_base, - qk_gain_init, - mlp_hidden=mlp_hidden, - ) - for i in range(num_layers) - ] - ) - self.final_norm = RMSNorm() - self.lm_head = None if tie_embeddings else CastedLinear(model_dim, vocab_size, bias=False) - if self.lm_head is not None: - self.lm_head._zero_init = True - self._init_weights() - - def _init_weights(self) -> None: - if self.tie_embeddings: - nn.init.normal_(self.tok_emb.weight, mean=0.0, std=self.tied_embed_init_std) - for module in self.modules(): - if isinstance(module, nn.Linear) and getattr(module, "_zero_init", False): - nn.init.zeros_(module.weight) - - def forward(self, input_ids: Tensor, target_ids: Tensor) -> Tensor: - x = self.tok_emb(input_ids) - x = F.rms_norm(x, (x.size(-1),)) - x0 = x - skips: list[Tensor] = [] - - # First half stores skips; second half reuses them in reverse order. - for i in range(self.num_encoder_layers): - x = self.blocks[i](x, x0) - skips.append(x) - for i in range(self.num_decoder_layers): - if skips: - x = x + self.skip_weights[i].to(dtype=x.dtype)[None, None, :] * skips.pop() - x = self.blocks[self.num_encoder_layers + i](x, x0) - - x = self.final_norm(x).reshape(-1, x.size(-1)) - targets = target_ids.reshape(-1) - if self.tie_embeddings: - logits_proj = F.linear(x, self.tok_emb.weight) - else: - if self.lm_head is None: - raise RuntimeError("lm_head is required when tie_embeddings=False") - logits_proj = self.lm_head(x) - logits = self.logit_softcap * torch.tanh(logits_proj / self.logit_softcap) - return F.cross_entropy(logits.float(), targets, reduction="mean") - - -# ----------------------------- -# TRAINING -# ----------------------------- - -def main() -> None: - global zeropower_via_newtonschulz5 - - code = Path(__file__).read_text(encoding="utf-8") - args = Hyperparameters() - zeropower_via_newtonschulz5 = torch.compile(zeropower_via_newtonschulz5) - - # ----------------------------- - # DISTRIBUTED + CUDA SETUP - # ----------------------------- - - distributed = "RANK" in os.environ and "WORLD_SIZE" in os.environ - rank = int(os.environ.get("RANK", "0")) - world_size = int(os.environ.get("WORLD_SIZE", "1")) - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - if world_size <= 0: - raise ValueError(f"WORLD_SIZE must be positive, got {world_size}") - if 8 % world_size != 0: - raise ValueError(f"WORLD_SIZE={world_size} must divide 8 so grad_accum_steps stays integral") - grad_accum_steps = 8 // world_size - grad_scale = 1.0 / grad_accum_steps - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required") - device = torch.device("cuda", local_rank) - torch.cuda.set_device(device) - if distributed: - dist.init_process_group(backend="nccl", device_id=device) - dist.barrier() - master_process = rank == 0 - - # Fast math knobs - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True - from torch.backends.cuda import enable_cudnn_sdp, enable_flash_sdp, enable_math_sdp, enable_mem_efficient_sdp - - enable_cudnn_sdp(False) - enable_flash_sdp(True) - enable_mem_efficient_sdp(False) - enable_math_sdp(False) - - logfile = None - if master_process: - os.makedirs("logs", exist_ok=True) - logfile = f"logs/{args.run_id}.txt" - print(logfile) - - def log0(msg: str, console: bool = True) -> None: - if not master_process: - return - if console: - print(msg) - if logfile is not None: - with open(logfile, "a", encoding="utf-8") as f: - print(msg, file=f) - - log0(code, console=False) - log0("=" * 100, console=False) - log0(f"Running Python {sys.version}", console=False) - log0(f"Running PyTorch {torch.__version__}", console=False) - log0( - subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False).stdout, - console=False, - ) - log0("=" * 100, console=False) - - # ----------------------------- - # TOKENIZER + VALIDATION METRIC SETUP - # ----------------------------- - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed_all(args.seed) - - if not args.tokenizer_path.endswith(".model"): - raise ValueError(f"Script only setup for SentencePiece .model file: {args.tokenizer_path}") - sp = spm.SentencePieceProcessor(model_file=args.tokenizer_path) - if int(sp.vocab_size()) != args.vocab_size: - raise ValueError( - f"VOCAB_SIZE={args.vocab_size} does not match tokenizer vocab_size={int(sp.vocab_size())}" - ) - dataset_dir = Path(args.data_path).resolve() - actual_train_files = len(list(dataset_dir.glob("fineweb_train_*.bin"))) - val_tokens = load_validation_tokens(args.val_files, args.train_seq_len) - base_bytes_lut, has_leading_space_lut, is_boundary_token_lut = build_sentencepiece_luts( - sp, args.vocab_size, device - ) - log0(f"val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path={args.tokenizer_path}") - log0(f"train_loader:dataset:{dataset_dir.name} train_shards:{actual_train_files}") - log0(f"val_loader:shards pattern={args.val_files} tokens:{val_tokens.numel() - 1}") - - # ----------------------------- - # MODEL + OPTIMIZER SETUP - # ----------------------------- - - base_model = GPT( - vocab_size=args.vocab_size, - num_layers=args.num_layers, - model_dim=args.model_dim, - num_heads=args.num_heads, - num_kv_heads=args.num_kv_heads, - mlp_mult=args.mlp_mult, - tie_embeddings=args.tie_embeddings, - tied_embed_init_std=args.tied_embed_init_std, - logit_softcap=args.logit_softcap, - rope_base=args.rope_base, - qk_gain_init=args.qk_gain_init, - mlp_hidden=args.mlp_hidden, - ).to(device).bfloat16() - for module in base_model.modules(): - if isinstance(module, CastedLinear): - module.float() - restore_low_dim_params_to_fp32(base_model) - compiled_model = torch.compile(base_model, dynamic=False, fullgraph=True) - model: nn.Module = DDP(compiled_model, device_ids=[local_rank], broadcast_buffers=False) if distributed else compiled_model - - # Optimizer split: - # - token embedding (Adam) uses EMBED_LR - # - untied lm_head (Adam) uses HEAD_LR - # - matrix params in transformer blocks use MATRIX_LR via Muon - # - vectors/scalars use SCALAR_LR via Adam - block_named_params = list(base_model.blocks.named_parameters()) - matrix_params = [ - p - for name, p in block_named_params - if p.ndim == 2 and not any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - scalar_params = [ - p - for name, p in block_named_params - if p.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - if base_model.skip_weights.numel() > 0: - scalar_params.append(base_model.skip_weights) - token_lr = args.tied_embed_lr if args.tie_embeddings else args.embed_lr - optimizer_tok = torch.optim.Adam( - [{"params": [base_model.tok_emb.weight], "lr": token_lr, "base_lr": token_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizer_muon = Muon( - matrix_params, - lr=args.matrix_lr, - momentum=args.muon_momentum, - backend_steps=args.muon_backend_steps, - ) - for group in optimizer_muon.param_groups: - group["base_lr"] = args.matrix_lr - optimizer_scalar = torch.optim.Adam( - [{"params": scalar_params, "lr": args.scalar_lr, "base_lr": args.scalar_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers: list[torch.optim.Optimizer] = [optimizer_tok, optimizer_muon, optimizer_scalar] - if base_model.lm_head is not None: - optimizer_head = torch.optim.Adam( - [{"params": [base_model.lm_head.weight], "lr": args.head_lr, "base_lr": args.head_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers.insert(1, optimizer_head) - - n_params = sum(p.numel() for p in base_model.parameters()) - log0(f"model_params:{n_params}") - log0(f"world_size:{world_size} grad_accum_steps:{grad_accum_steps}") - log0("sdp_backends:cudnn=False flash=True mem_efficient=False math=False") - log0(f"attention_mode:gqa num_heads:{args.num_heads} num_kv_heads:{args.num_kv_heads}") - log0( - f"tie_embeddings:{args.tie_embeddings} embed_lr:{token_lr} " - f"head_lr:{args.head_lr if base_model.lm_head is not None else 0.0} " - f"matrix_lr:{args.matrix_lr} scalar_lr:{args.scalar_lr}" - ) - log0( - f"train_batch_tokens:{args.train_batch_tokens} train_seq_len:{args.train_seq_len} " - f"iterations:{args.iterations} warmup_steps:{args.warmup_steps} " - f"max_wallclock_seconds:{args.max_wallclock_seconds:.3f}" - ) - log0(f"seed:{args.seed}") - - # ----------------------------- - # DATA LOADER & MODEL WARMUP - # ----------------------------- - - train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device) - - def zero_grad_all() -> None: - for opt in optimizers: - opt.zero_grad(set_to_none=True) - - max_wallclock_ms = 1000.0 * args.max_wallclock_seconds if args.max_wallclock_seconds > 0 else None - - def lr_mul(step: int, elapsed_ms: float) -> float: - if args.warmdown_iters <= 0: - return 1.0 - if max_wallclock_ms is None: - warmdown_start = max(args.iterations - args.warmdown_iters, 0) - return max((args.iterations - step) / max(args.warmdown_iters, 1), 0.0) if warmdown_start <= step < args.iterations else 1.0 - step_ms = elapsed_ms / max(step, 1) - warmdown_ms = args.warmdown_iters * step_ms - remaining_ms = max(max_wallclock_ms - elapsed_ms, 0.0) - return remaining_ms / max(warmdown_ms, 1e-9) if remaining_ms <= warmdown_ms else 1.0 - - # Warmup primes the compiled forward/backward/optimizer paths, then we restore the - # initial weights/optimizer state so measured training starts from the true init. - if args.warmup_steps > 0: - initial_model_state = {name: tensor.detach().cpu().clone() for name, tensor in base_model.state_dict().items()} - initial_optimizer_states = [copy.deepcopy(opt.state_dict()) for opt in optimizers] - model.train() - for warmup_step in range(args.warmup_steps): - zero_grad_all() - for micro_step in range(grad_accum_steps): - if distributed: - model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 - x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - warmup_loss = model(x, y) - (warmup_loss * grad_scale).backward() - for opt in optimizers: - opt.step() - zero_grad_all() - if args.warmup_steps <= 20 or (warmup_step + 1) % 10 == 0 or warmup_step + 1 == args.warmup_steps: - log0(f"warmup_step:{warmup_step + 1}/{args.warmup_steps}") - base_model.load_state_dict(initial_model_state, strict=True) - for opt, state in zip(optimizers, initial_optimizer_states, strict=True): - opt.load_state_dict(state) - zero_grad_all() - if distributed: - model.require_backward_grad_sync = True - train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device) - - # ----------------------------- - # MAIN TRAINING LOOP - # ----------------------------- - - training_time_ms = 0.0 - stop_after_step: int | None = None - torch.cuda.synchronize() - t0 = time.perf_counter() - - step = 0 - while True: - last_step = step == args.iterations or (stop_after_step is not None and step >= stop_after_step) - - should_validate = last_step or (args.val_loss_every > 0 and step % args.val_loss_every == 0) - if should_validate: - torch.cuda.synchronize() - training_time_ms += 1000.0 * (time.perf_counter() - t0) - val_loss, val_bpb = eval_val( - args, - model, - rank, - world_size, - device, - grad_accum_steps, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - ) - log0( - f"step:{step}/{args.iterations} val_loss:{val_loss:.4f} val_bpb:{val_bpb:.4f} " - f"train_time:{training_time_ms:.0f}ms step_avg:{training_time_ms / max(step, 1):.2f}ms" - ) - torch.cuda.synchronize() - t0 = time.perf_counter() - - if last_step: - if stop_after_step is not None and step < args.iterations: - log0( - f"stopping_early: wallclock_cap train_time:{training_time_ms:.0f}ms " - f"step:{step}/{args.iterations}" - ) - break - - elapsed_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) - scale = lr_mul(step, elapsed_ms) - zero_grad_all() - train_loss = torch.zeros((), device=device) - for micro_step in range(grad_accum_steps): - if distributed: - model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 - x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - loss = model(x, y) - train_loss += loss.detach() - (loss * grad_scale).backward() - train_loss /= grad_accum_steps - - frac = min(step / args.muon_momentum_warmup_steps, 1.0) if args.muon_momentum_warmup_steps > 0 else 1.0 - muon_momentum = (1 - frac) * args.muon_momentum_warmup_start + frac * args.muon_momentum - for group in optimizer_muon.param_groups: - group["momentum"] = muon_momentum - - for opt in optimizers: - for group in opt.param_groups: - group["lr"] = group["base_lr"] * scale - - if args.grad_clip_norm > 0: - torch.nn.utils.clip_grad_norm_(base_model.parameters(), args.grad_clip_norm) - for opt in optimizers: - opt.step() - zero_grad_all() - - step += 1 - approx_training_time_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) - should_log_train = ( - args.train_log_every > 0 - and (step <= 10 or step % args.train_log_every == 0 or stop_after_step is not None) - ) - if should_log_train: - log0( - f"step:{step}/{args.iterations} train_loss:{train_loss.item():.4f} " - f"train_time:{approx_training_time_ms:.0f}ms step_avg:{approx_training_time_ms / step:.2f}ms" - ) - - # Needed to sync whether we've reached the wallclock cap. - reached_cap = max_wallclock_ms is not None and approx_training_time_ms >= max_wallclock_ms - if distributed and max_wallclock_ms is not None: - reached_cap_tensor = torch.tensor(int(reached_cap), device=device) - dist.all_reduce(reached_cap_tensor, op=dist.ReduceOp.MAX) - reached_cap = bool(reached_cap_tensor.item()) - if stop_after_step is None and reached_cap: - stop_after_step = step - - log0( - f"peak memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB " - f"reserved: {torch.cuda.max_memory_reserved() // 1024 // 1024} MiB" - ) - - # ----------------------------- - # SERIALIZATION + ROUNDTRIP VALIDATION - # ----------------------------- - # Save the raw state (useful for debugging/loading in PyTorch directly), then always produce - # the compressed int8+zlib artifact and validate the round-tripped weights. - - if master_process: - torch.save(base_model.state_dict(), "final_model.pt") - model_bytes = os.path.getsize("final_model.pt") - code_bytes = len(code.encode("utf-8")) - log0(f"Serialized model: {model_bytes} bytes") - log0(f"Code size: {code_bytes} bytes") - log0(f"Total submission size: {model_bytes + code_bytes} bytes") - - quant_obj, quant_stats = quantize_state_dict_int8(base_model.state_dict()) - quant_buf = io.BytesIO() - torch.save(quant_obj, quant_buf) - quant_raw = quant_buf.getvalue() - quant_blob = zlib.compress(quant_raw, level=9) - quant_raw_bytes = len(quant_raw) - if master_process: - with open("final_model.int8.ptz", "wb") as f: - f.write(quant_blob) - quant_file_bytes = os.path.getsize("final_model.int8.ptz") - code_bytes = len(code.encode("utf-8")) - ratio = quant_stats["baseline_tensor_bytes"] / max(quant_stats["int8_payload_bytes"], 1) - log0( - f"Serialized model int8+zlib: {quant_file_bytes} bytes " - f"(payload:{quant_stats['int8_payload_bytes']} raw_torch:{quant_raw_bytes} payload_ratio:{ratio:.2f}x)" - ) - log0(f"Total submission size int8+zlib: {quant_file_bytes + code_bytes} bytes") - - if distributed: - dist.barrier() - with open("final_model.int8.ptz", "rb") as f: - quant_blob_disk = f.read() - quant_state = torch.load(io.BytesIO(zlib.decompress(quant_blob_disk)), map_location="cpu") - base_model.load_state_dict(dequantize_state_dict_int8(quant_state), strict=True) - torch.cuda.synchronize() - t_qeval = time.perf_counter() - q_val_loss, q_val_bpb = eval_val( - args, - model, - rank, - world_size, - device, - grad_accum_steps, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - ) - torch.cuda.synchronize() - log0( - f"final_int8_zlib_roundtrip val_loss:{q_val_loss:.4f} val_bpb:{q_val_bpb:.4f} " - f"eval_time:{1000.0 * (time.perf_counter() - t_qeval):.0f}ms" - ) - log0(f"final_int8_zlib_roundtrip_exact val_loss:{q_val_loss:.8f} val_bpb:{q_val_bpb:.8f}") - - if distributed: - dist.destroy_process_group() - - -if __name__ == "__main__": - main() - -==================================================================================================== -Running Python 3.12.3 (main, Nov 6 2025, 13:44:16) [GCC 13.3.0] -Running PyTorch 2.9.1+cu128 -Thu Mar 19 00:32:54 2026 -+-----------------------------------------------------------------------------------------+ -| NVIDIA-SMI 565.57.01 Driver Version: 565.57.01 CUDA Version: 12.7 | -|-----------------------------------------+------------------------+----------------------+ -| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC | -| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. | -| | | MIG M. | -|=========================================+========================+======================| -| 0 NVIDIA H100 80GB HBM3 On | 00000000:18:00.0 Off | 0 | -| N/A 33C P0 147W / 700W | 1518MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 1 NVIDIA H100 80GB HBM3 On | 00000000:2A:00.0 Off | 0 | -| N/A 37C P0 152W / 700W | 1518MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 2 NVIDIA H100 80GB HBM3 On | 00000000:3A:00.0 Off | 0 | -| N/A 37C P0 153W / 700W | 1518MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 3 NVIDIA H100 80GB HBM3 On | 00000000:5D:00.0 Off | 0 | -| N/A 34C P0 151W / 700W | 1518MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 4 NVIDIA H100 80GB HBM3 On | 00000000:9A:00.0 Off | 0 | -| N/A 35C P0 151W / 700W | 1518MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 5 NVIDIA H100 80GB HBM3 On | 00000000:AB:00.0 Off | 0 | -| N/A 38C P0 150W / 700W | 1518MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 6 NVIDIA H100 80GB HBM3 On | 00000000:BA:00.0 Off | 0 | -| N/A 36C P0 150W / 700W | 1518MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 7 NVIDIA H100 80GB HBM3 On | 00000000:DB:00.0 Off | 0 | -| N/A 34C P0 150W / 700W | 1518MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ - -+-----------------------------------------------------------------------------------------+ -| Processes: | -| GPU GI CI PID Type Process name GPU Memory | -| ID ID Usage | -|=========================================================================================| -+-----------------------------------------------------------------------------------------+ - -==================================================================================================== -val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path=./data/tokenizers/fineweb_1024_bpe.model -train_loader:dataset:fineweb10B_sp1024 train_shards:80 -val_loader:shards pattern=./data/datasets/fineweb10B_sp1024/fineweb_val_*.bin tokens:62021632 -model_params:16765000 -world_size:8 grad_accum_steps:1 -sdp_backends:cudnn=False flash=True mem_efficient=False math=False -attention_mode:gqa num_heads:8 num_kv_heads:4 -tie_embeddings:True embed_lr:0.05 head_lr:0.0 matrix_lr:0.06 scalar_lr:0.04 -train_batch_tokens:524288 train_seq_len:1024 iterations:20000 warmup_steps:20 max_wallclock_seconds:600.000 -seed:1337 -warmup_step:1/20 -warmup_step:2/20 -warmup_step:3/20 -warmup_step:4/20 -warmup_step:5/20 -warmup_step:6/20 -warmup_step:7/20 -warmup_step:8/20 -warmup_step:9/20 -warmup_step:10/20 -warmup_step:11/20 -warmup_step:12/20 -warmup_step:13/20 -warmup_step:14/20 -warmup_step:15/20 -warmup_step:16/20 -warmup_step:17/20 -warmup_step:18/20 -warmup_step:19/20 -warmup_step:20/20 -step:1/20000 train_loss:6.9362 train_time:68ms step_avg:67.67ms -step:2/20000 train_loss:16.8730 train_time:108ms step_avg:54.09ms -step:3/20000 train_loss:8.5225 train_time:160ms step_avg:53.29ms -step:4/20000 train_loss:6.6152 train_time:206ms step_avg:51.59ms -step:5/20000 train_loss:6.8377 train_time:261ms step_avg:52.14ms -step:6/20000 train_loss:7.5396 train_time:310ms step_avg:51.68ms -step:7/20000 train_loss:6.4540 train_time:365ms step_avg:52.08ms -step:8/20000 train_loss:6.1992 train_time:416ms step_avg:52.03ms -step:9/20000 train_loss:6.0263 train_time:472ms step_avg:52.42ms -step:10/20000 train_loss:5.9363 train_time:527ms step_avg:52.68ms -step:500/20000 train_loss:2.5003 train_time:21847ms step_avg:43.69ms -step:1000/20000 train_loss:2.3731 train_time:43708ms step_avg:43.71ms -step:1500/20000 train_loss:2.2249 train_time:65435ms step_avg:43.62ms -step:2000/20000 train_loss:2.2537 train_time:87313ms step_avg:43.66ms -step:2500/20000 train_loss:2.1287 train_time:109175ms step_avg:43.67ms -step:3000/20000 train_loss:2.2325 train_time:130926ms step_avg:43.64ms -step:3500/20000 train_loss:2.2540 train_time:152759ms step_avg:43.65ms -step:4000/20000 train_loss:2.1636 train_time:174500ms step_avg:43.62ms -step:4500/20000 train_loss:2.2599 train_time:196331ms step_avg:43.63ms -step:5000/20000 train_loss:2.0318 train_time:218162ms step_avg:43.63ms -step:5500/20000 train_loss:2.1445 train_time:239884ms step_avg:43.62ms -step:6000/20000 train_loss:2.2222 train_time:261726ms step_avg:43.62ms -step:6500/20000 train_loss:2.0817 train_time:283581ms step_avg:43.63ms -step:7000/20000 train_loss:2.2303 train_time:305317ms step_avg:43.62ms -step:7500/20000 train_loss:2.1147 train_time:327164ms step_avg:43.62ms -step:8000/20000 train_loss:2.1190 train_time:348888ms step_avg:43.61ms -step:8500/20000 train_loss:2.1278 train_time:370717ms step_avg:43.61ms -step:9000/20000 train_loss:2.0327 train_time:392536ms step_avg:43.62ms -step:9500/20000 train_loss:2.2631 train_time:414362ms step_avg:43.62ms -step:10000/20000 train_loss:2.1185 train_time:436186ms step_avg:43.62ms -step:10500/20000 train_loss:1.9820 train_time:458034ms step_avg:43.62ms -step:11000/20000 train_loss:2.1087 train_time:479776ms step_avg:43.62ms -step:11500/20000 train_loss:2.0948 train_time:501635ms step_avg:43.62ms -step:12000/20000 train_loss:2.0525 train_time:523376ms step_avg:43.61ms -step:12500/20000 train_loss:2.3395 train_time:545215ms step_avg:43.62ms -step:13000/20000 train_loss:2.1638 train_time:567073ms step_avg:43.62ms -step:13500/20000 train_loss:2.0154 train_time:588838ms step_avg:43.62ms -step:13692/20000 val_loss:2.0586 val_bpb:1.2192 train_time:600008ms step_avg:43.82ms -stopping_early: wallclock_cap train_time:600008ms step:13692/20000 -peak memory allocated: 10111 MiB reserved: 10302 MiB -Serialized model: 66045335 bytes -Code size: 48125 bytes -Total submission size: 66093460 bytes -Serialized model int8+zlib: 15848097 bytes (payload:17405664 raw_torch:17450459 payload_ratio:3.79x) -Total submission size int8+zlib: 15896222 bytes -final_int8_zlib_roundtrip val_loss:2.0595 val_bpb:1.2197 eval_time:1399ms -final_int8_zlib_roundtrip_exact val_loss:2.05945460 val_bpb:1.21972502 diff --git a/records/track_10min_16mb/2026-03-18_FP16Embed_WD3600/train_gpt.py b/records/track_10min_16mb/2026-03-18_FP16Embed_WD3600/train_gpt.py deleted file mode 100644 index fd6c415533..0000000000 --- a/records/track_10min_16mb/2026-03-18_FP16Embed_WD3600/train_gpt.py +++ /dev/null @@ -1,1133 +0,0 @@ -""" -The `train_gpt.py` and `train_gpt_mlx.py` scripts are intended as good launching-off points for new participants, not SOTA configs. We'll accept PRs that tune, improve, or simplify these scripts without significantly increasing complexity, but competitive submissions should stay in the `/records` folder. - -Hard stop: `train_gpt.py` and `train_gpt_mlx.py` must never be longer than 1500 lines. -""" - -from __future__ import annotations - -import copy -import glob -import io -import math -import os -import random -import subprocess -import sys -import time -import uuid -import zlib -from pathlib import Path - -import numpy as np -import sentencepiece as spm -import torch -import torch.distributed as dist -import torch.nn.functional as F -from torch import Tensor, nn -from torch.nn.parallel import DistributedDataParallel as DDP - -# ----------------------------- -# HYPERPARAMETERS -# ----------------------------- -# Default Simple Baseline run: -# - 9 transformer blocks at width 512 -# - 8 attention heads with 4 KV heads (GQA) and 2x MLP expansion -# - vocab size 1024, sequence length 1024, tied embeddings -# - 524,288 train tokens per step for 20,000 iterations with a ~10 minute cap - -class Hyperparameters: - # Data paths are shard globs produced by the existing preprocessing pipeline. - data_path = os.environ.get("DATA_PATH", "./data/datasets/fineweb10B_sp1024") - train_files = os.path.join(data_path, "fineweb_train_*.bin") - val_files = os.path.join(data_path, "fineweb_val_*.bin") - tokenizer_path = os.environ.get("TOKENIZER_PATH", "./data/tokenizers/fineweb_1024_bpe.model") - run_id = os.environ.get("RUN_ID", str(uuid.uuid4())) - seed = int(os.environ.get("SEED", 1337)) - - # Validation cadence and batch size. Validation always uses the full fineweb_val split. - val_batch_size = int(os.environ.get("VAL_BATCH_SIZE", 524_288)) - val_loss_every = int(os.environ.get("VAL_LOSS_EVERY", 1000)) - train_log_every = int(os.environ.get("TRAIN_LOG_EVERY", 200)) - - # Training length. - iterations = int(os.environ.get("ITERATIONS", 20000)) - warmdown_iters = int(os.environ.get("WARMDOWN_ITERS", 1200)) - warmup_steps = int(os.environ.get("WARMUP_STEPS", 20)) - train_batch_tokens = int(os.environ.get("TRAIN_BATCH_TOKENS", 524_288)) - train_seq_len = int(os.environ.get("TRAIN_SEQ_LEN", 1024)) - max_wallclock_seconds = float(os.environ.get("MAX_WALLCLOCK_SECONDS", 600.0)) - qk_gain_init = float(os.environ.get("QK_GAIN_INIT", 1.5)) - - # Model shape. - vocab_size = int(os.environ.get("VOCAB_SIZE", 1024)) - num_layers = int(os.environ.get("NUM_LAYERS", 9)) - num_kv_heads = int(os.environ.get("NUM_KV_HEADS", 4)) - model_dim = int(os.environ.get("MODEL_DIM", 512)) - num_heads = int(os.environ.get("NUM_HEADS", 8)) - mlp_mult = int(os.environ.get("MLP_MULT", 2)) - mlp_hidden = int(os.environ.get("MLP_HIDDEN", 0)) # override mlp_mult if > 0 - tie_embeddings = bool(int(os.environ.get("TIE_EMBEDDINGS", "1"))) - rope_base = float(os.environ.get("ROPE_BASE", 10000.0)) - logit_softcap = float(os.environ.get("LOGIT_SOFTCAP", 30.0)) - - # Optimizer hyperparameters. - embed_lr = float(os.environ.get("EMBED_LR", 0.6)) - head_lr = float(os.environ.get("HEAD_LR", 0.008)) - tied_embed_lr = float(os.environ.get("TIED_EMBED_LR", 0.05)) - tied_embed_init_std = float(os.environ.get("TIED_EMBED_INIT_STD", 0.005)) - matrix_lr = float(os.environ.get("MATRIX_LR", 0.04)) - scalar_lr = float(os.environ.get("SCALAR_LR", 0.04)) - muon_momentum = float(os.environ.get("MUON_MOMENTUM", 0.95)) - muon_backend_steps = int(os.environ.get("MUON_BACKEND_STEPS", 5)) - muon_momentum_warmup_start = float(os.environ.get("MUON_MOMENTUM_WARMUP_START", 0.85)) - muon_momentum_warmup_steps = int(os.environ.get("MUON_MOMENTUM_WARMUP_STEPS", 500)) - beta1 = float(os.environ.get("BETA1", 0.9)) - beta2 = float(os.environ.get("BETA2", 0.95)) - adam_eps = float(os.environ.get("ADAM_EPS", 1e-8)) - grad_clip_norm = float(os.environ.get("GRAD_CLIP_NORM", 0.0)) - -# ----------------------------- -# MUON OPTIMIZER -# ----------------------------- -# -# As borrowed from modded-nanogpt -# Background on Muon: https://kellerjordan.github.io/posts/muon/ - -def zeropower_via_newtonschulz5(G: Tensor, steps: int = 10, eps: float = 1e-7) -> Tensor: - # Orthogonalize a 2D update matrix with a fast Newton-Schulz iteration. - # Muon uses this to normalize matrix-shaped gradients before applying them. - a, b, c = (3.4445, -4.7750, 2.0315) - X = G.bfloat16() - X /= X.norm() + eps - transposed = G.size(0) > G.size(1) - if transposed: - X = X.T - for _ in range(steps): - A = X @ X.T - B = b * A + c * A @ A - X = a * X + B @ X - return X.T if transposed else X - - -class Muon(torch.optim.Optimizer): - def __init__(self, params, lr: float, momentum: float, backend_steps: int, nesterov: bool = True): - super().__init__( - params, - dict(lr=lr, momentum=momentum, backend_steps=backend_steps, nesterov=nesterov), - ) - - @torch.no_grad() - def step(self, closure=None): - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - distributed = dist.is_available() and dist.is_initialized() - world_size = dist.get_world_size() if distributed else 1 - rank = dist.get_rank() if distributed else 0 - - for group in self.param_groups: - params = group["params"] - if not params: - continue - lr = group["lr"] - momentum = group["momentum"] - backend_steps = group["backend_steps"] - nesterov = group["nesterov"] - - total_params = sum(int(p.numel()) for p in params) - updates_flat = torch.zeros(total_params, device=params[0].device, dtype=torch.bfloat16) - - curr = 0 - for i, p in enumerate(params): - if i % world_size == rank and p.grad is not None: - g = p.grad - state = self.state[p] - if "momentum_buffer" not in state: - state["momentum_buffer"] = torch.zeros_like(g) - buf = state["momentum_buffer"] - buf.mul_(momentum).add_(g) - if nesterov: - g = g.add(buf, alpha=momentum) - g = zeropower_via_newtonschulz5(g, steps=backend_steps) - # Scale correction from Muon reference implementations. - g *= max(1, g.size(0) / g.size(1)) ** 0.5 - updates_flat[curr : curr + p.numel()] = g.reshape(-1) - curr += p.numel() - - if distributed: - dist.all_reduce(updates_flat, op=dist.ReduceOp.SUM) - - curr = 0 - for p in params: - g = updates_flat[curr : curr + p.numel()].view_as(p).to(dtype=p.dtype) - p.add_(g, alpha=-lr) - curr += p.numel() - - return loss - - -# ----------------------------- -# TOKENIZER-AGNOSTIC EVALUATION SETUP -# ----------------------------- -# -# It's common for small models have a large fraction of their parameters be embeddings, since the 2 * d_model * d_vocab vectors can be gigantic. -# Instead of locking the tokenizer, we let you bring your own and calculate our validation metrics on the average compression of the validation set. -# We calculate BPB (bits-per-byte) instead of validation loss, so we need methods to count the number of bits per token in the tokenizer. -# Note: Submissions that edit the tokenizer will be examined more carefully, since screwing this up might unjustly improve your score. - -def build_sentencepiece_luts( - sp: spm.SentencePieceProcessor, vocab_size: int, device: torch.device -) -> tuple[Tensor, Tensor, Tensor]: - sp_vocab_size = int(sp.vocab_size()) - table_size = max(sp_vocab_size, vocab_size) - base_bytes_np = np.zeros((table_size,), dtype=np.int16) - has_leading_space_np = np.zeros((table_size,), dtype=np.bool_) - is_boundary_token_np = np.ones((table_size,), dtype=np.bool_) - for token_id in range(sp_vocab_size): - if sp.is_control(token_id) or sp.is_unknown(token_id) or sp.is_unused(token_id): - continue - is_boundary_token_np[token_id] = False - if sp.is_byte(token_id): - base_bytes_np[token_id] = 1 - continue - piece = sp.id_to_piece(token_id) - if piece.startswith("▁"): - has_leading_space_np[token_id] = True - piece = piece[1:] - base_bytes_np[token_id] = len(piece.encode("utf-8")) - return ( - torch.tensor(base_bytes_np, dtype=torch.int16, device=device), - torch.tensor(has_leading_space_np, dtype=torch.bool, device=device), - torch.tensor(is_boundary_token_np, dtype=torch.bool, device=device), - ) - - -def load_validation_tokens(pattern: str, seq_len: int) -> Tensor: - files = [Path(p) for p in sorted(glob.glob(pattern))] - if not files: - raise FileNotFoundError(f"No files found for pattern: {pattern}") - # The export pipeline writes the fixed first-50k-doc validation set to fineweb_val_*. - tokens = torch.cat([load_data_shard(file) for file in files]).contiguous() - usable = ((tokens.numel() - 1) // seq_len) * seq_len - if usable <= 0: - raise ValueError(f"Validation split is too short for TRAIN_SEQ_LEN={seq_len}") - return tokens[: usable + 1] - - -def eval_val( - args: Hyperparameters, - model: nn.Module, - rank: int, - world_size: int, - device: torch.device, - grad_accum_steps: int, - val_tokens: Tensor, - base_bytes_lut: Tensor, - has_leading_space_lut: Tensor, - is_boundary_token_lut: Tensor, -) -> tuple[float, float]: - # Validation computes two metrics: - # - val_loss: token cross-entropy (natural log) - # - val_bpb: tokenizer-agnostic compression metric used by the challenge - local_batch_tokens = args.val_batch_size // (world_size * grad_accum_steps) - if local_batch_tokens < args.train_seq_len: - raise ValueError( - "VAL_BATCH_SIZE must provide at least one sequence per rank; " - f"got VAL_BATCH_SIZE={args.val_batch_size}, WORLD_SIZE={world_size}, " - f"GRAD_ACCUM_STEPS={grad_accum_steps}, TRAIN_SEQ_LEN={args.train_seq_len}" - ) - local_batch_seqs = local_batch_tokens // args.train_seq_len - total_seqs = (val_tokens.numel() - 1) // args.train_seq_len - seq_start = (total_seqs * rank) // world_size - seq_end = (total_seqs * (rank + 1)) // world_size - val_loss_sum = torch.zeros((), device=device, dtype=torch.float64) - val_token_count = torch.zeros((), device=device, dtype=torch.float64) - val_byte_count = torch.zeros((), device=device, dtype=torch.float64) - - model.eval() - with torch.inference_mode(): - for batch_seq_start in range(seq_start, seq_end, local_batch_seqs): - batch_seq_end = min(batch_seq_start + local_batch_seqs, seq_end) - raw_start = batch_seq_start * args.train_seq_len - raw_end = batch_seq_end * args.train_seq_len + 1 - local = val_tokens[raw_start:raw_end].to(device=device, dtype=torch.int64, non_blocking=True) - x = local[:-1].reshape(-1, args.train_seq_len) - y = local[1:].reshape(-1, args.train_seq_len) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - batch_loss = model(x, y).detach() - batch_token_count = float(y.numel()) - val_loss_sum += batch_loss.to(torch.float64) * batch_token_count - val_token_count += batch_token_count - prev_ids = x.reshape(-1) - tgt_ids = y.reshape(-1) - token_bytes = base_bytes_lut[tgt_ids].to(dtype=torch.int16) - token_bytes += (has_leading_space_lut[tgt_ids] & ~is_boundary_token_lut[prev_ids]).to(dtype=torch.int16) - val_byte_count += token_bytes.to(torch.float64).sum() - - if dist.is_available() and dist.is_initialized(): - dist.all_reduce(val_loss_sum, op=dist.ReduceOp.SUM) - dist.all_reduce(val_token_count, op=dist.ReduceOp.SUM) - dist.all_reduce(val_byte_count, op=dist.ReduceOp.SUM) - - val_loss = val_loss_sum / val_token_count - bits_per_token = val_loss.item() / math.log(2.0) - tokens_per_byte = val_token_count.item() / val_byte_count.item() - model.train() - return float(val_loss.item()), float(bits_per_token * tokens_per_byte) - -# ----------------------------- -# POST-TRAINING QUANTIZATION -# ----------------------------- -# -# It's silly to export our model, which is trained in bf16 and fp32, at that same precision. -# Instead, we get approximately the same model (with a small hit) by quantizing the model to int8 & zlib compressing. -# We can then decompress the model and run in higher precision for evaluation, after closing in under the size limit. - -CONTROL_TENSOR_NAME_PATTERNS = tuple( - pattern - for pattern in os.environ.get( - "CONTROL_TENSOR_NAME_PATTERNS", - "attn_scale,attn_scales,mlp_scale,mlp_scales,resid_mix,resid_mixes,q_gain,skip_weight,skip_weights", - ).split(",") - if pattern -) -INT8_KEEP_FLOAT_FP32_NAME_PATTERNS = tuple( - pattern - for pattern in os.environ.get( - "INT8_KEEP_FLOAT_FP32_NAME_PATTERNS", - ",".join(CONTROL_TENSOR_NAME_PATTERNS), - ).split(",") - if pattern -) -INT8_KEEP_FLOAT_MAX_NUMEL = 65_536 -INT8_KEEP_FLOAT_STORE_DTYPE = torch.float16 -INT8_PER_ROW_SCALE_DTYPE = torch.float16 -INT8_CLIP_PERCENTILE = 99.99984 -INT8_CLIP_Q = INT8_CLIP_PERCENTILE / 100.0 - -def tensor_nbytes(t: Tensor) -> int: - return int(t.numel()) * int(t.element_size()) - -def keep_float_tensor(name: str, t: Tensor, passthrough_orig_dtypes: dict[str, str]) -> Tensor: - if any(pattern in name for pattern in INT8_KEEP_FLOAT_FP32_NAME_PATTERNS): - return t.float().contiguous() - if t.dtype in {torch.float32, torch.bfloat16}: - passthrough_orig_dtypes[name] = str(t.dtype).removeprefix("torch.") - return t.to(dtype=INT8_KEEP_FLOAT_STORE_DTYPE).contiguous() - return t - -def quantize_float_tensor(t: Tensor) -> tuple[Tensor, Tensor]: - t32 = t.float() - if t32.ndim == 2: - # Matrices get one scale per row, which usually tracks output-channel - # ranges much better than a single tensor-wide scale. - clip_abs = ( - torch.quantile(t32.abs(), INT8_CLIP_Q, dim=1) - if t32.numel() - else torch.empty((t32.shape[0],), dtype=torch.float32) - ) - clipped = torch.maximum(torch.minimum(t32, clip_abs[:, None]), -clip_abs[:, None]) - scale = (clip_abs / 127.0).clamp_min(1.0 / 127.0) - q = torch.clamp(torch.round(clipped / scale[:, None]), -127, 127).to(torch.int8).contiguous() - return q, scale.to(dtype=INT8_PER_ROW_SCALE_DTYPE).contiguous() - - # Vectors / scalars use a simpler per-tensor scale. - clip_abs = float(torch.quantile(t32.abs().flatten(), INT8_CLIP_Q).item()) if t32.numel() else 0.0 - scale = torch.tensor(clip_abs / 127.0 if clip_abs > 0 else 1.0, dtype=torch.float32) - q = torch.clamp(torch.round(torch.clamp(t32, -clip_abs, clip_abs) / scale), -127, 127).to(torch.int8).contiguous() - return q, scale - -def quantize_state_dict_int8(state_dict: dict[str, Tensor]): - # Single supported clean-script export format: - # - per-row int8 for 2D float tensors - # - per-tensor int8 for other float tensors - # - exact passthrough for non-floats - # - passthrough for small float tensors, stored as fp16 to save bytes - quantized: dict[str, Tensor] = {} - scales: dict[str, Tensor] = {} - dtypes: dict[str, str] = {} - passthrough: dict[str, Tensor] = {} - passthrough_orig_dtypes: dict[str, str] = {} - qmeta: dict[str, dict[str, object]] = {} - stats = dict.fromkeys( - ("param_count", "num_tensors", "num_float_tensors", "num_nonfloat_tensors", "baseline_tensor_bytes", "int8_payload_bytes"), - 0, - ) - - for name, tensor in state_dict.items(): - t = tensor.detach().to("cpu").contiguous() - stats["param_count"] += int(t.numel()) - stats["num_tensors"] += 1 - stats["baseline_tensor_bytes"] += tensor_nbytes(t) - - if not t.is_floating_point(): - stats["num_nonfloat_tensors"] += 1 - passthrough[name] = t - stats["int8_payload_bytes"] += tensor_nbytes(t) - continue - - # Small float tensors are cheap enough to keep directly. We still downcast - # fp32/bf16 passthrough tensors to fp16 so metadata does not dominate size. - # Also keep the tied embedding in fp16 to minimize quantization degradation - # on the output head, which is the most sensitive tensor for BPB. - if t.numel() <= INT8_KEEP_FLOAT_MAX_NUMEL or name == "tok_emb.weight": - kept = keep_float_tensor(name, t, passthrough_orig_dtypes) - passthrough[name] = kept - stats["int8_payload_bytes"] += tensor_nbytes(kept) - continue - - stats["num_float_tensors"] += 1 - q, s = quantize_float_tensor(t) - if s.ndim > 0: - qmeta[name] = {"scheme": "per_row", "axis": 0} - quantized[name] = q - scales[name] = s - dtypes[name] = str(t.dtype).removeprefix("torch.") - stats["int8_payload_bytes"] += tensor_nbytes(q) + tensor_nbytes(s) - - obj: dict[str, object] = { - "__quant_format__": "int8_clean_per_row_v1", - "quantized": quantized, - "scales": scales, - "dtypes": dtypes, - "passthrough": passthrough, - } - if qmeta: - obj["qmeta"] = qmeta - if passthrough_orig_dtypes: - obj["passthrough_orig_dtypes"] = passthrough_orig_dtypes - return obj, stats - -def dequantize_state_dict_int8(obj: dict[str, object]) -> dict[str, Tensor]: - out: dict[str, Tensor] = {} - qmeta = obj.get("qmeta", {}) - passthrough_orig_dtypes = obj.get("passthrough_orig_dtypes", {}) - for name, q in obj["quantized"].items(): - dtype = getattr(torch, obj["dtypes"][name]) - s = obj["scales"][name] - if qmeta.get(name, {}).get("scheme") == "per_row" or s.ndim > 0: - s = s.to(dtype=torch.float32) - # Broadcast the saved row scale back across trailing dimensions. - out[name] = (q.float() * s.view(q.shape[0], *([1] * (q.ndim - 1)))).to(dtype=dtype).contiguous() - else: - scale = float(s.item()) - out[name] = (q.float() * scale).to(dtype=dtype).contiguous() - for name, t in obj["passthrough"].items(): - # Restore small tensors, undoing the temporary fp16 storage cast if needed. - out_t = t.detach().to("cpu").contiguous() - orig_dtype = passthrough_orig_dtypes.get(name) - if isinstance(orig_dtype, str): - out_t = out_t.to(dtype=getattr(torch, orig_dtype)).contiguous() - out[name] = out_t - return out - - -# ----------------------------- -# DATA LOADING -# ----------------------------- - -def load_data_shard(file: Path) -> Tensor: - header_bytes = 256 * np.dtype(" None: - self.file_idx = (self.file_idx + 1) % len(self.files) - self.tokens = load_data_shard(self.files[self.file_idx]) - self.pos = 0 - - def take(self, n: int) -> Tensor: - chunks: list[Tensor] = [] - remaining = n - while remaining > 0: - avail = self.tokens.numel() - self.pos - if avail <= 0: - self._advance_file() - continue - k = min(remaining, avail) - chunks.append(self.tokens[self.pos : self.pos + k]) - self.pos += k - remaining -= k - return chunks[0] if len(chunks) == 1 else torch.cat(chunks) - - -class DistributedTokenLoader: - # Each call consumes a contiguous chunk from the shared token stream, then slices out - # one disjoint span per rank. The extra "+1" token lets us build (x, y) by shifting. - def __init__(self, pattern: str, rank: int, world_size: int, device: torch.device): - self.rank = rank - self.world_size = world_size - self.device = device - self.stream = TokenStream(pattern) - - def next_batch(self, global_tokens: int, seq_len: int, grad_accum_steps: int) -> tuple[Tensor, Tensor]: - local_tokens = global_tokens // (self.world_size * grad_accum_steps) - per_rank_span = local_tokens + 1 - chunk = self.stream.take(per_rank_span * self.world_size) - start = self.rank * per_rank_span - local = chunk[start : start + per_rank_span].to(dtype=torch.int64) - x = local[:-1].reshape(-1, seq_len) - y = local[1:].reshape(-1, seq_len) - return x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True) - -# ----------------------------- -# TRANSFORMER MODULES -# ----------------------------- - -class RMSNorm(nn.Module): - def __init__(self, eps: float | None = None): - super().__init__() - self.eps = eps - - def forward(self, x: Tensor) -> Tensor: - return F.rms_norm(x, (x.size(-1),), eps=self.eps) - - -class CastedLinear(nn.Linear): - # Keep weights in fp32 for optimizer/state quality, cast at matmul time for bf16 compute. - def forward(self, x: Tensor) -> Tensor: - bias = self.bias.to(x.dtype) if self.bias is not None else None - return F.linear(x, self.weight.to(x.dtype), bias) - - -def restore_low_dim_params_to_fp32(module: nn.Module) -> None: - # Keep small/control parameters in fp32 even when the model body runs in bf16. - with torch.no_grad(): - for name, param in module.named_parameters(): - if (param.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS)) and param.dtype != torch.float32: - param.data = param.data.float() - - -class Rotary(nn.Module): - # Caches cos/sin tables per sequence length on the current device. - def __init__(self, dim: int, base: float = 10000.0): - super().__init__() - inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self._seq_len_cached = 0 - self._cos_cached: Tensor | None = None - self._sin_cached: Tensor | None = None - - def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> tuple[Tensor, Tensor]: - if ( - self._cos_cached is None - or self._sin_cached is None - or self._seq_len_cached != seq_len - or self._cos_cached.device != device - ): - t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) - freqs = torch.outer(t, self.inv_freq.to(device)) - self._cos_cached = freqs.cos()[None, None, :, :] - self._sin_cached = freqs.sin()[None, None, :, :] - self._seq_len_cached = seq_len - return self._cos_cached.to(dtype=dtype), self._sin_cached.to(dtype=dtype) - - -def apply_rotary_emb(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: - half = x.size(-1) // 2 - x1, x2 = x[..., :half], x[..., half:] - return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1) - - -class CausalSelfAttention(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - num_kv_heads: int, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - if dim % num_heads != 0: - raise ValueError("model_dim must be divisible by num_heads") - if num_heads % num_kv_heads != 0: - raise ValueError("num_heads must be divisible by num_kv_heads") - self.num_heads = num_heads - self.num_kv_heads = num_kv_heads - self.head_dim = dim // num_heads - if self.head_dim % 2 != 0: - raise ValueError("head_dim must be even for RoPE") - kv_dim = self.num_kv_heads * self.head_dim - self.c_q = CastedLinear(dim, dim, bias=False) - self.c_k = CastedLinear(dim, kv_dim, bias=False) - self.c_v = CastedLinear(dim, kv_dim, bias=False) - self.proj = CastedLinear(dim, dim, bias=False) - self.proj._zero_init = True - self.q_gain = nn.Parameter(torch.full((num_heads,), qk_gain_init, dtype=torch.float32)) - self.rotary = Rotary(self.head_dim, base=rope_base) - - def forward(self, x: Tensor) -> Tensor: - bsz, seqlen, dim = x.shape - q = self.c_q(x).reshape(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2) - k = self.c_k(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - v = self.c_v(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - q = F.rms_norm(q, (q.size(-1),)) - k = F.rms_norm(k, (k.size(-1),)) - cos, sin = self.rotary(seqlen, x.device, q.dtype) - q = apply_rotary_emb(q, cos, sin) - k = apply_rotary_emb(k, cos, sin) - q = q * self.q_gain.to(dtype=q.dtype)[None, :, None, None] - y = F.scaled_dot_product_attention( - q, - k, - v, - attn_mask=None, - is_causal=True, - enable_gqa=(self.num_kv_heads != self.num_heads), - ) - y = y.transpose(1, 2).contiguous().reshape(bsz, seqlen, dim) - return self.proj(y) - - -class MLP(nn.Module): - # relu^2 MLP from the original modded-nanogpt setup - def __init__(self, dim: int, mlp_mult: int, mlp_hidden: int = 0): - super().__init__() - hidden = mlp_hidden if mlp_hidden > 0 else mlp_mult * dim - self.fc = CastedLinear(dim, hidden, bias=False) - self.proj = CastedLinear(hidden, dim, bias=False) - self.proj._zero_init = True - - def forward(self, x: Tensor) -> Tensor: - x = torch.relu(self.fc(x)) - return self.proj(x.square()) - - -class Block(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - num_kv_heads: int, - mlp_mult: int, - rope_base: float, - qk_gain_init: float, - mlp_hidden: int = 0, - ): - super().__init__() - self.attn_norm = RMSNorm() - self.mlp_norm = RMSNorm() - self.attn = CausalSelfAttention(dim, num_heads, num_kv_heads, rope_base, qk_gain_init) - self.mlp = MLP(dim, mlp_mult, mlp_hidden=mlp_hidden) - self.attn_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) - self.mlp_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) - self.resid_mix = nn.Parameter(torch.stack((torch.ones(dim), torch.zeros(dim))).float()) - - def forward(self, x: Tensor, x0: Tensor) -> Tensor: - mix = self.resid_mix.to(dtype=x.dtype) - x = mix[0][None, None, :] * x + mix[1][None, None, :] * x0 - attn_out = self.attn(self.attn_norm(x)) - x = x + self.attn_scale.to(dtype=x.dtype)[None, None, :] * attn_out - x = x + self.mlp_scale.to(dtype=x.dtype)[None, None, :] * self.mlp(self.mlp_norm(x)) - return x - - -class GPT(nn.Module): - def __init__( - self, - vocab_size: int, - num_layers: int, - model_dim: int, - num_heads: int, - num_kv_heads: int, - mlp_mult: int, - tie_embeddings: bool, - tied_embed_init_std: float, - logit_softcap: float, - rope_base: float, - qk_gain_init: float, - mlp_hidden: int = 0, - ): - super().__init__() - if logit_softcap <= 0.0: - raise ValueError(f"logit_softcap must be positive, got {logit_softcap}") - self.tie_embeddings = tie_embeddings - self.tied_embed_init_std = tied_embed_init_std - self.logit_softcap = logit_softcap - self.tok_emb = nn.Embedding(vocab_size, model_dim) - self.num_encoder_layers = num_layers // 2 - self.num_decoder_layers = num_layers - self.num_encoder_layers - self.num_skip_weights = min(self.num_encoder_layers, self.num_decoder_layers) - self.skip_weights = nn.Parameter(torch.ones(self.num_skip_weights, model_dim, dtype=torch.float32)) - self.blocks = nn.ModuleList( - [ - Block( - model_dim, - num_heads, - num_kv_heads, - mlp_mult, - rope_base, - qk_gain_init, - mlp_hidden=mlp_hidden, - ) - for i in range(num_layers) - ] - ) - self.final_norm = RMSNorm() - self.lm_head = None if tie_embeddings else CastedLinear(model_dim, vocab_size, bias=False) - if self.lm_head is not None: - self.lm_head._zero_init = True - self._init_weights() - - def _init_weights(self) -> None: - if self.tie_embeddings: - nn.init.normal_(self.tok_emb.weight, mean=0.0, std=self.tied_embed_init_std) - for module in self.modules(): - if isinstance(module, nn.Linear) and getattr(module, "_zero_init", False): - nn.init.zeros_(module.weight) - - def forward(self, input_ids: Tensor, target_ids: Tensor) -> Tensor: - x = self.tok_emb(input_ids) - x = F.rms_norm(x, (x.size(-1),)) - x0 = x - skips: list[Tensor] = [] - - # First half stores skips; second half reuses them in reverse order. - for i in range(self.num_encoder_layers): - x = self.blocks[i](x, x0) - skips.append(x) - for i in range(self.num_decoder_layers): - if skips: - x = x + self.skip_weights[i].to(dtype=x.dtype)[None, None, :] * skips.pop() - x = self.blocks[self.num_encoder_layers + i](x, x0) - - x = self.final_norm(x).reshape(-1, x.size(-1)) - targets = target_ids.reshape(-1) - if self.tie_embeddings: - logits_proj = F.linear(x, self.tok_emb.weight) - else: - if self.lm_head is None: - raise RuntimeError("lm_head is required when tie_embeddings=False") - logits_proj = self.lm_head(x) - logits = self.logit_softcap * torch.tanh(logits_proj / self.logit_softcap) - return F.cross_entropy(logits.float(), targets, reduction="mean") - - -# ----------------------------- -# TRAINING -# ----------------------------- - -def main() -> None: - global zeropower_via_newtonschulz5 - - code = Path(__file__).read_text(encoding="utf-8") - args = Hyperparameters() - zeropower_via_newtonschulz5 = torch.compile(zeropower_via_newtonschulz5) - - # ----------------------------- - # DISTRIBUTED + CUDA SETUP - # ----------------------------- - - distributed = "RANK" in os.environ and "WORLD_SIZE" in os.environ - rank = int(os.environ.get("RANK", "0")) - world_size = int(os.environ.get("WORLD_SIZE", "1")) - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - if world_size <= 0: - raise ValueError(f"WORLD_SIZE must be positive, got {world_size}") - if 8 % world_size != 0: - raise ValueError(f"WORLD_SIZE={world_size} must divide 8 so grad_accum_steps stays integral") - grad_accum_steps = 8 // world_size - grad_scale = 1.0 / grad_accum_steps - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required") - device = torch.device("cuda", local_rank) - torch.cuda.set_device(device) - if distributed: - dist.init_process_group(backend="nccl", device_id=device) - dist.barrier() - master_process = rank == 0 - - # Fast math knobs - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True - from torch.backends.cuda import enable_cudnn_sdp, enable_flash_sdp, enable_math_sdp, enable_mem_efficient_sdp - - enable_cudnn_sdp(False) - enable_flash_sdp(True) - enable_mem_efficient_sdp(False) - enable_math_sdp(False) - - logfile = None - if master_process: - os.makedirs("logs", exist_ok=True) - logfile = f"logs/{args.run_id}.txt" - print(logfile) - - def log0(msg: str, console: bool = True) -> None: - if not master_process: - return - if console: - print(msg) - if logfile is not None: - with open(logfile, "a", encoding="utf-8") as f: - print(msg, file=f) - - log0(code, console=False) - log0("=" * 100, console=False) - log0(f"Running Python {sys.version}", console=False) - log0(f"Running PyTorch {torch.__version__}", console=False) - log0( - subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False).stdout, - console=False, - ) - log0("=" * 100, console=False) - - # ----------------------------- - # TOKENIZER + VALIDATION METRIC SETUP - # ----------------------------- - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed_all(args.seed) - - if not args.tokenizer_path.endswith(".model"): - raise ValueError(f"Script only setup for SentencePiece .model file: {args.tokenizer_path}") - sp = spm.SentencePieceProcessor(model_file=args.tokenizer_path) - if int(sp.vocab_size()) != args.vocab_size: - raise ValueError( - f"VOCAB_SIZE={args.vocab_size} does not match tokenizer vocab_size={int(sp.vocab_size())}" - ) - dataset_dir = Path(args.data_path).resolve() - actual_train_files = len(list(dataset_dir.glob("fineweb_train_*.bin"))) - val_tokens = load_validation_tokens(args.val_files, args.train_seq_len) - base_bytes_lut, has_leading_space_lut, is_boundary_token_lut = build_sentencepiece_luts( - sp, args.vocab_size, device - ) - log0(f"val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path={args.tokenizer_path}") - log0(f"train_loader:dataset:{dataset_dir.name} train_shards:{actual_train_files}") - log0(f"val_loader:shards pattern={args.val_files} tokens:{val_tokens.numel() - 1}") - - # ----------------------------- - # MODEL + OPTIMIZER SETUP - # ----------------------------- - - base_model = GPT( - vocab_size=args.vocab_size, - num_layers=args.num_layers, - model_dim=args.model_dim, - num_heads=args.num_heads, - num_kv_heads=args.num_kv_heads, - mlp_mult=args.mlp_mult, - tie_embeddings=args.tie_embeddings, - tied_embed_init_std=args.tied_embed_init_std, - logit_softcap=args.logit_softcap, - rope_base=args.rope_base, - qk_gain_init=args.qk_gain_init, - mlp_hidden=args.mlp_hidden, - ).to(device).bfloat16() - for module in base_model.modules(): - if isinstance(module, CastedLinear): - module.float() - restore_low_dim_params_to_fp32(base_model) - compiled_model = torch.compile(base_model, dynamic=False, fullgraph=True) - model: nn.Module = DDP(compiled_model, device_ids=[local_rank], broadcast_buffers=False) if distributed else compiled_model - - # Optimizer split: - # - token embedding (Adam) uses EMBED_LR - # - untied lm_head (Adam) uses HEAD_LR - # - matrix params in transformer blocks use MATRIX_LR via Muon - # - vectors/scalars use SCALAR_LR via Adam - block_named_params = list(base_model.blocks.named_parameters()) - matrix_params = [ - p - for name, p in block_named_params - if p.ndim == 2 and not any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - scalar_params = [ - p - for name, p in block_named_params - if p.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - if base_model.skip_weights.numel() > 0: - scalar_params.append(base_model.skip_weights) - token_lr = args.tied_embed_lr if args.tie_embeddings else args.embed_lr - optimizer_tok = torch.optim.Adam( - [{"params": [base_model.tok_emb.weight], "lr": token_lr, "base_lr": token_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizer_muon = Muon( - matrix_params, - lr=args.matrix_lr, - momentum=args.muon_momentum, - backend_steps=args.muon_backend_steps, - ) - for group in optimizer_muon.param_groups: - group["base_lr"] = args.matrix_lr - optimizer_scalar = torch.optim.Adam( - [{"params": scalar_params, "lr": args.scalar_lr, "base_lr": args.scalar_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers: list[torch.optim.Optimizer] = [optimizer_tok, optimizer_muon, optimizer_scalar] - if base_model.lm_head is not None: - optimizer_head = torch.optim.Adam( - [{"params": [base_model.lm_head.weight], "lr": args.head_lr, "base_lr": args.head_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers.insert(1, optimizer_head) - - n_params = sum(p.numel() for p in base_model.parameters()) - log0(f"model_params:{n_params}") - log0(f"world_size:{world_size} grad_accum_steps:{grad_accum_steps}") - log0("sdp_backends:cudnn=False flash=True mem_efficient=False math=False") - log0(f"attention_mode:gqa num_heads:{args.num_heads} num_kv_heads:{args.num_kv_heads}") - log0( - f"tie_embeddings:{args.tie_embeddings} embed_lr:{token_lr} " - f"head_lr:{args.head_lr if base_model.lm_head is not None else 0.0} " - f"matrix_lr:{args.matrix_lr} scalar_lr:{args.scalar_lr}" - ) - log0( - f"train_batch_tokens:{args.train_batch_tokens} train_seq_len:{args.train_seq_len} " - f"iterations:{args.iterations} warmup_steps:{args.warmup_steps} " - f"max_wallclock_seconds:{args.max_wallclock_seconds:.3f}" - ) - log0(f"seed:{args.seed}") - - # ----------------------------- - # DATA LOADER & MODEL WARMUP - # ----------------------------- - - train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device) - - def zero_grad_all() -> None: - for opt in optimizers: - opt.zero_grad(set_to_none=True) - - max_wallclock_ms = 1000.0 * args.max_wallclock_seconds if args.max_wallclock_seconds > 0 else None - - def lr_mul(step: int, elapsed_ms: float) -> float: - if args.warmdown_iters <= 0: - return 1.0 - if max_wallclock_ms is None: - warmdown_start = max(args.iterations - args.warmdown_iters, 0) - return max((args.iterations - step) / max(args.warmdown_iters, 1), 0.0) if warmdown_start <= step < args.iterations else 1.0 - step_ms = elapsed_ms / max(step, 1) - warmdown_ms = args.warmdown_iters * step_ms - remaining_ms = max(max_wallclock_ms - elapsed_ms, 0.0) - return remaining_ms / max(warmdown_ms, 1e-9) if remaining_ms <= warmdown_ms else 1.0 - - # Warmup primes the compiled forward/backward/optimizer paths, then we restore the - # initial weights/optimizer state so measured training starts from the true init. - if args.warmup_steps > 0: - initial_model_state = {name: tensor.detach().cpu().clone() for name, tensor in base_model.state_dict().items()} - initial_optimizer_states = [copy.deepcopy(opt.state_dict()) for opt in optimizers] - model.train() - for warmup_step in range(args.warmup_steps): - zero_grad_all() - for micro_step in range(grad_accum_steps): - if distributed: - model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 - x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - warmup_loss = model(x, y) - (warmup_loss * grad_scale).backward() - for opt in optimizers: - opt.step() - zero_grad_all() - if args.warmup_steps <= 20 or (warmup_step + 1) % 10 == 0 or warmup_step + 1 == args.warmup_steps: - log0(f"warmup_step:{warmup_step + 1}/{args.warmup_steps}") - base_model.load_state_dict(initial_model_state, strict=True) - for opt, state in zip(optimizers, initial_optimizer_states, strict=True): - opt.load_state_dict(state) - zero_grad_all() - if distributed: - model.require_backward_grad_sync = True - train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device) - - # ----------------------------- - # MAIN TRAINING LOOP - # ----------------------------- - - training_time_ms = 0.0 - stop_after_step: int | None = None - torch.cuda.synchronize() - t0 = time.perf_counter() - - step = 0 - while True: - last_step = step == args.iterations or (stop_after_step is not None and step >= stop_after_step) - - should_validate = last_step or (args.val_loss_every > 0 and step % args.val_loss_every == 0) - if should_validate: - torch.cuda.synchronize() - training_time_ms += 1000.0 * (time.perf_counter() - t0) - val_loss, val_bpb = eval_val( - args, - model, - rank, - world_size, - device, - grad_accum_steps, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - ) - log0( - f"step:{step}/{args.iterations} val_loss:{val_loss:.4f} val_bpb:{val_bpb:.4f} " - f"train_time:{training_time_ms:.0f}ms step_avg:{training_time_ms / max(step, 1):.2f}ms" - ) - torch.cuda.synchronize() - t0 = time.perf_counter() - - if last_step: - if stop_after_step is not None and step < args.iterations: - log0( - f"stopping_early: wallclock_cap train_time:{training_time_ms:.0f}ms " - f"step:{step}/{args.iterations}" - ) - break - - elapsed_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) - scale = lr_mul(step, elapsed_ms) - zero_grad_all() - train_loss = torch.zeros((), device=device) - for micro_step in range(grad_accum_steps): - if distributed: - model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 - x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - loss = model(x, y) - train_loss += loss.detach() - (loss * grad_scale).backward() - train_loss /= grad_accum_steps - - frac = min(step / args.muon_momentum_warmup_steps, 1.0) if args.muon_momentum_warmup_steps > 0 else 1.0 - muon_momentum = (1 - frac) * args.muon_momentum_warmup_start + frac * args.muon_momentum - for group in optimizer_muon.param_groups: - group["momentum"] = muon_momentum - - for opt in optimizers: - for group in opt.param_groups: - group["lr"] = group["base_lr"] * scale - - if args.grad_clip_norm > 0: - torch.nn.utils.clip_grad_norm_(base_model.parameters(), args.grad_clip_norm) - for opt in optimizers: - opt.step() - zero_grad_all() - - step += 1 - approx_training_time_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) - should_log_train = ( - args.train_log_every > 0 - and (step <= 10 or step % args.train_log_every == 0 or stop_after_step is not None) - ) - if should_log_train: - log0( - f"step:{step}/{args.iterations} train_loss:{train_loss.item():.4f} " - f"train_time:{approx_training_time_ms:.0f}ms step_avg:{approx_training_time_ms / step:.2f}ms" - ) - - # Needed to sync whether we've reached the wallclock cap. - reached_cap = max_wallclock_ms is not None and approx_training_time_ms >= max_wallclock_ms - if distributed and max_wallclock_ms is not None: - reached_cap_tensor = torch.tensor(int(reached_cap), device=device) - dist.all_reduce(reached_cap_tensor, op=dist.ReduceOp.MAX) - reached_cap = bool(reached_cap_tensor.item()) - if stop_after_step is None and reached_cap: - stop_after_step = step - - log0( - f"peak memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB " - f"reserved: {torch.cuda.max_memory_reserved() // 1024 // 1024} MiB" - ) - - # ----------------------------- - # SERIALIZATION + ROUNDTRIP VALIDATION - # ----------------------------- - # Save the raw state (useful for debugging/loading in PyTorch directly), then always produce - # the compressed int8+zlib artifact and validate the round-tripped weights. - - if master_process: - torch.save(base_model.state_dict(), "final_model.pt") - model_bytes = os.path.getsize("final_model.pt") - code_bytes = len(code.encode("utf-8")) - log0(f"Serialized model: {model_bytes} bytes") - log0(f"Code size: {code_bytes} bytes") - log0(f"Total submission size: {model_bytes + code_bytes} bytes") - - quant_obj, quant_stats = quantize_state_dict_int8(base_model.state_dict()) - quant_buf = io.BytesIO() - torch.save(quant_obj, quant_buf) - quant_raw = quant_buf.getvalue() - quant_blob = zlib.compress(quant_raw, level=9) - quant_raw_bytes = len(quant_raw) - if master_process: - with open("final_model.int8.ptz", "wb") as f: - f.write(quant_blob) - quant_file_bytes = os.path.getsize("final_model.int8.ptz") - code_bytes = len(code.encode("utf-8")) - ratio = quant_stats["baseline_tensor_bytes"] / max(quant_stats["int8_payload_bytes"], 1) - log0( - f"Serialized model int8+zlib: {quant_file_bytes} bytes " - f"(payload:{quant_stats['int8_payload_bytes']} raw_torch:{quant_raw_bytes} payload_ratio:{ratio:.2f}x)" - ) - log0(f"Total submission size int8+zlib: {quant_file_bytes + code_bytes} bytes") - - if distributed: - dist.barrier() - with open("final_model.int8.ptz", "rb") as f: - quant_blob_disk = f.read() - quant_state = torch.load(io.BytesIO(zlib.decompress(quant_blob_disk)), map_location="cpu") - base_model.load_state_dict(dequantize_state_dict_int8(quant_state), strict=True) - torch.cuda.synchronize() - t_qeval = time.perf_counter() - q_val_loss, q_val_bpb = eval_val( - args, - model, - rank, - world_size, - device, - grad_accum_steps, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - ) - torch.cuda.synchronize() - log0( - f"final_int8_zlib_roundtrip val_loss:{q_val_loss:.4f} val_bpb:{q_val_bpb:.4f} " - f"eval_time:{1000.0 * (time.perf_counter() - t_qeval):.0f}ms" - ) - log0(f"final_int8_zlib_roundtrip_exact val_loss:{q_val_loss:.8f} val_bpb:{q_val_bpb:.8f}") - - if distributed: - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/records/track_10min_16mb/2026-03-18_FP16Embed_WD3600/train_seed42.log b/records/track_10min_16mb/2026-03-18_FP16Embed_WD3600/train_seed42.log deleted file mode 100644 index 6b74aefaef..0000000000 --- a/records/track_10min_16mb/2026-03-18_FP16Embed_WD3600/train_seed42.log +++ /dev/null @@ -1,1311 +0,0 @@ -""" -The `train_gpt.py` and `train_gpt_mlx.py` scripts are intended as good launching-off points for new participants, not SOTA configs. We'll accept PRs that tune, improve, or simplify these scripts without significantly increasing complexity, but competitive submissions should stay in the `/records` folder. - -Hard stop: `train_gpt.py` and `train_gpt_mlx.py` must never be longer than 1500 lines. -""" - -from __future__ import annotations - -import copy -import glob -import io -import math -import os -import random -import subprocess -import sys -import time -import uuid -import zlib -from pathlib import Path - -import numpy as np -import sentencepiece as spm -import torch -import torch.distributed as dist -import torch.nn.functional as F -from torch import Tensor, nn -from torch.nn.parallel import DistributedDataParallel as DDP - -# ----------------------------- -# HYPERPARAMETERS -# ----------------------------- -# Default Simple Baseline run: -# - 9 transformer blocks at width 512 -# - 8 attention heads with 4 KV heads (GQA) and 2x MLP expansion -# - vocab size 1024, sequence length 1024, tied embeddings -# - 524,288 train tokens per step for 20,000 iterations with a ~10 minute cap - -class Hyperparameters: - # Data paths are shard globs produced by the existing preprocessing pipeline. - data_path = os.environ.get("DATA_PATH", "./data/datasets/fineweb10B_sp1024") - train_files = os.path.join(data_path, "fineweb_train_*.bin") - val_files = os.path.join(data_path, "fineweb_val_*.bin") - tokenizer_path = os.environ.get("TOKENIZER_PATH", "./data/tokenizers/fineweb_1024_bpe.model") - run_id = os.environ.get("RUN_ID", str(uuid.uuid4())) - seed = int(os.environ.get("SEED", 1337)) - - # Validation cadence and batch size. Validation always uses the full fineweb_val split. - val_batch_size = int(os.environ.get("VAL_BATCH_SIZE", 524_288)) - val_loss_every = int(os.environ.get("VAL_LOSS_EVERY", 1000)) - train_log_every = int(os.environ.get("TRAIN_LOG_EVERY", 200)) - - # Training length. - iterations = int(os.environ.get("ITERATIONS", 20000)) - warmdown_iters = int(os.environ.get("WARMDOWN_ITERS", 1200)) - warmup_steps = int(os.environ.get("WARMUP_STEPS", 20)) - train_batch_tokens = int(os.environ.get("TRAIN_BATCH_TOKENS", 524_288)) - train_seq_len = int(os.environ.get("TRAIN_SEQ_LEN", 1024)) - max_wallclock_seconds = float(os.environ.get("MAX_WALLCLOCK_SECONDS", 600.0)) - qk_gain_init = float(os.environ.get("QK_GAIN_INIT", 1.5)) - - # Model shape. - vocab_size = int(os.environ.get("VOCAB_SIZE", 1024)) - num_layers = int(os.environ.get("NUM_LAYERS", 9)) - num_kv_heads = int(os.environ.get("NUM_KV_HEADS", 4)) - model_dim = int(os.environ.get("MODEL_DIM", 512)) - num_heads = int(os.environ.get("NUM_HEADS", 8)) - mlp_mult = int(os.environ.get("MLP_MULT", 2)) - mlp_hidden = int(os.environ.get("MLP_HIDDEN", 0)) # override mlp_mult if > 0 - tie_embeddings = bool(int(os.environ.get("TIE_EMBEDDINGS", "1"))) - rope_base = float(os.environ.get("ROPE_BASE", 10000.0)) - logit_softcap = float(os.environ.get("LOGIT_SOFTCAP", 30.0)) - - # Optimizer hyperparameters. - embed_lr = float(os.environ.get("EMBED_LR", 0.6)) - head_lr = float(os.environ.get("HEAD_LR", 0.008)) - tied_embed_lr = float(os.environ.get("TIED_EMBED_LR", 0.05)) - tied_embed_init_std = float(os.environ.get("TIED_EMBED_INIT_STD", 0.005)) - matrix_lr = float(os.environ.get("MATRIX_LR", 0.04)) - scalar_lr = float(os.environ.get("SCALAR_LR", 0.04)) - muon_momentum = float(os.environ.get("MUON_MOMENTUM", 0.95)) - muon_backend_steps = int(os.environ.get("MUON_BACKEND_STEPS", 5)) - muon_momentum_warmup_start = float(os.environ.get("MUON_MOMENTUM_WARMUP_START", 0.85)) - muon_momentum_warmup_steps = int(os.environ.get("MUON_MOMENTUM_WARMUP_STEPS", 500)) - beta1 = float(os.environ.get("BETA1", 0.9)) - beta2 = float(os.environ.get("BETA2", 0.95)) - adam_eps = float(os.environ.get("ADAM_EPS", 1e-8)) - grad_clip_norm = float(os.environ.get("GRAD_CLIP_NORM", 0.0)) - -# ----------------------------- -# MUON OPTIMIZER -# ----------------------------- -# -# As borrowed from modded-nanogpt -# Background on Muon: https://kellerjordan.github.io/posts/muon/ - -def zeropower_via_newtonschulz5(G: Tensor, steps: int = 10, eps: float = 1e-7) -> Tensor: - # Orthogonalize a 2D update matrix with a fast Newton-Schulz iteration. - # Muon uses this to normalize matrix-shaped gradients before applying them. - a, b, c = (3.4445, -4.7750, 2.0315) - X = G.bfloat16() - X /= X.norm() + eps - transposed = G.size(0) > G.size(1) - if transposed: - X = X.T - for _ in range(steps): - A = X @ X.T - B = b * A + c * A @ A - X = a * X + B @ X - return X.T if transposed else X - - -class Muon(torch.optim.Optimizer): - def __init__(self, params, lr: float, momentum: float, backend_steps: int, nesterov: bool = True): - super().__init__( - params, - dict(lr=lr, momentum=momentum, backend_steps=backend_steps, nesterov=nesterov), - ) - - @torch.no_grad() - def step(self, closure=None): - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - distributed = dist.is_available() and dist.is_initialized() - world_size = dist.get_world_size() if distributed else 1 - rank = dist.get_rank() if distributed else 0 - - for group in self.param_groups: - params = group["params"] - if not params: - continue - lr = group["lr"] - momentum = group["momentum"] - backend_steps = group["backend_steps"] - nesterov = group["nesterov"] - - total_params = sum(int(p.numel()) for p in params) - updates_flat = torch.zeros(total_params, device=params[0].device, dtype=torch.bfloat16) - - curr = 0 - for i, p in enumerate(params): - if i % world_size == rank and p.grad is not None: - g = p.grad - state = self.state[p] - if "momentum_buffer" not in state: - state["momentum_buffer"] = torch.zeros_like(g) - buf = state["momentum_buffer"] - buf.mul_(momentum).add_(g) - if nesterov: - g = g.add(buf, alpha=momentum) - g = zeropower_via_newtonschulz5(g, steps=backend_steps) - # Scale correction from Muon reference implementations. - g *= max(1, g.size(0) / g.size(1)) ** 0.5 - updates_flat[curr : curr + p.numel()] = g.reshape(-1) - curr += p.numel() - - if distributed: - dist.all_reduce(updates_flat, op=dist.ReduceOp.SUM) - - curr = 0 - for p in params: - g = updates_flat[curr : curr + p.numel()].view_as(p).to(dtype=p.dtype) - p.add_(g, alpha=-lr) - curr += p.numel() - - return loss - - -# ----------------------------- -# TOKENIZER-AGNOSTIC EVALUATION SETUP -# ----------------------------- -# -# It's common for small models have a large fraction of their parameters be embeddings, since the 2 * d_model * d_vocab vectors can be gigantic. -# Instead of locking the tokenizer, we let you bring your own and calculate our validation metrics on the average compression of the validation set. -# We calculate BPB (bits-per-byte) instead of validation loss, so we need methods to count the number of bits per token in the tokenizer. -# Note: Submissions that edit the tokenizer will be examined more carefully, since screwing this up might unjustly improve your score. - -def build_sentencepiece_luts( - sp: spm.SentencePieceProcessor, vocab_size: int, device: torch.device -) -> tuple[Tensor, Tensor, Tensor]: - sp_vocab_size = int(sp.vocab_size()) - table_size = max(sp_vocab_size, vocab_size) - base_bytes_np = np.zeros((table_size,), dtype=np.int16) - has_leading_space_np = np.zeros((table_size,), dtype=np.bool_) - is_boundary_token_np = np.ones((table_size,), dtype=np.bool_) - for token_id in range(sp_vocab_size): - if sp.is_control(token_id) or sp.is_unknown(token_id) or sp.is_unused(token_id): - continue - is_boundary_token_np[token_id] = False - if sp.is_byte(token_id): - base_bytes_np[token_id] = 1 - continue - piece = sp.id_to_piece(token_id) - if piece.startswith("▁"): - has_leading_space_np[token_id] = True - piece = piece[1:] - base_bytes_np[token_id] = len(piece.encode("utf-8")) - return ( - torch.tensor(base_bytes_np, dtype=torch.int16, device=device), - torch.tensor(has_leading_space_np, dtype=torch.bool, device=device), - torch.tensor(is_boundary_token_np, dtype=torch.bool, device=device), - ) - - -def load_validation_tokens(pattern: str, seq_len: int) -> Tensor: - files = [Path(p) for p in sorted(glob.glob(pattern))] - if not files: - raise FileNotFoundError(f"No files found for pattern: {pattern}") - # The export pipeline writes the fixed first-50k-doc validation set to fineweb_val_*. - tokens = torch.cat([load_data_shard(file) for file in files]).contiguous() - usable = ((tokens.numel() - 1) // seq_len) * seq_len - if usable <= 0: - raise ValueError(f"Validation split is too short for TRAIN_SEQ_LEN={seq_len}") - return tokens[: usable + 1] - - -def eval_val( - args: Hyperparameters, - model: nn.Module, - rank: int, - world_size: int, - device: torch.device, - grad_accum_steps: int, - val_tokens: Tensor, - base_bytes_lut: Tensor, - has_leading_space_lut: Tensor, - is_boundary_token_lut: Tensor, -) -> tuple[float, float]: - # Validation computes two metrics: - # - val_loss: token cross-entropy (natural log) - # - val_bpb: tokenizer-agnostic compression metric used by the challenge - local_batch_tokens = args.val_batch_size // (world_size * grad_accum_steps) - if local_batch_tokens < args.train_seq_len: - raise ValueError( - "VAL_BATCH_SIZE must provide at least one sequence per rank; " - f"got VAL_BATCH_SIZE={args.val_batch_size}, WORLD_SIZE={world_size}, " - f"GRAD_ACCUM_STEPS={grad_accum_steps}, TRAIN_SEQ_LEN={args.train_seq_len}" - ) - local_batch_seqs = local_batch_tokens // args.train_seq_len - total_seqs = (val_tokens.numel() - 1) // args.train_seq_len - seq_start = (total_seqs * rank) // world_size - seq_end = (total_seqs * (rank + 1)) // world_size - val_loss_sum = torch.zeros((), device=device, dtype=torch.float64) - val_token_count = torch.zeros((), device=device, dtype=torch.float64) - val_byte_count = torch.zeros((), device=device, dtype=torch.float64) - - model.eval() - with torch.inference_mode(): - for batch_seq_start in range(seq_start, seq_end, local_batch_seqs): - batch_seq_end = min(batch_seq_start + local_batch_seqs, seq_end) - raw_start = batch_seq_start * args.train_seq_len - raw_end = batch_seq_end * args.train_seq_len + 1 - local = val_tokens[raw_start:raw_end].to(device=device, dtype=torch.int64, non_blocking=True) - x = local[:-1].reshape(-1, args.train_seq_len) - y = local[1:].reshape(-1, args.train_seq_len) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - batch_loss = model(x, y).detach() - batch_token_count = float(y.numel()) - val_loss_sum += batch_loss.to(torch.float64) * batch_token_count - val_token_count += batch_token_count - prev_ids = x.reshape(-1) - tgt_ids = y.reshape(-1) - token_bytes = base_bytes_lut[tgt_ids].to(dtype=torch.int16) - token_bytes += (has_leading_space_lut[tgt_ids] & ~is_boundary_token_lut[prev_ids]).to(dtype=torch.int16) - val_byte_count += token_bytes.to(torch.float64).sum() - - if dist.is_available() and dist.is_initialized(): - dist.all_reduce(val_loss_sum, op=dist.ReduceOp.SUM) - dist.all_reduce(val_token_count, op=dist.ReduceOp.SUM) - dist.all_reduce(val_byte_count, op=dist.ReduceOp.SUM) - - val_loss = val_loss_sum / val_token_count - bits_per_token = val_loss.item() / math.log(2.0) - tokens_per_byte = val_token_count.item() / val_byte_count.item() - model.train() - return float(val_loss.item()), float(bits_per_token * tokens_per_byte) - -# ----------------------------- -# POST-TRAINING QUANTIZATION -# ----------------------------- -# -# It's silly to export our model, which is trained in bf16 and fp32, at that same precision. -# Instead, we get approximately the same model (with a small hit) by quantizing the model to int8 & zlib compressing. -# We can then decompress the model and run in higher precision for evaluation, after closing in under the size limit. - -CONTROL_TENSOR_NAME_PATTERNS = tuple( - pattern - for pattern in os.environ.get( - "CONTROL_TENSOR_NAME_PATTERNS", - "attn_scale,attn_scales,mlp_scale,mlp_scales,resid_mix,resid_mixes,q_gain,skip_weight,skip_weights", - ).split(",") - if pattern -) -INT8_KEEP_FLOAT_FP32_NAME_PATTERNS = tuple( - pattern - for pattern in os.environ.get( - "INT8_KEEP_FLOAT_FP32_NAME_PATTERNS", - ",".join(CONTROL_TENSOR_NAME_PATTERNS), - ).split(",") - if pattern -) -INT8_KEEP_FLOAT_MAX_NUMEL = 65_536 -INT8_KEEP_FLOAT_STORE_DTYPE = torch.float16 -INT8_PER_ROW_SCALE_DTYPE = torch.float16 -INT8_CLIP_PERCENTILE = 99.99984 -INT8_CLIP_Q = INT8_CLIP_PERCENTILE / 100.0 - -def tensor_nbytes(t: Tensor) -> int: - return int(t.numel()) * int(t.element_size()) - -def keep_float_tensor(name: str, t: Tensor, passthrough_orig_dtypes: dict[str, str]) -> Tensor: - if any(pattern in name for pattern in INT8_KEEP_FLOAT_FP32_NAME_PATTERNS): - return t.float().contiguous() - if t.dtype in {torch.float32, torch.bfloat16}: - passthrough_orig_dtypes[name] = str(t.dtype).removeprefix("torch.") - return t.to(dtype=INT8_KEEP_FLOAT_STORE_DTYPE).contiguous() - return t - -def quantize_float_tensor(t: Tensor) -> tuple[Tensor, Tensor]: - t32 = t.float() - if t32.ndim == 2: - # Matrices get one scale per row, which usually tracks output-channel - # ranges much better than a single tensor-wide scale. - clip_abs = ( - torch.quantile(t32.abs(), INT8_CLIP_Q, dim=1) - if t32.numel() - else torch.empty((t32.shape[0],), dtype=torch.float32) - ) - clipped = torch.maximum(torch.minimum(t32, clip_abs[:, None]), -clip_abs[:, None]) - scale = (clip_abs / 127.0).clamp_min(1.0 / 127.0) - q = torch.clamp(torch.round(clipped / scale[:, None]), -127, 127).to(torch.int8).contiguous() - return q, scale.to(dtype=INT8_PER_ROW_SCALE_DTYPE).contiguous() - - # Vectors / scalars use a simpler per-tensor scale. - clip_abs = float(torch.quantile(t32.abs().flatten(), INT8_CLIP_Q).item()) if t32.numel() else 0.0 - scale = torch.tensor(clip_abs / 127.0 if clip_abs > 0 else 1.0, dtype=torch.float32) - q = torch.clamp(torch.round(torch.clamp(t32, -clip_abs, clip_abs) / scale), -127, 127).to(torch.int8).contiguous() - return q, scale - -def quantize_state_dict_int8(state_dict: dict[str, Tensor]): - # Single supported clean-script export format: - # - per-row int8 for 2D float tensors - # - per-tensor int8 for other float tensors - # - exact passthrough for non-floats - # - passthrough for small float tensors, stored as fp16 to save bytes - quantized: dict[str, Tensor] = {} - scales: dict[str, Tensor] = {} - dtypes: dict[str, str] = {} - passthrough: dict[str, Tensor] = {} - passthrough_orig_dtypes: dict[str, str] = {} - qmeta: dict[str, dict[str, object]] = {} - stats = dict.fromkeys( - ("param_count", "num_tensors", "num_float_tensors", "num_nonfloat_tensors", "baseline_tensor_bytes", "int8_payload_bytes"), - 0, - ) - - for name, tensor in state_dict.items(): - t = tensor.detach().to("cpu").contiguous() - stats["param_count"] += int(t.numel()) - stats["num_tensors"] += 1 - stats["baseline_tensor_bytes"] += tensor_nbytes(t) - - if not t.is_floating_point(): - stats["num_nonfloat_tensors"] += 1 - passthrough[name] = t - stats["int8_payload_bytes"] += tensor_nbytes(t) - continue - - # Small float tensors are cheap enough to keep directly. We still downcast - # fp32/bf16 passthrough tensors to fp16 so metadata does not dominate size. - # Also keep the tied embedding in fp16 to minimize quantization degradation - # on the output head, which is the most sensitive tensor for BPB. - if t.numel() <= INT8_KEEP_FLOAT_MAX_NUMEL or name == "tok_emb.weight": - kept = keep_float_tensor(name, t, passthrough_orig_dtypes) - passthrough[name] = kept - stats["int8_payload_bytes"] += tensor_nbytes(kept) - continue - - stats["num_float_tensors"] += 1 - q, s = quantize_float_tensor(t) - if s.ndim > 0: - qmeta[name] = {"scheme": "per_row", "axis": 0} - quantized[name] = q - scales[name] = s - dtypes[name] = str(t.dtype).removeprefix("torch.") - stats["int8_payload_bytes"] += tensor_nbytes(q) + tensor_nbytes(s) - - obj: dict[str, object] = { - "__quant_format__": "int8_clean_per_row_v1", - "quantized": quantized, - "scales": scales, - "dtypes": dtypes, - "passthrough": passthrough, - } - if qmeta: - obj["qmeta"] = qmeta - if passthrough_orig_dtypes: - obj["passthrough_orig_dtypes"] = passthrough_orig_dtypes - return obj, stats - -def dequantize_state_dict_int8(obj: dict[str, object]) -> dict[str, Tensor]: - out: dict[str, Tensor] = {} - qmeta = obj.get("qmeta", {}) - passthrough_orig_dtypes = obj.get("passthrough_orig_dtypes", {}) - for name, q in obj["quantized"].items(): - dtype = getattr(torch, obj["dtypes"][name]) - s = obj["scales"][name] - if qmeta.get(name, {}).get("scheme") == "per_row" or s.ndim > 0: - s = s.to(dtype=torch.float32) - # Broadcast the saved row scale back across trailing dimensions. - out[name] = (q.float() * s.view(q.shape[0], *([1] * (q.ndim - 1)))).to(dtype=dtype).contiguous() - else: - scale = float(s.item()) - out[name] = (q.float() * scale).to(dtype=dtype).contiguous() - for name, t in obj["passthrough"].items(): - # Restore small tensors, undoing the temporary fp16 storage cast if needed. - out_t = t.detach().to("cpu").contiguous() - orig_dtype = passthrough_orig_dtypes.get(name) - if isinstance(orig_dtype, str): - out_t = out_t.to(dtype=getattr(torch, orig_dtype)).contiguous() - out[name] = out_t - return out - - -# ----------------------------- -# DATA LOADING -# ----------------------------- - -def load_data_shard(file: Path) -> Tensor: - header_bytes = 256 * np.dtype(" None: - self.file_idx = (self.file_idx + 1) % len(self.files) - self.tokens = load_data_shard(self.files[self.file_idx]) - self.pos = 0 - - def take(self, n: int) -> Tensor: - chunks: list[Tensor] = [] - remaining = n - while remaining > 0: - avail = self.tokens.numel() - self.pos - if avail <= 0: - self._advance_file() - continue - k = min(remaining, avail) - chunks.append(self.tokens[self.pos : self.pos + k]) - self.pos += k - remaining -= k - return chunks[0] if len(chunks) == 1 else torch.cat(chunks) - - -class DistributedTokenLoader: - # Each call consumes a contiguous chunk from the shared token stream, then slices out - # one disjoint span per rank. The extra "+1" token lets us build (x, y) by shifting. - def __init__(self, pattern: str, rank: int, world_size: int, device: torch.device): - self.rank = rank - self.world_size = world_size - self.device = device - self.stream = TokenStream(pattern) - - def next_batch(self, global_tokens: int, seq_len: int, grad_accum_steps: int) -> tuple[Tensor, Tensor]: - local_tokens = global_tokens // (self.world_size * grad_accum_steps) - per_rank_span = local_tokens + 1 - chunk = self.stream.take(per_rank_span * self.world_size) - start = self.rank * per_rank_span - local = chunk[start : start + per_rank_span].to(dtype=torch.int64) - x = local[:-1].reshape(-1, seq_len) - y = local[1:].reshape(-1, seq_len) - return x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True) - -# ----------------------------- -# TRANSFORMER MODULES -# ----------------------------- - -class RMSNorm(nn.Module): - def __init__(self, eps: float | None = None): - super().__init__() - self.eps = eps - - def forward(self, x: Tensor) -> Tensor: - return F.rms_norm(x, (x.size(-1),), eps=self.eps) - - -class CastedLinear(nn.Linear): - # Keep weights in fp32 for optimizer/state quality, cast at matmul time for bf16 compute. - def forward(self, x: Tensor) -> Tensor: - bias = self.bias.to(x.dtype) if self.bias is not None else None - return F.linear(x, self.weight.to(x.dtype), bias) - - -def restore_low_dim_params_to_fp32(module: nn.Module) -> None: - # Keep small/control parameters in fp32 even when the model body runs in bf16. - with torch.no_grad(): - for name, param in module.named_parameters(): - if (param.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS)) and param.dtype != torch.float32: - param.data = param.data.float() - - -class Rotary(nn.Module): - # Caches cos/sin tables per sequence length on the current device. - def __init__(self, dim: int, base: float = 10000.0): - super().__init__() - inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self._seq_len_cached = 0 - self._cos_cached: Tensor | None = None - self._sin_cached: Tensor | None = None - - def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> tuple[Tensor, Tensor]: - if ( - self._cos_cached is None - or self._sin_cached is None - or self._seq_len_cached != seq_len - or self._cos_cached.device != device - ): - t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) - freqs = torch.outer(t, self.inv_freq.to(device)) - self._cos_cached = freqs.cos()[None, None, :, :] - self._sin_cached = freqs.sin()[None, None, :, :] - self._seq_len_cached = seq_len - return self._cos_cached.to(dtype=dtype), self._sin_cached.to(dtype=dtype) - - -def apply_rotary_emb(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: - half = x.size(-1) // 2 - x1, x2 = x[..., :half], x[..., half:] - return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1) - - -class CausalSelfAttention(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - num_kv_heads: int, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - if dim % num_heads != 0: - raise ValueError("model_dim must be divisible by num_heads") - if num_heads % num_kv_heads != 0: - raise ValueError("num_heads must be divisible by num_kv_heads") - self.num_heads = num_heads - self.num_kv_heads = num_kv_heads - self.head_dim = dim // num_heads - if self.head_dim % 2 != 0: - raise ValueError("head_dim must be even for RoPE") - kv_dim = self.num_kv_heads * self.head_dim - self.c_q = CastedLinear(dim, dim, bias=False) - self.c_k = CastedLinear(dim, kv_dim, bias=False) - self.c_v = CastedLinear(dim, kv_dim, bias=False) - self.proj = CastedLinear(dim, dim, bias=False) - self.proj._zero_init = True - self.q_gain = nn.Parameter(torch.full((num_heads,), qk_gain_init, dtype=torch.float32)) - self.rotary = Rotary(self.head_dim, base=rope_base) - - def forward(self, x: Tensor) -> Tensor: - bsz, seqlen, dim = x.shape - q = self.c_q(x).reshape(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2) - k = self.c_k(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - v = self.c_v(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - q = F.rms_norm(q, (q.size(-1),)) - k = F.rms_norm(k, (k.size(-1),)) - cos, sin = self.rotary(seqlen, x.device, q.dtype) - q = apply_rotary_emb(q, cos, sin) - k = apply_rotary_emb(k, cos, sin) - q = q * self.q_gain.to(dtype=q.dtype)[None, :, None, None] - y = F.scaled_dot_product_attention( - q, - k, - v, - attn_mask=None, - is_causal=True, - enable_gqa=(self.num_kv_heads != self.num_heads), - ) - y = y.transpose(1, 2).contiguous().reshape(bsz, seqlen, dim) - return self.proj(y) - - -class MLP(nn.Module): - # relu^2 MLP from the original modded-nanogpt setup - def __init__(self, dim: int, mlp_mult: int, mlp_hidden: int = 0): - super().__init__() - hidden = mlp_hidden if mlp_hidden > 0 else mlp_mult * dim - self.fc = CastedLinear(dim, hidden, bias=False) - self.proj = CastedLinear(hidden, dim, bias=False) - self.proj._zero_init = True - - def forward(self, x: Tensor) -> Tensor: - x = torch.relu(self.fc(x)) - return self.proj(x.square()) - - -class Block(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - num_kv_heads: int, - mlp_mult: int, - rope_base: float, - qk_gain_init: float, - mlp_hidden: int = 0, - ): - super().__init__() - self.attn_norm = RMSNorm() - self.mlp_norm = RMSNorm() - self.attn = CausalSelfAttention(dim, num_heads, num_kv_heads, rope_base, qk_gain_init) - self.mlp = MLP(dim, mlp_mult, mlp_hidden=mlp_hidden) - self.attn_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) - self.mlp_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) - self.resid_mix = nn.Parameter(torch.stack((torch.ones(dim), torch.zeros(dim))).float()) - - def forward(self, x: Tensor, x0: Tensor) -> Tensor: - mix = self.resid_mix.to(dtype=x.dtype) - x = mix[0][None, None, :] * x + mix[1][None, None, :] * x0 - attn_out = self.attn(self.attn_norm(x)) - x = x + self.attn_scale.to(dtype=x.dtype)[None, None, :] * attn_out - x = x + self.mlp_scale.to(dtype=x.dtype)[None, None, :] * self.mlp(self.mlp_norm(x)) - return x - - -class GPT(nn.Module): - def __init__( - self, - vocab_size: int, - num_layers: int, - model_dim: int, - num_heads: int, - num_kv_heads: int, - mlp_mult: int, - tie_embeddings: bool, - tied_embed_init_std: float, - logit_softcap: float, - rope_base: float, - qk_gain_init: float, - mlp_hidden: int = 0, - ): - super().__init__() - if logit_softcap <= 0.0: - raise ValueError(f"logit_softcap must be positive, got {logit_softcap}") - self.tie_embeddings = tie_embeddings - self.tied_embed_init_std = tied_embed_init_std - self.logit_softcap = logit_softcap - self.tok_emb = nn.Embedding(vocab_size, model_dim) - self.num_encoder_layers = num_layers // 2 - self.num_decoder_layers = num_layers - self.num_encoder_layers - self.num_skip_weights = min(self.num_encoder_layers, self.num_decoder_layers) - self.skip_weights = nn.Parameter(torch.ones(self.num_skip_weights, model_dim, dtype=torch.float32)) - self.blocks = nn.ModuleList( - [ - Block( - model_dim, - num_heads, - num_kv_heads, - mlp_mult, - rope_base, - qk_gain_init, - mlp_hidden=mlp_hidden, - ) - for i in range(num_layers) - ] - ) - self.final_norm = RMSNorm() - self.lm_head = None if tie_embeddings else CastedLinear(model_dim, vocab_size, bias=False) - if self.lm_head is not None: - self.lm_head._zero_init = True - self._init_weights() - - def _init_weights(self) -> None: - if self.tie_embeddings: - nn.init.normal_(self.tok_emb.weight, mean=0.0, std=self.tied_embed_init_std) - for module in self.modules(): - if isinstance(module, nn.Linear) and getattr(module, "_zero_init", False): - nn.init.zeros_(module.weight) - - def forward(self, input_ids: Tensor, target_ids: Tensor) -> Tensor: - x = self.tok_emb(input_ids) - x = F.rms_norm(x, (x.size(-1),)) - x0 = x - skips: list[Tensor] = [] - - # First half stores skips; second half reuses them in reverse order. - for i in range(self.num_encoder_layers): - x = self.blocks[i](x, x0) - skips.append(x) - for i in range(self.num_decoder_layers): - if skips: - x = x + self.skip_weights[i].to(dtype=x.dtype)[None, None, :] * skips.pop() - x = self.blocks[self.num_encoder_layers + i](x, x0) - - x = self.final_norm(x).reshape(-1, x.size(-1)) - targets = target_ids.reshape(-1) - if self.tie_embeddings: - logits_proj = F.linear(x, self.tok_emb.weight) - else: - if self.lm_head is None: - raise RuntimeError("lm_head is required when tie_embeddings=False") - logits_proj = self.lm_head(x) - logits = self.logit_softcap * torch.tanh(logits_proj / self.logit_softcap) - return F.cross_entropy(logits.float(), targets, reduction="mean") - - -# ----------------------------- -# TRAINING -# ----------------------------- - -def main() -> None: - global zeropower_via_newtonschulz5 - - code = Path(__file__).read_text(encoding="utf-8") - args = Hyperparameters() - zeropower_via_newtonschulz5 = torch.compile(zeropower_via_newtonschulz5) - - # ----------------------------- - # DISTRIBUTED + CUDA SETUP - # ----------------------------- - - distributed = "RANK" in os.environ and "WORLD_SIZE" in os.environ - rank = int(os.environ.get("RANK", "0")) - world_size = int(os.environ.get("WORLD_SIZE", "1")) - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - if world_size <= 0: - raise ValueError(f"WORLD_SIZE must be positive, got {world_size}") - if 8 % world_size != 0: - raise ValueError(f"WORLD_SIZE={world_size} must divide 8 so grad_accum_steps stays integral") - grad_accum_steps = 8 // world_size - grad_scale = 1.0 / grad_accum_steps - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required") - device = torch.device("cuda", local_rank) - torch.cuda.set_device(device) - if distributed: - dist.init_process_group(backend="nccl", device_id=device) - dist.barrier() - master_process = rank == 0 - - # Fast math knobs - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True - from torch.backends.cuda import enable_cudnn_sdp, enable_flash_sdp, enable_math_sdp, enable_mem_efficient_sdp - - enable_cudnn_sdp(False) - enable_flash_sdp(True) - enable_mem_efficient_sdp(False) - enable_math_sdp(False) - - logfile = None - if master_process: - os.makedirs("logs", exist_ok=True) - logfile = f"logs/{args.run_id}.txt" - print(logfile) - - def log0(msg: str, console: bool = True) -> None: - if not master_process: - return - if console: - print(msg) - if logfile is not None: - with open(logfile, "a", encoding="utf-8") as f: - print(msg, file=f) - - log0(code, console=False) - log0("=" * 100, console=False) - log0(f"Running Python {sys.version}", console=False) - log0(f"Running PyTorch {torch.__version__}", console=False) - log0( - subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False).stdout, - console=False, - ) - log0("=" * 100, console=False) - - # ----------------------------- - # TOKENIZER + VALIDATION METRIC SETUP - # ----------------------------- - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed_all(args.seed) - - if not args.tokenizer_path.endswith(".model"): - raise ValueError(f"Script only setup for SentencePiece .model file: {args.tokenizer_path}") - sp = spm.SentencePieceProcessor(model_file=args.tokenizer_path) - if int(sp.vocab_size()) != args.vocab_size: - raise ValueError( - f"VOCAB_SIZE={args.vocab_size} does not match tokenizer vocab_size={int(sp.vocab_size())}" - ) - dataset_dir = Path(args.data_path).resolve() - actual_train_files = len(list(dataset_dir.glob("fineweb_train_*.bin"))) - val_tokens = load_validation_tokens(args.val_files, args.train_seq_len) - base_bytes_lut, has_leading_space_lut, is_boundary_token_lut = build_sentencepiece_luts( - sp, args.vocab_size, device - ) - log0(f"val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path={args.tokenizer_path}") - log0(f"train_loader:dataset:{dataset_dir.name} train_shards:{actual_train_files}") - log0(f"val_loader:shards pattern={args.val_files} tokens:{val_tokens.numel() - 1}") - - # ----------------------------- - # MODEL + OPTIMIZER SETUP - # ----------------------------- - - base_model = GPT( - vocab_size=args.vocab_size, - num_layers=args.num_layers, - model_dim=args.model_dim, - num_heads=args.num_heads, - num_kv_heads=args.num_kv_heads, - mlp_mult=args.mlp_mult, - tie_embeddings=args.tie_embeddings, - tied_embed_init_std=args.tied_embed_init_std, - logit_softcap=args.logit_softcap, - rope_base=args.rope_base, - qk_gain_init=args.qk_gain_init, - mlp_hidden=args.mlp_hidden, - ).to(device).bfloat16() - for module in base_model.modules(): - if isinstance(module, CastedLinear): - module.float() - restore_low_dim_params_to_fp32(base_model) - compiled_model = torch.compile(base_model, dynamic=False, fullgraph=True) - model: nn.Module = DDP(compiled_model, device_ids=[local_rank], broadcast_buffers=False) if distributed else compiled_model - - # Optimizer split: - # - token embedding (Adam) uses EMBED_LR - # - untied lm_head (Adam) uses HEAD_LR - # - matrix params in transformer blocks use MATRIX_LR via Muon - # - vectors/scalars use SCALAR_LR via Adam - block_named_params = list(base_model.blocks.named_parameters()) - matrix_params = [ - p - for name, p in block_named_params - if p.ndim == 2 and not any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - scalar_params = [ - p - for name, p in block_named_params - if p.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - if base_model.skip_weights.numel() > 0: - scalar_params.append(base_model.skip_weights) - token_lr = args.tied_embed_lr if args.tie_embeddings else args.embed_lr - optimizer_tok = torch.optim.Adam( - [{"params": [base_model.tok_emb.weight], "lr": token_lr, "base_lr": token_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizer_muon = Muon( - matrix_params, - lr=args.matrix_lr, - momentum=args.muon_momentum, - backend_steps=args.muon_backend_steps, - ) - for group in optimizer_muon.param_groups: - group["base_lr"] = args.matrix_lr - optimizer_scalar = torch.optim.Adam( - [{"params": scalar_params, "lr": args.scalar_lr, "base_lr": args.scalar_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers: list[torch.optim.Optimizer] = [optimizer_tok, optimizer_muon, optimizer_scalar] - if base_model.lm_head is not None: - optimizer_head = torch.optim.Adam( - [{"params": [base_model.lm_head.weight], "lr": args.head_lr, "base_lr": args.head_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers.insert(1, optimizer_head) - - n_params = sum(p.numel() for p in base_model.parameters()) - log0(f"model_params:{n_params}") - log0(f"world_size:{world_size} grad_accum_steps:{grad_accum_steps}") - log0("sdp_backends:cudnn=False flash=True mem_efficient=False math=False") - log0(f"attention_mode:gqa num_heads:{args.num_heads} num_kv_heads:{args.num_kv_heads}") - log0( - f"tie_embeddings:{args.tie_embeddings} embed_lr:{token_lr} " - f"head_lr:{args.head_lr if base_model.lm_head is not None else 0.0} " - f"matrix_lr:{args.matrix_lr} scalar_lr:{args.scalar_lr}" - ) - log0( - f"train_batch_tokens:{args.train_batch_tokens} train_seq_len:{args.train_seq_len} " - f"iterations:{args.iterations} warmup_steps:{args.warmup_steps} " - f"max_wallclock_seconds:{args.max_wallclock_seconds:.3f}" - ) - log0(f"seed:{args.seed}") - - # ----------------------------- - # DATA LOADER & MODEL WARMUP - # ----------------------------- - - train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device) - - def zero_grad_all() -> None: - for opt in optimizers: - opt.zero_grad(set_to_none=True) - - max_wallclock_ms = 1000.0 * args.max_wallclock_seconds if args.max_wallclock_seconds > 0 else None - - def lr_mul(step: int, elapsed_ms: float) -> float: - if args.warmdown_iters <= 0: - return 1.0 - if max_wallclock_ms is None: - warmdown_start = max(args.iterations - args.warmdown_iters, 0) - return max((args.iterations - step) / max(args.warmdown_iters, 1), 0.0) if warmdown_start <= step < args.iterations else 1.0 - step_ms = elapsed_ms / max(step, 1) - warmdown_ms = args.warmdown_iters * step_ms - remaining_ms = max(max_wallclock_ms - elapsed_ms, 0.0) - return remaining_ms / max(warmdown_ms, 1e-9) if remaining_ms <= warmdown_ms else 1.0 - - # Warmup primes the compiled forward/backward/optimizer paths, then we restore the - # initial weights/optimizer state so measured training starts from the true init. - if args.warmup_steps > 0: - initial_model_state = {name: tensor.detach().cpu().clone() for name, tensor in base_model.state_dict().items()} - initial_optimizer_states = [copy.deepcopy(opt.state_dict()) for opt in optimizers] - model.train() - for warmup_step in range(args.warmup_steps): - zero_grad_all() - for micro_step in range(grad_accum_steps): - if distributed: - model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 - x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - warmup_loss = model(x, y) - (warmup_loss * grad_scale).backward() - for opt in optimizers: - opt.step() - zero_grad_all() - if args.warmup_steps <= 20 or (warmup_step + 1) % 10 == 0 or warmup_step + 1 == args.warmup_steps: - log0(f"warmup_step:{warmup_step + 1}/{args.warmup_steps}") - base_model.load_state_dict(initial_model_state, strict=True) - for opt, state in zip(optimizers, initial_optimizer_states, strict=True): - opt.load_state_dict(state) - zero_grad_all() - if distributed: - model.require_backward_grad_sync = True - train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device) - - # ----------------------------- - # MAIN TRAINING LOOP - # ----------------------------- - - training_time_ms = 0.0 - stop_after_step: int | None = None - torch.cuda.synchronize() - t0 = time.perf_counter() - - step = 0 - while True: - last_step = step == args.iterations or (stop_after_step is not None and step >= stop_after_step) - - should_validate = last_step or (args.val_loss_every > 0 and step % args.val_loss_every == 0) - if should_validate: - torch.cuda.synchronize() - training_time_ms += 1000.0 * (time.perf_counter() - t0) - val_loss, val_bpb = eval_val( - args, - model, - rank, - world_size, - device, - grad_accum_steps, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - ) - log0( - f"step:{step}/{args.iterations} val_loss:{val_loss:.4f} val_bpb:{val_bpb:.4f} " - f"train_time:{training_time_ms:.0f}ms step_avg:{training_time_ms / max(step, 1):.2f}ms" - ) - torch.cuda.synchronize() - t0 = time.perf_counter() - - if last_step: - if stop_after_step is not None and step < args.iterations: - log0( - f"stopping_early: wallclock_cap train_time:{training_time_ms:.0f}ms " - f"step:{step}/{args.iterations}" - ) - break - - elapsed_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) - scale = lr_mul(step, elapsed_ms) - zero_grad_all() - train_loss = torch.zeros((), device=device) - for micro_step in range(grad_accum_steps): - if distributed: - model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 - x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - loss = model(x, y) - train_loss += loss.detach() - (loss * grad_scale).backward() - train_loss /= grad_accum_steps - - frac = min(step / args.muon_momentum_warmup_steps, 1.0) if args.muon_momentum_warmup_steps > 0 else 1.0 - muon_momentum = (1 - frac) * args.muon_momentum_warmup_start + frac * args.muon_momentum - for group in optimizer_muon.param_groups: - group["momentum"] = muon_momentum - - for opt in optimizers: - for group in opt.param_groups: - group["lr"] = group["base_lr"] * scale - - if args.grad_clip_norm > 0: - torch.nn.utils.clip_grad_norm_(base_model.parameters(), args.grad_clip_norm) - for opt in optimizers: - opt.step() - zero_grad_all() - - step += 1 - approx_training_time_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) - should_log_train = ( - args.train_log_every > 0 - and (step <= 10 or step % args.train_log_every == 0 or stop_after_step is not None) - ) - if should_log_train: - log0( - f"step:{step}/{args.iterations} train_loss:{train_loss.item():.4f} " - f"train_time:{approx_training_time_ms:.0f}ms step_avg:{approx_training_time_ms / step:.2f}ms" - ) - - # Needed to sync whether we've reached the wallclock cap. - reached_cap = max_wallclock_ms is not None and approx_training_time_ms >= max_wallclock_ms - if distributed and max_wallclock_ms is not None: - reached_cap_tensor = torch.tensor(int(reached_cap), device=device) - dist.all_reduce(reached_cap_tensor, op=dist.ReduceOp.MAX) - reached_cap = bool(reached_cap_tensor.item()) - if stop_after_step is None and reached_cap: - stop_after_step = step - - log0( - f"peak memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB " - f"reserved: {torch.cuda.max_memory_reserved() // 1024 // 1024} MiB" - ) - - # ----------------------------- - # SERIALIZATION + ROUNDTRIP VALIDATION - # ----------------------------- - # Save the raw state (useful for debugging/loading in PyTorch directly), then always produce - # the compressed int8+zlib artifact and validate the round-tripped weights. - - if master_process: - torch.save(base_model.state_dict(), "final_model.pt") - model_bytes = os.path.getsize("final_model.pt") - code_bytes = len(code.encode("utf-8")) - log0(f"Serialized model: {model_bytes} bytes") - log0(f"Code size: {code_bytes} bytes") - log0(f"Total submission size: {model_bytes + code_bytes} bytes") - - quant_obj, quant_stats = quantize_state_dict_int8(base_model.state_dict()) - quant_buf = io.BytesIO() - torch.save(quant_obj, quant_buf) - quant_raw = quant_buf.getvalue() - quant_blob = zlib.compress(quant_raw, level=9) - quant_raw_bytes = len(quant_raw) - if master_process: - with open("final_model.int8.ptz", "wb") as f: - f.write(quant_blob) - quant_file_bytes = os.path.getsize("final_model.int8.ptz") - code_bytes = len(code.encode("utf-8")) - ratio = quant_stats["baseline_tensor_bytes"] / max(quant_stats["int8_payload_bytes"], 1) - log0( - f"Serialized model int8+zlib: {quant_file_bytes} bytes " - f"(payload:{quant_stats['int8_payload_bytes']} raw_torch:{quant_raw_bytes} payload_ratio:{ratio:.2f}x)" - ) - log0(f"Total submission size int8+zlib: {quant_file_bytes + code_bytes} bytes") - - if distributed: - dist.barrier() - with open("final_model.int8.ptz", "rb") as f: - quant_blob_disk = f.read() - quant_state = torch.load(io.BytesIO(zlib.decompress(quant_blob_disk)), map_location="cpu") - base_model.load_state_dict(dequantize_state_dict_int8(quant_state), strict=True) - torch.cuda.synchronize() - t_qeval = time.perf_counter() - q_val_loss, q_val_bpb = eval_val( - args, - model, - rank, - world_size, - device, - grad_accum_steps, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - ) - torch.cuda.synchronize() - log0( - f"final_int8_zlib_roundtrip val_loss:{q_val_loss:.4f} val_bpb:{q_val_bpb:.4f} " - f"eval_time:{1000.0 * (time.perf_counter() - t_qeval):.0f}ms" - ) - log0(f"final_int8_zlib_roundtrip_exact val_loss:{q_val_loss:.8f} val_bpb:{q_val_bpb:.8f}") - - if distributed: - dist.destroy_process_group() - - -if __name__ == "__main__": - main() - -==================================================================================================== -Running Python 3.12.3 (main, Nov 6 2025, 13:44:16) [GCC 13.3.0] -Running PyTorch 2.9.1+cu128 -Thu Mar 19 00:57:34 2026 -+-----------------------------------------------------------------------------------------+ -| NVIDIA-SMI 565.57.01 Driver Version: 565.57.01 CUDA Version: 12.7 | -|-----------------------------------------+------------------------+----------------------+ -| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC | -| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. | -| | | MIG M. | -|=========================================+========================+======================| -| 0 NVIDIA H100 80GB HBM3 On | 00000000:18:00.0 Off | 0 | -| N/A 33C P0 147W / 700W | 1518MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 1 NVIDIA H100 80GB HBM3 On | 00000000:2A:00.0 Off | 0 | -| N/A 36C P0 151W / 700W | 1518MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 2 NVIDIA H100 80GB HBM3 On | 00000000:3A:00.0 Off | 0 | -| N/A 35C P0 149W / 700W | 1518MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 3 NVIDIA H100 80GB HBM3 On | 00000000:5D:00.0 Off | 0 | -| N/A 33C P0 152W / 700W | 1518MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 4 NVIDIA H100 80GB HBM3 On | 00000000:9A:00.0 Off | 0 | -| N/A 35C P0 151W / 700W | 1518MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 5 NVIDIA H100 80GB HBM3 On | 00000000:AB:00.0 Off | 0 | -| N/A 37C P0 148W / 700W | 1518MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 6 NVIDIA H100 80GB HBM3 On | 00000000:BA:00.0 Off | 0 | -| N/A 35C P0 149W / 700W | 1518MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 7 NVIDIA H100 80GB HBM3 On | 00000000:DB:00.0 Off | 0 | -| N/A 33C P0 147W / 700W | 1518MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ - -+-----------------------------------------------------------------------------------------+ -| Processes: | -| GPU GI CI PID Type Process name GPU Memory | -| ID ID Usage | -|=========================================================================================| -+-----------------------------------------------------------------------------------------+ - -==================================================================================================== -val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path=./data/tokenizers/fineweb_1024_bpe.model -train_loader:dataset:fineweb10B_sp1024 train_shards:80 -val_loader:shards pattern=./data/datasets/fineweb10B_sp1024/fineweb_val_*.bin tokens:62021632 -model_params:16765000 -world_size:8 grad_accum_steps:1 -sdp_backends:cudnn=False flash=True mem_efficient=False math=False -attention_mode:gqa num_heads:8 num_kv_heads:4 -tie_embeddings:True embed_lr:0.05 head_lr:0.0 matrix_lr:0.06 scalar_lr:0.04 -train_batch_tokens:524288 train_seq_len:1024 iterations:20000 warmup_steps:20 max_wallclock_seconds:600.000 -seed:42 -warmup_step:1/20 -warmup_step:2/20 -warmup_step:3/20 -warmup_step:4/20 -warmup_step:5/20 -warmup_step:6/20 -warmup_step:7/20 -warmup_step:8/20 -warmup_step:9/20 -warmup_step:10/20 -warmup_step:11/20 -warmup_step:12/20 -warmup_step:13/20 -warmup_step:14/20 -warmup_step:15/20 -warmup_step:16/20 -warmup_step:17/20 -warmup_step:18/20 -warmup_step:19/20 -warmup_step:20/20 -step:0/20000 val_loss:6.9386 val_bpb:4.1094 train_time:0ms step_avg:0.02ms -step:1/20000 train_loss:6.9374 train_time:32ms step_avg:31.73ms -step:2/20000 train_loss:16.7845 train_time:82ms step_avg:41.03ms -step:3/20000 train_loss:8.3491 train_time:133ms step_avg:44.45ms -step:4/20000 train_loss:6.7040 train_time:185ms step_avg:46.20ms -step:5/20000 train_loss:6.8139 train_time:235ms step_avg:47.07ms -step:6/20000 train_loss:7.4557 train_time:291ms step_avg:48.43ms -step:7/20000 train_loss:6.4028 train_time:343ms step_avg:49.03ms -step:8/20000 train_loss:6.1309 train_time:398ms step_avg:49.75ms -step:9/20000 train_loss:6.0002 train_time:455ms step_avg:50.51ms -step:10/20000 train_loss:5.8936 train_time:513ms step_avg:51.26ms -step:200/20000 train_loss:2.8638 train_time:8830ms step_avg:44.15ms -step:400/20000 train_loss:2.3624 train_time:17564ms step_avg:43.91ms -step:600/20000 train_loss:2.5469 train_time:26310ms step_avg:43.85ms -step:800/20000 train_loss:2.2999 train_time:35076ms step_avg:43.85ms -step:1000/20000 train_loss:2.3767 train_time:43826ms step_avg:43.83ms -step:1200/20000 train_loss:2.3883 train_time:52566ms step_avg:43.81ms -step:1400/20000 train_loss:2.4343 train_time:61328ms step_avg:43.81ms -step:1600/20000 train_loss:2.1019 train_time:70068ms step_avg:43.79ms -step:1800/20000 train_loss:2.2065 train_time:78798ms step_avg:43.78ms -step:2000/20000 train_loss:2.2557 train_time:87543ms step_avg:43.77ms -step:2000/20000 val_loss:2.2388 val_bpb:1.3259 train_time:87562ms step_avg:43.78ms -step:2200/20000 train_loss:2.0773 train_time:96276ms step_avg:43.76ms -step:2400/20000 train_loss:2.2039 train_time:105005ms step_avg:43.75ms -step:2600/20000 train_loss:2.4142 train_time:113743ms step_avg:43.75ms -step:2800/20000 train_loss:2.2384 train_time:122552ms step_avg:43.77ms -step:3000/20000 train_loss:2.2291 train_time:131285ms step_avg:43.76ms -step:3200/20000 train_loss:2.1880 train_time:140018ms step_avg:43.76ms -step:3400/20000 train_loss:2.1604 train_time:148771ms step_avg:43.76ms -step:3600/20000 train_loss:2.1172 train_time:157501ms step_avg:43.75ms -step:3800/20000 train_loss:2.2261 train_time:166232ms step_avg:43.75ms -step:4000/20000 train_loss:2.1641 train_time:174959ms step_avg:43.74ms -step:4000/20000 val_loss:2.1722 val_bpb:1.2865 train_time:174975ms step_avg:43.74ms -step:4200/20000 train_loss:2.1771 train_time:183777ms step_avg:43.76ms -step:4400/20000 train_loss:2.1164 train_time:192514ms step_avg:43.75ms -step:4600/20000 train_loss:1.9734 train_time:201249ms step_avg:43.75ms -step:4800/20000 train_loss:2.2695 train_time:209978ms step_avg:43.75ms -step:5000/20000 train_loss:2.0314 train_time:218701ms step_avg:43.74ms -step:5200/20000 train_loss:2.1773 train_time:227437ms step_avg:43.74ms -step:5400/20000 train_loss:2.1858 train_time:236176ms step_avg:43.74ms -step:5600/20000 train_loss:2.1886 train_time:244909ms step_avg:43.73ms -step:5800/20000 train_loss:2.1449 train_time:253638ms step_avg:43.73ms -step:6000/20000 train_loss:2.2212 train_time:262378ms step_avg:43.73ms -step:6000/20000 val_loss:2.1458 val_bpb:1.2709 train_time:262394ms step_avg:43.73ms -step:6200/20000 train_loss:2.0900 train_time:271098ms step_avg:43.73ms -step:6400/20000 train_loss:2.1637 train_time:279831ms step_avg:43.72ms -step:6600/20000 train_loss:2.1237 train_time:288565ms step_avg:43.72ms -step:6800/20000 train_loss:2.1928 train_time:297295ms step_avg:43.72ms -step:7000/20000 train_loss:2.2285 train_time:306032ms step_avg:43.72ms -step:7200/20000 train_loss:2.1982 train_time:314770ms step_avg:43.72ms -step:7400/20000 train_loss:2.1229 train_time:323525ms step_avg:43.72ms -step:7600/20000 train_loss:1.9980 train_time:332277ms step_avg:43.72ms -step:7800/20000 train_loss:2.1451 train_time:341024ms step_avg:43.72ms -step:8000/20000 train_loss:2.1144 train_time:349774ms step_avg:43.72ms -step:8000/20000 val_loss:2.1252 val_bpb:1.2587 train_time:349793ms step_avg:43.72ms -step:8200/20000 train_loss:2.1890 train_time:358502ms step_avg:43.72ms -step:8400/20000 train_loss:2.1378 train_time:367330ms step_avg:43.73ms -step:8600/20000 train_loss:2.1407 train_time:376094ms step_avg:43.73ms -step:8800/20000 train_loss:2.0997 train_time:384842ms step_avg:43.73ms -step:9000/20000 train_loss:2.0303 train_time:393596ms step_avg:43.73ms -step:9200/20000 train_loss:2.0892 train_time:402333ms step_avg:43.73ms -step:9400/20000 train_loss:2.1323 train_time:411086ms step_avg:43.73ms -step:9600/20000 train_loss:2.1527 train_time:419832ms step_avg:43.73ms -step:9800/20000 train_loss:2.0751 train_time:428571ms step_avg:43.73ms -step:10000/20000 train_loss:2.1173 train_time:437330ms step_avg:43.73ms -step:10000/20000 val_loss:2.1148 val_bpb:1.2525 train_time:437347ms step_avg:43.73ms -step:10200/20000 train_loss:2.0678 train_time:446065ms step_avg:43.73ms -step:10400/20000 train_loss:2.1034 train_time:454794ms step_avg:43.73ms -step:10600/20000 train_loss:1.9767 train_time:463528ms step_avg:43.73ms -step:10800/20000 train_loss:2.1827 train_time:472270ms step_avg:43.73ms -step:11000/20000 train_loss:2.1106 train_time:481007ms step_avg:43.73ms -step:11200/20000 train_loss:2.0630 train_time:489740ms step_avg:43.73ms -step:11400/20000 train_loss:2.0444 train_time:498486ms step_avg:43.73ms -step:11600/20000 train_loss:2.0506 train_time:507218ms step_avg:43.73ms -step:11800/20000 train_loss:2.0765 train_time:515969ms step_avg:43.73ms -step:12000/20000 train_loss:2.0517 train_time:524706ms step_avg:43.73ms -step:12000/20000 val_loss:2.0813 val_bpb:1.2326 train_time:524722ms step_avg:43.73ms -step:12200/20000 train_loss:2.1971 train_time:533452ms step_avg:43.73ms -step:12400/20000 train_loss:1.8441 train_time:542306ms step_avg:43.73ms -step:12600/20000 train_loss:2.0663 train_time:551034ms step_avg:43.73ms -step:12800/20000 train_loss:2.0870 train_time:559771ms step_avg:43.73ms -step:13000/20000 train_loss:2.1597 train_time:568510ms step_avg:43.73ms -step:13200/20000 train_loss:2.1700 train_time:577237ms step_avg:43.73ms -step:13400/20000 train_loss:2.0473 train_time:585968ms step_avg:43.73ms -step:13600/20000 train_loss:1.9135 train_time:594715ms step_avg:43.73ms -step:13722/20000 val_loss:2.0576 val_bpb:1.2186 train_time:599992ms step_avg:43.72ms -stopping_early: wallclock_cap train_time:599992ms step:13722/20000 -peak memory allocated: 10111 MiB reserved: 10302 MiB -Serialized model: 66045335 bytes -Code size: 48125 bytes -Total submission size: 66093460 bytes -Serialized model int8+zlib: 15849048 bytes (payload:17405664 raw_torch:17450459 payload_ratio:3.79x) -Total submission size int8+zlib: 15897173 bytes -final_int8_zlib_roundtrip val_loss:2.0600 val_bpb:1.2201 eval_time:1408ms -final_int8_zlib_roundtrip_exact val_loss:2.06002794 val_bpb:1.22006458 diff --git a/records/track_10min_16mb/2026-03-18_LongContextSeq2048/README.md b/records/track_10min_16mb/2026-03-18_LongContextSeq2048/README.md deleted file mode 100644 index 735cc0ccc6..0000000000 --- a/records/track_10min_16mb/2026-03-18_LongContextSeq2048/README.md +++ /dev/null @@ -1,65 +0,0 @@ -This record submission is called `Long Context Seq2048 v2`. - -Configuration: -- Layout: `VOCAB_SIZE=1024 NUM_LAYERS=9 MODEL_DIM=512 NUM_HEADS=8 NUM_KV_HEADS=4 MLP_MULT=2` -- Tied output/input embeddings: `TIE_EMBEDDINGS=1` -- Sequence length: `TRAIN_SEQ_LEN=2048` -- Batching: `TRAIN_BATCH_TOKENS=524288` -- Learning rates: `TIED_EMBED_LR=0.04 MATRIX_LR=0.032 SCALAR_LR=0.032` - -Command: -```bash -NCCL_IB_DISABLE=1 \ -RUN_ID=seq2048_sxm28_full_20260319a \ -DATA_PATH=./data/datasets/fineweb10B_sp1024 \ -TOKENIZER_PATH=./data/tokenizers/fineweb_1024_bpe.model \ -VOCAB_SIZE=1024 \ -MAX_WALLCLOCK_SECONDS=600 \ -torchrun --standalone --nproc_per_node=8 \ - records/track_10min_16mb/2026-03-18_LongContextSeq2048/train_gpt.py -``` - -Verification environment: -- `8x H100 80GB HBM3` -- all-to-all `NV18` topology -- `torch 2.8.0+cu128` - -Key metrics (from `train.log` in this folder, rerun on the target SXM-class box): -- Timed training stopped at `11564/20000` steps due to the wallclock cap. -- Pre-quant eval at stop: `val_loss:2.0269`, `val_bpb:1.2005` -- Post-quant roundtrip eval: `val_loss:2.0359`, `val_bpb:1.2058` -- Exact printed metric: `final_int8_zlib_roundtrip_exact val_bpb:1.20576485` -- Train time: `600038ms` (`step_avg:51.89ms`) -- Peak memory: `10247 MiB allocated`, `10488 MiB reserved` -- Serialized model int8+zlib: `15819554 bytes` -- Code size for this standalone record script: `47716 bytes` -- Total submission size int8+zlib: `15867270 bytes` - -Additional full-run reproducibility logs included in this folder: -- `train.log`: canonical SXM rerun, `SEED=1337`, `val_bpb=1.20576485` -- `train_seed1338.log`: SXM rerun, `SEED=1338`, `val_bpb=1.20617460` -- `train_seed1339.log`: SXM rerun, `SEED=1339`, `val_bpb=1.20715923` - -Record-track significance note: -- The public repo state for this submission has `Naive Baseline` at `1.2243657`. -- The challenge therefore requires beating `1.2193657` to claim a new record. -- All three included SXM full runs clear that threshold: - - `SEED=1337`: `1.20576485` - - `SEED=1338`: `1.20617460` - - `SEED=1339`: `1.20715923` -- Sample mean across the three runs: `1.20636623` -- Sample standard deviation: `0.00071667` -- One-sided one-sample t-test against `1.2193657`: `t=31.42` with `df=2`, which gives `p=0.00051` - -Why this folder is standalone: -- `train_gpt.py` compiles from inside this record folder and was used for the canonical rerun whose output is saved as `train.log`. -- No extra Python source files are required for the training path. -- The only inputs expected at runtime are the cached dataset and tokenizer paths described in the main repo README. - -Included files: -- `train_gpt.py` (standalone winning recipe with defaults baked in) -- `README.md` (this file) -- `submission.json` (leaderboard metadata) -- `train.log` (canonical full log from the standalone record script) -- `train_seed1338.log`, `train_seed1339.log` (extra full reruns for reproducibility) -- `logs/seq2048_sxm28_*` (raw per-run tee output and trainer text logs from the SXM verification box) diff --git a/records/track_10min_16mb/2026-03-18_LongContextSeq2048/submission.json b/records/track_10min_16mb/2026-03-18_LongContextSeq2048/submission.json deleted file mode 100644 index 049a8b74ad..0000000000 --- a/records/track_10min_16mb/2026-03-18_LongContextSeq2048/submission.json +++ /dev/null @@ -1,11 +0,0 @@ -{ - "author": "Spokane Way", - "github_id": "spokane-way", - "name": "Long Context Seq2048 v2", - "blurb": "SP-1024 9x512 KV4 run at TRAIN_SEQ_LEN=2048 with tuned seq2048 learning rates (0.040/0.032/0.032). This standalone record script reproduces the SXM-verified 10-minute artifact under the 16,000,000-byte cap.", - "date": "2026-03-19T04:50:00Z", - "val_loss": 2.03588345, - "val_bpb": 1.20576485, - "bytes_total": 15867270, - "bytes_code": 47716 -} diff --git a/records/track_10min_16mb/2026-03-18_LongContextSeq2048/train.log b/records/track_10min_16mb/2026-03-18_LongContextSeq2048/train.log deleted file mode 100644 index 15a6f65b12..0000000000 --- a/records/track_10min_16mb/2026-03-18_LongContextSeq2048/train.log +++ /dev/null @@ -1,124 +0,0 @@ - -***************************************** -Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. -***************************************** -logs/seq2048_sxm28_full_20260319a.txt -val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path=/root/parameter-golf-sxm28/data/tokenizers/fineweb_1024_bpe.model -train_loader:dataset:fineweb10B_sp1024 train_shards:80 -val_loader:shards pattern=/root/parameter-golf-sxm28/data/datasets/fineweb10B_sp1024/fineweb_val_*.bin tokens:62021632 -model_params:17059912 -world_size:8 grad_accum_steps:1 -sdp_backends:cudnn=False flash=True mem_efficient=False math=False -attention_mode:gqa num_heads:8 num_kv_heads:4 -tie_embeddings:True embed_lr:0.04 head_lr:0.0 matrix_lr:0.032 scalar_lr:0.032 -train_batch_tokens:524288 train_seq_len:2048 iterations:20000 warmup_steps:20 max_wallclock_seconds:600.000 -seed:1337 -warmup_step:1/20 -warmup_step:2/20 -warmup_step:3/20 -warmup_step:4/20 -warmup_step:5/20 -warmup_step:6/20 -warmup_step:7/20 -warmup_step:8/20 -warmup_step:9/20 -warmup_step:10/20 -warmup_step:11/20 -warmup_step:12/20 -warmup_step:13/20 -warmup_step:14/20 -warmup_step:15/20 -warmup_step:16/20 -warmup_step:17/20 -warmup_step:18/20 -warmup_step:19/20 -warmup_step:20/20 -step:0/20000 val_loss:6.9357 val_bpb:4.1077 train_time:0ms step_avg:0.02ms -step:1/20000 train_loss:6.9370 train_time:27ms step_avg:27.23ms -step:2/20000 train_loss:14.7712 train_time:74ms step_avg:36.88ms -step:3/20000 train_loss:8.1324 train_time:125ms step_avg:41.59ms -step:4/20000 train_loss:6.6083 train_time:176ms step_avg:44.01ms -step:5/20000 train_loss:6.9060 train_time:227ms step_avg:45.47ms -step:6/20000 train_loss:7.6667 train_time:279ms step_avg:46.44ms -step:7/20000 train_loss:6.6546 train_time:330ms step_avg:47.13ms -step:8/20000 train_loss:6.3864 train_time:381ms step_avg:47.66ms -step:9/20000 train_loss:6.2280 train_time:433ms step_avg:48.07ms -step:10/20000 train_loss:6.1411 train_time:484ms step_avg:48.42ms -step:200/20000 train_loss:2.7753 train_time:10282ms step_avg:51.41ms -step:400/20000 train_loss:2.2990 train_time:20615ms step_avg:51.54ms -step:600/20000 train_loss:2.5004 train_time:30958ms step_avg:51.60ms -step:800/20000 train_loss:2.2435 train_time:41311ms step_avg:51.64ms -step:1000/20000 train_loss:2.3383 train_time:51684ms step_avg:51.68ms -step:1000/20000 val_loss:2.2909 val_bpb:1.3568 train_time:51717ms step_avg:51.72ms -step:1200/20000 train_loss:2.3520 train_time:62063ms step_avg:51.72ms -step:1400/20000 train_loss:2.3778 train_time:72454ms step_avg:51.75ms -step:1600/20000 train_loss:2.0422 train_time:82841ms step_avg:51.78ms -step:1800/20000 train_loss:2.1630 train_time:93248ms step_avg:51.80ms -step:2000/20000 train_loss:2.2122 train_time:103654ms step_avg:51.83ms -step:2000/20000 val_loss:2.1924 val_bpb:1.2984 train_time:103687ms step_avg:51.84ms -step:2200/20000 train_loss:2.0339 train_time:114067ms step_avg:51.85ms -step:2400/20000 train_loss:2.1666 train_time:124488ms step_avg:51.87ms -step:2600/20000 train_loss:2.3803 train_time:134904ms step_avg:51.89ms -step:2800/20000 train_loss:2.1944 train_time:145315ms step_avg:51.90ms -step:3000/20000 train_loss:2.1889 train_time:155728ms step_avg:51.91ms -step:3000/20000 val_loss:2.1524 val_bpb:1.2748 train_time:155761ms step_avg:51.92ms -step:3200/20000 train_loss:2.1507 train_time:166139ms step_avg:51.92ms -step:3400/20000 train_loss:2.1186 train_time:176537ms step_avg:51.92ms -step:3600/20000 train_loss:2.0636 train_time:186950ms step_avg:51.93ms -step:3800/20000 train_loss:2.1715 train_time:197346ms step_avg:51.93ms -step:4000/20000 train_loss:2.1326 train_time:207738ms step_avg:51.93ms -step:4000/20000 val_loss:2.1285 val_bpb:1.2606 train_time:207770ms step_avg:51.94ms -step:4200/20000 train_loss:2.1300 train_time:218180ms step_avg:51.95ms -step:4400/20000 train_loss:2.0635 train_time:228563ms step_avg:51.95ms -step:4600/20000 train_loss:1.9340 train_time:238947ms step_avg:51.95ms -step:4800/20000 train_loss:2.2169 train_time:249326ms step_avg:51.94ms -step:5000/20000 train_loss:1.9728 train_time:259712ms step_avg:51.94ms -step:5000/20000 val_loss:2.1118 val_bpb:1.2507 train_time:259745ms step_avg:51.95ms -step:5200/20000 train_loss:2.1346 train_time:270102ms step_avg:51.94ms -step:5400/20000 train_loss:2.1480 train_time:280489ms step_avg:51.94ms -step:5600/20000 train_loss:2.1403 train_time:290858ms step_avg:51.94ms -step:5800/20000 train_loss:2.0939 train_time:301230ms step_avg:51.94ms -step:6000/20000 train_loss:2.1745 train_time:311608ms step_avg:51.93ms -step:6000/20000 val_loss:2.1015 val_bpb:1.2446 train_time:311642ms step_avg:51.94ms -step:6200/20000 train_loss:2.0438 train_time:321983ms step_avg:51.93ms -step:6400/20000 train_loss:2.1272 train_time:332352ms step_avg:51.93ms -step:6600/20000 train_loss:2.0825 train_time:342718ms step_avg:51.93ms -step:6800/20000 train_loss:2.1436 train_time:353087ms step_avg:51.92ms -step:7000/20000 train_loss:2.1914 train_time:363453ms step_avg:51.92ms -step:7000/20000 val_loss:2.0907 val_bpb:1.2382 train_time:363485ms step_avg:51.93ms -step:7200/20000 train_loss:2.1618 train_time:373813ms step_avg:51.92ms -step:7400/20000 train_loss:2.0806 train_time:384181ms step_avg:51.92ms -step:7600/20000 train_loss:1.9643 train_time:394550ms step_avg:51.91ms -step:7800/20000 train_loss:2.1069 train_time:404903ms step_avg:51.91ms -step:8000/20000 train_loss:2.0808 train_time:415270ms step_avg:51.91ms -step:8000/20000 val_loss:2.0816 val_bpb:1.2328 train_time:415302ms step_avg:51.91ms -step:8200/20000 train_loss:2.1517 train_time:425628ms step_avg:51.91ms -step:8400/20000 train_loss:2.0958 train_time:436033ms step_avg:51.91ms -step:8600/20000 train_loss:2.1052 train_time:446388ms step_avg:51.91ms -step:8800/20000 train_loss:2.0699 train_time:456752ms step_avg:51.90ms -step:9000/20000 train_loss:1.9858 train_time:467109ms step_avg:51.90ms -step:9000/20000 val_loss:2.0765 val_bpb:1.2298 train_time:467142ms step_avg:51.90ms -step:9200/20000 train_loss:2.0473 train_time:477468ms step_avg:51.90ms -step:9400/20000 train_loss:2.0934 train_time:487824ms step_avg:51.90ms -step:9600/20000 train_loss:2.1151 train_time:498188ms step_avg:51.89ms -step:9800/20000 train_loss:2.0174 train_time:508551ms step_avg:51.89ms -step:10000/20000 train_loss:2.0742 train_time:518903ms step_avg:51.89ms -step:10000/20000 val_loss:2.0715 val_bpb:1.2268 train_time:518936ms step_avg:51.89ms -step:10200/20000 train_loss:2.0357 train_time:529265ms step_avg:51.89ms -step:10400/20000 train_loss:2.0548 train_time:539622ms step_avg:51.89ms -step:10600/20000 train_loss:1.9345 train_time:549977ms step_avg:51.88ms -step:10800/20000 train_loss:2.1369 train_time:560331ms step_avg:51.88ms -step:11000/20000 train_loss:2.0578 train_time:570691ms step_avg:51.88ms -step:11000/20000 val_loss:2.0447 val_bpb:1.2110 train_time:570724ms step_avg:51.88ms -step:11200/20000 train_loss:2.0111 train_time:581136ms step_avg:51.89ms -step:11400/20000 train_loss:1.9882 train_time:591500ms step_avg:51.89ms -step:11564/20000 val_loss:2.0269 val_bpb:1.2005 train_time:600038ms step_avg:51.89ms -stopping_early: wallclock_cap train_time:600038ms step:11564/20000 -peak memory allocated: 10247 MiB reserved: 10488 MiB -Serialized model: 67224983 bytes -Code size: 47716 bytes -Total submission size: 67272699 bytes -Serialized model int8+zlib: 15819554 bytes (payload:17178912 raw_torch:17224025 payload_ratio:3.91x) -Total submission size int8+zlib: 15867270 bytes -final_int8_zlib_roundtrip val_loss:2.0359 val_bpb:1.2058 eval_time:1639ms -final_int8_zlib_roundtrip_exact val_loss:2.03588345 val_bpb:1.20576485 diff --git a/records/track_10min_16mb/2026-03-18_LongContextSeq2048/train_gpt.py b/records/track_10min_16mb/2026-03-18_LongContextSeq2048/train_gpt.py deleted file mode 100644 index a9ceb044e5..0000000000 --- a/records/track_10min_16mb/2026-03-18_LongContextSeq2048/train_gpt.py +++ /dev/null @@ -1,1127 +0,0 @@ -""" -The `train_gpt.py` and `train_gpt_mlx.py` scripts are intended as good launching-off points for new participants, not SOTA configs. We'll accept PRs that tune, improve, or simplify these scripts without significantly increasing complexity, but competitive submissions should stay in the `/records` folder. - -Hard stop: `train_gpt.py` and `train_gpt_mlx.py` must never be longer than 1500 lines. -""" - -from __future__ import annotations - -import copy -import glob -import io -import math -import os -import random -import subprocess -import sys -import time -import uuid -import zlib -from pathlib import Path - -import numpy as np -import sentencepiece as spm -import torch -import torch.distributed as dist -import torch.nn.functional as F -from torch import Tensor, nn -from torch.nn.parallel import DistributedDataParallel as DDP - -# ----------------------------- -# HYPERPARAMETERS -# ----------------------------- -# Default Long Context Seq2048 v2 run: -# - 9 transformer blocks at width 512 -# - 8 attention heads with 4 KV heads (GQA) and 2x MLP expansion -# - vocab size 1024, sequence length 2048, tied embeddings -# - 524,288 train tokens per step for 20,000 iterations with a ~10 minute cap -# - tuned seq2048 learning rates: 0.040 / 0.032 / 0.032 - -class Hyperparameters: - # Data paths are shard globs produced by the existing preprocessing pipeline. - data_path = os.environ.get("DATA_PATH", "./data/datasets/fineweb10B_sp1024") - train_files = os.path.join(data_path, "fineweb_train_*.bin") - val_files = os.path.join(data_path, "fineweb_val_*.bin") - tokenizer_path = os.environ.get("TOKENIZER_PATH", "./data/tokenizers/fineweb_1024_bpe.model") - run_id = os.environ.get("RUN_ID", "long_context_seq2048_v2") - seed = int(os.environ.get("SEED", 1337)) - - # Validation cadence and batch size. Validation always uses the full fineweb_val split. - val_batch_size = int(os.environ.get("VAL_BATCH_SIZE", 524_288)) - val_loss_every = int(os.environ.get("VAL_LOSS_EVERY", 1000)) - train_log_every = int(os.environ.get("TRAIN_LOG_EVERY", 200)) - - # Training length. - iterations = int(os.environ.get("ITERATIONS", 20000)) - warmdown_iters = int(os.environ.get("WARMDOWN_ITERS", 1200)) - warmup_steps = int(os.environ.get("WARMUP_STEPS", 20)) - train_batch_tokens = int(os.environ.get("TRAIN_BATCH_TOKENS", 524_288)) - train_seq_len = int(os.environ.get("TRAIN_SEQ_LEN", 2048)) - max_wallclock_seconds = float(os.environ.get("MAX_WALLCLOCK_SECONDS", 600.0)) - qk_gain_init = float(os.environ.get("QK_GAIN_INIT", 1.5)) - - # Model shape. - vocab_size = int(os.environ.get("VOCAB_SIZE", 1024)) - num_layers = int(os.environ.get("NUM_LAYERS", 9)) - num_kv_heads = int(os.environ.get("NUM_KV_HEADS", 4)) - model_dim = int(os.environ.get("MODEL_DIM", 512)) - num_heads = int(os.environ.get("NUM_HEADS", 8)) - mlp_mult = int(os.environ.get("MLP_MULT", 2)) - tie_embeddings = bool(int(os.environ.get("TIE_EMBEDDINGS", "1"))) - rope_base = float(os.environ.get("ROPE_BASE", 10000.0)) - logit_softcap = float(os.environ.get("LOGIT_SOFTCAP", 30.0)) - - # Optimizer hyperparameters. - embed_lr = float(os.environ.get("EMBED_LR", 0.6)) - head_lr = float(os.environ.get("HEAD_LR", 0.008)) - tied_embed_lr = float(os.environ.get("TIED_EMBED_LR", 0.04)) - tied_embed_init_std = float(os.environ.get("TIED_EMBED_INIT_STD", 0.005)) - matrix_lr = float(os.environ.get("MATRIX_LR", 0.032)) - scalar_lr = float(os.environ.get("SCALAR_LR", 0.032)) - muon_momentum = float(os.environ.get("MUON_MOMENTUM", 0.95)) - muon_backend_steps = int(os.environ.get("MUON_BACKEND_STEPS", 5)) - muon_momentum_warmup_start = float(os.environ.get("MUON_MOMENTUM_WARMUP_START", 0.85)) - muon_momentum_warmup_steps = int(os.environ.get("MUON_MOMENTUM_WARMUP_STEPS", 500)) - beta1 = float(os.environ.get("BETA1", 0.9)) - beta2 = float(os.environ.get("BETA2", 0.95)) - adam_eps = float(os.environ.get("ADAM_EPS", 1e-8)) - grad_clip_norm = float(os.environ.get("GRAD_CLIP_NORM", 0.0)) - -# ----------------------------- -# MUON OPTIMIZER -# ----------------------------- -# -# As borrowed from modded-nanogpt -# Background on Muon: https://kellerjordan.github.io/posts/muon/ - -def zeropower_via_newtonschulz5(G: Tensor, steps: int = 10, eps: float = 1e-7) -> Tensor: - # Orthogonalize a 2D update matrix with a fast Newton-Schulz iteration. - # Muon uses this to normalize matrix-shaped gradients before applying them. - a, b, c = (3.4445, -4.7750, 2.0315) - X = G.bfloat16() - X /= X.norm() + eps - transposed = G.size(0) > G.size(1) - if transposed: - X = X.T - for _ in range(steps): - A = X @ X.T - B = b * A + c * A @ A - X = a * X + B @ X - return X.T if transposed else X - - -class Muon(torch.optim.Optimizer): - def __init__(self, params, lr: float, momentum: float, backend_steps: int, nesterov: bool = True): - super().__init__( - params, - dict(lr=lr, momentum=momentum, backend_steps=backend_steps, nesterov=nesterov), - ) - - @torch.no_grad() - def step(self, closure=None): - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - distributed = dist.is_available() and dist.is_initialized() - world_size = dist.get_world_size() if distributed else 1 - rank = dist.get_rank() if distributed else 0 - - for group in self.param_groups: - params = group["params"] - if not params: - continue - lr = group["lr"] - momentum = group["momentum"] - backend_steps = group["backend_steps"] - nesterov = group["nesterov"] - - total_params = sum(int(p.numel()) for p in params) - updates_flat = torch.zeros(total_params, device=params[0].device, dtype=torch.bfloat16) - - curr = 0 - for i, p in enumerate(params): - if i % world_size == rank and p.grad is not None: - g = p.grad - state = self.state[p] - if "momentum_buffer" not in state: - state["momentum_buffer"] = torch.zeros_like(g) - buf = state["momentum_buffer"] - buf.mul_(momentum).add_(g) - if nesterov: - g = g.add(buf, alpha=momentum) - g = zeropower_via_newtonschulz5(g, steps=backend_steps) - # Scale correction from Muon reference implementations. - g *= max(1, g.size(0) / g.size(1)) ** 0.5 - updates_flat[curr : curr + p.numel()] = g.reshape(-1) - curr += p.numel() - - if distributed: - dist.all_reduce(updates_flat, op=dist.ReduceOp.SUM) - - curr = 0 - for p in params: - g = updates_flat[curr : curr + p.numel()].view_as(p).to(dtype=p.dtype) - p.add_(g, alpha=-lr) - curr += p.numel() - - return loss - - -# ----------------------------- -# TOKENIZER-AGNOSTIC EVALUATION SETUP -# ----------------------------- -# -# It's common for small models have a large fraction of their parameters be embeddings, since the 2 * d_model * d_vocab vectors can be gigantic. -# Instead of locking the tokenizer, we let you bring your own and calculate our validation metrics on the average compression of the validation set. -# We calculate BPB (bits-per-byte) instead of validation loss, so we need methods to count the number of bits per token in the tokenizer. -# Note: Submissions that edit the tokenizer will be examined more carefully, since screwing this up might unjustly improve your score. - -def build_sentencepiece_luts( - sp: spm.SentencePieceProcessor, vocab_size: int, device: torch.device -) -> tuple[Tensor, Tensor, Tensor]: - sp_vocab_size = int(sp.vocab_size()) - table_size = max(sp_vocab_size, vocab_size) - base_bytes_np = np.zeros((table_size,), dtype=np.int16) - has_leading_space_np = np.zeros((table_size,), dtype=np.bool_) - is_boundary_token_np = np.ones((table_size,), dtype=np.bool_) - for token_id in range(sp_vocab_size): - if sp.is_control(token_id) or sp.is_unknown(token_id) or sp.is_unused(token_id): - continue - is_boundary_token_np[token_id] = False - if sp.is_byte(token_id): - base_bytes_np[token_id] = 1 - continue - piece = sp.id_to_piece(token_id) - if piece.startswith("▁"): - has_leading_space_np[token_id] = True - piece = piece[1:] - base_bytes_np[token_id] = len(piece.encode("utf-8")) - return ( - torch.tensor(base_bytes_np, dtype=torch.int16, device=device), - torch.tensor(has_leading_space_np, dtype=torch.bool, device=device), - torch.tensor(is_boundary_token_np, dtype=torch.bool, device=device), - ) - - -def load_validation_tokens(pattern: str, seq_len: int) -> Tensor: - files = [Path(p) for p in sorted(glob.glob(pattern))] - if not files: - raise FileNotFoundError(f"No files found for pattern: {pattern}") - # The export pipeline writes the fixed first-50k-doc validation set to fineweb_val_*. - tokens = torch.cat([load_data_shard(file) for file in files]).contiguous() - usable = ((tokens.numel() - 1) // seq_len) * seq_len - if usable <= 0: - raise ValueError(f"Validation split is too short for TRAIN_SEQ_LEN={seq_len}") - return tokens[: usable + 1] - - -def eval_val( - args: Hyperparameters, - model: nn.Module, - rank: int, - world_size: int, - device: torch.device, - grad_accum_steps: int, - val_tokens: Tensor, - base_bytes_lut: Tensor, - has_leading_space_lut: Tensor, - is_boundary_token_lut: Tensor, -) -> tuple[float, float]: - # Validation computes two metrics: - # - val_loss: token cross-entropy (natural log) - # - val_bpb: tokenizer-agnostic compression metric used by the challenge - local_batch_tokens = args.val_batch_size // (world_size * grad_accum_steps) - if local_batch_tokens < args.train_seq_len: - raise ValueError( - "VAL_BATCH_SIZE must provide at least one sequence per rank; " - f"got VAL_BATCH_SIZE={args.val_batch_size}, WORLD_SIZE={world_size}, " - f"GRAD_ACCUM_STEPS={grad_accum_steps}, TRAIN_SEQ_LEN={args.train_seq_len}" - ) - local_batch_seqs = local_batch_tokens // args.train_seq_len - total_seqs = (val_tokens.numel() - 1) // args.train_seq_len - seq_start = (total_seqs * rank) // world_size - seq_end = (total_seqs * (rank + 1)) // world_size - val_loss_sum = torch.zeros((), device=device, dtype=torch.float64) - val_token_count = torch.zeros((), device=device, dtype=torch.float64) - val_byte_count = torch.zeros((), device=device, dtype=torch.float64) - - model.eval() - with torch.inference_mode(): - for batch_seq_start in range(seq_start, seq_end, local_batch_seqs): - batch_seq_end = min(batch_seq_start + local_batch_seqs, seq_end) - raw_start = batch_seq_start * args.train_seq_len - raw_end = batch_seq_end * args.train_seq_len + 1 - local = val_tokens[raw_start:raw_end].to(device=device, dtype=torch.int64, non_blocking=True) - x = local[:-1].reshape(-1, args.train_seq_len) - y = local[1:].reshape(-1, args.train_seq_len) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - batch_loss = model(x, y).detach() - batch_token_count = float(y.numel()) - val_loss_sum += batch_loss.to(torch.float64) * batch_token_count - val_token_count += batch_token_count - prev_ids = x.reshape(-1) - tgt_ids = y.reshape(-1) - token_bytes = base_bytes_lut[tgt_ids].to(dtype=torch.int16) - token_bytes += (has_leading_space_lut[tgt_ids] & ~is_boundary_token_lut[prev_ids]).to(dtype=torch.int16) - val_byte_count += token_bytes.to(torch.float64).sum() - - if dist.is_available() and dist.is_initialized(): - dist.all_reduce(val_loss_sum, op=dist.ReduceOp.SUM) - dist.all_reduce(val_token_count, op=dist.ReduceOp.SUM) - dist.all_reduce(val_byte_count, op=dist.ReduceOp.SUM) - - val_loss = val_loss_sum / val_token_count - bits_per_token = val_loss.item() / math.log(2.0) - tokens_per_byte = val_token_count.item() / val_byte_count.item() - model.train() - return float(val_loss.item()), float(bits_per_token * tokens_per_byte) - -# ----------------------------- -# POST-TRAINING QUANTIZATION -# ----------------------------- -# -# It's silly to export our model, which is trained in bf16 and fp32, at that same precision. -# Instead, we get approximately the same model (with a small hit) by quantizing the model to int8 & zlib compressing. -# We can then decompress the model and run in higher precision for evaluation, after closing in under the size limit. - -CONTROL_TENSOR_NAME_PATTERNS = tuple( - pattern - for pattern in os.environ.get( - "CONTROL_TENSOR_NAME_PATTERNS", - "attn_scale,attn_scales,mlp_scale,mlp_scales,resid_mix,resid_mixes,q_gain,skip_weight,skip_weights", - ).split(",") - if pattern -) -INT8_KEEP_FLOAT_FP32_NAME_PATTERNS = tuple( - pattern - for pattern in os.environ.get( - "INT8_KEEP_FLOAT_FP32_NAME_PATTERNS", - ",".join(CONTROL_TENSOR_NAME_PATTERNS), - ).split(",") - if pattern -) -INT8_KEEP_FLOAT_MAX_NUMEL = 65_536 -INT8_KEEP_FLOAT_STORE_DTYPE = torch.float16 -INT8_PER_ROW_SCALE_DTYPE = torch.float16 -INT8_CLIP_PERCENTILE = 99.99984 -INT8_CLIP_Q = INT8_CLIP_PERCENTILE / 100.0 - -def tensor_nbytes(t: Tensor) -> int: - return int(t.numel()) * int(t.element_size()) - -def keep_float_tensor(name: str, t: Tensor, passthrough_orig_dtypes: dict[str, str]) -> Tensor: - if any(pattern in name for pattern in INT8_KEEP_FLOAT_FP32_NAME_PATTERNS): - return t.float().contiguous() - if t.dtype in {torch.float32, torch.bfloat16}: - passthrough_orig_dtypes[name] = str(t.dtype).removeprefix("torch.") - return t.to(dtype=INT8_KEEP_FLOAT_STORE_DTYPE).contiguous() - return t - -def quantize_float_tensor(t: Tensor) -> tuple[Tensor, Tensor]: - t32 = t.float() - if t32.ndim == 2: - # Matrices get one scale per row, which usually tracks output-channel - # ranges much better than a single tensor-wide scale. - clip_abs = ( - torch.quantile(t32.abs(), INT8_CLIP_Q, dim=1) - if t32.numel() - else torch.empty((t32.shape[0],), dtype=torch.float32) - ) - clipped = torch.maximum(torch.minimum(t32, clip_abs[:, None]), -clip_abs[:, None]) - scale = (clip_abs / 127.0).clamp_min(1.0 / 127.0) - q = torch.clamp(torch.round(clipped / scale[:, None]), -127, 127).to(torch.int8).contiguous() - return q, scale.to(dtype=INT8_PER_ROW_SCALE_DTYPE).contiguous() - - # Vectors / scalars use a simpler per-tensor scale. - clip_abs = float(torch.quantile(t32.abs().flatten(), INT8_CLIP_Q).item()) if t32.numel() else 0.0 - scale = torch.tensor(clip_abs / 127.0 if clip_abs > 0 else 1.0, dtype=torch.float32) - q = torch.clamp(torch.round(torch.clamp(t32, -clip_abs, clip_abs) / scale), -127, 127).to(torch.int8).contiguous() - return q, scale - -def quantize_state_dict_int8(state_dict: dict[str, Tensor]): - # Single supported clean-script export format: - # - per-row int8 for 2D float tensors - # - per-tensor int8 for other float tensors - # - exact passthrough for non-floats - # - passthrough for small float tensors, stored as fp16 to save bytes - quantized: dict[str, Tensor] = {} - scales: dict[str, Tensor] = {} - dtypes: dict[str, str] = {} - passthrough: dict[str, Tensor] = {} - passthrough_orig_dtypes: dict[str, str] = {} - qmeta: dict[str, dict[str, object]] = {} - stats = dict.fromkeys( - ("param_count", "num_tensors", "num_float_tensors", "num_nonfloat_tensors", "baseline_tensor_bytes", "int8_payload_bytes"), - 0, - ) - - for name, tensor in state_dict.items(): - t = tensor.detach().to("cpu").contiguous() - stats["param_count"] += int(t.numel()) - stats["num_tensors"] += 1 - stats["baseline_tensor_bytes"] += tensor_nbytes(t) - - if not t.is_floating_point(): - stats["num_nonfloat_tensors"] += 1 - passthrough[name] = t - stats["int8_payload_bytes"] += tensor_nbytes(t) - continue - - # Small float tensors are cheap enough to keep directly. We still downcast - # fp32/bf16 passthrough tensors to fp16 so metadata does not dominate size. - if t.numel() <= INT8_KEEP_FLOAT_MAX_NUMEL: - kept = keep_float_tensor(name, t, passthrough_orig_dtypes) - passthrough[name] = kept - stats["int8_payload_bytes"] += tensor_nbytes(kept) - continue - - stats["num_float_tensors"] += 1 - q, s = quantize_float_tensor(t) - if s.ndim > 0: - qmeta[name] = {"scheme": "per_row", "axis": 0} - quantized[name] = q - scales[name] = s - dtypes[name] = str(t.dtype).removeprefix("torch.") - stats["int8_payload_bytes"] += tensor_nbytes(q) + tensor_nbytes(s) - - obj: dict[str, object] = { - "__quant_format__": "int8_clean_per_row_v1", - "quantized": quantized, - "scales": scales, - "dtypes": dtypes, - "passthrough": passthrough, - } - if qmeta: - obj["qmeta"] = qmeta - if passthrough_orig_dtypes: - obj["passthrough_orig_dtypes"] = passthrough_orig_dtypes - return obj, stats - -def dequantize_state_dict_int8(obj: dict[str, object]) -> dict[str, Tensor]: - out: dict[str, Tensor] = {} - qmeta = obj.get("qmeta", {}) - passthrough_orig_dtypes = obj.get("passthrough_orig_dtypes", {}) - for name, q in obj["quantized"].items(): - dtype = getattr(torch, obj["dtypes"][name]) - s = obj["scales"][name] - if qmeta.get(name, {}).get("scheme") == "per_row" or s.ndim > 0: - s = s.to(dtype=torch.float32) - # Broadcast the saved row scale back across trailing dimensions. - out[name] = (q.float() * s.view(q.shape[0], *([1] * (q.ndim - 1)))).to(dtype=dtype).contiguous() - else: - scale = float(s.item()) - out[name] = (q.float() * scale).to(dtype=dtype).contiguous() - for name, t in obj["passthrough"].items(): - # Restore small tensors, undoing the temporary fp16 storage cast if needed. - out_t = t.detach().to("cpu").contiguous() - orig_dtype = passthrough_orig_dtypes.get(name) - if isinstance(orig_dtype, str): - out_t = out_t.to(dtype=getattr(torch, orig_dtype)).contiguous() - out[name] = out_t - return out - - -# ----------------------------- -# DATA LOADING -# ----------------------------- - -def load_data_shard(file: Path) -> Tensor: - header_bytes = 256 * np.dtype(" None: - self.file_idx = (self.file_idx + 1) % len(self.files) - self.tokens = load_data_shard(self.files[self.file_idx]) - self.pos = 0 - - def take(self, n: int) -> Tensor: - chunks: list[Tensor] = [] - remaining = n - while remaining > 0: - avail = self.tokens.numel() - self.pos - if avail <= 0: - self._advance_file() - continue - k = min(remaining, avail) - chunks.append(self.tokens[self.pos : self.pos + k]) - self.pos += k - remaining -= k - return chunks[0] if len(chunks) == 1 else torch.cat(chunks) - - -class DistributedTokenLoader: - # Each call consumes a contiguous chunk from the shared token stream, then slices out - # one disjoint span per rank. The extra "+1" token lets us build (x, y) by shifting. - def __init__(self, pattern: str, rank: int, world_size: int, device: torch.device): - self.rank = rank - self.world_size = world_size - self.device = device - self.stream = TokenStream(pattern) - - def next_batch(self, global_tokens: int, seq_len: int, grad_accum_steps: int) -> tuple[Tensor, Tensor]: - local_tokens = global_tokens // (self.world_size * grad_accum_steps) - per_rank_span = local_tokens + 1 - chunk = self.stream.take(per_rank_span * self.world_size) - start = self.rank * per_rank_span - local = chunk[start : start + per_rank_span].to(dtype=torch.int64) - x = local[:-1].reshape(-1, seq_len) - y = local[1:].reshape(-1, seq_len) - return x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True) - -# ----------------------------- -# TRANSFORMER MODULES -# ----------------------------- - -class RMSNorm(nn.Module): - def __init__(self, eps: float | None = None): - super().__init__() - self.eps = eps - - def forward(self, x: Tensor) -> Tensor: - return F.rms_norm(x, (x.size(-1),), eps=self.eps) - - -class CastedLinear(nn.Linear): - # Keep weights in fp32 for optimizer/state quality, cast at matmul time for bf16 compute. - def forward(self, x: Tensor) -> Tensor: - bias = self.bias.to(x.dtype) if self.bias is not None else None - return F.linear(x, self.weight.to(x.dtype), bias) - - -def restore_low_dim_params_to_fp32(module: nn.Module) -> None: - # Keep small/control parameters in fp32 even when the model body runs in bf16. - with torch.no_grad(): - for name, param in module.named_parameters(): - if (param.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS)) and param.dtype != torch.float32: - param.data = param.data.float() - - -class Rotary(nn.Module): - # Caches cos/sin tables per sequence length on the current device. - def __init__(self, dim: int, base: float = 10000.0): - super().__init__() - inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self._seq_len_cached = 0 - self._cos_cached: Tensor | None = None - self._sin_cached: Tensor | None = None - - def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> tuple[Tensor, Tensor]: - if ( - self._cos_cached is None - or self._sin_cached is None - or self._seq_len_cached != seq_len - or self._cos_cached.device != device - ): - t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) - freqs = torch.outer(t, self.inv_freq.to(device)) - self._cos_cached = freqs.cos()[None, None, :, :] - self._sin_cached = freqs.sin()[None, None, :, :] - self._seq_len_cached = seq_len - return self._cos_cached.to(dtype=dtype), self._sin_cached.to(dtype=dtype) - - -def apply_rotary_emb(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: - half = x.size(-1) // 2 - x1, x2 = x[..., :half], x[..., half:] - return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1) - - -class CausalSelfAttention(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - num_kv_heads: int, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - if dim % num_heads != 0: - raise ValueError("model_dim must be divisible by num_heads") - if num_heads % num_kv_heads != 0: - raise ValueError("num_heads must be divisible by num_kv_heads") - self.num_heads = num_heads - self.num_kv_heads = num_kv_heads - self.head_dim = dim // num_heads - if self.head_dim % 2 != 0: - raise ValueError("head_dim must be even for RoPE") - kv_dim = self.num_kv_heads * self.head_dim - self.c_q = CastedLinear(dim, dim, bias=False) - self.c_k = CastedLinear(dim, kv_dim, bias=False) - self.c_v = CastedLinear(dim, kv_dim, bias=False) - self.proj = CastedLinear(dim, dim, bias=False) - self.proj._zero_init = True - self.q_gain = nn.Parameter(torch.full((num_heads,), qk_gain_init, dtype=torch.float32)) - self.rotary = Rotary(self.head_dim, base=rope_base) - - def forward(self, x: Tensor) -> Tensor: - bsz, seqlen, dim = x.shape - q = self.c_q(x).reshape(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2) - k = self.c_k(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - v = self.c_v(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - q = F.rms_norm(q, (q.size(-1),)) - k = F.rms_norm(k, (k.size(-1),)) - cos, sin = self.rotary(seqlen, x.device, q.dtype) - q = apply_rotary_emb(q, cos, sin) - k = apply_rotary_emb(k, cos, sin) - q = q * self.q_gain.to(dtype=q.dtype)[None, :, None, None] - y = F.scaled_dot_product_attention( - q, - k, - v, - attn_mask=None, - is_causal=True, - enable_gqa=(self.num_kv_heads != self.num_heads), - ) - y = y.transpose(1, 2).contiguous().reshape(bsz, seqlen, dim) - return self.proj(y) - - -class MLP(nn.Module): - # relu^2 MLP from the original modded-nanogpt setup - def __init__(self, dim: int, mlp_mult: int): - super().__init__() - hidden = mlp_mult * dim - self.fc = CastedLinear(dim, hidden, bias=False) - self.proj = CastedLinear(hidden, dim, bias=False) - self.proj._zero_init = True - - def forward(self, x: Tensor) -> Tensor: - x = torch.relu(self.fc(x)) - return self.proj(x.square()) - - -class Block(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - num_kv_heads: int, - mlp_mult: int, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - self.attn_norm = RMSNorm() - self.mlp_norm = RMSNorm() - self.attn = CausalSelfAttention(dim, num_heads, num_kv_heads, rope_base, qk_gain_init) - self.mlp = MLP(dim, mlp_mult) - self.attn_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) - self.mlp_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) - self.resid_mix = nn.Parameter(torch.stack((torch.ones(dim), torch.zeros(dim))).float()) - - def forward(self, x: Tensor, x0: Tensor) -> Tensor: - mix = self.resid_mix.to(dtype=x.dtype) - x = mix[0][None, None, :] * x + mix[1][None, None, :] * x0 - attn_out = self.attn(self.attn_norm(x)) - x = x + self.attn_scale.to(dtype=x.dtype)[None, None, :] * attn_out - x = x + self.mlp_scale.to(dtype=x.dtype)[None, None, :] * self.mlp(self.mlp_norm(x)) - return x - - -class GPT(nn.Module): - def __init__( - self, - vocab_size: int, - num_layers: int, - model_dim: int, - num_heads: int, - num_kv_heads: int, - mlp_mult: int, - tie_embeddings: bool, - tied_embed_init_std: float, - logit_softcap: float, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - if logit_softcap <= 0.0: - raise ValueError(f"logit_softcap must be positive, got {logit_softcap}") - self.tie_embeddings = tie_embeddings - self.tied_embed_init_std = tied_embed_init_std - self.logit_softcap = logit_softcap - self.tok_emb = nn.Embedding(vocab_size, model_dim) - self.num_encoder_layers = num_layers // 2 - self.num_decoder_layers = num_layers - self.num_encoder_layers - self.num_skip_weights = min(self.num_encoder_layers, self.num_decoder_layers) - self.skip_weights = nn.Parameter(torch.ones(self.num_skip_weights, model_dim, dtype=torch.float32)) - self.blocks = nn.ModuleList( - [ - Block( - model_dim, - num_heads, - num_kv_heads, - mlp_mult, - rope_base, - qk_gain_init, - ) - for i in range(num_layers) - ] - ) - self.final_norm = RMSNorm() - self.lm_head = None if tie_embeddings else CastedLinear(model_dim, vocab_size, bias=False) - if self.lm_head is not None: - self.lm_head._zero_init = True - self._init_weights() - - def _init_weights(self) -> None: - if self.tie_embeddings: - nn.init.normal_(self.tok_emb.weight, mean=0.0, std=self.tied_embed_init_std) - for module in self.modules(): - if isinstance(module, nn.Linear) and getattr(module, "_zero_init", False): - nn.init.zeros_(module.weight) - - def forward(self, input_ids: Tensor, target_ids: Tensor) -> Tensor: - x = self.tok_emb(input_ids) - x = F.rms_norm(x, (x.size(-1),)) - x0 = x - skips: list[Tensor] = [] - - # First half stores skips; second half reuses them in reverse order. - for i in range(self.num_encoder_layers): - x = self.blocks[i](x, x0) - skips.append(x) - for i in range(self.num_decoder_layers): - if skips: - x = x + self.skip_weights[i].to(dtype=x.dtype)[None, None, :] * skips.pop() - x = self.blocks[self.num_encoder_layers + i](x, x0) - - x = self.final_norm(x).reshape(-1, x.size(-1)) - targets = target_ids.reshape(-1) - if self.tie_embeddings: - logits_proj = F.linear(x, self.tok_emb.weight) - else: - if self.lm_head is None: - raise RuntimeError("lm_head is required when tie_embeddings=False") - logits_proj = self.lm_head(x) - logits = self.logit_softcap * torch.tanh(logits_proj / self.logit_softcap) - return F.cross_entropy(logits.float(), targets, reduction="mean") - - -# ----------------------------- -# TRAINING -# ----------------------------- - -def main() -> None: - global zeropower_via_newtonschulz5 - - code = Path(__file__).read_text(encoding="utf-8") - args = Hyperparameters() - zeropower_via_newtonschulz5 = torch.compile(zeropower_via_newtonschulz5) - - # ----------------------------- - # DISTRIBUTED + CUDA SETUP - # ----------------------------- - - distributed = "RANK" in os.environ and "WORLD_SIZE" in os.environ - rank = int(os.environ.get("RANK", "0")) - world_size = int(os.environ.get("WORLD_SIZE", "1")) - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - if world_size <= 0: - raise ValueError(f"WORLD_SIZE must be positive, got {world_size}") - if 8 % world_size != 0: - raise ValueError(f"WORLD_SIZE={world_size} must divide 8 so grad_accum_steps stays integral") - grad_accum_steps = 8 // world_size - grad_scale = 1.0 / grad_accum_steps - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required") - device = torch.device("cuda", local_rank) - torch.cuda.set_device(device) - if distributed: - dist.init_process_group(backend="nccl", device_id=device) - dist.barrier() - master_process = rank == 0 - - # Fast math knobs - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True - from torch.backends.cuda import enable_cudnn_sdp, enable_flash_sdp, enable_math_sdp, enable_mem_efficient_sdp - - enable_cudnn_sdp(False) - enable_flash_sdp(True) - enable_mem_efficient_sdp(False) - enable_math_sdp(False) - - logfile = None - if master_process: - os.makedirs("logs", exist_ok=True) - logfile = f"logs/{args.run_id}.txt" - print(logfile) - - def log0(msg: str, console: bool = True) -> None: - if not master_process: - return - if console: - print(msg) - if logfile is not None: - with open(logfile, "a", encoding="utf-8") as f: - print(msg, file=f) - - log0(code, console=False) - log0("=" * 100, console=False) - log0(f"Running Python {sys.version}", console=False) - log0(f"Running PyTorch {torch.__version__}", console=False) - log0( - subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False).stdout, - console=False, - ) - log0("=" * 100, console=False) - - # ----------------------------- - # TOKENIZER + VALIDATION METRIC SETUP - # ----------------------------- - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed_all(args.seed) - - if not args.tokenizer_path.endswith(".model"): - raise ValueError(f"Script only setup for SentencePiece .model file: {args.tokenizer_path}") - sp = spm.SentencePieceProcessor(model_file=args.tokenizer_path) - if int(sp.vocab_size()) != args.vocab_size: - raise ValueError( - f"VOCAB_SIZE={args.vocab_size} does not match tokenizer vocab_size={int(sp.vocab_size())}" - ) - dataset_dir = Path(args.data_path).resolve() - actual_train_files = len(list(dataset_dir.glob("fineweb_train_*.bin"))) - val_tokens = load_validation_tokens(args.val_files, args.train_seq_len) - base_bytes_lut, has_leading_space_lut, is_boundary_token_lut = build_sentencepiece_luts( - sp, args.vocab_size, device - ) - log0(f"val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path={args.tokenizer_path}") - log0(f"train_loader:dataset:{dataset_dir.name} train_shards:{actual_train_files}") - log0(f"val_loader:shards pattern={args.val_files} tokens:{val_tokens.numel() - 1}") - - # ----------------------------- - # MODEL + OPTIMIZER SETUP - # ----------------------------- - - base_model = GPT( - vocab_size=args.vocab_size, - num_layers=args.num_layers, - model_dim=args.model_dim, - num_heads=args.num_heads, - num_kv_heads=args.num_kv_heads, - mlp_mult=args.mlp_mult, - tie_embeddings=args.tie_embeddings, - tied_embed_init_std=args.tied_embed_init_std, - logit_softcap=args.logit_softcap, - rope_base=args.rope_base, - qk_gain_init=args.qk_gain_init, - ).to(device).bfloat16() - for module in base_model.modules(): - if isinstance(module, CastedLinear): - module.float() - restore_low_dim_params_to_fp32(base_model) - compiled_model = torch.compile(base_model, dynamic=False, fullgraph=True) - model: nn.Module = DDP(compiled_model, device_ids=[local_rank], broadcast_buffers=False) if distributed else compiled_model - - # Optimizer split: - # - token embedding (Adam) uses EMBED_LR - # - untied lm_head (Adam) uses HEAD_LR - # - matrix params in transformer blocks use MATRIX_LR via Muon - # - vectors/scalars use SCALAR_LR via Adam - block_named_params = list(base_model.blocks.named_parameters()) - matrix_params = [ - p - for name, p in block_named_params - if p.ndim == 2 and not any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - scalar_params = [ - p - for name, p in block_named_params - if p.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - if base_model.skip_weights.numel() > 0: - scalar_params.append(base_model.skip_weights) - token_lr = args.tied_embed_lr if args.tie_embeddings else args.embed_lr - optimizer_tok = torch.optim.Adam( - [{"params": [base_model.tok_emb.weight], "lr": token_lr, "base_lr": token_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizer_muon = Muon( - matrix_params, - lr=args.matrix_lr, - momentum=args.muon_momentum, - backend_steps=args.muon_backend_steps, - ) - for group in optimizer_muon.param_groups: - group["base_lr"] = args.matrix_lr - optimizer_scalar = torch.optim.Adam( - [{"params": scalar_params, "lr": args.scalar_lr, "base_lr": args.scalar_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers: list[torch.optim.Optimizer] = [optimizer_tok, optimizer_muon, optimizer_scalar] - if base_model.lm_head is not None: - optimizer_head = torch.optim.Adam( - [{"params": [base_model.lm_head.weight], "lr": args.head_lr, "base_lr": args.head_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers.insert(1, optimizer_head) - - n_params = sum(p.numel() for p in base_model.parameters()) - log0(f"model_params:{n_params}") - log0(f"world_size:{world_size} grad_accum_steps:{grad_accum_steps}") - log0("sdp_backends:cudnn=False flash=True mem_efficient=False math=False") - log0(f"attention_mode:gqa num_heads:{args.num_heads} num_kv_heads:{args.num_kv_heads}") - log0( - f"tie_embeddings:{args.tie_embeddings} embed_lr:{token_lr} " - f"head_lr:{args.head_lr if base_model.lm_head is not None else 0.0} " - f"matrix_lr:{args.matrix_lr} scalar_lr:{args.scalar_lr}" - ) - log0( - f"train_batch_tokens:{args.train_batch_tokens} train_seq_len:{args.train_seq_len} " - f"iterations:{args.iterations} warmup_steps:{args.warmup_steps} " - f"max_wallclock_seconds:{args.max_wallclock_seconds:.3f}" - ) - log0(f"seed:{args.seed}") - - # ----------------------------- - # DATA LOADER & MODEL WARMUP - # ----------------------------- - - train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device) - - def zero_grad_all() -> None: - for opt in optimizers: - opt.zero_grad(set_to_none=True) - - max_wallclock_ms = 1000.0 * args.max_wallclock_seconds if args.max_wallclock_seconds > 0 else None - - def lr_mul(step: int, elapsed_ms: float) -> float: - if args.warmdown_iters <= 0: - return 1.0 - if max_wallclock_ms is None: - warmdown_start = max(args.iterations - args.warmdown_iters, 0) - return max((args.iterations - step) / max(args.warmdown_iters, 1), 0.0) if warmdown_start <= step < args.iterations else 1.0 - step_ms = elapsed_ms / max(step, 1) - warmdown_ms = args.warmdown_iters * step_ms - remaining_ms = max(max_wallclock_ms - elapsed_ms, 0.0) - return remaining_ms / max(warmdown_ms, 1e-9) if remaining_ms <= warmdown_ms else 1.0 - - # Warmup primes the compiled forward/backward/optimizer paths, then we restore the - # initial weights/optimizer state so measured training starts from the true init. - if args.warmup_steps > 0: - initial_model_state = {name: tensor.detach().cpu().clone() for name, tensor in base_model.state_dict().items()} - initial_optimizer_states = [copy.deepcopy(opt.state_dict()) for opt in optimizers] - model.train() - for warmup_step in range(args.warmup_steps): - zero_grad_all() - for micro_step in range(grad_accum_steps): - if distributed: - model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 - x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - warmup_loss = model(x, y) - (warmup_loss * grad_scale).backward() - for opt in optimizers: - opt.step() - zero_grad_all() - if args.warmup_steps <= 20 or (warmup_step + 1) % 10 == 0 or warmup_step + 1 == args.warmup_steps: - log0(f"warmup_step:{warmup_step + 1}/{args.warmup_steps}") - base_model.load_state_dict(initial_model_state, strict=True) - for opt, state in zip(optimizers, initial_optimizer_states, strict=True): - opt.load_state_dict(state) - zero_grad_all() - if distributed: - model.require_backward_grad_sync = True - train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device) - - # ----------------------------- - # MAIN TRAINING LOOP - # ----------------------------- - - training_time_ms = 0.0 - stop_after_step: int | None = None - torch.cuda.synchronize() - t0 = time.perf_counter() - - step = 0 - while True: - last_step = step == args.iterations or (stop_after_step is not None and step >= stop_after_step) - - should_validate = last_step or (args.val_loss_every > 0 and step % args.val_loss_every == 0) - if should_validate: - torch.cuda.synchronize() - training_time_ms += 1000.0 * (time.perf_counter() - t0) - val_loss, val_bpb = eval_val( - args, - model, - rank, - world_size, - device, - grad_accum_steps, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - ) - log0( - f"step:{step}/{args.iterations} val_loss:{val_loss:.4f} val_bpb:{val_bpb:.4f} " - f"train_time:{training_time_ms:.0f}ms step_avg:{training_time_ms / max(step, 1):.2f}ms" - ) - torch.cuda.synchronize() - t0 = time.perf_counter() - - if last_step: - if stop_after_step is not None and step < args.iterations: - log0( - f"stopping_early: wallclock_cap train_time:{training_time_ms:.0f}ms " - f"step:{step}/{args.iterations}" - ) - break - - elapsed_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) - scale = lr_mul(step, elapsed_ms) - zero_grad_all() - train_loss = torch.zeros((), device=device) - for micro_step in range(grad_accum_steps): - if distributed: - model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 - x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - loss = model(x, y) - train_loss += loss.detach() - (loss * grad_scale).backward() - train_loss /= grad_accum_steps - - frac = min(step / args.muon_momentum_warmup_steps, 1.0) if args.muon_momentum_warmup_steps > 0 else 1.0 - muon_momentum = (1 - frac) * args.muon_momentum_warmup_start + frac * args.muon_momentum - for group in optimizer_muon.param_groups: - group["momentum"] = muon_momentum - - for opt in optimizers: - for group in opt.param_groups: - group["lr"] = group["base_lr"] * scale - - if args.grad_clip_norm > 0: - torch.nn.utils.clip_grad_norm_(base_model.parameters(), args.grad_clip_norm) - for opt in optimizers: - opt.step() - zero_grad_all() - - step += 1 - approx_training_time_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) - should_log_train = ( - args.train_log_every > 0 - and (step <= 10 or step % args.train_log_every == 0 or stop_after_step is not None) - ) - if should_log_train: - log0( - f"step:{step}/{args.iterations} train_loss:{train_loss.item():.4f} " - f"train_time:{approx_training_time_ms:.0f}ms step_avg:{approx_training_time_ms / step:.2f}ms" - ) - - # Needed to sync whether we've reached the wallclock cap. - reached_cap = max_wallclock_ms is not None and approx_training_time_ms >= max_wallclock_ms - if distributed and max_wallclock_ms is not None: - reached_cap_tensor = torch.tensor(int(reached_cap), device=device) - dist.all_reduce(reached_cap_tensor, op=dist.ReduceOp.MAX) - reached_cap = bool(reached_cap_tensor.item()) - if stop_after_step is None and reached_cap: - stop_after_step = step - - log0( - f"peak memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB " - f"reserved: {torch.cuda.max_memory_reserved() // 1024 // 1024} MiB" - ) - - # ----------------------------- - # SERIALIZATION + ROUNDTRIP VALIDATION - # ----------------------------- - # Save the raw state (useful for debugging/loading in PyTorch directly), then always produce - # the compressed int8+zlib artifact and validate the round-tripped weights. - - if master_process: - torch.save(base_model.state_dict(), "final_model.pt") - model_bytes = os.path.getsize("final_model.pt") - code_bytes = len(code.encode("utf-8")) - log0(f"Serialized model: {model_bytes} bytes") - log0(f"Code size: {code_bytes} bytes") - log0(f"Total submission size: {model_bytes + code_bytes} bytes") - - quant_obj, quant_stats = quantize_state_dict_int8(base_model.state_dict()) - quant_buf = io.BytesIO() - torch.save(quant_obj, quant_buf) - quant_raw = quant_buf.getvalue() - quant_blob = zlib.compress(quant_raw, level=9) - quant_raw_bytes = len(quant_raw) - if master_process: - with open("final_model.int8.ptz", "wb") as f: - f.write(quant_blob) - quant_file_bytes = os.path.getsize("final_model.int8.ptz") - code_bytes = len(code.encode("utf-8")) - ratio = quant_stats["baseline_tensor_bytes"] / max(quant_stats["int8_payload_bytes"], 1) - log0( - f"Serialized model int8+zlib: {quant_file_bytes} bytes " - f"(payload:{quant_stats['int8_payload_bytes']} raw_torch:{quant_raw_bytes} payload_ratio:{ratio:.2f}x)" - ) - log0(f"Total submission size int8+zlib: {quant_file_bytes + code_bytes} bytes") - - if distributed: - dist.barrier() - with open("final_model.int8.ptz", "rb") as f: - quant_blob_disk = f.read() - quant_state = torch.load(io.BytesIO(zlib.decompress(quant_blob_disk)), map_location="cpu") - base_model.load_state_dict(dequantize_state_dict_int8(quant_state), strict=True) - torch.cuda.synchronize() - t_qeval = time.perf_counter() - q_val_loss, q_val_bpb = eval_val( - args, - model, - rank, - world_size, - device, - grad_accum_steps, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - ) - torch.cuda.synchronize() - log0( - f"final_int8_zlib_roundtrip val_loss:{q_val_loss:.4f} val_bpb:{q_val_bpb:.4f} " - f"eval_time:{1000.0 * (time.perf_counter() - t_qeval):.0f}ms" - ) - log0(f"final_int8_zlib_roundtrip_exact val_loss:{q_val_loss:.8f} val_bpb:{q_val_bpb:.8f}") - - if distributed: - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/records/track_10min_16mb/2026-03-18_LongContextSeq2048/train_seed1338.log b/records/track_10min_16mb/2026-03-18_LongContextSeq2048/train_seed1338.log deleted file mode 100644 index 5851fe0a3c..0000000000 --- a/records/track_10min_16mb/2026-03-18_LongContextSeq2048/train_seed1338.log +++ /dev/null @@ -1,124 +0,0 @@ - -***************************************** -Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. -***************************************** -logs/seq2048_sxm28_seed1338_20260319a.txt -val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path=/root/parameter-golf-sxm28/data/tokenizers/fineweb_1024_bpe.model -train_loader:dataset:fineweb10B_sp1024 train_shards:80 -val_loader:shards pattern=/root/parameter-golf-sxm28/data/datasets/fineweb10B_sp1024/fineweb_val_*.bin tokens:62021632 -model_params:17059912 -world_size:8 grad_accum_steps:1 -sdp_backends:cudnn=False flash=True mem_efficient=False math=False -attention_mode:gqa num_heads:8 num_kv_heads:4 -tie_embeddings:True embed_lr:0.04 head_lr:0.0 matrix_lr:0.032 scalar_lr:0.032 -train_batch_tokens:524288 train_seq_len:2048 iterations:20000 warmup_steps:20 max_wallclock_seconds:600.000 -seed:1338 -warmup_step:1/20 -warmup_step:2/20 -warmup_step:3/20 -warmup_step:4/20 -warmup_step:5/20 -warmup_step:6/20 -warmup_step:7/20 -warmup_step:8/20 -warmup_step:9/20 -warmup_step:10/20 -warmup_step:11/20 -warmup_step:12/20 -warmup_step:13/20 -warmup_step:14/20 -warmup_step:15/20 -warmup_step:16/20 -warmup_step:17/20 -warmup_step:18/20 -warmup_step:19/20 -warmup_step:20/20 -step:0/20000 val_loss:6.9373 val_bpb:4.1086 train_time:0ms step_avg:0.02ms -step:1/20000 train_loss:6.9363 train_time:27ms step_avg:27.47ms -step:2/20000 train_loss:14.7813 train_time:73ms step_avg:36.26ms -step:3/20000 train_loss:8.1303 train_time:124ms step_avg:41.23ms -step:4/20000 train_loss:6.6459 train_time:175ms step_avg:43.65ms -step:5/20000 train_loss:6.9413 train_time:227ms step_avg:45.34ms -step:6/20000 train_loss:7.8068 train_time:278ms step_avg:46.34ms -step:7/20000 train_loss:6.8490 train_time:330ms step_avg:47.09ms -step:8/20000 train_loss:6.4809 train_time:381ms step_avg:47.61ms -step:9/20000 train_loss:6.2006 train_time:434ms step_avg:48.24ms -step:10/20000 train_loss:6.0502 train_time:483ms step_avg:48.35ms -step:200/20000 train_loss:2.7815 train_time:10297ms step_avg:51.48ms -step:400/20000 train_loss:2.3049 train_time:20646ms step_avg:51.62ms -step:600/20000 train_loss:2.4980 train_time:30990ms step_avg:51.65ms -step:800/20000 train_loss:2.2491 train_time:41345ms step_avg:51.68ms -step:1000/20000 train_loss:2.3361 train_time:51721ms step_avg:51.72ms -step:1000/20000 val_loss:2.2898 val_bpb:1.3562 train_time:51754ms step_avg:51.75ms -step:1200/20000 train_loss:2.3575 train_time:62108ms step_avg:51.76ms -step:1400/20000 train_loss:2.3813 train_time:72504ms step_avg:51.79ms -step:1600/20000 train_loss:2.0496 train_time:82894ms step_avg:51.81ms -step:1800/20000 train_loss:2.1676 train_time:93289ms step_avg:51.83ms -step:2000/20000 train_loss:2.2119 train_time:103684ms step_avg:51.84ms -step:2000/20000 val_loss:2.1937 val_bpb:1.2993 train_time:103717ms step_avg:51.86ms -step:2200/20000 train_loss:2.0320 train_time:114081ms step_avg:51.86ms -step:2400/20000 train_loss:2.1652 train_time:124471ms step_avg:51.86ms -step:2600/20000 train_loss:2.3821 train_time:134872ms step_avg:51.87ms -step:2800/20000 train_loss:2.1987 train_time:145264ms step_avg:51.88ms -step:3000/20000 train_loss:2.1912 train_time:155647ms step_avg:51.88ms -step:3000/20000 val_loss:2.1534 val_bpb:1.2754 train_time:155680ms step_avg:51.89ms -step:3200/20000 train_loss:2.1516 train_time:166037ms step_avg:51.89ms -step:3400/20000 train_loss:2.1199 train_time:176417ms step_avg:51.89ms -step:3600/20000 train_loss:2.0651 train_time:186795ms step_avg:51.89ms -step:3800/20000 train_loss:2.1694 train_time:197173ms step_avg:51.89ms -step:4000/20000 train_loss:2.1330 train_time:207554ms step_avg:51.89ms -step:4000/20000 val_loss:2.1292 val_bpb:1.2610 train_time:207587ms step_avg:51.90ms -step:4200/20000 train_loss:2.1284 train_time:217976ms step_avg:51.90ms -step:4400/20000 train_loss:2.0686 train_time:228351ms step_avg:51.90ms -step:4600/20000 train_loss:1.9371 train_time:238738ms step_avg:51.90ms -step:4800/20000 train_loss:2.2171 train_time:249109ms step_avg:51.90ms -step:5000/20000 train_loss:1.9744 train_time:259476ms step_avg:51.90ms -step:5000/20000 val_loss:2.1127 val_bpb:1.2512 train_time:259509ms step_avg:51.90ms -step:5200/20000 train_loss:2.1356 train_time:269848ms step_avg:51.89ms -step:5400/20000 train_loss:2.1527 train_time:280217ms step_avg:51.89ms -step:5600/20000 train_loss:2.1390 train_time:290578ms step_avg:51.89ms -step:5800/20000 train_loss:2.0944 train_time:300933ms step_avg:51.89ms -step:6000/20000 train_loss:2.1752 train_time:311294ms step_avg:51.88ms -step:6000/20000 val_loss:2.1026 val_bpb:1.2453 train_time:311327ms step_avg:51.89ms -step:6200/20000 train_loss:2.0458 train_time:321653ms step_avg:51.88ms -step:6400/20000 train_loss:2.1240 train_time:332019ms step_avg:51.88ms -step:6600/20000 train_loss:2.0830 train_time:342381ms step_avg:51.88ms -step:6800/20000 train_loss:2.1434 train_time:352738ms step_avg:51.87ms -step:7000/20000 train_loss:2.1907 train_time:363096ms step_avg:51.87ms -step:7000/20000 val_loss:2.0916 val_bpb:1.2388 train_time:363129ms step_avg:51.88ms -step:7200/20000 train_loss:2.1672 train_time:373450ms step_avg:51.87ms -step:7400/20000 train_loss:2.0842 train_time:383806ms step_avg:51.87ms -step:7600/20000 train_loss:1.9615 train_time:394163ms step_avg:51.86ms -step:7800/20000 train_loss:2.1113 train_time:404518ms step_avg:51.86ms -step:8000/20000 train_loss:2.0788 train_time:414870ms step_avg:51.86ms -step:8000/20000 val_loss:2.0826 val_bpb:1.2334 train_time:414903ms step_avg:51.86ms -step:8200/20000 train_loss:2.1505 train_time:425230ms step_avg:51.86ms -step:8400/20000 train_loss:2.0933 train_time:435626ms step_avg:51.86ms -step:8600/20000 train_loss:2.1070 train_time:445977ms step_avg:51.86ms -step:8800/20000 train_loss:2.0708 train_time:456329ms step_avg:51.86ms -step:9000/20000 train_loss:1.9882 train_time:466685ms step_avg:51.85ms -step:9000/20000 val_loss:2.0772 val_bpb:1.2302 train_time:466718ms step_avg:51.86ms -step:9200/20000 train_loss:2.0470 train_time:477091ms step_avg:51.86ms -step:9400/20000 train_loss:2.0941 train_time:487469ms step_avg:51.86ms -step:9600/20000 train_loss:2.1116 train_time:497817ms step_avg:51.86ms -step:9800/20000 train_loss:2.0202 train_time:508164ms step_avg:51.85ms -step:10000/20000 train_loss:2.0783 train_time:518510ms step_avg:51.85ms -step:10000/20000 val_loss:2.0723 val_bpb:1.2274 train_time:518543ms step_avg:51.85ms -step:10200/20000 train_loss:2.0342 train_time:528861ms step_avg:51.85ms -step:10400/20000 train_loss:2.0587 train_time:539208ms step_avg:51.85ms -step:10600/20000 train_loss:1.9323 train_time:549552ms step_avg:51.84ms -step:10800/20000 train_loss:2.1371 train_time:559906ms step_avg:51.84ms -step:11000/20000 train_loss:2.0567 train_time:570254ms step_avg:51.84ms -step:11000/20000 val_loss:2.0458 val_bpb:1.2116 train_time:570286ms step_avg:51.84ms -step:11200/20000 train_loss:2.0119 train_time:580608ms step_avg:51.84ms -step:11400/20000 train_loss:1.9920 train_time:590954ms step_avg:51.84ms -step:11575/20000 val_loss:2.0274 val_bpb:1.2007 train_time:600043ms step_avg:51.84ms -stopping_early: wallclock_cap train_time:600043ms step:11575/20000 -peak memory allocated: 10247 MiB reserved: 10312 MiB -Serialized model: 67224983 bytes -Code size: 47716 bytes -Total submission size: 67272699 bytes -Serialized model int8+zlib: 15813523 bytes (payload:17178912 raw_torch:17224025 payload_ratio:3.91x) -Total submission size int8+zlib: 15861239 bytes -final_int8_zlib_roundtrip val_loss:2.0366 val_bpb:1.2062 eval_time:1638ms -final_int8_zlib_roundtrip_exact val_loss:2.03657529 val_bpb:1.20617460 diff --git a/records/track_10min_16mb/2026-03-18_LongContextSeq2048/train_seed1339.log b/records/track_10min_16mb/2026-03-18_LongContextSeq2048/train_seed1339.log deleted file mode 100644 index 527d535a46..0000000000 --- a/records/track_10min_16mb/2026-03-18_LongContextSeq2048/train_seed1339.log +++ /dev/null @@ -1,124 +0,0 @@ - -***************************************** -Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. -***************************************** -logs/seq2048_sxm28_seed1339_20260319a.txt -val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path=/root/parameter-golf-sxm28/data/tokenizers/fineweb_1024_bpe.model -train_loader:dataset:fineweb10B_sp1024 train_shards:80 -val_loader:shards pattern=/root/parameter-golf-sxm28/data/datasets/fineweb10B_sp1024/fineweb_val_*.bin tokens:62021632 -model_params:17059912 -world_size:8 grad_accum_steps:1 -sdp_backends:cudnn=False flash=True mem_efficient=False math=False -attention_mode:gqa num_heads:8 num_kv_heads:4 -tie_embeddings:True embed_lr:0.04 head_lr:0.0 matrix_lr:0.032 scalar_lr:0.032 -train_batch_tokens:524288 train_seq_len:2048 iterations:20000 warmup_steps:20 max_wallclock_seconds:600.000 -seed:1339 -warmup_step:1/20 -warmup_step:2/20 -warmup_step:3/20 -warmup_step:4/20 -warmup_step:5/20 -warmup_step:6/20 -warmup_step:7/20 -warmup_step:8/20 -warmup_step:9/20 -warmup_step:10/20 -warmup_step:11/20 -warmup_step:12/20 -warmup_step:13/20 -warmup_step:14/20 -warmup_step:15/20 -warmup_step:16/20 -warmup_step:17/20 -warmup_step:18/20 -warmup_step:19/20 -warmup_step:20/20 -step:0/20000 val_loss:6.9372 val_bpb:4.1086 train_time:0ms step_avg:0.02ms -step:1/20000 train_loss:6.9375 train_time:27ms step_avg:27.41ms -step:2/20000 train_loss:14.8195 train_time:75ms step_avg:37.30ms -step:3/20000 train_loss:8.0864 train_time:126ms step_avg:41.96ms -step:4/20000 train_loss:6.5534 train_time:177ms step_avg:44.18ms -step:5/20000 train_loss:6.8329 train_time:228ms step_avg:45.62ms -step:6/20000 train_loss:7.6678 train_time:279ms step_avg:46.55ms -step:7/20000 train_loss:6.6834 train_time:330ms step_avg:47.21ms -step:8/20000 train_loss:6.3318 train_time:382ms step_avg:47.72ms -step:9/20000 train_loss:6.1684 train_time:433ms step_avg:48.12ms -step:10/20000 train_loss:6.0971 train_time:484ms step_avg:48.42ms -step:200/20000 train_loss:2.7828 train_time:10300ms step_avg:51.50ms -step:400/20000 train_loss:2.2900 train_time:20640ms step_avg:51.60ms -step:600/20000 train_loss:2.5016 train_time:30985ms step_avg:51.64ms -step:800/20000 train_loss:2.2388 train_time:41342ms step_avg:51.68ms -step:1000/20000 train_loss:2.3326 train_time:51710ms step_avg:51.71ms -step:1000/20000 val_loss:2.2851 val_bpb:1.3534 train_time:51743ms step_avg:51.74ms -step:1200/20000 train_loss:2.3498 train_time:62082ms step_avg:51.74ms -step:1400/20000 train_loss:2.3834 train_time:72461ms step_avg:51.76ms -step:1600/20000 train_loss:2.0438 train_time:82831ms step_avg:51.77ms -step:1800/20000 train_loss:2.1610 train_time:93204ms step_avg:51.78ms -step:2000/20000 train_loss:2.2097 train_time:103580ms step_avg:51.79ms -step:2000/20000 val_loss:2.1921 val_bpb:1.2983 train_time:103613ms step_avg:51.81ms -step:2200/20000 train_loss:2.0283 train_time:113950ms step_avg:51.80ms -step:2400/20000 train_loss:2.1558 train_time:124322ms step_avg:51.80ms -step:2600/20000 train_loss:2.3815 train_time:134697ms step_avg:51.81ms -step:2800/20000 train_loss:2.1964 train_time:145067ms step_avg:51.81ms -step:3000/20000 train_loss:2.1879 train_time:155433ms step_avg:51.81ms -step:3000/20000 val_loss:2.1535 val_bpb:1.2754 train_time:155467ms step_avg:51.82ms -step:3200/20000 train_loss:2.1499 train_time:165801ms step_avg:51.81ms -step:3400/20000 train_loss:2.1196 train_time:176179ms step_avg:51.82ms -step:3600/20000 train_loss:2.0663 train_time:186549ms step_avg:51.82ms -step:3800/20000 train_loss:2.1720 train_time:196916ms step_avg:51.82ms -step:4000/20000 train_loss:2.1340 train_time:207282ms step_avg:51.82ms -step:4000/20000 val_loss:2.1294 val_bpb:1.2611 train_time:207315ms step_avg:51.83ms -step:4200/20000 train_loss:2.1284 train_time:217691ms step_avg:51.83ms -step:4400/20000 train_loss:2.0666 train_time:228047ms step_avg:51.83ms -step:4600/20000 train_loss:1.9387 train_time:238412ms step_avg:51.83ms -step:4800/20000 train_loss:2.2201 train_time:248762ms step_avg:51.83ms -step:5000/20000 train_loss:1.9748 train_time:259127ms step_avg:51.83ms -step:5000/20000 val_loss:2.1132 val_bpb:1.2516 train_time:259160ms step_avg:51.83ms -step:5200/20000 train_loss:2.1342 train_time:269493ms step_avg:51.83ms -step:5400/20000 train_loss:2.1527 train_time:279859ms step_avg:51.83ms -step:5600/20000 train_loss:2.1413 train_time:290221ms step_avg:51.83ms -step:5800/20000 train_loss:2.0992 train_time:300579ms step_avg:51.82ms -step:6000/20000 train_loss:2.1790 train_time:310939ms step_avg:51.82ms -step:6000/20000 val_loss:2.1039 val_bpb:1.2460 train_time:310972ms step_avg:51.83ms -step:6200/20000 train_loss:2.0485 train_time:321291ms step_avg:51.82ms -step:6400/20000 train_loss:2.1251 train_time:331655ms step_avg:51.82ms -step:6600/20000 train_loss:2.0805 train_time:342011ms step_avg:51.82ms -step:6800/20000 train_loss:2.1480 train_time:352372ms step_avg:51.82ms -step:7000/20000 train_loss:2.1942 train_time:362734ms step_avg:51.82ms -step:7000/20000 val_loss:2.0925 val_bpb:1.2393 train_time:362767ms step_avg:51.82ms -step:7200/20000 train_loss:2.1633 train_time:373093ms step_avg:51.82ms -step:7400/20000 train_loss:2.0834 train_time:383453ms step_avg:51.82ms -step:7600/20000 train_loss:1.9632 train_time:393812ms step_avg:51.82ms -step:7800/20000 train_loss:2.1106 train_time:404251ms step_avg:51.83ms -step:8000/20000 train_loss:2.0791 train_time:414608ms step_avg:51.83ms -step:8000/20000 val_loss:2.0836 val_bpb:1.2340 train_time:414641ms step_avg:51.83ms -step:8200/20000 train_loss:2.1540 train_time:424958ms step_avg:51.82ms -step:8400/20000 train_loss:2.0970 train_time:435353ms step_avg:51.83ms -step:8600/20000 train_loss:2.1104 train_time:445701ms step_avg:51.83ms -step:8800/20000 train_loss:2.0694 train_time:456052ms step_avg:51.82ms -step:9000/20000 train_loss:1.9874 train_time:466402ms step_avg:51.82ms -step:9000/20000 val_loss:2.0785 val_bpb:1.2310 train_time:466435ms step_avg:51.83ms -step:9200/20000 train_loss:2.0479 train_time:476785ms step_avg:51.82ms -step:9400/20000 train_loss:2.0944 train_time:487193ms step_avg:51.83ms -step:9600/20000 train_loss:2.1114 train_time:497576ms step_avg:51.83ms -step:9800/20000 train_loss:2.0204 train_time:507942ms step_avg:51.83ms -step:10000/20000 train_loss:2.0788 train_time:518312ms step_avg:51.83ms -step:10000/20000 val_loss:2.0729 val_bpb:1.2277 train_time:518345ms step_avg:51.83ms -step:10200/20000 train_loss:2.0372 train_time:528681ms step_avg:51.83ms -step:10400/20000 train_loss:2.0577 train_time:539056ms step_avg:51.83ms -step:10600/20000 train_loss:1.9348 train_time:549410ms step_avg:51.83ms -step:10800/20000 train_loss:2.1410 train_time:559767ms step_avg:51.83ms -step:11000/20000 train_loss:2.0562 train_time:570117ms step_avg:51.83ms -step:11000/20000 val_loss:2.0472 val_bpb:1.2125 train_time:570149ms step_avg:51.83ms -step:11200/20000 train_loss:2.0166 train_time:580461ms step_avg:51.83ms -step:11400/20000 train_loss:1.9939 train_time:590808ms step_avg:51.83ms -step:11578/20000 val_loss:2.0286 val_bpb:1.2015 train_time:600051ms step_avg:51.83ms -stopping_early: wallclock_cap train_time:600051ms step:11578/20000 -peak memory allocated: 10247 MiB reserved: 10312 MiB -Serialized model: 67224983 bytes -Code size: 47716 bytes -Total submission size: 67272699 bytes -Serialized model int8+zlib: 15814036 bytes (payload:17178912 raw_torch:17224025 payload_ratio:3.91x) -Total submission size int8+zlib: 15861752 bytes -final_int8_zlib_roundtrip val_loss:2.0382 val_bpb:1.2072 eval_time:1640ms -final_int8_zlib_roundtrip_exact val_loss:2.03823779 val_bpb:1.20715923 diff --git a/records/track_10min_16mb/2026-03-19_SlidingWindowEval/README.md b/records/track_10min_16mb/2026-03-19_SlidingWindowEval/README.md deleted file mode 100644 index e4c8ef6fd4..0000000000 --- a/records/track_10min_16mb/2026-03-19_SlidingWindowEval/README.md +++ /dev/null @@ -1,78 +0,0 @@ -This record implements sliding window evaluation, showing that eval strategies alone can provide significant improvements. - -**Note on `train_gpt.py`:** The included script contains some unused experimental code paths (QAT, looped architectures) that are **all disabled by default** and were not active during this run. Only the sliding window evaluation code (`eval_val_sliding`, `forward_logits`, `EVAL_STRIDE`, `EVAL_BATCH_SEQS`) is used. The command below shows the exact invocation. - -## Key Idea: Sliding Window Evaluation - -The baseline evaluates by chopping the validation set into non-overlapping 1024-token chunks. The problem is that the first token in each chunk has zero context. On average, each token gets ~512 tokens of context. - -Sliding window evaluation uses overlapping windows with a configurable stride. With `EVAL_STRIDE=64` and `TRAIN_SEQ_LEN=1024`, each window advances by 64 tokens, but only the rightmost 64 tokens (which have 960+ tokens of context) are scored. Every token in the validation set is scored exactly once, but with near-maximum context. - -## Results - -| Metric | Naive Baseline | This Submission | -|---|---|---| -| Pre-quant val_bpb | 1.2172 | 1.2196 | -| **Post-quant val_bpb** | **1.2244** | **1.1925** | -| **Improvement** | — | **-0.0319** | -| Training steps | 13,780 | 13,450 | -| Eval time (8xH100) | ~16s | 70s | -| Artifact size | 15,863,489 bytes | 15,874,829 bytes | - -The pre-quant BPB is nearly identical (training is the same). The 0.032 improvement comes entirely from scoring tokens with richer context during evaluation. - -## Configuration - -Architecture and training are identical to the Naive Baseline: -- Layout: `VOCAB_SIZE=1024 NUM_LAYERS=9 MODEL_DIM=512 NUM_HEADS=8 NUM_KV_HEADS=4 MLP_MULT=2` -- Tied output/input embeddings: `TIE_EMBEDDINGS=1` -- Tied embedding LR: `TIED_EMBED_LR=0.05` -- Batching: `TRAIN_BATCH_TOKENS=524288 TRAIN_SEQ_LEN=1024` - -Evaluation-specific parameters: -- `EVAL_STRIDE=64` (sliding window stride; baseline uses non-overlapping = stride 1024) -- `EVAL_BATCH_SEQS=1024` (number of windows per forward pass for GPU utilization) - -## Command - -```bash -RUN_ID=8xh100_slide64_v2 \ -DATA_PATH=./data/datasets/fineweb10B_sp1024/ \ -TOKENIZER_PATH=./data/tokenizers/fineweb_1024_bpe.model \ -VOCAB_SIZE=1024 \ -NUM_LOOPS=1 \ -LORA_RANK=0 \ -QAT=0 \ -EVAL_STRIDE=64 \ -EVAL_BATCH_SEQS=1024 \ -MAX_WALLCLOCK_SECONDS=600 \ -TRAIN_LOG_EVERY=200 \ -VAL_LOSS_EVERY=1000 \ -torchrun --standalone --nproc_per_node=8 train_gpt.py -``` - -The `NUM_LOOPS=1 LORA_RANK=0 QAT=0` flags explicitly disable all unused code paths (these are also the defaults). - -## Key Metrics (from `train.log`) - -- Timed training stopped at `13450/20000` steps due to the wallclock cap. -- Pre-quant eval at stop: `val_loss:2.0592`, `val_bpb:1.2196` -- Post-quant sliding window eval: `val_loss:2.0135`, `val_bpb:1.1925` -- Exact printed metric: `final_int8_zlib_roundtrip_exact val_bpb:1.19250007` -- Train time: `600028ms` (`step_avg:44.61ms`) -- Peak memory: `10119 MiB allocated`, `10294 MiB reserved` -- Eval time: `69881ms` (sliding window, stride=64, batch_seqs=1024) -- Serialized model int8+zlib: `15816489 bytes` -- Code size: `58340 bytes` -- Total submission size int8+zlib: `15874829 bytes` - -## Training Volume - -- Global batch: `524288` tokens/step -- Total train tokens seen: `7,055,769,600` - -## Included Files - -- `train_gpt.py` (code snapshot used for the run, includes `eval_val_sliding` function) -- `train.log` (exact remote training log) -- `submission.json` (leaderboard metadata) diff --git a/records/track_10min_16mb/2026-03-19_SlidingWindowEval/submission.json b/records/track_10min_16mb/2026-03-19_SlidingWindowEval/submission.json deleted file mode 100644 index d25b325f97..0000000000 --- a/records/track_10min_16mb/2026-03-19_SlidingWindowEval/submission.json +++ /dev/null @@ -1,17 +0,0 @@ -{ - "author": "Matthew Li", - "github_id": "mattqlf", - "name": "Sliding Window Eval (stride=64)", - "blurb": "Baseline 9x512 SP-1024 architecture with sliding window evaluation at stride=64. Each token is scored with 960+ tokens of context instead of the baseline's 0-1023. Training is identical to the naive baseline; the improvement comes entirely from the evaluation strategy. Post-quant int8+zlib roundtrip under the 16,000,000-byte cap.", - "date": "2026-03-19T04:48:00Z", - "val_loss": 2.01348383, - "val_bpb": 1.19250007, - "pre_quant_val_loss": 2.0592, - "pre_quant_val_bpb": 1.2196, - "step_stop": 13450, - "wallclock_seconds": 600.028, - "eval_time_seconds": 69.881, - "bytes_total": 15874829, - "bytes_model_int8_zlib": 15816489, - "bytes_code": 58340 -} \ No newline at end of file diff --git a/records/track_10min_16mb/2026-03-19_SlidingWindowEval/train.log b/records/track_10min_16mb/2026-03-19_SlidingWindowEval/train.log deleted file mode 100644 index 8bd9edc7d8..0000000000 --- a/records/track_10min_16mb/2026-03-19_SlidingWindowEval/train.log +++ /dev/null @@ -1,133 +0,0 @@ -val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path=./data/tokenizers/fineweb_1024_bpe.model -train_loader:dataset:fineweb10B_sp1024 train_shards:80 -val_loader:shards pattern=./data/datasets/fineweb10B_sp1024/fineweb_val_*.bin tokens:62021632 -qat:False -model_params:17059912 (unique_layers:9 loops:1 effective_depth:9 lora_rank:0 lora_params:0) -world_size:8 grad_accum_steps:1 -sdp_backends:cudnn=False flash=True mem_efficient=False math=False -attention_mode:gqa num_heads:8 num_kv_heads:4 -tie_embeddings:True embed_lr:0.05 head_lr:0.0 matrix_lr:0.04 scalar_lr:0.04 -train_batch_tokens:524288 train_seq_len:1024 iterations:20000 warmup_steps:20 max_wallclock_seconds:600.000 -seed:1337 -warmup_step:1/20 -warmup_step:2/20 -warmup_step:3/20 -warmup_step:4/20 -warmup_step:5/20 -warmup_step:6/20 -warmup_step:7/20 -warmup_step:8/20 -warmup_step:9/20 -warmup_step:10/20 -warmup_step:11/20 -warmup_step:12/20 -warmup_step:13/20 -warmup_step:14/20 -warmup_step:15/20 -warmup_step:16/20 -warmup_step:17/20 -warmup_step:18/20 -warmup_step:19/20 -warmup_step:20/20 -step:0/20000 val_loss:6.9357 val_bpb:4.1077 train_time:0ms step_avg:0.01ms -step:1/20000 train_loss:6.9370 train_time:24ms step_avg:23.88ms -step:2/20000 train_loss:16.8366 train_time:62ms step_avg:31.18ms -step:3/20000 train_loss:8.7608 train_time:105ms step_avg:35.06ms -step:4/20000 train_loss:6.6385 train_time:148ms step_avg:37.01ms -step:5/20000 train_loss:6.6114 train_time:191ms step_avg:38.20ms -step:6/20000 train_loss:7.4220 train_time:234ms step_avg:39.05ms -step:7/20000 train_loss:6.3508 train_time:277ms step_avg:39.61ms -step:8/20000 train_loss:6.1582 train_time:320ms step_avg:40.05ms -step:9/20000 train_loss:6.0678 train_time:364ms step_avg:40.39ms -step:10/20000 train_loss:5.9747 train_time:407ms step_avg:40.66ms -step:200/20000 train_loss:2.8545 train_time:8724ms step_avg:43.62ms -step:400/20000 train_loss:2.3579 train_time:17484ms step_avg:43.71ms -step:600/20000 train_loss:2.5468 train_time:26272ms step_avg:43.79ms -step:800/20000 train_loss:2.2933 train_time:35060ms step_avg:43.83ms -step:1000/20000 train_loss:2.3741 train_time:43870ms step_avg:43.87ms -step:1000/20000 val_loss:2.3339 val_bpb:1.3823 train_time:43898ms step_avg:43.90ms -step:1200/20000 train_loss:2.3859 train_time:52691ms step_avg:43.91ms -step:1400/20000 train_loss:2.4313 train_time:61546ms step_avg:43.96ms -step:1600/20000 train_loss:2.0990 train_time:70406ms step_avg:44.00ms -step:1800/20000 train_loss:2.1989 train_time:79268ms step_avg:44.04ms -step:2000/20000 train_loss:2.2537 train_time:88144ms step_avg:44.07ms -step:2000/20000 val_loss:2.2349 val_bpb:1.3236 train_time:88173ms step_avg:44.09ms -step:2200/20000 train_loss:2.0705 train_time:97057ms step_avg:44.12ms -step:2400/20000 train_loss:2.2003 train_time:105955ms step_avg:44.15ms -step:2600/20000 train_loss:2.4100 train_time:114848ms step_avg:44.17ms -step:2800/20000 train_loss:2.2339 train_time:123759ms step_avg:44.20ms -step:3000/20000 train_loss:2.2271 train_time:132674ms step_avg:44.22ms -step:3000/20000 val_loss:2.1940 val_bpb:1.2994 train_time:132702ms step_avg:44.23ms -step:3200/20000 train_loss:2.1853 train_time:141596ms step_avg:44.25ms -step:3400/20000 train_loss:2.1579 train_time:150510ms step_avg:44.27ms -step:3600/20000 train_loss:2.1150 train_time:159433ms step_avg:44.29ms -step:3800/20000 train_loss:2.2207 train_time:168353ms step_avg:44.30ms -step:4000/20000 train_loss:2.1629 train_time:177281ms step_avg:44.32ms -step:4000/20000 val_loss:2.1691 val_bpb:1.2846 train_time:177309ms step_avg:44.33ms -step:4200/20000 train_loss:2.1755 train_time:186254ms step_avg:44.35ms -step:4400/20000 train_loss:2.1075 train_time:195164ms step_avg:44.36ms -step:4600/20000 train_loss:1.9721 train_time:204095ms step_avg:44.37ms -step:4800/20000 train_loss:2.2620 train_time:213026ms step_avg:44.38ms -step:5000/20000 train_loss:2.0261 train_time:221961ms step_avg:44.39ms -step:5000/20000 val_loss:2.1527 val_bpb:1.2749 train_time:221991ms step_avg:44.40ms -step:5200/20000 train_loss:2.1734 train_time:230894ms step_avg:44.40ms -step:5400/20000 train_loss:2.1832 train_time:239840ms step_avg:44.41ms -step:5600/20000 train_loss:2.1834 train_time:248772ms step_avg:44.42ms -step:5800/20000 train_loss:2.1438 train_time:257705ms step_avg:44.43ms -step:6000/20000 train_loss:2.2213 train_time:266645ms step_avg:44.44ms -step:6000/20000 val_loss:2.1428 val_bpb:1.2691 train_time:266673ms step_avg:44.45ms -step:6200/20000 train_loss:2.0903 train_time:275590ms step_avg:44.45ms -step:6400/20000 train_loss:2.1614 train_time:284523ms step_avg:44.46ms -step:6600/20000 train_loss:2.1233 train_time:293461ms step_avg:44.46ms -step:6800/20000 train_loss:2.1883 train_time:302396ms step_avg:44.47ms -step:7000/20000 train_loss:2.2269 train_time:311350ms step_avg:44.48ms -step:7000/20000 val_loss:2.1319 val_bpb:1.2626 train_time:311378ms step_avg:44.48ms -step:7200/20000 train_loss:2.1985 train_time:320283ms step_avg:44.48ms -step:7400/20000 train_loss:2.1159 train_time:329218ms step_avg:44.49ms -step:7600/20000 train_loss:2.0015 train_time:338182ms step_avg:44.50ms -step:7800/20000 train_loss:2.1457 train_time:347121ms step_avg:44.50ms -step:8000/20000 train_loss:2.1162 train_time:356081ms step_avg:44.51ms -step:8000/20000 val_loss:2.1223 val_bpb:1.2570 train_time:356110ms step_avg:44.51ms -step:8200/20000 train_loss:2.1840 train_time:365027ms step_avg:44.52ms -step:8400/20000 train_loss:2.1384 train_time:374085ms step_avg:44.53ms -step:8600/20000 train_loss:2.1382 train_time:383022ms step_avg:44.54ms -step:8800/20000 train_loss:2.1010 train_time:391971ms step_avg:44.54ms -step:9000/20000 train_loss:2.0244 train_time:400928ms step_avg:44.55ms -step:9000/20000 val_loss:2.1174 val_bpb:1.2540 train_time:400957ms step_avg:44.55ms -step:9200/20000 train_loss:2.0847 train_time:409874ms step_avg:44.55ms -step:9400/20000 train_loss:2.1341 train_time:418805ms step_avg:44.55ms -step:9600/20000 train_loss:2.1481 train_time:427753ms step_avg:44.56ms -step:9800/20000 train_loss:2.0727 train_time:436682ms step_avg:44.56ms -step:10000/20000 train_loss:2.1143 train_time:445623ms step_avg:44.56ms -step:10000/20000 val_loss:2.1124 val_bpb:1.2511 train_time:445652ms step_avg:44.57ms -step:10200/20000 train_loss:2.0665 train_time:454563ms step_avg:44.57ms -step:10400/20000 train_loss:2.0990 train_time:463504ms step_avg:44.57ms -step:10600/20000 train_loss:1.9760 train_time:472458ms step_avg:44.57ms -step:10800/20000 train_loss:2.1863 train_time:481398ms step_avg:44.57ms -step:11000/20000 train_loss:2.1152 train_time:490335ms step_avg:44.58ms -step:11000/20000 val_loss:2.1058 val_bpb:1.2472 train_time:490363ms step_avg:44.58ms -step:11200/20000 train_loss:2.0681 train_time:499305ms step_avg:44.58ms -step:11400/20000 train_loss:2.0572 train_time:508232ms step_avg:44.58ms -step:11600/20000 train_loss:2.0625 train_time:517178ms step_avg:44.58ms -step:11800/20000 train_loss:2.0980 train_time:526122ms step_avg:44.59ms -step:12000/20000 train_loss:2.0710 train_time:535066ms step_avg:44.59ms -step:12000/20000 val_loss:2.1003 val_bpb:1.2439 train_time:535094ms step_avg:44.59ms -step:12200/20000 train_loss:2.2155 train_time:544026ms step_avg:44.59ms -step:12400/20000 train_loss:1.8595 train_time:553021ms step_avg:44.60ms -step:12600/20000 train_loss:2.0846 train_time:561982ms step_avg:44.60ms -step:12800/20000 train_loss:2.0964 train_time:570913ms step_avg:44.60ms -step:13000/20000 train_loss:2.1690 train_time:579870ms step_avg:44.61ms -step:13000/20000 val_loss:2.0744 val_bpb:1.2286 train_time:579898ms step_avg:44.61ms -step:13200/20000 train_loss:2.1741 train_time:588820ms step_avg:44.61ms -step:13400/20000 train_loss:2.0456 train_time:597778ms step_avg:44.61ms -step:13450/20000 val_loss:2.0592 val_bpb:1.2196 train_time:600028ms step_avg:44.61ms -stopping_early: wallclock_cap train_time:600028ms step:13450/20000 -peak memory allocated: 10119 MiB reserved: 10294 MiB -Serialized model: 67224983 bytes -Code size: 58340 bytes -Total submission size: 67283323 bytes -Serialized model int8+zlib: 15816489 bytes (payload:17178912 raw_torch:17224025 payload_ratio:3.91x) -Total submission size int8+zlib: 15874829 bytes -final_eval_mode:sliding_window stride:64 batch_seqs:1024 -final_int8_zlib_roundtrip val_loss:2.0135 val_bpb:1.1925 eval_time:69881ms -final_int8_zlib_roundtrip_exact val_loss:2.01348383 val_bpb:1.19250007 diff --git a/records/track_10min_16mb/2026-03-19_SlidingWindowEval/train_gpt.py b/records/track_10min_16mb/2026-03-19_SlidingWindowEval/train_gpt.py deleted file mode 100644 index 8555de625c..0000000000 --- a/records/track_10min_16mb/2026-03-19_SlidingWindowEval/train_gpt.py +++ /dev/null @@ -1,1366 +0,0 @@ -""" -The `train_gpt.py` and `train_gpt_mlx.py` scripts are intended as good launching-off points for new participants, not SOTA configs. We'll accept PRs that tune, improve, or simplify these scripts without significantly increasing complexity, but competitive submissions should stay in the `/records` folder. - -Hard stop: To keep readable for newcomers, let's make sure `train_gpt.py` and `train_gpt_mlx.py` never are longer than 1500 lines. -""" - -from __future__ import annotations - -import copy -import glob -import io -import math -import os -import random -import subprocess -import sys -import time -import uuid -import zlib -from pathlib import Path - -import numpy as np -import sentencepiece as spm -import torch -import torch.distributed as dist -import torch.nn.functional as F -from torch import Tensor, nn -from torch.nn.parallel import DistributedDataParallel as DDP - -# ----------------------------- -# HYPERPARAMETERS -# ----------------------------- -# Default Simple Baseline run: -# - 9 transformer blocks at width 512 -# - 8 attention heads with 4 KV heads (GQA) and 2x MLP expansion -# - vocab size 1024, sequence length 1024, tied embeddings -# - 524,288 train tokens per step for 20,000 iterations with a ~10 minute cap - -class Hyperparameters: - # Data paths are shard globs produced by the existing preprocessing pipeline. - data_path = os.environ.get("DATA_PATH", "./data/datasets/fineweb10B_sp1024") - train_files = os.path.join(data_path, "fineweb_train_*.bin") - val_files = os.path.join(data_path, "fineweb_val_*.bin") - tokenizer_path = os.environ.get("TOKENIZER_PATH", "./data/tokenizers/fineweb_1024_bpe.model") - run_id = os.environ.get("RUN_ID", str(uuid.uuid4())) - seed = int(os.environ.get("SEED", 1337)) - - # Validation cadence and batch size. Validation always uses the full fineweb_val split. - val_batch_size = int(os.environ.get("VAL_BATCH_SIZE", 524_288)) - val_loss_every = int(os.environ.get("VAL_LOSS_EVERY", 1000)) - train_log_every = int(os.environ.get("TRAIN_LOG_EVERY", 200)) - - # Training length. - iterations = int(os.environ.get("ITERATIONS", 20000)) - warmdown_iters = int(os.environ.get("WARMDOWN_ITERS", 1200)) - warmup_steps = int(os.environ.get("WARMUP_STEPS", 20)) - train_batch_tokens = int(os.environ.get("TRAIN_BATCH_TOKENS", 524_288)) - train_seq_len = int(os.environ.get("TRAIN_SEQ_LEN", 1024)) - max_wallclock_seconds = float(os.environ.get("MAX_WALLCLOCK_SECONDS", 600.0)) - qk_gain_init = float(os.environ.get("QK_GAIN_INIT", 1.5)) - - # Model shape. - vocab_size = int(os.environ.get("VOCAB_SIZE", 1024)) - num_layers = int(os.environ.get("NUM_LAYERS", 9)) - num_kv_heads = int(os.environ.get("NUM_KV_HEADS", 4)) - model_dim = int(os.environ.get("MODEL_DIM", 512)) - num_heads = int(os.environ.get("NUM_HEADS", 8)) - mlp_mult = int(os.environ.get("MLP_MULT", 2)) - num_loops = int(os.environ.get("NUM_LOOPS", 1)) - lora_rank = int(os.environ.get("LORA_RANK", 0)) - qat = bool(int(os.environ.get("QAT", "1"))) - eval_stride = int(os.environ.get("EVAL_STRIDE", 64)) - eval_batch_seqs = int(os.environ.get("EVAL_BATCH_SEQS", 32)) - tie_embeddings = bool(int(os.environ.get("TIE_EMBEDDINGS", "1"))) - rope_base = float(os.environ.get("ROPE_BASE", 10000.0)) - logit_softcap = float(os.environ.get("LOGIT_SOFTCAP", 30.0)) - - # Optimizer hyperparameters. - embed_lr = float(os.environ.get("EMBED_LR", 0.6)) - head_lr = float(os.environ.get("HEAD_LR", 0.008)) - tied_embed_lr = float(os.environ.get("TIED_EMBED_LR", 0.05)) - tied_embed_init_std = float(os.environ.get("TIED_EMBED_INIT_STD", 0.005)) - matrix_lr = float(os.environ.get("MATRIX_LR", 0.04)) - scalar_lr = float(os.environ.get("SCALAR_LR", 0.04)) - muon_momentum = float(os.environ.get("MUON_MOMENTUM", 0.95)) - muon_backend_steps = int(os.environ.get("MUON_BACKEND_STEPS", 5)) - muon_momentum_warmup_start = float(os.environ.get("MUON_MOMENTUM_WARMUP_START", 0.85)) - muon_momentum_warmup_steps = int(os.environ.get("MUON_MOMENTUM_WARMUP_STEPS", 500)) - beta1 = float(os.environ.get("BETA1", 0.9)) - beta2 = float(os.environ.get("BETA2", 0.95)) - adam_eps = float(os.environ.get("ADAM_EPS", 1e-8)) - lora_lr = float(os.environ.get("LORA_LR", 0.01)) - grad_clip_norm = float(os.environ.get("GRAD_CLIP_NORM", 0.0)) - -# ----------------------------- -# MUON OPTIMIZER -# ----------------------------- -# -# As borrowed from modded-nanogpt -# Background on Muon: https://kellerjordan.github.io/posts/muon/ - -def zeropower_via_newtonschulz5(G: Tensor, steps: int = 10, eps: float = 1e-7) -> Tensor: - # Orthogonalize a 2D update matrix with a fast Newton-Schulz iteration. - # Muon uses this to normalize matrix-shaped gradients before applying them. - a, b, c = (3.4445, -4.7750, 2.0315) - X = G.bfloat16() - X /= X.norm() + eps - transposed = G.size(0) > G.size(1) - if transposed: - X = X.T - for _ in range(steps): - A = X @ X.T - B = b * A + c * A @ A - X = a * X + B @ X - return X.T if transposed else X - - -class Muon(torch.optim.Optimizer): - def __init__(self, params, lr: float, momentum: float, backend_steps: int, nesterov: bool = True): - super().__init__( - params, - dict(lr=lr, momentum=momentum, backend_steps=backend_steps, nesterov=nesterov), - ) - - @torch.no_grad() - def step(self, closure=None): - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - distributed = dist.is_available() and dist.is_initialized() - world_size = dist.get_world_size() if distributed else 1 - rank = dist.get_rank() if distributed else 0 - - for group in self.param_groups: - params = group["params"] - if not params: - continue - lr = group["lr"] - momentum = group["momentum"] - backend_steps = group["backend_steps"] - nesterov = group["nesterov"] - - total_params = sum(int(p.numel()) for p in params) - updates_flat = torch.zeros(total_params, device=params[0].device, dtype=torch.bfloat16) - - curr = 0 - for i, p in enumerate(params): - if i % world_size == rank and p.grad is not None: - g = p.grad - state = self.state[p] - if "momentum_buffer" not in state: - state["momentum_buffer"] = torch.zeros_like(g) - buf = state["momentum_buffer"] - buf.mul_(momentum).add_(g) - if nesterov: - g = g.add(buf, alpha=momentum) - g = zeropower_via_newtonschulz5(g, steps=backend_steps) - # Scale correction from Muon reference implementations. - g *= max(1, g.size(0) / g.size(1)) ** 0.5 - updates_flat[curr : curr + p.numel()] = g.reshape(-1) - curr += p.numel() - - if distributed: - dist.all_reduce(updates_flat, op=dist.ReduceOp.SUM) - - curr = 0 - for p in params: - g = updates_flat[curr : curr + p.numel()].view_as(p).to(dtype=p.dtype) - p.add_(g, alpha=-lr) - curr += p.numel() - - return loss - - -# ----------------------------- -# TOKENIZER-AGNOSTIC EVALUATION SETUP -# ----------------------------- -# -# It's common for small models have a large fraction of their parameters be embeddings, since the 2 * d_model * d_vocab vectors can be gigantic. -# Instead of locking the tokenizer, we let you bring your own and calculate our validation metrics on the average compression of the validation set. -# We calculate BPB (bits-per-byte) instead of validation loss, so we need methods to count the number of bits per token in the tokenizer. -# Note: Submissions that edit the tokenizer will be examined more carefully, since screwing this up might unjustly improve your score. - -def build_sentencepiece_luts( - sp: spm.SentencePieceProcessor, vocab_size: int, device: torch.device -) -> tuple[Tensor, Tensor, Tensor]: - sp_vocab_size = int(sp.vocab_size()) - table_size = max(sp_vocab_size, vocab_size) - base_bytes_np = np.zeros((table_size,), dtype=np.int16) - has_leading_space_np = np.zeros((table_size,), dtype=np.bool_) - is_boundary_token_np = np.ones((table_size,), dtype=np.bool_) - for token_id in range(sp_vocab_size): - if sp.is_control(token_id) or sp.is_unknown(token_id) or sp.is_unused(token_id): - continue - is_boundary_token_np[token_id] = False - if sp.is_byte(token_id): - base_bytes_np[token_id] = 1 - continue - piece = sp.id_to_piece(token_id) - if piece.startswith("▁"): - has_leading_space_np[token_id] = True - piece = piece[1:] - base_bytes_np[token_id] = len(piece.encode("utf-8")) - return ( - torch.tensor(base_bytes_np, dtype=torch.int16, device=device), - torch.tensor(has_leading_space_np, dtype=torch.bool, device=device), - torch.tensor(is_boundary_token_np, dtype=torch.bool, device=device), - ) - - -def load_validation_tokens(pattern: str, seq_len: int) -> Tensor: - files = [Path(p) for p in sorted(glob.glob(pattern))] - if not files: - raise FileNotFoundError(f"No files found for pattern: {pattern}") - # The export pipeline writes the fixed first-50k-doc validation set to fineweb_val_*. - tokens = torch.cat([load_data_shard(file) for file in files]).contiguous() - usable = ((tokens.numel() - 1) // seq_len) * seq_len - if usable <= 0: - raise ValueError(f"Validation split is too short for TRAIN_SEQ_LEN={seq_len}") - return tokens[: usable + 1] - - -def eval_val( - args: Hyperparameters, - model: nn.Module, - rank: int, - world_size: int, - device: torch.device, - grad_accum_steps: int, - val_tokens: Tensor, - base_bytes_lut: Tensor, - has_leading_space_lut: Tensor, - is_boundary_token_lut: Tensor, -) -> tuple[float, float]: - # Validation computes two metrics: - # - val_loss: token cross-entropy (natural log) - # - val_bpb: tokenizer-agnostic compression metric used by the challenge - local_batch_tokens = args.val_batch_size // (world_size * grad_accum_steps) - if local_batch_tokens < args.train_seq_len: - raise ValueError( - "VAL_BATCH_SIZE must provide at least one sequence per rank; " - f"got VAL_BATCH_SIZE={args.val_batch_size}, WORLD_SIZE={world_size}, " - f"GRAD_ACCUM_STEPS={grad_accum_steps}, TRAIN_SEQ_LEN={args.train_seq_len}" - ) - local_batch_seqs = local_batch_tokens // args.train_seq_len - total_seqs = (val_tokens.numel() - 1) // args.train_seq_len - seq_start = (total_seqs * rank) // world_size - seq_end = (total_seqs * (rank + 1)) // world_size - val_loss_sum = torch.zeros((), device=device, dtype=torch.float64) - val_token_count = torch.zeros((), device=device, dtype=torch.float64) - val_byte_count = torch.zeros((), device=device, dtype=torch.float64) - - model.eval() - with torch.inference_mode(): - for batch_seq_start in range(seq_start, seq_end, local_batch_seqs): - batch_seq_end = min(batch_seq_start + local_batch_seqs, seq_end) - raw_start = batch_seq_start * args.train_seq_len - raw_end = batch_seq_end * args.train_seq_len + 1 - local = val_tokens[raw_start:raw_end].to(device=device, dtype=torch.int64, non_blocking=True) - x = local[:-1].reshape(-1, args.train_seq_len) - y = local[1:].reshape(-1, args.train_seq_len) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - batch_loss = model(x, y).detach() - batch_token_count = float(y.numel()) - val_loss_sum += batch_loss.to(torch.float64) * batch_token_count - val_token_count += batch_token_count - prev_ids = x.reshape(-1) - tgt_ids = y.reshape(-1) - token_bytes = base_bytes_lut[tgt_ids].to(dtype=torch.int16) - token_bytes += (has_leading_space_lut[tgt_ids] & ~is_boundary_token_lut[prev_ids]).to(dtype=torch.int16) - val_byte_count += token_bytes.to(torch.float64).sum() - - if dist.is_available() and dist.is_initialized(): - dist.all_reduce(val_loss_sum, op=dist.ReduceOp.SUM) - dist.all_reduce(val_token_count, op=dist.ReduceOp.SUM) - dist.all_reduce(val_byte_count, op=dist.ReduceOp.SUM) - - val_loss = val_loss_sum / val_token_count - bits_per_token = val_loss.item() / math.log(2.0) - tokens_per_byte = val_token_count.item() / val_byte_count.item() - model.train() - return float(val_loss.item()), float(bits_per_token * tokens_per_byte) - -# ----------------------------- -# POST-TRAINING QUANTIZATION -# ----------------------------- -# -# It's silly to export our model, which is trained in bf16 and fp32, at that same precision. -# Instead, we get approximately the same model (with a small hit) by quantizing the model to int8 & zlib compressing. -# We can then decompress the model and run in higher precision for evaluation, after closing in under the size limit. - -CONTROL_TENSOR_NAME_PATTERNS = tuple( - pattern - for pattern in os.environ.get( - "CONTROL_TENSOR_NAME_PATTERNS", - "attn_scale,attn_scales,mlp_scale,mlp_scales,resid_mix,resid_mixes,q_gain,skip_weight,skip_weights", - ).split(",") - if pattern -) -INT8_KEEP_FLOAT_FP32_NAME_PATTERNS = tuple( - pattern - for pattern in os.environ.get( - "INT8_KEEP_FLOAT_FP32_NAME_PATTERNS", - ",".join(CONTROL_TENSOR_NAME_PATTERNS), - ).split(",") - if pattern -) -INT8_KEEP_FLOAT_MAX_NUMEL = 65_536 -INT8_KEEP_FLOAT_STORE_DTYPE = torch.float16 -INT8_PER_ROW_SCALE_DTYPE = torch.float16 -INT8_CLIP_PERCENTILE = 99.99984 -INT8_CLIP_Q = INT8_CLIP_PERCENTILE / 100.0 - -def tensor_nbytes(t: Tensor) -> int: - return int(t.numel()) * int(t.element_size()) - -def keep_float_tensor(name: str, t: Tensor, passthrough_orig_dtypes: dict[str, str]) -> Tensor: - if any(pattern in name for pattern in INT8_KEEP_FLOAT_FP32_NAME_PATTERNS): - return t.float().contiguous() - if t.dtype in {torch.float32, torch.bfloat16}: - passthrough_orig_dtypes[name] = str(t.dtype).removeprefix("torch.") - return t.to(dtype=INT8_KEEP_FLOAT_STORE_DTYPE).contiguous() - return t - -def quantize_float_tensor(t: Tensor) -> tuple[Tensor, Tensor]: - t32 = t.float() - if t32.ndim == 2: - # Matrices get one scale per row, which usually tracks output-channel - # ranges much better than a single tensor-wide scale. - clip_abs = ( - torch.quantile(t32.abs(), INT8_CLIP_Q, dim=1) - if t32.numel() - else torch.empty((t32.shape[0],), dtype=torch.float32) - ) - clipped = torch.maximum(torch.minimum(t32, clip_abs[:, None]), -clip_abs[:, None]) - scale = (clip_abs / 127.0).clamp_min(1.0 / 127.0) - q = torch.clamp(torch.round(clipped / scale[:, None]), -127, 127).to(torch.int8).contiguous() - return q, scale.to(dtype=INT8_PER_ROW_SCALE_DTYPE).contiguous() - - # Vectors / scalars use a simpler per-tensor scale. - clip_abs = float(torch.quantile(t32.abs().flatten(), INT8_CLIP_Q).item()) if t32.numel() else 0.0 - scale = torch.tensor(clip_abs / 127.0 if clip_abs > 0 else 1.0, dtype=torch.float32) - q = torch.clamp(torch.round(torch.clamp(t32, -clip_abs, clip_abs) / scale), -127, 127).to(torch.int8).contiguous() - return q, scale - -def quantize_state_dict_int8(state_dict: dict[str, Tensor]): - # Single supported clean-script export format: - # - per-row int8 for 2D float tensors - # - per-tensor int8 for other float tensors - # - exact passthrough for non-floats - # - passthrough for small float tensors, stored as fp16 to save bytes - quantized: dict[str, Tensor] = {} - scales: dict[str, Tensor] = {} - dtypes: dict[str, str] = {} - passthrough: dict[str, Tensor] = {} - passthrough_orig_dtypes: dict[str, str] = {} - qmeta: dict[str, dict[str, object]] = {} - stats = dict.fromkeys( - ("param_count", "num_tensors", "num_float_tensors", "num_nonfloat_tensors", "baseline_tensor_bytes", "int8_payload_bytes"), - 0, - ) - - for name, tensor in state_dict.items(): - t = tensor.detach().to("cpu").contiguous() - stats["param_count"] += int(t.numel()) - stats["num_tensors"] += 1 - stats["baseline_tensor_bytes"] += tensor_nbytes(t) - - if not t.is_floating_point(): - stats["num_nonfloat_tensors"] += 1 - passthrough[name] = t - stats["int8_payload_bytes"] += tensor_nbytes(t) - continue - - # Small float tensors are cheap enough to keep directly. We still downcast - # fp32/bf16 passthrough tensors to fp16 so metadata does not dominate size. - if t.numel() <= INT8_KEEP_FLOAT_MAX_NUMEL: - kept = keep_float_tensor(name, t, passthrough_orig_dtypes) - passthrough[name] = kept - stats["int8_payload_bytes"] += tensor_nbytes(kept) - continue - - stats["num_float_tensors"] += 1 - q, s = quantize_float_tensor(t) - if s.ndim > 0: - qmeta[name] = {"scheme": "per_row", "axis": 0} - quantized[name] = q - scales[name] = s - dtypes[name] = str(t.dtype).removeprefix("torch.") - stats["int8_payload_bytes"] += tensor_nbytes(q) + tensor_nbytes(s) - - obj: dict[str, object] = { - "__quant_format__": "int8_clean_per_row_v1", - "quantized": quantized, - "scales": scales, - "dtypes": dtypes, - "passthrough": passthrough, - } - if qmeta: - obj["qmeta"] = qmeta - if passthrough_orig_dtypes: - obj["passthrough_orig_dtypes"] = passthrough_orig_dtypes - return obj, stats - -def dequantize_state_dict_int8(obj: dict[str, object]) -> dict[str, Tensor]: - out: dict[str, Tensor] = {} - qmeta = obj.get("qmeta", {}) - passthrough_orig_dtypes = obj.get("passthrough_orig_dtypes", {}) - for name, q in obj["quantized"].items(): - dtype = getattr(torch, obj["dtypes"][name]) - s = obj["scales"][name] - if qmeta.get(name, {}).get("scheme") == "per_row" or s.ndim > 0: - s = s.to(dtype=torch.float32) - # Broadcast the saved row scale back across trailing dimensions. - out[name] = (q.float() * s.view(q.shape[0], *([1] * (q.ndim - 1)))).to(dtype=dtype).contiguous() - else: - scale = float(s.item()) - out[name] = (q.float() * scale).to(dtype=dtype).contiguous() - for name, t in obj["passthrough"].items(): - # Restore small tensors, undoing the temporary fp16 storage cast if needed. - out_t = t.detach().to("cpu").contiguous() - orig_dtype = passthrough_orig_dtypes.get(name) - if isinstance(orig_dtype, str): - out_t = out_t.to(dtype=getattr(torch, orig_dtype)).contiguous() - out[name] = out_t - return out - - -# ----------------------------- -# DATA LOADING -# ----------------------------- - -def load_data_shard(file: Path) -> Tensor: - header_bytes = 256 * np.dtype(" None: - self.file_idx = (self.file_idx + 1) % len(self.files) - self.tokens = load_data_shard(self.files[self.file_idx]) - self.pos = 0 - - def take(self, n: int) -> Tensor: - chunks: list[Tensor] = [] - remaining = n - while remaining > 0: - avail = self.tokens.numel() - self.pos - if avail <= 0: - self._advance_file() - continue - k = min(remaining, avail) - chunks.append(self.tokens[self.pos : self.pos + k]) - self.pos += k - remaining -= k - return chunks[0] if len(chunks) == 1 else torch.cat(chunks) - - -class DistributedTokenLoader: - # Each call consumes a contiguous chunk from the shared token stream, then slices out - # one disjoint span per rank. The extra "+1" token lets us build (x, y) by shifting. - def __init__(self, pattern: str, rank: int, world_size: int, device: torch.device): - self.rank = rank - self.world_size = world_size - self.device = device - self.stream = TokenStream(pattern) - - def next_batch(self, global_tokens: int, seq_len: int, grad_accum_steps: int) -> tuple[Tensor, Tensor]: - local_tokens = global_tokens // (self.world_size * grad_accum_steps) - per_rank_span = local_tokens + 1 - chunk = self.stream.take(per_rank_span * self.world_size) - start = self.rank * per_rank_span - local = chunk[start : start + per_rank_span].to(dtype=torch.int64) - x = local[:-1].reshape(-1, seq_len) - y = local[1:].reshape(-1, seq_len) - return x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True) - -# ----------------------------- -# TRANSFORMER MODULES -# ----------------------------- - -class RMSNorm(nn.Module): - def __init__(self, eps: float | None = None): - super().__init__() - self.eps = eps - - def forward(self, x: Tensor) -> Tensor: - return F.rms_norm(x, (x.size(-1),), eps=self.eps) - - -def fake_quantize_int8_per_row(w: Tensor) -> Tensor: - """Simulate per-row int8 quantization with straight-through estimator. - - Forward: uses quantized-then-dequantized weights (same rounding as post-training). - Backward: gradients pass through as if no quantization happened (STE). - """ - scale = w.detach().abs().amax(dim=-1, keepdim=True).div_(127.0).clamp_(min=1.0 / 127.0) - w_deq = (w / scale).round().clamp_(-127, 127) * scale - return w + (w_deq - w).detach() - - -class CastedLinear(nn.Linear): - # Keep weights in fp32 for optimizer/state quality, cast at matmul time for bf16 compute. - _qat: bool = False - - def forward(self, x: Tensor) -> Tensor: - w = self.weight - if self._qat and self.training: - w = fake_quantize_int8_per_row(w) - bias = self.bias.to(x.dtype) if self.bias is not None else None - return F.linear(x, w.to(x.dtype), bias) - - -class AttentionLoRA(nn.Module): - """Per-iteration LoRA adapters for attention Q, K, V, and output projections. - - Initialized so that the LoRA contribution is zero at the start of training - (B matrices are zeros). During training, the optimizer learns per-iteration - specialization while the base attention weights remain shared across loops. - """ - def __init__(self, dim: int, kv_dim: int, rank: int): - super().__init__() - self.q_A = nn.Parameter(torch.empty(dim, rank)) - self.q_B = nn.Parameter(torch.zeros(rank, dim)) - self.k_A = nn.Parameter(torch.empty(dim, rank)) - self.k_B = nn.Parameter(torch.zeros(rank, kv_dim)) - self.v_A = nn.Parameter(torch.empty(dim, rank)) - self.v_B = nn.Parameter(torch.zeros(rank, kv_dim)) - self.proj_A = nn.Parameter(torch.empty(dim, rank)) - self.proj_B = nn.Parameter(torch.zeros(rank, dim)) - self._init_lora() - - def _init_lora(self) -> None: - for name in ("q_A", "k_A", "v_A", "proj_A"): - nn.init.kaiming_uniform_(getattr(self, name), a=math.sqrt(5)) - - -def restore_low_dim_params_to_fp32(module: nn.Module) -> None: - # Keep small/control parameters in fp32 even when the model body runs in bf16. - with torch.no_grad(): - for name, param in module.named_parameters(): - if (param.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS)) and param.dtype != torch.float32: - param.data = param.data.float() - - -class Rotary(nn.Module): - # Caches cos/sin tables per sequence length on the current device. - def __init__(self, dim: int, base: float = 10000.0): - super().__init__() - inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self._seq_len_cached = 0 - self._cos_cached: Tensor | None = None - self._sin_cached: Tensor | None = None - - def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> tuple[Tensor, Tensor]: - if ( - self._cos_cached is None - or self._sin_cached is None - or self._seq_len_cached != seq_len - or self._cos_cached.device != device - ): - t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) - freqs = torch.outer(t, self.inv_freq.to(device)) - self._cos_cached = freqs.cos()[None, None, :, :] - self._sin_cached = freqs.sin()[None, None, :, :] - self._seq_len_cached = seq_len - return self._cos_cached.to(dtype=dtype), self._sin_cached.to(dtype=dtype) - - -def apply_rotary_emb(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: - half = x.size(-1) // 2 - x1, x2 = x[..., :half], x[..., half:] - return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1) - - -class CausalSelfAttention(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - num_kv_heads: int, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - if dim % num_heads != 0: - raise ValueError("model_dim must be divisible by num_heads") - if num_heads % num_kv_heads != 0: - raise ValueError("num_heads must be divisible by num_kv_heads") - self.num_heads = num_heads - self.num_kv_heads = num_kv_heads - self.head_dim = dim // num_heads - if self.head_dim % 2 != 0: - raise ValueError("head_dim must be even for RoPE") - kv_dim = self.num_kv_heads * self.head_dim - self.c_q = CastedLinear(dim, dim, bias=False) - self.c_k = CastedLinear(dim, kv_dim, bias=False) - self.c_v = CastedLinear(dim, kv_dim, bias=False) - self.proj = CastedLinear(dim, dim, bias=False) - self.proj._zero_init = True - self.q_gain = nn.Parameter(torch.full((num_heads,), qk_gain_init, dtype=torch.float32)) - self.rotary = Rotary(self.head_dim, base=rope_base) - - def forward(self, x: Tensor, lora: AttentionLoRA | None = None) -> Tensor: - bsz, seqlen, dim = x.shape - q = self.c_q(x) - k = self.c_k(x) - v = self.c_v(x) - if lora is not None: - # LoRA delta: (bsz, seqlen, dim) @ (dim, rank) @ (rank, out_dim) - # autocast handles fp32->bf16 cast of LoRA params automatically - q = q + (x @ lora.q_A) @ lora.q_B - k = k + (x @ lora.k_A) @ lora.k_B - v = v + (x @ lora.v_A) @ lora.v_B - q = q.reshape(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2) - k = k.reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - v = v.reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - q = F.rms_norm(q, (q.size(-1),)) - k = F.rms_norm(k, (k.size(-1),)) - cos, sin = self.rotary(seqlen, x.device, q.dtype) - q = apply_rotary_emb(q, cos, sin) - k = apply_rotary_emb(k, cos, sin) - q = q * self.q_gain.to(dtype=q.dtype)[None, :, None, None] - y = F.scaled_dot_product_attention( - q, - k, - v, - attn_mask=None, - is_causal=True, - enable_gqa=(self.num_kv_heads != self.num_heads), - ) - y = y.transpose(1, 2).contiguous().reshape(bsz, seqlen, dim) - out = self.proj(y) - if lora is not None: - out = out + (y @ lora.proj_A) @ lora.proj_B - return out - - -class MLP(nn.Module): - # relu^2 MLP from the original modded-nanogpt setup - def __init__(self, dim: int, mlp_mult: int): - super().__init__() - hidden = mlp_mult * dim - self.fc = CastedLinear(dim, hidden, bias=False) - self.proj = CastedLinear(hidden, dim, bias=False) - self.proj._zero_init = True - - def forward(self, x: Tensor) -> Tensor: - x = torch.relu(self.fc(x)) - return self.proj(x.square()) - - -class Block(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - num_kv_heads: int, - mlp_mult: int, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - self.attn_norm = RMSNorm() - self.mlp_norm = RMSNorm() - self.attn = CausalSelfAttention(dim, num_heads, num_kv_heads, rope_base, qk_gain_init) - self.mlp = MLP(dim, mlp_mult) - self.attn_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) - self.mlp_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) - self.resid_mix = nn.Parameter(torch.stack((torch.ones(dim), torch.zeros(dim))).float()) - - def forward(self, x: Tensor, x0: Tensor, lora: AttentionLoRA | None = None) -> Tensor: - mix = self.resid_mix.to(dtype=x.dtype) - x = mix[0][None, None, :] * x + mix[1][None, None, :] * x0 - attn_out = self.attn(self.attn_norm(x), lora=lora) - x = x + self.attn_scale.to(dtype=x.dtype)[None, None, :] * attn_out - x = x + self.mlp_scale.to(dtype=x.dtype)[None, None, :] * self.mlp(self.mlp_norm(x)) - return x - - -class GPT(nn.Module): - def __init__( - self, - vocab_size: int, - num_layers: int, - model_dim: int, - num_heads: int, - num_kv_heads: int, - mlp_mult: int, - tie_embeddings: bool, - tied_embed_init_std: float, - logit_softcap: float, - rope_base: float, - qk_gain_init: float, - num_loops: int = 1, - lora_rank: int = 0, - ): - super().__init__() - if logit_softcap <= 0.0: - raise ValueError(f"logit_softcap must be positive, got {logit_softcap}") - self.tie_embeddings = tie_embeddings - self.tied_embed_init_std = tied_embed_init_std - self.logit_softcap = logit_softcap - self.num_unique_layers = num_layers - self.num_loops = num_loops - effective_depth = num_layers * num_loops - self.tok_emb = nn.Embedding(vocab_size, model_dim) - self.num_encoder_layers = effective_depth // 2 - self.num_decoder_layers = effective_depth - self.num_encoder_layers - self.num_skip_weights = min(self.num_encoder_layers, self.num_decoder_layers) - self.skip_weights = nn.Parameter(torch.ones(self.num_skip_weights, model_dim, dtype=torch.float32)) - self.blocks = nn.ModuleList( - [ - Block( - model_dim, - num_heads, - num_kv_heads, - mlp_mult, - rope_base, - qk_gain_init, - ) - for i in range(num_layers) - ] - ) - # Per-(loop, block) LoRA adapters for attention projections. - # Only created when num_loops > 1 and lora_rank > 0. - kv_dim = num_kv_heads * (model_dim // num_heads) - if lora_rank > 0 and num_loops > 1: - self.lora_adapters = nn.ModuleList( - [ - nn.ModuleList( - [AttentionLoRA(model_dim, kv_dim, lora_rank) for _ in range(num_layers)] - ) - for _ in range(num_loops) - ] - ) - else: - self.lora_adapters = None - self.final_norm = RMSNorm() - self.lm_head = None if tie_embeddings else CastedLinear(model_dim, vocab_size, bias=False) - if self.lm_head is not None: - self.lm_head._zero_init = True - self._init_weights() - - def _init_weights(self) -> None: - if self.tie_embeddings: - nn.init.normal_(self.tok_emb.weight, mean=0.0, std=self.tied_embed_init_std) - for module in self.modules(): - if isinstance(module, nn.Linear) and getattr(module, "_zero_init", False): - nn.init.zeros_(module.weight) - - def forward(self, input_ids: Tensor, target_ids: Tensor) -> Tensor: - x = self.tok_emb(input_ids) - x = F.rms_norm(x, (x.size(-1),)) - x0 = x - skips: list[Tensor] = [] - - # Iterate through effective layers: each unique block is reused across loops. - # First half (encoder) stores skip connections; second half (decoder) pops them. - eff_idx = 0 - for loop_idx in range(self.num_loops): - for block_idx in range(self.num_unique_layers): - lora = self.lora_adapters[loop_idx][block_idx] if self.lora_adapters is not None else None - if eff_idx < self.num_encoder_layers: - x = self.blocks[block_idx](x, x0, lora=lora) - skips.append(x) - else: - dec_idx = eff_idx - self.num_encoder_layers - if dec_idx < self.num_skip_weights and skips: - x = x + self.skip_weights[dec_idx].to(dtype=x.dtype)[None, None, :] * skips.pop() - x = self.blocks[block_idx](x, x0, lora=lora) - eff_idx += 1 - - x = self.final_norm(x).reshape(-1, x.size(-1)) - targets = target_ids.reshape(-1) - if self.tie_embeddings: - logits_proj = F.linear(x, self.tok_emb.weight) - else: - if self.lm_head is None: - raise RuntimeError("lm_head is required when tie_embeddings=False") - logits_proj = self.lm_head(x) - logits = self.logit_softcap * torch.tanh(logits_proj / self.logit_softcap) - return F.cross_entropy(logits.float(), targets, reduction="mean") - - def forward_logits(self, input_ids: Tensor) -> Tensor: - """Return logits (bsz, seq_len, vocab) without computing loss.""" - x = self.tok_emb(input_ids) - x = F.rms_norm(x, (x.size(-1),)) - x0 = x - skips: list[Tensor] = [] - eff_idx = 0 - for loop_idx in range(self.num_loops): - for block_idx in range(self.num_unique_layers): - lora = self.lora_adapters[loop_idx][block_idx] if self.lora_adapters is not None else None - if eff_idx < self.num_encoder_layers: - x = self.blocks[block_idx](x, x0, lora=lora) - skips.append(x) - else: - dec_idx = eff_idx - self.num_encoder_layers - if dec_idx < self.num_skip_weights and skips: - x = x + self.skip_weights[dec_idx].to(dtype=x.dtype)[None, None, :] * skips.pop() - x = self.blocks[block_idx](x, x0, lora=lora) - eff_idx += 1 - x = self.final_norm(x) - if self.tie_embeddings: - logits_proj = F.linear(x, self.tok_emb.weight) - else: - logits_proj = self.lm_head(x) - return self.logit_softcap * torch.tanh(logits_proj / self.logit_softcap) - - -def eval_val_sliding( - args: Hyperparameters, - base_model: nn.Module, - rank: int, - world_size: int, - device: torch.device, - val_tokens: Tensor, - base_bytes_lut: Tensor, - has_leading_space_lut: Tensor, - is_boundary_token_lut: Tensor, - stride: int, - batch_seqs: int = 32, -) -> tuple[float, float]: - """Sliding window evaluation: each token scored with maximum context. - - Windows of train_seq_len advance by `stride`. Only the last `stride` tokens - per window contribute to the score (first window scores all). Windows are - batched and distributed across ranks. - """ - seq_len = args.train_seq_len - total_tokens = val_tokens.numel() - 1 - - # Build windows; skip any too short to score a full stride - window_starts = [ws for ws in range(0, total_tokens, stride) - if min(ws + seq_len, total_tokens) - ws >= stride] - total_windows = len(window_starts) - - # Distribute across ranks - my_s = (total_windows * rank) // world_size - my_e = (total_windows * (rank + 1)) // world_size - my_windows = window_starts[my_s:my_e] - - loss_sum = torch.zeros((), device=device, dtype=torch.float64) - token_count = torch.zeros((), device=device, dtype=torch.float64) - byte_count = torch.zeros((), device=device, dtype=torch.float64) - - base_model.eval() - with torch.inference_mode(): - for bi in range(0, len(my_windows), batch_seqs): - batch_ws = my_windows[bi:bi + batch_seqs] - bsz = len(batch_ws) - - x_batch = torch.zeros(bsz, seq_len, dtype=torch.int64, device=device) - y_batch = torch.zeros(bsz, seq_len, dtype=torch.int64, device=device) - wlens: list[int] = [] - - for i, ws in enumerate(batch_ws): - end = min(ws + seq_len, total_tokens) - wlen = end - ws - wlens.append(wlen) - chunk = val_tokens[ws:end + 1].to(dtype=torch.int64, device=device) - x_batch[i, :wlen] = chunk[:-1] - y_batch[i, :wlen] = chunk[1:] - - with torch.autocast(device_type="cuda", dtype=torch.bfloat16): - logits = base_model.forward_logits(x_batch) - - nll = F.cross_entropy( - logits.reshape(-1, logits.size(-1)).float(), - y_batch.reshape(-1), - reduction="none", - ).reshape(bsz, seq_len) - - for i, ws in enumerate(batch_ws): - wlen = wlens[i] - s = 0 if ws == 0 else wlen - stride - scored_nll = nll[i, s:wlen].to(torch.float64) - loss_sum += scored_nll.sum() - token_count += float(wlen - s) - tgt = y_batch[i, s:wlen] - prev = x_batch[i, s:wlen] - tb = base_bytes_lut[tgt].to(torch.float64) - tb += (has_leading_space_lut[tgt] & ~is_boundary_token_lut[prev]).to(torch.float64) - byte_count += tb.sum() - - # Progress (rank 0 only) - if rank == 0 and (bi // batch_seqs) % 50 == 0: - done = min(bi + batch_seqs, len(my_windows)) - pct = done / len(my_windows) * 100 - running_bpb = 0.0 - if token_count.item() > 0: - rl = (loss_sum / token_count).item() - running_bpb = rl / math.log(2.0) * (token_count.item() / byte_count.item()) - print(f" sliding_eval [{pct:5.1f}%] {done}/{len(my_windows)} windows running_bpb={running_bpb:.6f}", flush=True) - - if dist.is_available() and dist.is_initialized(): - dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM) - dist.all_reduce(token_count, op=dist.ReduceOp.SUM) - dist.all_reduce(byte_count, op=dist.ReduceOp.SUM) - - val_loss = (loss_sum / token_count).item() - bits_per_token = val_loss / math.log(2.0) - tokens_per_byte = token_count.item() / byte_count.item() - base_model.train() - return val_loss, bits_per_token * tokens_per_byte - - -# ----------------------------- -# TRAINING -# ----------------------------- - -def main() -> None: - global zeropower_via_newtonschulz5 - - code = Path(__file__).read_text(encoding="utf-8") - args = Hyperparameters() - zeropower_via_newtonschulz5 = torch.compile(zeropower_via_newtonschulz5) - - # ----------------------------- - # DISTRIBUTED + CUDA SETUP - # ----------------------------- - - distributed = "RANK" in os.environ and "WORLD_SIZE" in os.environ - rank = int(os.environ.get("RANK", "0")) - world_size = int(os.environ.get("WORLD_SIZE", "1")) - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - if world_size <= 0: - raise ValueError(f"WORLD_SIZE must be positive, got {world_size}") - if 8 % world_size != 0: - raise ValueError(f"WORLD_SIZE={world_size} must divide 8 so grad_accum_steps stays integral") - grad_accum_steps = 8 // world_size - grad_scale = 1.0 / grad_accum_steps - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required") - device = torch.device("cuda", local_rank) - torch.cuda.set_device(device) - if distributed: - dist.init_process_group(backend="nccl", device_id=device) - dist.barrier() - master_process = rank == 0 - - # Fast math knobs - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True - from torch.backends.cuda import enable_cudnn_sdp, enable_flash_sdp, enable_math_sdp, enable_mem_efficient_sdp - - enable_cudnn_sdp(False) - enable_flash_sdp(True) - enable_mem_efficient_sdp(False) - enable_math_sdp(False) - - logfile = None - if master_process: - os.makedirs("logs", exist_ok=True) - logfile = f"logs/{args.run_id}.txt" - print(logfile) - - def log0(msg: str, console: bool = True) -> None: - if not master_process: - return - if console: - print(msg) - if logfile is not None: - with open(logfile, "a", encoding="utf-8") as f: - print(msg, file=f) - - log0(code, console=False) - log0("=" * 100, console=False) - log0(f"Running Python {sys.version}", console=False) - log0(f"Running PyTorch {torch.__version__}", console=False) - log0( - subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False).stdout, - console=False, - ) - log0("=" * 100, console=False) - - # ----------------------------- - # TOKENIZER + VALIDATION METRIC SETUP - # ----------------------------- - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed_all(args.seed) - - if not args.tokenizer_path.endswith(".model"): - raise ValueError(f"Script only setup for SentencePiece .model file: {args.tokenizer_path}") - sp = spm.SentencePieceProcessor(model_file=args.tokenizer_path) - if int(sp.vocab_size()) != args.vocab_size: - raise ValueError( - f"VOCAB_SIZE={args.vocab_size} does not match tokenizer vocab_size={int(sp.vocab_size())}" - ) - dataset_dir = Path(args.data_path).resolve() - actual_train_files = len(list(dataset_dir.glob("fineweb_train_*.bin"))) - val_tokens = load_validation_tokens(args.val_files, args.train_seq_len) - base_bytes_lut, has_leading_space_lut, is_boundary_token_lut = build_sentencepiece_luts( - sp, args.vocab_size, device - ) - log0(f"val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path={args.tokenizer_path}") - log0(f"train_loader:dataset:{dataset_dir.name} train_shards:{actual_train_files}") - log0(f"val_loader:shards pattern={args.val_files} tokens:{val_tokens.numel() - 1}") - - # ----------------------------- - # MODEL + OPTIMIZER SETUP - # ----------------------------- - - base_model = GPT( - vocab_size=args.vocab_size, - num_layers=args.num_layers, - model_dim=args.model_dim, - num_heads=args.num_heads, - num_kv_heads=args.num_kv_heads, - mlp_mult=args.mlp_mult, - tie_embeddings=args.tie_embeddings, - tied_embed_init_std=args.tied_embed_init_std, - logit_softcap=args.logit_softcap, - rope_base=args.rope_base, - qk_gain_init=args.qk_gain_init, - num_loops=args.num_loops, - lora_rank=args.lora_rank, - ).to(device).bfloat16() - for module in base_model.modules(): - if isinstance(module, (CastedLinear, AttentionLoRA)): - module.float() - if isinstance(module, CastedLinear) and args.qat: - module._qat = True - restore_low_dim_params_to_fp32(base_model) - log0(f"qat:{args.qat}") - compiled_model = torch.compile(base_model, dynamic=False, fullgraph=True) - model: nn.Module = DDP(compiled_model, device_ids=[local_rank], broadcast_buffers=False) if distributed else compiled_model - - # Optimizer split: - # - token embedding (Adam) uses EMBED_LR - # - untied lm_head (Adam) uses HEAD_LR - # - matrix params in transformer blocks use MATRIX_LR via Muon - # - vectors/scalars use SCALAR_LR via Adam - block_named_params = list(base_model.blocks.named_parameters()) - matrix_params = [ - p - for name, p in block_named_params - if p.ndim == 2 and not any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - scalar_params = [ - p - for name, p in block_named_params - if p.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - if base_model.skip_weights.numel() > 0: - scalar_params.append(base_model.skip_weights) - token_lr = args.tied_embed_lr if args.tie_embeddings else args.embed_lr - optimizer_tok = torch.optim.Adam( - [{"params": [base_model.tok_emb.weight], "lr": token_lr, "base_lr": token_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizer_muon = Muon( - matrix_params, - lr=args.matrix_lr, - momentum=args.muon_momentum, - backend_steps=args.muon_backend_steps, - ) - for group in optimizer_muon.param_groups: - group["base_lr"] = args.matrix_lr - optimizer_scalar = torch.optim.Adam( - [{"params": scalar_params, "lr": args.scalar_lr, "base_lr": args.scalar_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers: list[torch.optim.Optimizer] = [optimizer_tok, optimizer_muon, optimizer_scalar] - if base_model.lora_adapters is not None: - lora_params = list(base_model.lora_adapters.parameters()) - optimizer_lora = torch.optim.Adam( - [{"params": lora_params, "lr": args.lora_lr, "base_lr": args.lora_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers.append(optimizer_lora) - if base_model.lm_head is not None: - optimizer_head = torch.optim.Adam( - [{"params": [base_model.lm_head.weight], "lr": args.head_lr, "base_lr": args.head_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers.insert(1, optimizer_head) - - n_params = sum(p.numel() for p in base_model.parameters()) - n_lora = sum(p.numel() for p in base_model.lora_adapters.parameters()) if base_model.lora_adapters is not None else 0 - effective_depth = args.num_layers * args.num_loops - log0(f"model_params:{n_params} (unique_layers:{args.num_layers} loops:{args.num_loops} effective_depth:{effective_depth} lora_rank:{args.lora_rank} lora_params:{n_lora})") - log0(f"world_size:{world_size} grad_accum_steps:{grad_accum_steps}") - log0("sdp_backends:cudnn=False flash=True mem_efficient=False math=False") - log0(f"attention_mode:gqa num_heads:{args.num_heads} num_kv_heads:{args.num_kv_heads}") - log0( - f"tie_embeddings:{args.tie_embeddings} embed_lr:{token_lr} " - f"head_lr:{args.head_lr if base_model.lm_head is not None else 0.0} " - f"matrix_lr:{args.matrix_lr} scalar_lr:{args.scalar_lr}" - ) - log0( - f"train_batch_tokens:{args.train_batch_tokens} train_seq_len:{args.train_seq_len} " - f"iterations:{args.iterations} warmup_steps:{args.warmup_steps} " - f"max_wallclock_seconds:{args.max_wallclock_seconds:.3f}" - ) - log0(f"seed:{args.seed}") - - # ----------------------------- - # DATA LOADER & MODEL WARMUP - # ----------------------------- - - train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device) - - def zero_grad_all() -> None: - for opt in optimizers: - opt.zero_grad(set_to_none=True) - - max_wallclock_ms = 1000.0 * args.max_wallclock_seconds if args.max_wallclock_seconds > 0 else None - - def lr_mul(step: int, elapsed_ms: float) -> float: - if args.warmdown_iters <= 0: - return 1.0 - if max_wallclock_ms is None: - warmdown_start = max(args.iterations - args.warmdown_iters, 0) - return max((args.iterations - step) / max(args.warmdown_iters, 1), 0.0) if warmdown_start <= step < args.iterations else 1.0 - step_ms = elapsed_ms / max(step, 1) - warmdown_ms = args.warmdown_iters * step_ms - remaining_ms = max(max_wallclock_ms - elapsed_ms, 0.0) - return remaining_ms / max(warmdown_ms, 1e-9) if remaining_ms <= warmdown_ms else 1.0 - - # Warmup primes the compiled forward/backward/optimizer paths, then we restore the - # initial weights/optimizer state so measured training starts from the true init. - if args.warmup_steps > 0: - initial_model_state = {name: tensor.detach().cpu().clone() for name, tensor in base_model.state_dict().items()} - initial_optimizer_states = [copy.deepcopy(opt.state_dict()) for opt in optimizers] - model.train() - for warmup_step in range(args.warmup_steps): - zero_grad_all() - for micro_step in range(grad_accum_steps): - if distributed: - model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 - x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - warmup_loss = model(x, y) - (warmup_loss * grad_scale).backward() - for opt in optimizers: - opt.step() - zero_grad_all() - if args.warmup_steps <= 20 or (warmup_step + 1) % 10 == 0 or warmup_step + 1 == args.warmup_steps: - log0(f"warmup_step:{warmup_step + 1}/{args.warmup_steps}") - base_model.load_state_dict(initial_model_state, strict=True) - for opt, state in zip(optimizers, initial_optimizer_states, strict=True): - opt.load_state_dict(state) - zero_grad_all() - if distributed: - model.require_backward_grad_sync = True - train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device) - - # ----------------------------- - # MAIN TRAINING LOOP - # ----------------------------- - - training_time_ms = 0.0 - stop_after_step: int | None = None - torch.cuda.synchronize() - t0 = time.perf_counter() - - step = 0 - while True: - last_step = step == args.iterations or (stop_after_step is not None and step >= stop_after_step) - - should_validate = last_step or (args.val_loss_every > 0 and step % args.val_loss_every == 0) - if should_validate: - torch.cuda.synchronize() - training_time_ms += 1000.0 * (time.perf_counter() - t0) - val_loss, val_bpb = eval_val( - args, - model, - rank, - world_size, - device, - grad_accum_steps, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - ) - log0( - f"step:{step}/{args.iterations} val_loss:{val_loss:.4f} val_bpb:{val_bpb:.4f} " - f"train_time:{training_time_ms:.0f}ms step_avg:{training_time_ms / max(step, 1):.2f}ms" - ) - torch.cuda.synchronize() - t0 = time.perf_counter() - - if last_step: - if stop_after_step is not None and step < args.iterations: - log0( - f"stopping_early: wallclock_cap train_time:{training_time_ms:.0f}ms " - f"step:{step}/{args.iterations}" - ) - break - - elapsed_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) - scale = lr_mul(step, elapsed_ms) - zero_grad_all() - train_loss = torch.zeros((), device=device) - for micro_step in range(grad_accum_steps): - if distributed: - model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 - x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - loss = model(x, y) - train_loss += loss.detach() - (loss * grad_scale).backward() - train_loss /= grad_accum_steps - - frac = min(step / args.muon_momentum_warmup_steps, 1.0) if args.muon_momentum_warmup_steps > 0 else 1.0 - muon_momentum = (1 - frac) * args.muon_momentum_warmup_start + frac * args.muon_momentum - for group in optimizer_muon.param_groups: - group["momentum"] = muon_momentum - - for opt in optimizers: - for group in opt.param_groups: - group["lr"] = group["base_lr"] * scale - - if args.grad_clip_norm > 0: - torch.nn.utils.clip_grad_norm_(base_model.parameters(), args.grad_clip_norm) - for opt in optimizers: - opt.step() - zero_grad_all() - - step += 1 - approx_training_time_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) - should_log_train = ( - args.train_log_every > 0 - and (step <= 10 or step % args.train_log_every == 0 or stop_after_step is not None) - ) - if should_log_train: - log0( - f"step:{step}/{args.iterations} train_loss:{train_loss.item():.4f} " - f"train_time:{approx_training_time_ms:.0f}ms step_avg:{approx_training_time_ms / step:.2f}ms" - ) - - # Needed to sync whether we've reached the wallclock cap. - reached_cap = max_wallclock_ms is not None and approx_training_time_ms >= max_wallclock_ms - if distributed and max_wallclock_ms is not None: - reached_cap_tensor = torch.tensor(int(reached_cap), device=device) - dist.all_reduce(reached_cap_tensor, op=dist.ReduceOp.MAX) - reached_cap = bool(reached_cap_tensor.item()) - if stop_after_step is None and reached_cap: - stop_after_step = step - - log0( - f"peak memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB " - f"reserved: {torch.cuda.max_memory_reserved() // 1024 // 1024} MiB" - ) - - # ----------------------------- - # SERIALIZATION + ROUNDTRIP VALIDATION - # ----------------------------- - # Save the raw state (useful for debugging/loading in PyTorch directly), then always produce - # the compressed int8+zlib artifact and validate the round-tripped weights. - - if master_process: - torch.save(base_model.state_dict(), "final_model.pt") - model_bytes = os.path.getsize("final_model.pt") - code_bytes = len(code.encode("utf-8")) - log0(f"Serialized model: {model_bytes} bytes") - log0(f"Code size: {code_bytes} bytes") - log0(f"Total submission size: {model_bytes + code_bytes} bytes") - - quant_obj, quant_stats = quantize_state_dict_int8(base_model.state_dict()) - quant_buf = io.BytesIO() - torch.save(quant_obj, quant_buf) - quant_raw = quant_buf.getvalue() - quant_blob = zlib.compress(quant_raw, level=9) - quant_raw_bytes = len(quant_raw) - if master_process: - with open("final_model.int8.ptz", "wb") as f: - f.write(quant_blob) - quant_file_bytes = os.path.getsize("final_model.int8.ptz") - code_bytes = len(code.encode("utf-8")) - ratio = quant_stats["baseline_tensor_bytes"] / max(quant_stats["int8_payload_bytes"], 1) - log0( - f"Serialized model int8+zlib: {quant_file_bytes} bytes " - f"(payload:{quant_stats['int8_payload_bytes']} raw_torch:{quant_raw_bytes} payload_ratio:{ratio:.2f}x)" - ) - log0(f"Total submission size int8+zlib: {quant_file_bytes + code_bytes} bytes") - - if distributed: - dist.barrier() - with open("final_model.int8.ptz", "rb") as f: - quant_blob_disk = f.read() - quant_state = torch.load(io.BytesIO(zlib.decompress(quant_blob_disk)), map_location="cpu") - base_model.load_state_dict(dequantize_state_dict_int8(quant_state), strict=True) - torch.cuda.synchronize() - t_qeval = time.perf_counter() - if args.eval_stride > 0 and args.eval_stride < args.train_seq_len: - log0(f"final_eval_mode:sliding_window stride:{args.eval_stride} batch_seqs:{args.eval_batch_seqs}") - q_val_loss, q_val_bpb = eval_val_sliding( - args, - base_model, - rank, - world_size, - device, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - stride=args.eval_stride, - batch_seqs=args.eval_batch_seqs, - ) - else: - log0("final_eval_mode:standard") - q_val_loss, q_val_bpb = eval_val( - args, - model, - rank, - world_size, - device, - grad_accum_steps, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - ) - torch.cuda.synchronize() - log0( - f"final_int8_zlib_roundtrip val_loss:{q_val_loss:.4f} val_bpb:{q_val_bpb:.4f} " - f"eval_time:{1000.0 * (time.perf_counter() - t_qeval):.0f}ms" - ) - log0(f"final_int8_zlib_roundtrip_exact val_loss:{q_val_loss:.8f} val_bpb:{q_val_bpb:.8f}") - - if distributed: - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/records/track_non_record_16mb/2026-03-18_Quasi10Bfrom50B_SP1024_9x512_KV4_4h_pgut3/README.md b/records/track_non_record_16mb/2026-03-18_Quasi10Bfrom50B_SP1024_9x512_KV4_4h_pgut3/README.md deleted file mode 100644 index c36c90b816..0000000000 --- a/records/track_non_record_16mb/2026-03-18_Quasi10Bfrom50B_SP1024_9x512_KV4_4h_pgut3/README.md +++ /dev/null @@ -1,56 +0,0 @@ -This record captures an unlimited-compute non-record submission built from the current root `train_gpt.py`. - -This run is not intended to satisfy the 10-minute cutoff for the main leaderboard. It uses the same 9x512 SP-1024 tied-embedding baseline layout, but extends training to a 4-hour wallclock cap on `pgut3` while evaluating against the full 50k-document validation split every 20k steps. - -Configuration: -- Track: `non-record`, unlimited compute, still under the `16,000,000` byte artifact cap -- Layout: `VOCAB_SIZE=1024 NUM_LAYERS=9 MODEL_DIM=512 NUM_HEADS=8 NUM_KV_HEADS=4 MLP_MULT=2` -- Tied output/input embeddings: `TIE_EMBEDDINGS=1` -- Tied embedding LR: `TIED_EMBED_LR=0.05` -- Batching: `TRAIN_BATCH_TOKENS=524288 TRAIN_SEQ_LEN=1024` -- Validation cadence: `VAL_LOSS_EVERY=20000` on the full `fineweb_val_*` split - -Command (track-relevant params): -```bash -OMP_NUM_THREADS=1 \ -TORCH_NCCL_ASYNC_ERROR_HANDLING=1 \ -RUN_ID=train_gpt_pgut3_quasi10b_sp1024_4h_20260318_075102 \ -DATA_PATH=/tmp/fineweb_quasi10Bfrom50B_50keval_sp1024_v0/datasets/fineweb10B_sp1024 \ -TOKENIZER_PATH=/tmp/fineweb_quasi10Bfrom50B_50keval_sp1024_v0/tokenizers/fineweb_1024_bpe.model \ -VOCAB_SIZE=1024 \ -NUM_LAYERS=9 \ -MODEL_DIM=512 \ -NUM_HEADS=8 \ -NUM_KV_HEADS=4 \ -MLP_MULT=2 \ -TIE_EMBEDDINGS=1 \ -TIED_EMBED_LR=0.05 \ -ITERATIONS=500000 \ -WARMUP_STEPS=20 \ -MAX_WALLCLOCK_SECONDS=14400 \ -TRAIN_BATCH_TOKENS=524288 \ -TRAIN_SEQ_LEN=1024 \ -TRAIN_LOG_EVERY=200 \ -VAL_LOSS_EVERY=20000 \ -torchrun --standalone --nproc_per_node=8 /root/code/parameter-golf/train_gpt.py -``` - -Key metrics (from `train.log`): -- Timed training stopped at `329430/500000` steps due to the 4-hour wallclock cap. -- Best pre-quant eval at stop: `val_loss:1.9837`, `val_bpb:1.1749` -- Post-quant roundtrip eval: `val_loss:2.0386`, `val_bpb:1.2074` -- Exact printed metric: `final_int8_zlib_roundtrip_exact val_bpb:1.20737944` -- Train time: `14400039ms` (`step_avg:43.71ms`) -- Peak memory: `10184 MiB allocated`, `10588 MiB reserved` -- Serialized model int8+zlib: `15762519 bytes` -- Code size: `47642 bytes` -- Total submission size int8+zlib: `15810161 bytes` - -Training volume: -- Global batch: `524288` tokens/step -- Total train tokens seen: `172716195840` - -Included files: -- `train_gpt.py` (code snapshot used for the run) -- `train.log` (exact remote training log) -- `submission.json` (leaderboard metadata) diff --git a/records/track_non_record_16mb/2026-03-18_Quasi10Bfrom50B_SP1024_9x512_KV4_4h_pgut3/submission.json b/records/track_non_record_16mb/2026-03-18_Quasi10Bfrom50B_SP1024_9x512_KV4_4h_pgut3/submission.json deleted file mode 100644 index 4192982daa..0000000000 --- a/records/track_non_record_16mb/2026-03-18_Quasi10Bfrom50B_SP1024_9x512_KV4_4h_pgut3/submission.json +++ /dev/null @@ -1,17 +0,0 @@ -{ - "author": "Will DePue", - "github_id": "williamd", - "name": "4-Hour Quasi-10B SP1024", - "blurb": "Unlimited compute track: SP-1024 9x512 KV4 run on pgut3 for 4 hours against the quasi10Bfrom50B 50k-eval export; pre-quant reached 1.1749 BPB at wallclock stop and final int8+zlib roundtrip scored 1.2074 under the 16,000,000-byte cap.", - "date": "2026-03-18T11:53:00Z", - "track": "non-record-unlimited-compute-16mb", - "val_loss": 2.03860961, - "val_bpb": 1.20737944, - "pre_quant_val_loss": 1.9837, - "pre_quant_val_bpb": 1.1749, - "step_stop": 329430, - "wallclock_seconds": 14400.039, - "bytes_total": 15810161, - "bytes_model_int8_zlib": 15762519, - "bytes_code": 47642 -} diff --git a/records/track_non_record_16mb/2026-03-18_Quasi10Bfrom50B_SP1024_9x512_KV4_4h_pgut3/train.log b/records/track_non_record_16mb/2026-03-18_Quasi10Bfrom50B_SP1024_9x512_KV4_4h_pgut3/train.log deleted file mode 100644 index fbc482360b..0000000000 --- a/records/track_non_record_16mb/2026-03-18_Quasi10Bfrom50B_SP1024_9x512_KV4_4h_pgut3/train.log +++ /dev/null @@ -1,2901 +0,0 @@ -""" -The `train_gpt.py` and `train_gpt_mlx.py` scripts are intended as good launching-off points for new participants, not SOTA configs. We'll accept PRs that tune, improve, or simplify these scripts without significantly increasing complexity, but competitive submissions should stay in the `/records` folder. - -Hard stop: `train_gpt.py` and `train_gpt_mlx.py` must never be longer than 1500 lines. -""" - -from __future__ import annotations - -import copy -import glob -import io -import math -import os -import random -import subprocess -import sys -import time -import uuid -import zlib -from pathlib import Path - -import numpy as np -import sentencepiece as spm -import torch -import torch.distributed as dist -import torch.nn.functional as F -from torch import Tensor, nn -from torch.nn.parallel import DistributedDataParallel as DDP - -# ----------------------------- -# HYPERPARAMETERS -# ----------------------------- -# Default Simple Baseline run: -# - 9 transformer blocks at width 512 -# - 8 attention heads with 4 KV heads (GQA) and 2x MLP expansion -# - vocab size 1024, sequence length 1024, tied embeddings -# - 524,288 train tokens per step for 20,000 iterations with a ~10 minute cap - -class Hyperparameters: - # Data paths are shard globs produced by the existing preprocessing pipeline. - data_path = os.environ.get("DATA_PATH", "./data/datasets/fineweb10B_sp1024") - train_files = os.path.join(data_path, "fineweb_train_*.bin") - val_files = os.path.join(data_path, "fineweb_val_*.bin") - tokenizer_path = os.environ.get("TOKENIZER_PATH", "./data/tokenizers/fineweb_1024_bpe.model") - run_id = os.environ.get("RUN_ID", str(uuid.uuid4())) - seed = int(os.environ.get("SEED", 1337)) - - # Validation cadence and batch size. Validation always uses the full fineweb_val split. - val_batch_size = int(os.environ.get("VAL_BATCH_SIZE", 524_288)) - val_loss_every = int(os.environ.get("VAL_LOSS_EVERY", 1000)) - train_log_every = int(os.environ.get("TRAIN_LOG_EVERY", 200)) - - # Training length. - iterations = int(os.environ.get("ITERATIONS", 20000)) - warmdown_iters = int(os.environ.get("WARMDOWN_ITERS", 1200)) - warmup_steps = int(os.environ.get("WARMUP_STEPS", 20)) - train_batch_tokens = int(os.environ.get("TRAIN_BATCH_TOKENS", 524_288)) - train_seq_len = int(os.environ.get("TRAIN_SEQ_LEN", 1024)) - max_wallclock_seconds = float(os.environ.get("MAX_WALLCLOCK_SECONDS", 600.0)) - qk_gain_init = float(os.environ.get("QK_GAIN_INIT", 1.5)) - - # Model shape. - vocab_size = int(os.environ.get("VOCAB_SIZE", 1024)) - num_layers = int(os.environ.get("NUM_LAYERS", 9)) - num_kv_heads = int(os.environ.get("NUM_KV_HEADS", 4)) - model_dim = int(os.environ.get("MODEL_DIM", 512)) - num_heads = int(os.environ.get("NUM_HEADS", 8)) - mlp_mult = int(os.environ.get("MLP_MULT", 2)) - tie_embeddings = bool(int(os.environ.get("TIE_EMBEDDINGS", "1"))) - rope_base = float(os.environ.get("ROPE_BASE", 10000.0)) - logit_softcap = float(os.environ.get("LOGIT_SOFTCAP", 30.0)) - - # Optimizer hyperparameters. - embed_lr = float(os.environ.get("EMBED_LR", 0.6)) - head_lr = float(os.environ.get("HEAD_LR", 0.008)) - tied_embed_lr = float(os.environ.get("TIED_EMBED_LR", 0.05)) - tied_embed_init_std = float(os.environ.get("TIED_EMBED_INIT_STD", 0.005)) - matrix_lr = float(os.environ.get("MATRIX_LR", 0.04)) - scalar_lr = float(os.environ.get("SCALAR_LR", 0.04)) - muon_momentum = float(os.environ.get("MUON_MOMENTUM", 0.95)) - muon_backend_steps = int(os.environ.get("MUON_BACKEND_STEPS", 5)) - muon_momentum_warmup_start = float(os.environ.get("MUON_MOMENTUM_WARMUP_START", 0.85)) - muon_momentum_warmup_steps = int(os.environ.get("MUON_MOMENTUM_WARMUP_STEPS", 500)) - beta1 = float(os.environ.get("BETA1", 0.9)) - beta2 = float(os.environ.get("BETA2", 0.95)) - adam_eps = float(os.environ.get("ADAM_EPS", 1e-8)) - grad_clip_norm = float(os.environ.get("GRAD_CLIP_NORM", 0.0)) - -# ----------------------------- -# MUON OPTIMIZER -# ----------------------------- -# -# As borrowed from modded-nanogpt -# Background on Muon: https://kellerjordan.github.io/posts/muon/ - -def zeropower_via_newtonschulz5(G: Tensor, steps: int = 10, eps: float = 1e-7) -> Tensor: - # Orthogonalize a 2D update matrix with a fast Newton-Schulz iteration. - # Muon uses this to normalize matrix-shaped gradients before applying them. - a, b, c = (3.4445, -4.7750, 2.0315) - X = G.bfloat16() - X /= X.norm() + eps - transposed = G.size(0) > G.size(1) - if transposed: - X = X.T - for _ in range(steps): - A = X @ X.T - B = b * A + c * A @ A - X = a * X + B @ X - return X.T if transposed else X - - -class Muon(torch.optim.Optimizer): - def __init__(self, params, lr: float, momentum: float, backend_steps: int, nesterov: bool = True): - super().__init__( - params, - dict(lr=lr, momentum=momentum, backend_steps=backend_steps, nesterov=nesterov), - ) - - @torch.no_grad() - def step(self, closure=None): - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - distributed = dist.is_available() and dist.is_initialized() - world_size = dist.get_world_size() if distributed else 1 - rank = dist.get_rank() if distributed else 0 - - for group in self.param_groups: - params = group["params"] - if not params: - continue - lr = group["lr"] - momentum = group["momentum"] - backend_steps = group["backend_steps"] - nesterov = group["nesterov"] - - total_params = sum(int(p.numel()) for p in params) - updates_flat = torch.zeros(total_params, device=params[0].device, dtype=torch.bfloat16) - - curr = 0 - for i, p in enumerate(params): - if i % world_size == rank and p.grad is not None: - g = p.grad - state = self.state[p] - if "momentum_buffer" not in state: - state["momentum_buffer"] = torch.zeros_like(g) - buf = state["momentum_buffer"] - buf.mul_(momentum).add_(g) - if nesterov: - g = g.add(buf, alpha=momentum) - g = zeropower_via_newtonschulz5(g, steps=backend_steps) - # Scale correction from Muon reference implementations. - g *= max(1, g.size(0) / g.size(1)) ** 0.5 - updates_flat[curr : curr + p.numel()] = g.reshape(-1) - curr += p.numel() - - if distributed: - dist.all_reduce(updates_flat, op=dist.ReduceOp.SUM) - - curr = 0 - for p in params: - g = updates_flat[curr : curr + p.numel()].view_as(p).to(dtype=p.dtype) - p.add_(g, alpha=-lr) - curr += p.numel() - - return loss - - -# ----------------------------- -# TOKENIZER-AGNOSTIC EVALUATION SETUP -# ----------------------------- -# -# It's common for small models have a large fraction of their parameters be embeddings, since the 2 * d_model * d_vocab vectors can be gigantic. -# Instead of locking the tokenizer, we let you bring your own and calculate our validation metrics on the average compression of the validation set. -# We calculate BPB (bits-per-byte) instead of validation loss, so we need methods to count the number of bits per token in the tokenizer. -# Note: Submissions that edit the tokenizer will be examined more carefully, since screwing this up might unjustly improve your score. - -def build_sentencepiece_luts( - sp: spm.SentencePieceProcessor, vocab_size: int, device: torch.device -) -> tuple[Tensor, Tensor, Tensor]: - sp_vocab_size = int(sp.vocab_size()) - table_size = max(sp_vocab_size, vocab_size) - base_bytes_np = np.zeros((table_size,), dtype=np.int16) - has_leading_space_np = np.zeros((table_size,), dtype=np.bool_) - is_boundary_token_np = np.ones((table_size,), dtype=np.bool_) - for token_id in range(sp_vocab_size): - if sp.is_control(token_id) or sp.is_unknown(token_id) or sp.is_unused(token_id): - continue - is_boundary_token_np[token_id] = False - if sp.is_byte(token_id): - base_bytes_np[token_id] = 1 - continue - piece = sp.id_to_piece(token_id) - if piece.startswith("▁"): - has_leading_space_np[token_id] = True - piece = piece[1:] - base_bytes_np[token_id] = len(piece.encode("utf-8")) - return ( - torch.tensor(base_bytes_np, dtype=torch.int16, device=device), - torch.tensor(has_leading_space_np, dtype=torch.bool, device=device), - torch.tensor(is_boundary_token_np, dtype=torch.bool, device=device), - ) - - -def load_validation_tokens(pattern: str, seq_len: int) -> Tensor: - files = [Path(p) for p in sorted(glob.glob(pattern))] - if not files: - raise FileNotFoundError(f"No files found for pattern: {pattern}") - # The export pipeline writes the fixed first-50k-doc validation set to fineweb_val_*. - tokens = torch.cat([load_data_shard(file) for file in files]).contiguous() - usable = ((tokens.numel() - 1) // seq_len) * seq_len - if usable <= 0: - raise ValueError(f"Validation split is too short for TRAIN_SEQ_LEN={seq_len}") - return tokens[: usable + 1] - - -def eval_val( - args: Hyperparameters, - model: nn.Module, - rank: int, - world_size: int, - device: torch.device, - grad_accum_steps: int, - val_tokens: Tensor, - base_bytes_lut: Tensor, - has_leading_space_lut: Tensor, - is_boundary_token_lut: Tensor, -) -> tuple[float, float]: - # Validation computes two metrics: - # - val_loss: token cross-entropy (natural log) - # - val_bpb: tokenizer-agnostic compression metric used by the challenge - local_batch_tokens = args.val_batch_size // (world_size * grad_accum_steps) - if local_batch_tokens < args.train_seq_len: - raise ValueError( - "VAL_BATCH_SIZE must provide at least one sequence per rank; " - f"got VAL_BATCH_SIZE={args.val_batch_size}, WORLD_SIZE={world_size}, " - f"GRAD_ACCUM_STEPS={grad_accum_steps}, TRAIN_SEQ_LEN={args.train_seq_len}" - ) - local_batch_seqs = local_batch_tokens // args.train_seq_len - total_seqs = (val_tokens.numel() - 1) // args.train_seq_len - seq_start = (total_seqs * rank) // world_size - seq_end = (total_seqs * (rank + 1)) // world_size - val_loss_sum = torch.zeros((), device=device, dtype=torch.float64) - val_token_count = torch.zeros((), device=device, dtype=torch.float64) - val_byte_count = torch.zeros((), device=device, dtype=torch.float64) - - model.eval() - with torch.inference_mode(): - for batch_seq_start in range(seq_start, seq_end, local_batch_seqs): - batch_seq_end = min(batch_seq_start + local_batch_seqs, seq_end) - raw_start = batch_seq_start * args.train_seq_len - raw_end = batch_seq_end * args.train_seq_len + 1 - local = val_tokens[raw_start:raw_end].to(device=device, dtype=torch.int64, non_blocking=True) - x = local[:-1].reshape(-1, args.train_seq_len) - y = local[1:].reshape(-1, args.train_seq_len) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - batch_loss = model(x, y).detach() - batch_token_count = float(y.numel()) - val_loss_sum += batch_loss.to(torch.float64) * batch_token_count - val_token_count += batch_token_count - prev_ids = x.reshape(-1) - tgt_ids = y.reshape(-1) - token_bytes = base_bytes_lut[tgt_ids].to(dtype=torch.int16) - token_bytes += (has_leading_space_lut[tgt_ids] & ~is_boundary_token_lut[prev_ids]).to(dtype=torch.int16) - val_byte_count += token_bytes.to(torch.float64).sum() - - if dist.is_available() and dist.is_initialized(): - dist.all_reduce(val_loss_sum, op=dist.ReduceOp.SUM) - dist.all_reduce(val_token_count, op=dist.ReduceOp.SUM) - dist.all_reduce(val_byte_count, op=dist.ReduceOp.SUM) - - val_loss = val_loss_sum / val_token_count - bits_per_token = val_loss.item() / math.log(2.0) - tokens_per_byte = val_token_count.item() / val_byte_count.item() - model.train() - return float(val_loss.item()), float(bits_per_token * tokens_per_byte) - -# ----------------------------- -# POST-TRAINING QUANTIZATION -# ----------------------------- -# -# It's silly to export our model, which is trained in bf16 and fp32, at that same precision. -# Instead, we get approximately the same model (with a small hit) by quantizing the model to int8 & zlib compressing. -# We can then decompress the model and run in higher precision for evaluation, after closing in under the size limit. - -CONTROL_TENSOR_NAME_PATTERNS = tuple( - pattern - for pattern in os.environ.get( - "CONTROL_TENSOR_NAME_PATTERNS", - "attn_scale,attn_scales,mlp_scale,mlp_scales,resid_mix,resid_mixes,q_gain,skip_weight,skip_weights", - ).split(",") - if pattern -) -INT8_KEEP_FLOAT_FP32_NAME_PATTERNS = tuple( - pattern - for pattern in os.environ.get( - "INT8_KEEP_FLOAT_FP32_NAME_PATTERNS", - ",".join(CONTROL_TENSOR_NAME_PATTERNS), - ).split(",") - if pattern -) -INT8_KEEP_FLOAT_MAX_NUMEL = 65_536 -INT8_KEEP_FLOAT_STORE_DTYPE = torch.float16 -INT8_PER_ROW_SCALE_DTYPE = torch.float16 -INT8_CLIP_PERCENTILE = 99.99984 -INT8_CLIP_Q = INT8_CLIP_PERCENTILE / 100.0 - -def tensor_nbytes(t: Tensor) -> int: - return int(t.numel()) * int(t.element_size()) - -def keep_float_tensor(name: str, t: Tensor, passthrough_orig_dtypes: dict[str, str]) -> Tensor: - if any(pattern in name for pattern in INT8_KEEP_FLOAT_FP32_NAME_PATTERNS): - return t.float().contiguous() - if t.dtype in {torch.float32, torch.bfloat16}: - passthrough_orig_dtypes[name] = str(t.dtype).removeprefix("torch.") - return t.to(dtype=INT8_KEEP_FLOAT_STORE_DTYPE).contiguous() - return t - -def quantize_float_tensor(t: Tensor) -> tuple[Tensor, Tensor]: - t32 = t.float() - if t32.ndim == 2: - # Matrices get one scale per row, which usually tracks output-channel - # ranges much better than a single tensor-wide scale. - clip_abs = ( - torch.quantile(t32.abs(), INT8_CLIP_Q, dim=1) - if t32.numel() - else torch.empty((t32.shape[0],), dtype=torch.float32) - ) - clipped = torch.maximum(torch.minimum(t32, clip_abs[:, None]), -clip_abs[:, None]) - scale = (clip_abs / 127.0).clamp_min(1.0 / 127.0) - q = torch.clamp(torch.round(clipped / scale[:, None]), -127, 127).to(torch.int8).contiguous() - return q, scale.to(dtype=INT8_PER_ROW_SCALE_DTYPE).contiguous() - - # Vectors / scalars use a simpler per-tensor scale. - clip_abs = float(torch.quantile(t32.abs().flatten(), INT8_CLIP_Q).item()) if t32.numel() else 0.0 - scale = torch.tensor(clip_abs / 127.0 if clip_abs > 0 else 1.0, dtype=torch.float32) - q = torch.clamp(torch.round(torch.clamp(t32, -clip_abs, clip_abs) / scale), -127, 127).to(torch.int8).contiguous() - return q, scale - -def quantize_state_dict_int8(state_dict: dict[str, Tensor]): - # Single supported clean-script export format: - # - per-row int8 for 2D float tensors - # - per-tensor int8 for other float tensors - # - exact passthrough for non-floats - # - passthrough for small float tensors, stored as fp16 to save bytes - quantized: dict[str, Tensor] = {} - scales: dict[str, Tensor] = {} - dtypes: dict[str, str] = {} - passthrough: dict[str, Tensor] = {} - passthrough_orig_dtypes: dict[str, str] = {} - qmeta: dict[str, dict[str, object]] = {} - stats = dict.fromkeys( - ("param_count", "num_tensors", "num_float_tensors", "num_nonfloat_tensors", "baseline_tensor_bytes", "int8_payload_bytes"), - 0, - ) - - for name, tensor in state_dict.items(): - t = tensor.detach().to("cpu").contiguous() - stats["param_count"] += int(t.numel()) - stats["num_tensors"] += 1 - stats["baseline_tensor_bytes"] += tensor_nbytes(t) - - if not t.is_floating_point(): - stats["num_nonfloat_tensors"] += 1 - passthrough[name] = t - stats["int8_payload_bytes"] += tensor_nbytes(t) - continue - - # Small float tensors are cheap enough to keep directly. We still downcast - # fp32/bf16 passthrough tensors to fp16 so metadata does not dominate size. - if t.numel() <= INT8_KEEP_FLOAT_MAX_NUMEL: - kept = keep_float_tensor(name, t, passthrough_orig_dtypes) - passthrough[name] = kept - stats["int8_payload_bytes"] += tensor_nbytes(kept) - continue - - stats["num_float_tensors"] += 1 - q, s = quantize_float_tensor(t) - if s.ndim > 0: - qmeta[name] = {"scheme": "per_row", "axis": 0} - quantized[name] = q - scales[name] = s - dtypes[name] = str(t.dtype).removeprefix("torch.") - stats["int8_payload_bytes"] += tensor_nbytes(q) + tensor_nbytes(s) - - obj: dict[str, object] = { - "__quant_format__": "int8_clean_per_row_v1", - "quantized": quantized, - "scales": scales, - "dtypes": dtypes, - "passthrough": passthrough, - } - if qmeta: - obj["qmeta"] = qmeta - if passthrough_orig_dtypes: - obj["passthrough_orig_dtypes"] = passthrough_orig_dtypes - return obj, stats - -def dequantize_state_dict_int8(obj: dict[str, object]) -> dict[str, Tensor]: - out: dict[str, Tensor] = {} - qmeta = obj.get("qmeta", {}) - passthrough_orig_dtypes = obj.get("passthrough_orig_dtypes", {}) - for name, q in obj["quantized"].items(): - dtype = getattr(torch, obj["dtypes"][name]) - s = obj["scales"][name] - if qmeta.get(name, {}).get("scheme") == "per_row" or s.ndim > 0: - s = s.to(dtype=torch.float32) - # Broadcast the saved row scale back across trailing dimensions. - out[name] = (q.float() * s.view(q.shape[0], *([1] * (q.ndim - 1)))).to(dtype=dtype).contiguous() - else: - scale = float(s.item()) - out[name] = (q.float() * scale).to(dtype=dtype).contiguous() - for name, t in obj["passthrough"].items(): - # Restore small tensors, undoing the temporary fp16 storage cast if needed. - out_t = t.detach().to("cpu").contiguous() - orig_dtype = passthrough_orig_dtypes.get(name) - if isinstance(orig_dtype, str): - out_t = out_t.to(dtype=getattr(torch, orig_dtype)).contiguous() - out[name] = out_t - return out - - -# ----------------------------- -# DATA LOADING -# ----------------------------- - -def load_data_shard(file: Path) -> Tensor: - header_bytes = 256 * np.dtype(" None: - self.file_idx = (self.file_idx + 1) % len(self.files) - self.tokens = load_data_shard(self.files[self.file_idx]) - self.pos = 0 - - def take(self, n: int) -> Tensor: - chunks: list[Tensor] = [] - remaining = n - while remaining > 0: - avail = self.tokens.numel() - self.pos - if avail <= 0: - self._advance_file() - continue - k = min(remaining, avail) - chunks.append(self.tokens[self.pos : self.pos + k]) - self.pos += k - remaining -= k - return chunks[0] if len(chunks) == 1 else torch.cat(chunks) - - -class DistributedTokenLoader: - # Each call consumes a contiguous chunk from the shared token stream, then slices out - # one disjoint span per rank. The extra "+1" token lets us build (x, y) by shifting. - def __init__(self, pattern: str, rank: int, world_size: int, device: torch.device): - self.rank = rank - self.world_size = world_size - self.device = device - self.stream = TokenStream(pattern) - - def next_batch(self, global_tokens: int, seq_len: int, grad_accum_steps: int) -> tuple[Tensor, Tensor]: - local_tokens = global_tokens // (self.world_size * grad_accum_steps) - per_rank_span = local_tokens + 1 - chunk = self.stream.take(per_rank_span * self.world_size) - start = self.rank * per_rank_span - local = chunk[start : start + per_rank_span].to(dtype=torch.int64) - x = local[:-1].reshape(-1, seq_len) - y = local[1:].reshape(-1, seq_len) - return x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True) - -# ----------------------------- -# TRANSFORMER MODULES -# ----------------------------- - -class RMSNorm(nn.Module): - def __init__(self, eps: float | None = None): - super().__init__() - self.eps = eps - - def forward(self, x: Tensor) -> Tensor: - return F.rms_norm(x, (x.size(-1),), eps=self.eps) - - -class CastedLinear(nn.Linear): - # Keep weights in fp32 for optimizer/state quality, cast at matmul time for bf16 compute. - def forward(self, x: Tensor) -> Tensor: - bias = self.bias.to(x.dtype) if self.bias is not None else None - return F.linear(x, self.weight.to(x.dtype), bias) - - -def restore_low_dim_params_to_fp32(module: nn.Module) -> None: - # Keep small/control parameters in fp32 even when the model body runs in bf16. - with torch.no_grad(): - for name, param in module.named_parameters(): - if (param.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS)) and param.dtype != torch.float32: - param.data = param.data.float() - - -class Rotary(nn.Module): - # Caches cos/sin tables per sequence length on the current device. - def __init__(self, dim: int, base: float = 10000.0): - super().__init__() - inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self._seq_len_cached = 0 - self._cos_cached: Tensor | None = None - self._sin_cached: Tensor | None = None - - def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> tuple[Tensor, Tensor]: - if ( - self._cos_cached is None - or self._sin_cached is None - or self._seq_len_cached != seq_len - or self._cos_cached.device != device - ): - t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) - freqs = torch.outer(t, self.inv_freq.to(device)) - self._cos_cached = freqs.cos()[None, None, :, :] - self._sin_cached = freqs.sin()[None, None, :, :] - self._seq_len_cached = seq_len - return self._cos_cached.to(dtype=dtype), self._sin_cached.to(dtype=dtype) - - -def apply_rotary_emb(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: - half = x.size(-1) // 2 - x1, x2 = x[..., :half], x[..., half:] - return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1) - - -class CausalSelfAttention(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - num_kv_heads: int, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - if dim % num_heads != 0: - raise ValueError("model_dim must be divisible by num_heads") - if num_heads % num_kv_heads != 0: - raise ValueError("num_heads must be divisible by num_kv_heads") - self.num_heads = num_heads - self.num_kv_heads = num_kv_heads - self.head_dim = dim // num_heads - if self.head_dim % 2 != 0: - raise ValueError("head_dim must be even for RoPE") - kv_dim = self.num_kv_heads * self.head_dim - self.c_q = CastedLinear(dim, dim, bias=False) - self.c_k = CastedLinear(dim, kv_dim, bias=False) - self.c_v = CastedLinear(dim, kv_dim, bias=False) - self.proj = CastedLinear(dim, dim, bias=False) - self.proj._zero_init = True - self.q_gain = nn.Parameter(torch.full((num_heads,), qk_gain_init, dtype=torch.float32)) - self.rotary = Rotary(self.head_dim, base=rope_base) - - def forward(self, x: Tensor) -> Tensor: - bsz, seqlen, dim = x.shape - q = self.c_q(x).reshape(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2) - k = self.c_k(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - v = self.c_v(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - q = F.rms_norm(q, (q.size(-1),)) - k = F.rms_norm(k, (k.size(-1),)) - cos, sin = self.rotary(seqlen, x.device, q.dtype) - q = apply_rotary_emb(q, cos, sin) - k = apply_rotary_emb(k, cos, sin) - q = q * self.q_gain.to(dtype=q.dtype)[None, :, None, None] - y = F.scaled_dot_product_attention( - q, - k, - v, - attn_mask=None, - is_causal=True, - enable_gqa=(self.num_kv_heads != self.num_heads), - ) - y = y.transpose(1, 2).contiguous().reshape(bsz, seqlen, dim) - return self.proj(y) - - -class MLP(nn.Module): - # relu^2 MLP from the original modded-nanogpt setup - def __init__(self, dim: int, mlp_mult: int): - super().__init__() - hidden = mlp_mult * dim - self.fc = CastedLinear(dim, hidden, bias=False) - self.proj = CastedLinear(hidden, dim, bias=False) - self.proj._zero_init = True - - def forward(self, x: Tensor) -> Tensor: - x = torch.relu(self.fc(x)) - return self.proj(x.square()) - - -class Block(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - num_kv_heads: int, - mlp_mult: int, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - self.attn_norm = RMSNorm() - self.mlp_norm = RMSNorm() - self.attn = CausalSelfAttention(dim, num_heads, num_kv_heads, rope_base, qk_gain_init) - self.mlp = MLP(dim, mlp_mult) - self.attn_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) - self.mlp_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) - self.resid_mix = nn.Parameter(torch.stack((torch.ones(dim), torch.zeros(dim))).float()) - - def forward(self, x: Tensor, x0: Tensor) -> Tensor: - mix = self.resid_mix.to(dtype=x.dtype) - x = mix[0][None, None, :] * x + mix[1][None, None, :] * x0 - attn_out = self.attn(self.attn_norm(x)) - x = x + self.attn_scale.to(dtype=x.dtype)[None, None, :] * attn_out - x = x + self.mlp_scale.to(dtype=x.dtype)[None, None, :] * self.mlp(self.mlp_norm(x)) - return x - - -class GPT(nn.Module): - def __init__( - self, - vocab_size: int, - num_layers: int, - model_dim: int, - num_heads: int, - num_kv_heads: int, - mlp_mult: int, - tie_embeddings: bool, - tied_embed_init_std: float, - logit_softcap: float, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - if logit_softcap <= 0.0: - raise ValueError(f"logit_softcap must be positive, got {logit_softcap}") - self.tie_embeddings = tie_embeddings - self.tied_embed_init_std = tied_embed_init_std - self.logit_softcap = logit_softcap - self.tok_emb = nn.Embedding(vocab_size, model_dim) - self.num_encoder_layers = num_layers // 2 - self.num_decoder_layers = num_layers - self.num_encoder_layers - self.num_skip_weights = min(self.num_encoder_layers, self.num_decoder_layers) - self.skip_weights = nn.Parameter(torch.ones(self.num_skip_weights, model_dim, dtype=torch.float32)) - self.blocks = nn.ModuleList( - [ - Block( - model_dim, - num_heads, - num_kv_heads, - mlp_mult, - rope_base, - qk_gain_init, - ) - for i in range(num_layers) - ] - ) - self.final_norm = RMSNorm() - self.lm_head = None if tie_embeddings else CastedLinear(model_dim, vocab_size, bias=False) - if self.lm_head is not None: - self.lm_head._zero_init = True - self._init_weights() - - def _init_weights(self) -> None: - if self.tie_embeddings: - nn.init.normal_(self.tok_emb.weight, mean=0.0, std=self.tied_embed_init_std) - for module in self.modules(): - if isinstance(module, nn.Linear) and getattr(module, "_zero_init", False): - nn.init.zeros_(module.weight) - - def forward(self, input_ids: Tensor, target_ids: Tensor) -> Tensor: - x = self.tok_emb(input_ids) - x = F.rms_norm(x, (x.size(-1),)) - x0 = x - skips: list[Tensor] = [] - - # First half stores skips; second half reuses them in reverse order. - for i in range(self.num_encoder_layers): - x = self.blocks[i](x, x0) - skips.append(x) - for i in range(self.num_decoder_layers): - if skips: - x = x + self.skip_weights[i].to(dtype=x.dtype)[None, None, :] * skips.pop() - x = self.blocks[self.num_encoder_layers + i](x, x0) - - x = self.final_norm(x).reshape(-1, x.size(-1)) - targets = target_ids.reshape(-1) - if self.tie_embeddings: - logits_proj = F.linear(x, self.tok_emb.weight) - else: - if self.lm_head is None: - raise RuntimeError("lm_head is required when tie_embeddings=False") - logits_proj = self.lm_head(x) - logits = self.logit_softcap * torch.tanh(logits_proj / self.logit_softcap) - return F.cross_entropy(logits.float(), targets, reduction="mean") - - -# ----------------------------- -# TRAINING -# ----------------------------- - -def main() -> None: - global zeropower_via_newtonschulz5 - - code = Path(__file__).read_text(encoding="utf-8") - args = Hyperparameters() - zeropower_via_newtonschulz5 = torch.compile(zeropower_via_newtonschulz5) - - # ----------------------------- - # DISTRIBUTED + CUDA SETUP - # ----------------------------- - - distributed = "RANK" in os.environ and "WORLD_SIZE" in os.environ - rank = int(os.environ.get("RANK", "0")) - world_size = int(os.environ.get("WORLD_SIZE", "1")) - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - if world_size <= 0: - raise ValueError(f"WORLD_SIZE must be positive, got {world_size}") - if 8 % world_size != 0: - raise ValueError(f"WORLD_SIZE={world_size} must divide 8 so grad_accum_steps stays integral") - grad_accum_steps = 8 // world_size - grad_scale = 1.0 / grad_accum_steps - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required") - device = torch.device("cuda", local_rank) - torch.cuda.set_device(device) - if distributed: - dist.init_process_group(backend="nccl", device_id=device) - dist.barrier() - master_process = rank == 0 - - # Fast math knobs - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True - from torch.backends.cuda import enable_cudnn_sdp, enable_flash_sdp, enable_math_sdp, enable_mem_efficient_sdp - - enable_cudnn_sdp(False) - enable_flash_sdp(True) - enable_mem_efficient_sdp(False) - enable_math_sdp(False) - - logfile = None - if master_process: - os.makedirs("logs", exist_ok=True) - logfile = f"logs/{args.run_id}.txt" - print(logfile) - - def log0(msg: str, console: bool = True) -> None: - if not master_process: - return - if console: - print(msg) - if logfile is not None: - with open(logfile, "a", encoding="utf-8") as f: - print(msg, file=f) - - log0(code, console=False) - log0("=" * 100, console=False) - log0(f"Running Python {sys.version}", console=False) - log0(f"Running PyTorch {torch.__version__}", console=False) - log0( - subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False).stdout, - console=False, - ) - log0("=" * 100, console=False) - - # ----------------------------- - # TOKENIZER + VALIDATION METRIC SETUP - # ----------------------------- - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed_all(args.seed) - - if not args.tokenizer_path.endswith(".model"): - raise ValueError(f"Script only setup for SentencePiece .model file: {args.tokenizer_path}") - sp = spm.SentencePieceProcessor(model_file=args.tokenizer_path) - if int(sp.vocab_size()) != args.vocab_size: - raise ValueError( - f"VOCAB_SIZE={args.vocab_size} does not match tokenizer vocab_size={int(sp.vocab_size())}" - ) - dataset_dir = Path(args.data_path).resolve() - actual_train_files = len(list(dataset_dir.glob("fineweb_train_*.bin"))) - val_tokens = load_validation_tokens(args.val_files, args.train_seq_len) - base_bytes_lut, has_leading_space_lut, is_boundary_token_lut = build_sentencepiece_luts( - sp, args.vocab_size, device - ) - log0(f"val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path={args.tokenizer_path}") - log0(f"train_loader:dataset:{dataset_dir.name} train_shards:{actual_train_files}") - log0(f"val_loader:shards pattern={args.val_files} tokens:{val_tokens.numel() - 1}") - - # ----------------------------- - # MODEL + OPTIMIZER SETUP - # ----------------------------- - - base_model = GPT( - vocab_size=args.vocab_size, - num_layers=args.num_layers, - model_dim=args.model_dim, - num_heads=args.num_heads, - num_kv_heads=args.num_kv_heads, - mlp_mult=args.mlp_mult, - tie_embeddings=args.tie_embeddings, - tied_embed_init_std=args.tied_embed_init_std, - logit_softcap=args.logit_softcap, - rope_base=args.rope_base, - qk_gain_init=args.qk_gain_init, - ).to(device).bfloat16() - for module in base_model.modules(): - if isinstance(module, CastedLinear): - module.float() - restore_low_dim_params_to_fp32(base_model) - compiled_model = torch.compile(base_model, dynamic=False, fullgraph=True) - model: nn.Module = DDP(compiled_model, device_ids=[local_rank], broadcast_buffers=False) if distributed else compiled_model - - # Optimizer split: - # - token embedding (Adam) uses EMBED_LR - # - untied lm_head (Adam) uses HEAD_LR - # - matrix params in transformer blocks use MATRIX_LR via Muon - # - vectors/scalars use SCALAR_LR via Adam - block_named_params = list(base_model.blocks.named_parameters()) - matrix_params = [ - p - for name, p in block_named_params - if p.ndim == 2 and not any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - scalar_params = [ - p - for name, p in block_named_params - if p.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - if base_model.skip_weights.numel() > 0: - scalar_params.append(base_model.skip_weights) - token_lr = args.tied_embed_lr if args.tie_embeddings else args.embed_lr - optimizer_tok = torch.optim.Adam( - [{"params": [base_model.tok_emb.weight], "lr": token_lr, "base_lr": token_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizer_muon = Muon( - matrix_params, - lr=args.matrix_lr, - momentum=args.muon_momentum, - backend_steps=args.muon_backend_steps, - ) - for group in optimizer_muon.param_groups: - group["base_lr"] = args.matrix_lr - optimizer_scalar = torch.optim.Adam( - [{"params": scalar_params, "lr": args.scalar_lr, "base_lr": args.scalar_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers: list[torch.optim.Optimizer] = [optimizer_tok, optimizer_muon, optimizer_scalar] - if base_model.lm_head is not None: - optimizer_head = torch.optim.Adam( - [{"params": [base_model.lm_head.weight], "lr": args.head_lr, "base_lr": args.head_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers.insert(1, optimizer_head) - - n_params = sum(p.numel() for p in base_model.parameters()) - log0(f"model_params:{n_params}") - log0(f"world_size:{world_size} grad_accum_steps:{grad_accum_steps}") - log0("sdp_backends:cudnn=False flash=True mem_efficient=False math=False") - log0(f"attention_mode:gqa num_heads:{args.num_heads} num_kv_heads:{args.num_kv_heads}") - log0( - f"tie_embeddings:{args.tie_embeddings} embed_lr:{token_lr} " - f"head_lr:{args.head_lr if base_model.lm_head is not None else 0.0} " - f"matrix_lr:{args.matrix_lr} scalar_lr:{args.scalar_lr}" - ) - log0( - f"train_batch_tokens:{args.train_batch_tokens} train_seq_len:{args.train_seq_len} " - f"iterations:{args.iterations} warmup_steps:{args.warmup_steps} " - f"max_wallclock_seconds:{args.max_wallclock_seconds:.3f}" - ) - log0(f"seed:{args.seed}") - - # ----------------------------- - # DATA LOADER & MODEL WARMUP - # ----------------------------- - - train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device) - - def zero_grad_all() -> None: - for opt in optimizers: - opt.zero_grad(set_to_none=True) - - max_wallclock_ms = 1000.0 * args.max_wallclock_seconds if args.max_wallclock_seconds > 0 else None - - def lr_mul(step: int, elapsed_ms: float) -> float: - if args.warmdown_iters <= 0: - return 1.0 - if max_wallclock_ms is None: - warmdown_start = max(args.iterations - args.warmdown_iters, 0) - return max((args.iterations - step) / max(args.warmdown_iters, 1), 0.0) if warmdown_start <= step < args.iterations else 1.0 - step_ms = elapsed_ms / max(step, 1) - warmdown_ms = args.warmdown_iters * step_ms - remaining_ms = max(max_wallclock_ms - elapsed_ms, 0.0) - return remaining_ms / max(warmdown_ms, 1e-9) if remaining_ms <= warmdown_ms else 1.0 - - # Warmup primes the compiled forward/backward/optimizer paths, then we restore the - # initial weights/optimizer state so measured training starts from the true init. - if args.warmup_steps > 0: - initial_model_state = {name: tensor.detach().cpu().clone() for name, tensor in base_model.state_dict().items()} - initial_optimizer_states = [copy.deepcopy(opt.state_dict()) for opt in optimizers] - model.train() - for warmup_step in range(args.warmup_steps): - zero_grad_all() - for micro_step in range(grad_accum_steps): - if distributed: - model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 - x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - warmup_loss = model(x, y) - (warmup_loss * grad_scale).backward() - for opt in optimizers: - opt.step() - zero_grad_all() - if args.warmup_steps <= 20 or (warmup_step + 1) % 10 == 0 or warmup_step + 1 == args.warmup_steps: - log0(f"warmup_step:{warmup_step + 1}/{args.warmup_steps}") - base_model.load_state_dict(initial_model_state, strict=True) - for opt, state in zip(optimizers, initial_optimizer_states, strict=True): - opt.load_state_dict(state) - zero_grad_all() - if distributed: - model.require_backward_grad_sync = True - train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device) - - # ----------------------------- - # MAIN TRAINING LOOP - # ----------------------------- - - training_time_ms = 0.0 - stop_after_step: int | None = None - torch.cuda.synchronize() - t0 = time.perf_counter() - - step = 0 - while True: - last_step = step == args.iterations or (stop_after_step is not None and step >= stop_after_step) - - should_validate = last_step or (args.val_loss_every > 0 and step % args.val_loss_every == 0) - if should_validate: - torch.cuda.synchronize() - training_time_ms += 1000.0 * (time.perf_counter() - t0) - val_loss, val_bpb = eval_val( - args, - model, - rank, - world_size, - device, - grad_accum_steps, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - ) - log0( - f"step:{step}/{args.iterations} val_loss:{val_loss:.4f} val_bpb:{val_bpb:.4f} " - f"train_time:{training_time_ms:.0f}ms step_avg:{training_time_ms / max(step, 1):.2f}ms" - ) - torch.cuda.synchronize() - t0 = time.perf_counter() - - if last_step: - if stop_after_step is not None and step < args.iterations: - log0( - f"stopping_early: wallclock_cap train_time:{training_time_ms:.0f}ms " - f"step:{step}/{args.iterations}" - ) - break - - elapsed_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) - scale = lr_mul(step, elapsed_ms) - zero_grad_all() - train_loss = torch.zeros((), device=device) - for micro_step in range(grad_accum_steps): - if distributed: - model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 - x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - loss = model(x, y) - train_loss += loss.detach() - (loss * grad_scale).backward() - train_loss /= grad_accum_steps - - frac = min(step / args.muon_momentum_warmup_steps, 1.0) if args.muon_momentum_warmup_steps > 0 else 1.0 - muon_momentum = (1 - frac) * args.muon_momentum_warmup_start + frac * args.muon_momentum - for group in optimizer_muon.param_groups: - group["momentum"] = muon_momentum - - for opt in optimizers: - for group in opt.param_groups: - group["lr"] = group["base_lr"] * scale - - if args.grad_clip_norm > 0: - torch.nn.utils.clip_grad_norm_(base_model.parameters(), args.grad_clip_norm) - for opt in optimizers: - opt.step() - zero_grad_all() - - step += 1 - approx_training_time_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) - should_log_train = ( - args.train_log_every > 0 - and (step <= 10 or step % args.train_log_every == 0 or stop_after_step is not None) - ) - if should_log_train: - log0( - f"step:{step}/{args.iterations} train_loss:{train_loss.item():.4f} " - f"train_time:{approx_training_time_ms:.0f}ms step_avg:{approx_training_time_ms / step:.2f}ms" - ) - - # Needed to sync whether we've reached the wallclock cap. - reached_cap = max_wallclock_ms is not None and approx_training_time_ms >= max_wallclock_ms - if distributed and max_wallclock_ms is not None: - reached_cap_tensor = torch.tensor(int(reached_cap), device=device) - dist.all_reduce(reached_cap_tensor, op=dist.ReduceOp.MAX) - reached_cap = bool(reached_cap_tensor.item()) - if stop_after_step is None and reached_cap: - stop_after_step = step - - log0( - f"peak memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB " - f"reserved: {torch.cuda.max_memory_reserved() // 1024 // 1024} MiB" - ) - - # ----------------------------- - # SERIALIZATION + ROUNDTRIP VALIDATION - # ----------------------------- - # Save the raw state (useful for debugging/loading in PyTorch directly), then always produce - # the compressed int8+zlib artifact and validate the round-tripped weights. - - if master_process: - torch.save(base_model.state_dict(), "final_model.pt") - model_bytes = os.path.getsize("final_model.pt") - code_bytes = len(code.encode("utf-8")) - log0(f"Serialized model: {model_bytes} bytes") - log0(f"Code size: {code_bytes} bytes") - log0(f"Total submission size: {model_bytes + code_bytes} bytes") - - quant_obj, quant_stats = quantize_state_dict_int8(base_model.state_dict()) - quant_buf = io.BytesIO() - torch.save(quant_obj, quant_buf) - quant_raw = quant_buf.getvalue() - quant_blob = zlib.compress(quant_raw, level=9) - quant_raw_bytes = len(quant_raw) - if master_process: - with open("final_model.int8.ptz", "wb") as f: - f.write(quant_blob) - quant_file_bytes = os.path.getsize("final_model.int8.ptz") - code_bytes = len(code.encode("utf-8")) - ratio = quant_stats["baseline_tensor_bytes"] / max(quant_stats["int8_payload_bytes"], 1) - log0( - f"Serialized model int8+zlib: {quant_file_bytes} bytes " - f"(payload:{quant_stats['int8_payload_bytes']} raw_torch:{quant_raw_bytes} payload_ratio:{ratio:.2f}x)" - ) - log0(f"Total submission size int8+zlib: {quant_file_bytes + code_bytes} bytes") - - if distributed: - dist.barrier() - with open("final_model.int8.ptz", "rb") as f: - quant_blob_disk = f.read() - quant_state = torch.load(io.BytesIO(zlib.decompress(quant_blob_disk)), map_location="cpu") - base_model.load_state_dict(dequantize_state_dict_int8(quant_state), strict=True) - torch.cuda.synchronize() - t_qeval = time.perf_counter() - q_val_loss, q_val_bpb = eval_val( - args, - model, - rank, - world_size, - device, - grad_accum_steps, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - ) - torch.cuda.synchronize() - log0( - f"final_int8_zlib_roundtrip val_loss:{q_val_loss:.4f} val_bpb:{q_val_bpb:.4f} " - f"eval_time:{1000.0 * (time.perf_counter() - t_qeval):.0f}ms" - ) - log0(f"final_int8_zlib_roundtrip_exact val_loss:{q_val_loss:.8f} val_bpb:{q_val_bpb:.8f}") - - if distributed: - dist.destroy_process_group() - - -if __name__ == "__main__": - main() - -==================================================================================================== -Running Python 3.12.9 (main, Jan 22 2026, 18:37:37) [GCC 13.3.0] -Running PyTorch 2.10.0+cu128 -Wed Mar 18 07:51:20 2026 -+-----------------------------------------------------------------------------------------+ -| NVIDIA-SMI 565.57.01 Driver Version: 565.57.01 CUDA Version: 13.0 | -|-----------------------------------------+------------------------+----------------------+ -| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC | -| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. | -| | | MIG M. | -|=========================================+========================+======================| -| 0 NVIDIA H100 80GB HBM3 On | 00000001:00:00.0 Off | 0 | -| N/A 31C P0 119W / 700W | 1774MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 1 NVIDIA H100 80GB HBM3 On | 00000002:00:00.0 Off | 0 | -| N/A 31C P0 120W / 700W | 1822MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 2 NVIDIA H100 80GB HBM3 On | 00000003:00:00.0 Off | 0 | -| N/A 30C P0 117W / 700W | 1822MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 3 NVIDIA H100 80GB HBM3 On | 00000008:00:00.0 Off | 0 | -| N/A 30C P0 119W / 700W | 1822MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 4 NVIDIA H100 80GB HBM3 On | 00000009:00:00.0 Off | 0 | -| N/A 29C P0 120W / 700W | 1822MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 5 NVIDIA H100 80GB HBM3 On | 0000000A:00:00.0 Off | 0 | -| N/A 31C P0 119W / 700W | 1822MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 6 NVIDIA H100 80GB HBM3 On | 0000000B:00:00.0 Off | 0 | -| N/A 30C P0 116W / 700W | 1822MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ -| 7 NVIDIA H100 80GB HBM3 On | 0000000C:00:00.0 Off | 0 | -| N/A 30C P0 118W / 700W | 1582MiB / 81559MiB | 0% Default | -| | | Disabled | -+-----------------------------------------+------------------------+----------------------+ - -+-----------------------------------------------------------------------------------------+ -| Processes: | -| GPU GI CI PID Type Process name GPU Memory | -| ID ID Usage | -|=========================================================================================| -| 0 N/A N/A 987032 C ...t/.pyenv/versions/3.12.9/bin/python 1764MiB | -| 1 N/A N/A 987033 C ...t/.pyenv/versions/3.12.9/bin/python 1812MiB | -| 2 N/A N/A 987034 C ...t/.pyenv/versions/3.12.9/bin/python 1812MiB | -| 3 N/A N/A 987035 C ...t/.pyenv/versions/3.12.9/bin/python 1812MiB | -| 4 N/A N/A 987036 C ...t/.pyenv/versions/3.12.9/bin/python 1812MiB | -| 5 N/A N/A 987037 C ...t/.pyenv/versions/3.12.9/bin/python 1812MiB | -| 6 N/A N/A 987038 C ...t/.pyenv/versions/3.12.9/bin/python 1812MiB | -| 7 N/A N/A 987039 C ...t/.pyenv/versions/3.12.9/bin/python 1572MiB | -+-----------------------------------------------------------------------------------------+ - -==================================================================================================== -val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path=/tmp/fineweb_quasi10Bfrom50B_50keval_sp1024_v0/tokenizers/fineweb_1024_bpe.model -train_loader:dataset:fineweb10B_sp1024 train_shards:195 -val_loader:shards pattern=/tmp/fineweb_quasi10Bfrom50B_50keval_sp1024_v0/datasets/fineweb10B_sp1024/fineweb_val_*.bin tokens:62021632 -model_params:17059912 -world_size:8 grad_accum_steps:1 -sdp_backends:cudnn=False flash=True mem_efficient=False math=False -attention_mode:gqa num_heads:8 num_kv_heads:4 -tie_embeddings:True embed_lr:0.05 head_lr:0.0 matrix_lr:0.04 scalar_lr:0.04 -train_batch_tokens:524288 train_seq_len:1024 iterations:500000 warmup_steps:20 max_wallclock_seconds:14400.000 -seed:1337 -warmup_step:1/20 -warmup_step:2/20 -warmup_step:3/20 -warmup_step:4/20 -warmup_step:5/20 -warmup_step:6/20 -warmup_step:7/20 -warmup_step:8/20 -warmup_step:9/20 -warmup_step:10/20 -warmup_step:11/20 -warmup_step:12/20 -warmup_step:13/20 -warmup_step:14/20 -warmup_step:15/20 -warmup_step:16/20 -warmup_step:17/20 -warmup_step:18/20 -warmup_step:19/20 -warmup_step:20/20 -step:0/500000 val_loss:6.9357 val_bpb:4.1077 train_time:0ms step_avg:0.02ms -step:1/500000 train_loss:6.9370 train_time:33ms step_avg:32.84ms -step:2/500000 train_loss:16.8366 train_time:77ms step_avg:38.37ms -step:3/500000 train_loss:8.7611 train_time:120ms step_avg:39.91ms -step:4/500000 train_loss:6.6384 train_time:162ms step_avg:40.39ms -step:5/500000 train_loss:6.6118 train_time:206ms step_avg:41.14ms -step:6/500000 train_loss:7.4217 train_time:249ms step_avg:41.44ms -step:7/500000 train_loss:6.3501 train_time:292ms step_avg:41.72ms -step:8/500000 train_loss:6.1581 train_time:336ms step_avg:41.97ms -step:9/500000 train_loss:6.0678 train_time:378ms step_avg:42.02ms -step:10/500000 train_loss:5.9746 train_time:422ms step_avg:42.15ms -step:200/500000 train_loss:2.8504 train_time:8738ms step_avg:43.69ms -step:400/500000 train_loss:2.3638 train_time:17484ms step_avg:43.71ms -step:600/500000 train_loss:2.5453 train_time:26216ms step_avg:43.69ms -step:800/500000 train_loss:2.2913 train_time:34960ms step_avg:43.70ms -step:1000/500000 train_loss:2.3707 train_time:43708ms step_avg:43.71ms -step:1200/500000 train_loss:2.3889 train_time:52441ms step_avg:43.70ms -step:1400/500000 train_loss:2.4331 train_time:61180ms step_avg:43.70ms -step:1600/500000 train_loss:2.0960 train_time:69912ms step_avg:43.70ms -step:1800/500000 train_loss:2.1998 train_time:78649ms step_avg:43.69ms -step:2000/500000 train_loss:2.2506 train_time:87392ms step_avg:43.70ms -step:2200/500000 train_loss:2.0768 train_time:96125ms step_avg:43.69ms -step:2400/500000 train_loss:2.2012 train_time:104860ms step_avg:43.69ms -step:2600/500000 train_loss:2.4155 train_time:113597ms step_avg:43.69ms -step:2800/500000 train_loss:2.2353 train_time:122349ms step_avg:43.70ms -step:3000/500000 train_loss:2.2275 train_time:131087ms step_avg:43.70ms -step:3200/500000 train_loss:2.1879 train_time:139825ms step_avg:43.70ms -step:3400/500000 train_loss:2.1611 train_time:148561ms step_avg:43.69ms -step:3600/500000 train_loss:2.1179 train_time:157297ms step_avg:43.69ms -step:3800/500000 train_loss:2.2211 train_time:166018ms step_avg:43.69ms -step:4000/500000 train_loss:2.1627 train_time:174758ms step_avg:43.69ms -step:4200/500000 train_loss:2.1757 train_time:183623ms step_avg:43.72ms -step:4400/500000 train_loss:2.1143 train_time:192357ms step_avg:43.72ms -step:4600/500000 train_loss:1.9730 train_time:201085ms step_avg:43.71ms -step:4800/500000 train_loss:2.2618 train_time:209818ms step_avg:43.71ms -step:5000/500000 train_loss:2.0275 train_time:218548ms step_avg:43.71ms -step:5200/500000 train_loss:2.1723 train_time:227276ms step_avg:43.71ms -step:5400/500000 train_loss:2.1842 train_time:235999ms step_avg:43.70ms -step:5600/500000 train_loss:2.1841 train_time:244724ms step_avg:43.70ms -step:5800/500000 train_loss:2.1435 train_time:253451ms step_avg:43.70ms -step:6000/500000 train_loss:2.2229 train_time:262182ms step_avg:43.70ms -step:6200/500000 train_loss:2.0860 train_time:270912ms step_avg:43.70ms -step:6400/500000 train_loss:2.1598 train_time:279652ms step_avg:43.70ms -step:6600/500000 train_loss:2.1228 train_time:288379ms step_avg:43.69ms -step:6800/500000 train_loss:2.1890 train_time:297103ms step_avg:43.69ms -step:7000/500000 train_loss:2.2279 train_time:305832ms step_avg:43.69ms -step:7200/500000 train_loss:2.1999 train_time:314562ms step_avg:43.69ms -step:7400/500000 train_loss:2.1173 train_time:323296ms step_avg:43.69ms -step:7600/500000 train_loss:1.9995 train_time:332032ms step_avg:43.69ms -step:7800/500000 train_loss:2.1500 train_time:340763ms step_avg:43.69ms -step:8000/500000 train_loss:2.1150 train_time:349492ms step_avg:43.69ms -step:8200/500000 train_loss:2.1884 train_time:358222ms step_avg:43.69ms -step:8400/500000 train_loss:2.1346 train_time:367069ms step_avg:43.70ms -step:8600/500000 train_loss:2.1364 train_time:375806ms step_avg:43.70ms -step:8800/500000 train_loss:2.0998 train_time:384530ms step_avg:43.70ms -step:9000/500000 train_loss:2.0268 train_time:393259ms step_avg:43.70ms -step:9200/500000 train_loss:2.0837 train_time:401988ms step_avg:43.69ms -step:9400/500000 train_loss:2.1327 train_time:410718ms step_avg:43.69ms -step:9600/500000 train_loss:2.1452 train_time:419449ms step_avg:43.69ms -step:9800/500000 train_loss:2.0746 train_time:428188ms step_avg:43.69ms -step:10000/500000 train_loss:2.1114 train_time:436922ms step_avg:43.69ms -step:10200/500000 train_loss:2.0700 train_time:445656ms step_avg:43.69ms -step:10400/500000 train_loss:2.0994 train_time:454390ms step_avg:43.69ms -step:10600/500000 train_loss:1.9771 train_time:463127ms step_avg:43.69ms -step:10800/500000 train_loss:2.1884 train_time:471870ms step_avg:43.69ms -step:11000/500000 train_loss:2.1136 train_time:480607ms step_avg:43.69ms -step:11200/500000 train_loss:2.0714 train_time:489337ms step_avg:43.69ms -step:11400/500000 train_loss:2.0570 train_time:498065ms step_avg:43.69ms -step:11600/500000 train_loss:2.0643 train_time:506798ms step_avg:43.69ms -step:11800/500000 train_loss:2.0966 train_time:515533ms step_avg:43.69ms -step:12000/500000 train_loss:2.0734 train_time:524265ms step_avg:43.69ms -step:12200/500000 train_loss:2.2214 train_time:532995ms step_avg:43.69ms -step:12400/500000 train_loss:1.8632 train_time:541847ms step_avg:43.70ms -step:12600/500000 train_loss:2.0944 train_time:550591ms step_avg:43.70ms -step:12800/500000 train_loss:2.1093 train_time:559322ms step_avg:43.70ms -step:13000/500000 train_loss:2.1906 train_time:568053ms step_avg:43.70ms -step:13200/500000 train_loss:2.2034 train_time:576786ms step_avg:43.70ms -step:13400/500000 train_loss:2.0818 train_time:585518ms step_avg:43.70ms -step:13600/500000 train_loss:1.9545 train_time:594253ms step_avg:43.70ms -step:13800/500000 train_loss:2.0350 train_time:602978ms step_avg:43.69ms -step:14000/500000 train_loss:2.0960 train_time:611714ms step_avg:43.69ms -step:14200/500000 train_loss:2.1828 train_time:620450ms step_avg:43.69ms -step:14400/500000 train_loss:2.0819 train_time:629189ms step_avg:43.69ms -step:14600/500000 train_loss:2.1389 train_time:637923ms step_avg:43.69ms -step:14800/500000 train_loss:1.9175 train_time:646657ms step_avg:43.69ms -step:15000/500000 train_loss:2.0332 train_time:655403ms step_avg:43.69ms -step:15200/500000 train_loss:2.1406 train_time:664139ms step_avg:43.69ms -step:15400/500000 train_loss:2.0916 train_time:672882ms step_avg:43.69ms -step:15600/500000 train_loss:2.1547 train_time:681620ms step_avg:43.69ms -step:15800/500000 train_loss:1.8204 train_time:690352ms step_avg:43.69ms -step:16000/500000 train_loss:2.1014 train_time:699084ms step_avg:43.69ms -step:16200/500000 train_loss:2.1821 train_time:707820ms step_avg:43.69ms -step:16400/500000 train_loss:2.1894 train_time:716552ms step_avg:43.69ms -step:16600/500000 train_loss:2.0465 train_time:725412ms step_avg:43.70ms -step:16800/500000 train_loss:1.9392 train_time:734145ms step_avg:43.70ms -step:17000/500000 train_loss:1.9603 train_time:742889ms step_avg:43.70ms -step:17200/500000 train_loss:2.0424 train_time:751641ms step_avg:43.70ms -step:17400/500000 train_loss:2.0173 train_time:760388ms step_avg:43.70ms -step:17600/500000 train_loss:2.0272 train_time:769119ms step_avg:43.70ms -step:17800/500000 train_loss:2.0577 train_time:777863ms step_avg:43.70ms -step:18000/500000 train_loss:2.1233 train_time:786590ms step_avg:43.70ms -step:18200/500000 train_loss:2.2190 train_time:795310ms step_avg:43.70ms -step:18400/500000 train_loss:2.0528 train_time:804037ms step_avg:43.70ms -step:18600/500000 train_loss:2.0409 train_time:812766ms step_avg:43.70ms -step:18800/500000 train_loss:2.1861 train_time:821488ms step_avg:43.70ms -step:19000/500000 train_loss:2.0094 train_time:830217ms step_avg:43.70ms -step:19200/500000 train_loss:2.1793 train_time:838942ms step_avg:43.69ms -step:19400/500000 train_loss:2.1423 train_time:847676ms step_avg:43.69ms -step:19600/500000 train_loss:1.9673 train_time:856408ms step_avg:43.69ms -step:19800/500000 train_loss:2.2609 train_time:865135ms step_avg:43.69ms -step:20000/500000 train_loss:1.9800 train_time:873911ms step_avg:43.70ms -step:20000/500000 val_loss:2.0779 val_bpb:1.2307 train_time:873923ms step_avg:43.70ms -step:20200/500000 train_loss:2.0109 train_time:882650ms step_avg:43.70ms -step:20400/500000 train_loss:2.0714 train_time:891387ms step_avg:43.70ms -step:20600/500000 train_loss:2.0559 train_time:900265ms step_avg:43.70ms -step:20800/500000 train_loss:2.0586 train_time:909008ms step_avg:43.70ms -step:21000/500000 train_loss:2.0829 train_time:917731ms step_avg:43.70ms -step:21200/500000 train_loss:2.0840 train_time:926484ms step_avg:43.70ms -step:21400/500000 train_loss:2.0728 train_time:935227ms step_avg:43.70ms -step:21600/500000 train_loss:1.9781 train_time:943962ms step_avg:43.70ms -step:21800/500000 train_loss:2.1248 train_time:952707ms step_avg:43.70ms -step:22000/500000 train_loss:2.0181 train_time:961450ms step_avg:43.70ms -step:22200/500000 train_loss:1.9460 train_time:970185ms step_avg:43.70ms -step:22400/500000 train_loss:2.1037 train_time:978935ms step_avg:43.70ms -step:22600/500000 train_loss:2.0755 train_time:987675ms step_avg:43.70ms -step:22800/500000 train_loss:2.0061 train_time:996415ms step_avg:43.70ms -step:23000/500000 train_loss:2.1035 train_time:1005151ms step_avg:43.70ms -step:23200/500000 train_loss:2.0248 train_time:1013891ms step_avg:43.70ms -step:23400/500000 train_loss:2.0755 train_time:1022624ms step_avg:43.70ms -step:23600/500000 train_loss:2.0633 train_time:1031363ms step_avg:43.70ms -step:23800/500000 train_loss:2.0454 train_time:1040107ms step_avg:43.70ms -step:24000/500000 train_loss:2.0450 train_time:1048848ms step_avg:43.70ms -step:24200/500000 train_loss:2.1111 train_time:1057577ms step_avg:43.70ms -step:24400/500000 train_loss:2.2038 train_time:1066306ms step_avg:43.70ms -step:24600/500000 train_loss:2.1040 train_time:1075040ms step_avg:43.70ms -step:24800/500000 train_loss:2.0724 train_time:1083908ms step_avg:43.71ms -step:25000/500000 train_loss:2.0719 train_time:1092641ms step_avg:43.71ms -step:25200/500000 train_loss:2.0243 train_time:1101373ms step_avg:43.71ms -step:25400/500000 train_loss:1.9389 train_time:1110103ms step_avg:43.70ms -step:25600/500000 train_loss:1.9308 train_time:1118832ms step_avg:43.70ms -step:25800/500000 train_loss:2.0562 train_time:1127565ms step_avg:43.70ms -step:26000/500000 train_loss:2.1770 train_time:1136291ms step_avg:43.70ms -step:26200/500000 train_loss:2.0391 train_time:1145021ms step_avg:43.70ms -step:26400/500000 train_loss:1.9735 train_time:1153756ms step_avg:43.70ms -step:26600/500000 train_loss:2.0604 train_time:1162477ms step_avg:43.70ms -step:26800/500000 train_loss:1.9815 train_time:1171216ms step_avg:43.70ms -step:27000/500000 train_loss:2.0145 train_time:1179957ms step_avg:43.70ms -step:27200/500000 train_loss:2.0304 train_time:1188694ms step_avg:43.70ms -step:27400/500000 train_loss:2.0649 train_time:1197422ms step_avg:43.70ms -step:27600/500000 train_loss:2.0879 train_time:1206157ms step_avg:43.70ms -step:27800/500000 train_loss:2.2113 train_time:1214902ms step_avg:43.70ms -step:28000/500000 train_loss:2.1648 train_time:1223631ms step_avg:43.70ms -step:28200/500000 train_loss:2.0987 train_time:1232368ms step_avg:43.70ms -step:28400/500000 train_loss:2.1582 train_time:1241115ms step_avg:43.70ms -step:28600/500000 train_loss:2.0779 train_time:1249862ms step_avg:43.70ms -step:28800/500000 train_loss:2.0425 train_time:1258604ms step_avg:43.70ms -step:29000/500000 train_loss:2.1468 train_time:1267458ms step_avg:43.71ms -step:29200/500000 train_loss:2.1703 train_time:1276186ms step_avg:43.70ms -step:29400/500000 train_loss:2.0395 train_time:1284915ms step_avg:43.70ms -step:29600/500000 train_loss:2.0624 train_time:1293655ms step_avg:43.70ms -step:29800/500000 train_loss:2.0755 train_time:1302387ms step_avg:43.70ms -step:30000/500000 train_loss:2.0016 train_time:1311121ms step_avg:43.70ms -step:30200/500000 train_loss:2.0520 train_time:1319857ms step_avg:43.70ms -step:30400/500000 train_loss:1.9604 train_time:1328595ms step_avg:43.70ms -step:30600/500000 train_loss:2.0195 train_time:1337334ms step_avg:43.70ms -step:30800/500000 train_loss:2.0386 train_time:1346076ms step_avg:43.70ms -step:31000/500000 train_loss:2.1222 train_time:1354809ms step_avg:43.70ms -step:31200/500000 train_loss:2.1388 train_time:1363547ms step_avg:43.70ms -step:31400/500000 train_loss:2.1097 train_time:1372279ms step_avg:43.70ms -step:31600/500000 train_loss:2.0476 train_time:1381026ms step_avg:43.70ms -step:31800/500000 train_loss:2.0700 train_time:1389780ms step_avg:43.70ms -step:32000/500000 train_loss:2.1175 train_time:1398523ms step_avg:43.70ms -step:32200/500000 train_loss:2.0473 train_time:1407258ms step_avg:43.70ms -step:32400/500000 train_loss:1.9902 train_time:1416002ms step_avg:43.70ms -step:32600/500000 train_loss:2.0894 train_time:1424736ms step_avg:43.70ms -step:32800/500000 train_loss:2.1487 train_time:1433470ms step_avg:43.70ms -step:33000/500000 train_loss:2.0049 train_time:1442332ms step_avg:43.71ms -step:33200/500000 train_loss:2.0618 train_time:1451061ms step_avg:43.71ms -step:33400/500000 train_loss:1.9522 train_time:1459784ms step_avg:43.71ms -step:33600/500000 train_loss:2.1620 train_time:1468519ms step_avg:43.71ms -step:33800/500000 train_loss:2.0664 train_time:1477245ms step_avg:43.71ms -step:34000/500000 train_loss:2.3050 train_time:1485969ms step_avg:43.70ms -step:34200/500000 train_loss:2.1240 train_time:1494699ms step_avg:43.70ms -step:34400/500000 train_loss:2.0971 train_time:1503428ms step_avg:43.70ms -step:34600/500000 train_loss:2.1775 train_time:1512165ms step_avg:43.70ms -step:34800/500000 train_loss:2.0770 train_time:1520903ms step_avg:43.70ms -step:35000/500000 train_loss:2.0156 train_time:1529639ms step_avg:43.70ms -step:35200/500000 train_loss:2.0252 train_time:1538380ms step_avg:43.70ms -step:35400/500000 train_loss:2.0857 train_time:1547117ms step_avg:43.70ms -step:35600/500000 train_loss:2.2012 train_time:1555857ms step_avg:43.70ms -step:35800/500000 train_loss:2.0054 train_time:1564595ms step_avg:43.70ms -step:36000/500000 train_loss:2.1126 train_time:1573347ms step_avg:43.70ms -step:36200/500000 train_loss:1.9196 train_time:1582070ms step_avg:43.70ms -step:36400/500000 train_loss:2.0816 train_time:1590801ms step_avg:43.70ms -step:36600/500000 train_loss:2.0577 train_time:1599527ms step_avg:43.70ms -step:36800/500000 train_loss:1.9859 train_time:1608254ms step_avg:43.70ms -step:37000/500000 train_loss:1.9921 train_time:1616984ms step_avg:43.70ms -step:37200/500000 train_loss:2.0478 train_time:1625821ms step_avg:43.70ms -step:37400/500000 train_loss:2.0508 train_time:1634571ms step_avg:43.71ms -step:37600/500000 train_loss:2.0218 train_time:1643314ms step_avg:43.71ms -step:37800/500000 train_loss:2.1115 train_time:1652053ms step_avg:43.71ms -step:38000/500000 train_loss:1.9950 train_time:1660788ms step_avg:43.70ms -step:38200/500000 train_loss:2.1599 train_time:1669527ms step_avg:43.70ms -step:38400/500000 train_loss:2.0854 train_time:1678274ms step_avg:43.71ms -step:38600/500000 train_loss:2.0777 train_time:1687012ms step_avg:43.70ms -step:38800/500000 train_loss:2.0381 train_time:1695754ms step_avg:43.71ms -step:39000/500000 train_loss:1.9603 train_time:1704494ms step_avg:43.70ms -step:39200/500000 train_loss:2.0654 train_time:1713230ms step_avg:43.70ms -step:39400/500000 train_loss:2.0960 train_time:1721966ms step_avg:43.70ms -step:39600/500000 train_loss:2.1160 train_time:1730705ms step_avg:43.70ms -step:39800/500000 train_loss:2.0706 train_time:1739449ms step_avg:43.70ms -step:40000/500000 train_loss:1.9931 train_time:1748179ms step_avg:43.70ms -step:40000/500000 val_loss:2.0513 val_bpb:1.2149 train_time:1748195ms step_avg:43.70ms -step:40200/500000 train_loss:2.0616 train_time:1757025ms step_avg:43.71ms -step:40400/500000 train_loss:2.1068 train_time:1765759ms step_avg:43.71ms -step:40600/500000 train_loss:2.1224 train_time:1774490ms step_avg:43.71ms -step:40800/500000 train_loss:1.9901 train_time:1783221ms step_avg:43.71ms -step:41000/500000 train_loss:2.0482 train_time:1791956ms step_avg:43.71ms -step:41200/500000 train_loss:2.0054 train_time:1800685ms step_avg:43.71ms -step:41400/500000 train_loss:2.1389 train_time:1809420ms step_avg:43.71ms -step:41600/500000 train_loss:1.9084 train_time:1818157ms step_avg:43.71ms -step:41800/500000 train_loss:1.9776 train_time:1826898ms step_avg:43.71ms -step:42000/500000 train_loss:2.1240 train_time:1835629ms step_avg:43.71ms -step:42200/500000 train_loss:2.0346 train_time:1844369ms step_avg:43.71ms -step:42400/500000 train_loss:1.8743 train_time:1853109ms step_avg:43.71ms -step:42600/500000 train_loss:2.0811 train_time:1861841ms step_avg:43.71ms -step:42800/500000 train_loss:1.9275 train_time:1870573ms step_avg:43.70ms -step:43000/500000 train_loss:2.0186 train_time:1879304ms step_avg:43.70ms -step:43200/500000 train_loss:1.9877 train_time:1888040ms step_avg:43.70ms -step:43400/500000 train_loss:2.0362 train_time:1896771ms step_avg:43.70ms -step:43600/500000 train_loss:2.0231 train_time:1905503ms step_avg:43.70ms -step:43800/500000 train_loss:2.0801 train_time:1914237ms step_avg:43.70ms -step:44000/500000 train_loss:2.0663 train_time:1922975ms step_avg:43.70ms -step:44200/500000 train_loss:2.0227 train_time:1931827ms step_avg:43.71ms -step:44400/500000 train_loss:1.9966 train_time:1940564ms step_avg:43.71ms -step:44600/500000 train_loss:2.0269 train_time:1949298ms step_avg:43.71ms -step:44800/500000 train_loss:2.4033 train_time:1958041ms step_avg:43.71ms -step:45000/500000 train_loss:2.1128 train_time:1966777ms step_avg:43.71ms -step:45200/500000 train_loss:2.1673 train_time:1975520ms step_avg:43.71ms -step:45400/500000 train_loss:2.1298 train_time:1984260ms step_avg:43.71ms -step:45600/500000 train_loss:1.8693 train_time:1992998ms step_avg:43.71ms -step:45800/500000 train_loss:2.0838 train_time:2001733ms step_avg:43.71ms -step:46000/500000 train_loss:2.0543 train_time:2010471ms step_avg:43.71ms -step:46200/500000 train_loss:2.1408 train_time:2019211ms step_avg:43.71ms -step:46400/500000 train_loss:2.0154 train_time:2027949ms step_avg:43.71ms -step:46600/500000 train_loss:2.0614 train_time:2036684ms step_avg:43.71ms -step:46800/500000 train_loss:2.1562 train_time:2045414ms step_avg:43.71ms -step:47000/500000 train_loss:2.0508 train_time:2054144ms step_avg:43.71ms -step:47200/500000 train_loss:2.0838 train_time:2062882ms step_avg:43.71ms -step:47400/500000 train_loss:2.0732 train_time:2071643ms step_avg:43.71ms -step:47600/500000 train_loss:2.0627 train_time:2080393ms step_avg:43.71ms -step:47800/500000 train_loss:2.0581 train_time:2089138ms step_avg:43.71ms -step:48000/500000 train_loss:2.1821 train_time:2097874ms step_avg:43.71ms -step:48200/500000 train_loss:2.0664 train_time:2106621ms step_avg:43.71ms -step:48400/500000 train_loss:2.0724 train_time:2115483ms step_avg:43.71ms -step:48600/500000 train_loss:1.9536 train_time:2124224ms step_avg:43.71ms -step:48800/500000 train_loss:2.0122 train_time:2132956ms step_avg:43.71ms -step:49000/500000 train_loss:2.1620 train_time:2141694ms step_avg:43.71ms -step:49200/500000 train_loss:1.9946 train_time:2150431ms step_avg:43.71ms -step:49400/500000 train_loss:2.0624 train_time:2159169ms step_avg:43.71ms -step:49600/500000 train_loss:2.0001 train_time:2167906ms step_avg:43.71ms -step:49800/500000 train_loss:2.0727 train_time:2176641ms step_avg:43.71ms -step:50000/500000 train_loss:2.0876 train_time:2185379ms step_avg:43.71ms -step:50200/500000 train_loss:2.0509 train_time:2194118ms step_avg:43.71ms -step:50400/500000 train_loss:2.1277 train_time:2202852ms step_avg:43.71ms -step:50600/500000 train_loss:2.0534 train_time:2211582ms step_avg:43.71ms -step:50800/500000 train_loss:1.9725 train_time:2220324ms step_avg:43.71ms -step:51000/500000 train_loss:1.9960 train_time:2229067ms step_avg:43.71ms -step:51200/500000 train_loss:2.0345 train_time:2237807ms step_avg:43.71ms -step:51400/500000 train_loss:2.0990 train_time:2246550ms step_avg:43.71ms -step:51600/500000 train_loss:2.0457 train_time:2255285ms step_avg:43.71ms -step:51800/500000 train_loss:2.0500 train_time:2264020ms step_avg:43.71ms -step:52000/500000 train_loss:1.9877 train_time:2272760ms step_avg:43.71ms -step:52200/500000 train_loss:1.9642 train_time:2281497ms step_avg:43.71ms -step:52400/500000 train_loss:2.0589 train_time:2290236ms step_avg:43.71ms -step:52600/500000 train_loss:2.0642 train_time:2299090ms step_avg:43.71ms -step:52800/500000 train_loss:2.0002 train_time:2307829ms step_avg:43.71ms -step:53000/500000 train_loss:2.1251 train_time:2316559ms step_avg:43.71ms -step:53200/500000 train_loss:2.2760 train_time:2325293ms step_avg:43.71ms -step:53400/500000 train_loss:1.9892 train_time:2334032ms step_avg:43.71ms -step:53600/500000 train_loss:2.0399 train_time:2342765ms step_avg:43.71ms -step:53800/500000 train_loss:2.0171 train_time:2351497ms step_avg:43.71ms -step:54000/500000 train_loss:1.9835 train_time:2360222ms step_avg:43.71ms -step:54200/500000 train_loss:1.9996 train_time:2368976ms step_avg:43.71ms -step:54400/500000 train_loss:2.1281 train_time:2377711ms step_avg:43.71ms -step:54600/500000 train_loss:2.0236 train_time:2386444ms step_avg:43.71ms -step:54800/500000 train_loss:1.9770 train_time:2395176ms step_avg:43.71ms -step:55000/500000 train_loss:2.1209 train_time:2403912ms step_avg:43.71ms -step:55200/500000 train_loss:2.0261 train_time:2412637ms step_avg:43.71ms -step:55400/500000 train_loss:2.0456 train_time:2421371ms step_avg:43.71ms -step:55600/500000 train_loss:2.0821 train_time:2430109ms step_avg:43.71ms -step:55800/500000 train_loss:1.9503 train_time:2438834ms step_avg:43.71ms -step:56000/500000 train_loss:2.0074 train_time:2447565ms step_avg:43.71ms -step:56200/500000 train_loss:2.0663 train_time:2456301ms step_avg:43.71ms -step:56400/500000 train_loss:2.0383 train_time:2465029ms step_avg:43.71ms -step:56600/500000 train_loss:2.0852 train_time:2473880ms step_avg:43.71ms -step:56800/500000 train_loss:1.9910 train_time:2482615ms step_avg:43.71ms -step:57000/500000 train_loss:2.0997 train_time:2491342ms step_avg:43.71ms -step:57200/500000 train_loss:2.0564 train_time:2500070ms step_avg:43.71ms -step:57400/500000 train_loss:1.9974 train_time:2508803ms step_avg:43.71ms -step:57600/500000 train_loss:2.0217 train_time:2517529ms step_avg:43.71ms -step:57800/500000 train_loss:1.9696 train_time:2526265ms step_avg:43.71ms -step:58000/500000 train_loss:2.0311 train_time:2534998ms step_avg:43.71ms -step:58200/500000 train_loss:2.0664 train_time:2543727ms step_avg:43.71ms -step:58400/500000 train_loss:2.0663 train_time:2552468ms step_avg:43.71ms -step:58600/500000 train_loss:1.8755 train_time:2561203ms step_avg:43.71ms -step:58800/500000 train_loss:2.0705 train_time:2569943ms step_avg:43.71ms -step:59000/500000 train_loss:2.1983 train_time:2578691ms step_avg:43.71ms -step:59200/500000 train_loss:2.0254 train_time:2587421ms step_avg:43.71ms -step:59400/500000 train_loss:1.9081 train_time:2596157ms step_avg:43.71ms -step:59600/500000 train_loss:1.8892 train_time:2604900ms step_avg:43.71ms -step:59800/500000 train_loss:2.0723 train_time:2613631ms step_avg:43.71ms -step:60000/500000 train_loss:1.9396 train_time:2622368ms step_avg:43.71ms -step:60000/500000 val_loss:2.0377 val_bpb:1.2069 train_time:2622385ms step_avg:43.71ms -step:60200/500000 train_loss:2.0829 train_time:2631107ms step_avg:43.71ms -step:60400/500000 train_loss:2.0247 train_time:2639847ms step_avg:43.71ms -step:60600/500000 train_loss:2.0155 train_time:2648581ms step_avg:43.71ms -step:60800/500000 train_loss:1.9590 train_time:2657437ms step_avg:43.71ms -step:61000/500000 train_loss:2.0764 train_time:2666173ms step_avg:43.71ms -step:61200/500000 train_loss:2.0025 train_time:2674908ms step_avg:43.71ms -step:61400/500000 train_loss:2.0852 train_time:2683660ms step_avg:43.71ms -step:61600/500000 train_loss:2.1964 train_time:2692401ms step_avg:43.71ms -step:61800/500000 train_loss:2.1862 train_time:2701152ms step_avg:43.71ms -step:62000/500000 train_loss:2.0954 train_time:2709889ms step_avg:43.71ms -step:62200/500000 train_loss:2.0075 train_time:2718626ms step_avg:43.71ms -step:62400/500000 train_loss:1.9852 train_time:2727373ms step_avg:43.71ms -step:62600/500000 train_loss:2.0035 train_time:2736123ms step_avg:43.71ms -step:62800/500000 train_loss:2.1441 train_time:2744854ms step_avg:43.71ms -step:63000/500000 train_loss:1.8316 train_time:2753594ms step_avg:43.71ms -step:63200/500000 train_loss:2.0102 train_time:2762341ms step_avg:43.71ms -step:63400/500000 train_loss:1.8111 train_time:2771087ms step_avg:43.71ms -step:63600/500000 train_loss:2.0566 train_time:2779840ms step_avg:43.71ms -step:63800/500000 train_loss:2.0424 train_time:2788589ms step_avg:43.71ms -step:64000/500000 train_loss:1.9626 train_time:2797329ms step_avg:43.71ms -step:64200/500000 train_loss:2.0900 train_time:2806076ms step_avg:43.71ms -step:64400/500000 train_loss:2.0560 train_time:2814809ms step_avg:43.71ms -step:64600/500000 train_loss:1.9803 train_time:2823559ms step_avg:43.71ms -step:64800/500000 train_loss:2.0888 train_time:2832415ms step_avg:43.71ms -step:65000/500000 train_loss:1.9832 train_time:2841158ms step_avg:43.71ms -step:65200/500000 train_loss:2.0433 train_time:2849898ms step_avg:43.71ms -step:65400/500000 train_loss:1.8724 train_time:2858640ms step_avg:43.71ms -step:65600/500000 train_loss:1.9251 train_time:2867381ms step_avg:43.71ms -step:65800/500000 train_loss:2.0526 train_time:2876123ms step_avg:43.71ms -step:66000/500000 train_loss:1.9819 train_time:2884855ms step_avg:43.71ms -step:66200/500000 train_loss:2.1081 train_time:2893587ms step_avg:43.71ms -step:66400/500000 train_loss:2.0156 train_time:2902325ms step_avg:43.71ms -step:66600/500000 train_loss:2.0231 train_time:2911066ms step_avg:43.71ms -step:66800/500000 train_loss:1.9735 train_time:2919799ms step_avg:43.71ms -step:67000/500000 train_loss:1.9554 train_time:2928542ms step_avg:43.71ms -step:67200/500000 train_loss:2.0098 train_time:2937282ms step_avg:43.71ms -step:67400/500000 train_loss:2.1122 train_time:2946019ms step_avg:43.71ms -step:67600/500000 train_loss:2.0330 train_time:2954754ms step_avg:43.71ms -step:67800/500000 train_loss:1.9299 train_time:2963497ms step_avg:43.71ms -step:68000/500000 train_loss:2.0881 train_time:2972236ms step_avg:43.71ms -step:68200/500000 train_loss:1.9947 train_time:2980982ms step_avg:43.71ms -step:68400/500000 train_loss:2.1097 train_time:2989720ms step_avg:43.71ms -step:68600/500000 train_loss:2.1249 train_time:2998464ms step_avg:43.71ms -step:68800/500000 train_loss:1.9755 train_time:3007204ms step_avg:43.71ms -step:69000/500000 train_loss:1.9303 train_time:3016067ms step_avg:43.71ms -step:69200/500000 train_loss:2.0941 train_time:3024806ms step_avg:43.71ms -step:69400/500000 train_loss:2.0425 train_time:3033540ms step_avg:43.71ms -step:69600/500000 train_loss:1.9208 train_time:3042275ms step_avg:43.71ms -step:69800/500000 train_loss:2.0736 train_time:3051002ms step_avg:43.71ms -step:70000/500000 train_loss:2.0506 train_time:3059740ms step_avg:43.71ms -step:70200/500000 train_loss:1.9177 train_time:3068480ms step_avg:43.71ms -step:70400/500000 train_loss:2.0552 train_time:3077209ms step_avg:43.71ms -step:70600/500000 train_loss:2.1593 train_time:3085950ms step_avg:43.71ms -step:70800/500000 train_loss:2.3957 train_time:3094684ms step_avg:43.71ms -step:71000/500000 train_loss:2.1222 train_time:3103419ms step_avg:43.71ms -step:71200/500000 train_loss:2.1032 train_time:3112149ms step_avg:43.71ms -step:71400/500000 train_loss:1.9591 train_time:3120894ms step_avg:43.71ms -step:71600/500000 train_loss:1.9592 train_time:3129636ms step_avg:43.71ms -step:71800/500000 train_loss:2.0074 train_time:3138374ms step_avg:43.71ms -step:72000/500000 train_loss:2.2693 train_time:3147117ms step_avg:43.71ms -step:72200/500000 train_loss:2.1045 train_time:3155857ms step_avg:43.71ms -step:72400/500000 train_loss:1.8356 train_time:3164598ms step_avg:43.71ms -step:72600/500000 train_loss:2.0390 train_time:3173344ms step_avg:43.71ms -step:72800/500000 train_loss:1.9867 train_time:3182090ms step_avg:43.71ms -step:73000/500000 train_loss:2.0488 train_time:3190957ms step_avg:43.71ms -step:73200/500000 train_loss:2.0355 train_time:3199692ms step_avg:43.71ms -step:73400/500000 train_loss:1.9839 train_time:3208438ms step_avg:43.71ms -step:73600/500000 train_loss:2.0317 train_time:3217184ms step_avg:43.71ms -step:73800/500000 train_loss:1.9925 train_time:3225921ms step_avg:43.71ms -step:74000/500000 train_loss:1.9522 train_time:3234658ms step_avg:43.71ms -step:74200/500000 train_loss:2.1049 train_time:3243364ms step_avg:43.71ms -step:74400/500000 train_loss:1.9785 train_time:3252099ms step_avg:43.71ms -step:74600/500000 train_loss:2.0592 train_time:3260837ms step_avg:43.71ms -step:74800/500000 train_loss:2.0489 train_time:3269587ms step_avg:43.71ms -step:75000/500000 train_loss:2.0631 train_time:3278321ms step_avg:43.71ms -step:75200/500000 train_loss:2.0235 train_time:3287054ms step_avg:43.71ms -step:75400/500000 train_loss:1.8544 train_time:3295787ms step_avg:43.71ms -step:75600/500000 train_loss:1.8574 train_time:3304514ms step_avg:43.71ms -step:75800/500000 train_loss:2.0504 train_time:3313237ms step_avg:43.71ms -step:76000/500000 train_loss:1.9967 train_time:3322083ms step_avg:43.71ms -step:76200/500000 train_loss:2.0440 train_time:3330827ms step_avg:43.71ms -step:76400/500000 train_loss:1.9953 train_time:3339574ms step_avg:43.71ms -step:76600/500000 train_loss:2.0072 train_time:3348313ms step_avg:43.71ms -step:76800/500000 train_loss:2.0488 train_time:3357061ms step_avg:43.71ms -step:77000/500000 train_loss:2.0404 train_time:3365811ms step_avg:43.71ms -step:77200/500000 train_loss:2.0290 train_time:3374553ms step_avg:43.71ms -step:77400/500000 train_loss:1.8698 train_time:3383295ms step_avg:43.71ms -step:77600/500000 train_loss:1.9732 train_time:3392031ms step_avg:43.71ms -step:77800/500000 train_loss:2.0140 train_time:3400774ms step_avg:43.71ms -step:78000/500000 train_loss:1.9584 train_time:3409520ms step_avg:43.71ms -step:78200/500000 train_loss:2.2657 train_time:3418280ms step_avg:43.71ms -step:78400/500000 train_loss:2.0041 train_time:3427037ms step_avg:43.71ms -step:78600/500000 train_loss:2.0953 train_time:3435793ms step_avg:43.71ms -step:78800/500000 train_loss:2.0426 train_time:3444542ms step_avg:43.71ms -step:79000/500000 train_loss:2.0681 train_time:3453286ms step_avg:43.71ms -step:79200/500000 train_loss:2.0154 train_time:3462037ms step_avg:43.71ms -step:79400/500000 train_loss:1.9769 train_time:3470781ms step_avg:43.71ms -step:79600/500000 train_loss:2.0938 train_time:3479530ms step_avg:43.71ms -step:79800/500000 train_loss:2.0043 train_time:3488267ms step_avg:43.71ms -step:80000/500000 train_loss:2.0123 train_time:3497014ms step_avg:43.71ms -step:80000/500000 val_loss:2.0306 val_bpb:1.2026 train_time:3497030ms step_avg:43.71ms -step:80200/500000 train_loss:1.9839 train_time:3505891ms step_avg:43.71ms -step:80400/500000 train_loss:2.0201 train_time:3514643ms step_avg:43.71ms -step:80600/500000 train_loss:1.9697 train_time:3523393ms step_avg:43.71ms -step:80800/500000 train_loss:1.8849 train_time:3532145ms step_avg:43.71ms -step:81000/500000 train_loss:2.0080 train_time:3540891ms step_avg:43.71ms -step:81200/500000 train_loss:2.0728 train_time:3549647ms step_avg:43.71ms -step:81400/500000 train_loss:2.1014 train_time:3558381ms step_avg:43.71ms -step:81600/500000 train_loss:2.0568 train_time:3567112ms step_avg:43.71ms -step:81800/500000 train_loss:2.0455 train_time:3575852ms step_avg:43.71ms -step:82000/500000 train_loss:2.0547 train_time:3584589ms step_avg:43.71ms -step:82200/500000 train_loss:2.0608 train_time:3593331ms step_avg:43.71ms -step:82400/500000 train_loss:2.3063 train_time:3602075ms step_avg:43.71ms -step:82600/500000 train_loss:2.0339 train_time:3610813ms step_avg:43.71ms -step:82800/500000 train_loss:2.1101 train_time:3619548ms step_avg:43.71ms -step:83000/500000 train_loss:1.9921 train_time:3628292ms step_avg:43.71ms -step:83200/500000 train_loss:1.9948 train_time:3637032ms step_avg:43.71ms -step:83400/500000 train_loss:2.0244 train_time:3645763ms step_avg:43.71ms -step:83600/500000 train_loss:1.9101 train_time:3654495ms step_avg:43.71ms -step:83800/500000 train_loss:2.0107 train_time:3663231ms step_avg:43.71ms -step:84000/500000 train_loss:1.9378 train_time:3671965ms step_avg:43.71ms -step:84200/500000 train_loss:2.0769 train_time:3680704ms step_avg:43.71ms -step:84400/500000 train_loss:1.9761 train_time:3689562ms step_avg:43.72ms -step:84600/500000 train_loss:2.0646 train_time:3698295ms step_avg:43.72ms -step:84800/500000 train_loss:2.0266 train_time:3707034ms step_avg:43.72ms -step:85000/500000 train_loss:2.0077 train_time:3715778ms step_avg:43.72ms -step:85200/500000 train_loss:2.1978 train_time:3724509ms step_avg:43.71ms -step:85400/500000 train_loss:2.1181 train_time:3733240ms step_avg:43.71ms -step:85600/500000 train_loss:1.9731 train_time:3741988ms step_avg:43.71ms -step:85800/500000 train_loss:1.9917 train_time:3750739ms step_avg:43.71ms -step:86000/500000 train_loss:1.9015 train_time:3759481ms step_avg:43.71ms -step:86200/500000 train_loss:1.9947 train_time:3768217ms step_avg:43.71ms -step:86400/500000 train_loss:2.0555 train_time:3776958ms step_avg:43.71ms -step:86600/500000 train_loss:1.9431 train_time:3785705ms step_avg:43.71ms -step:86800/500000 train_loss:1.9248 train_time:3794444ms step_avg:43.71ms -step:87000/500000 train_loss:2.2830 train_time:3803191ms step_avg:43.71ms -step:87200/500000 train_loss:1.9020 train_time:3811936ms step_avg:43.71ms -step:87400/500000 train_loss:2.0276 train_time:3820677ms step_avg:43.71ms -step:87600/500000 train_loss:1.9286 train_time:3829421ms step_avg:43.71ms -step:87800/500000 train_loss:2.0150 train_time:3838163ms step_avg:43.71ms -step:88000/500000 train_loss:2.1064 train_time:3846914ms step_avg:43.71ms -step:88200/500000 train_loss:1.8532 train_time:3855651ms step_avg:43.71ms -step:88400/500000 train_loss:2.1184 train_time:3864514ms step_avg:43.72ms -step:88600/500000 train_loss:1.9995 train_time:3873261ms step_avg:43.72ms -step:88800/500000 train_loss:2.0633 train_time:3882020ms step_avg:43.72ms -step:89000/500000 train_loss:2.0382 train_time:3890758ms step_avg:43.72ms -step:89200/500000 train_loss:1.9755 train_time:3899507ms step_avg:43.72ms -step:89400/500000 train_loss:1.9484 train_time:3908254ms step_avg:43.72ms -step:89600/500000 train_loss:2.0261 train_time:3916999ms step_avg:43.72ms -step:89800/500000 train_loss:1.9487 train_time:3925742ms step_avg:43.72ms -step:90000/500000 train_loss:2.0480 train_time:3934482ms step_avg:43.72ms -step:90200/500000 train_loss:2.1049 train_time:3943216ms step_avg:43.72ms -step:90400/500000 train_loss:2.0033 train_time:3951955ms step_avg:43.72ms -step:90600/500000 train_loss:2.1195 train_time:3960696ms step_avg:43.72ms -step:90800/500000 train_loss:2.0872 train_time:3969435ms step_avg:43.72ms -step:91000/500000 train_loss:2.0414 train_time:3978176ms step_avg:43.72ms -step:91200/500000 train_loss:1.9898 train_time:3986934ms step_avg:43.72ms -step:91400/500000 train_loss:2.1097 train_time:3995676ms step_avg:43.72ms -step:91600/500000 train_loss:2.0608 train_time:4004427ms step_avg:43.72ms -step:91800/500000 train_loss:2.0089 train_time:4013164ms step_avg:43.72ms -step:92000/500000 train_loss:2.0163 train_time:4021901ms step_avg:43.72ms -step:92200/500000 train_loss:2.0968 train_time:4030647ms step_avg:43.72ms -step:92400/500000 train_loss:1.9743 train_time:4039391ms step_avg:43.72ms -step:92600/500000 train_loss:2.0348 train_time:4048245ms step_avg:43.72ms -step:92800/500000 train_loss:2.0199 train_time:4056989ms step_avg:43.72ms -step:93000/500000 train_loss:2.0642 train_time:4065735ms step_avg:43.72ms -step:93200/500000 train_loss:2.1773 train_time:4074480ms step_avg:43.72ms -step:93400/500000 train_loss:2.0545 train_time:4083215ms step_avg:43.72ms -step:93600/500000 train_loss:1.9643 train_time:4091955ms step_avg:43.72ms -step:93800/500000 train_loss:2.0638 train_time:4100694ms step_avg:43.72ms -step:94000/500000 train_loss:2.0236 train_time:4109446ms step_avg:43.72ms -step:94200/500000 train_loss:2.0259 train_time:4118191ms step_avg:43.72ms -step:94400/500000 train_loss:1.9855 train_time:4126935ms step_avg:43.72ms -step:94600/500000 train_loss:2.0248 train_time:4135681ms step_avg:43.72ms -step:94800/500000 train_loss:2.0139 train_time:4144419ms step_avg:43.72ms -step:95000/500000 train_loss:2.1720 train_time:4153155ms step_avg:43.72ms -step:95200/500000 train_loss:2.0017 train_time:4161900ms step_avg:43.72ms -step:95400/500000 train_loss:2.0994 train_time:4170642ms step_avg:43.72ms -step:95600/500000 train_loss:2.0451 train_time:4179383ms step_avg:43.72ms -step:95800/500000 train_loss:1.9698 train_time:4188120ms step_avg:43.72ms -step:96000/500000 train_loss:2.0817 train_time:4196853ms step_avg:43.72ms -step:96200/500000 train_loss:2.1018 train_time:4205589ms step_avg:43.72ms -step:96400/500000 train_loss:2.1385 train_time:4214327ms step_avg:43.72ms -step:96600/500000 train_loss:2.0407 train_time:4223177ms step_avg:43.72ms -step:96800/500000 train_loss:1.9939 train_time:4231910ms step_avg:43.72ms -step:97000/500000 train_loss:2.0416 train_time:4240640ms step_avg:43.72ms -step:97200/500000 train_loss:2.0274 train_time:4249373ms step_avg:43.72ms -step:97400/500000 train_loss:1.8758 train_time:4258108ms step_avg:43.72ms -step:97600/500000 train_loss:1.9845 train_time:4266837ms step_avg:43.72ms -step:97800/500000 train_loss:2.1325 train_time:4275574ms step_avg:43.72ms -step:98000/500000 train_loss:1.9135 train_time:4284317ms step_avg:43.72ms -step:98200/500000 train_loss:2.1159 train_time:4293049ms step_avg:43.72ms -step:98400/500000 train_loss:1.9857 train_time:4301781ms step_avg:43.72ms -step:98600/500000 train_loss:1.9655 train_time:4310513ms step_avg:43.72ms -step:98800/500000 train_loss:1.9529 train_time:4319254ms step_avg:43.72ms -step:99000/500000 train_loss:2.0475 train_time:4327994ms step_avg:43.72ms -step:99200/500000 train_loss:2.0054 train_time:4336730ms step_avg:43.72ms -step:99400/500000 train_loss:1.9931 train_time:4345463ms step_avg:43.72ms -step:99600/500000 train_loss:1.9279 train_time:4354197ms step_avg:43.72ms -step:99800/500000 train_loss:1.9186 train_time:4362935ms step_avg:43.72ms -step:100000/500000 train_loss:1.7548 train_time:4371686ms step_avg:43.72ms -step:100000/500000 val_loss:2.0253 val_bpb:1.1995 train_time:4371696ms step_avg:43.72ms -step:100200/500000 train_loss:1.9653 train_time:4380428ms step_avg:43.72ms -step:100400/500000 train_loss:2.1490 train_time:4389172ms step_avg:43.72ms -step:100600/500000 train_loss:2.1338 train_time:4397923ms step_avg:43.72ms -step:100800/500000 train_loss:1.9365 train_time:4406793ms step_avg:43.72ms -step:101000/500000 train_loss:2.1717 train_time:4415535ms step_avg:43.72ms -step:101200/500000 train_loss:2.0397 train_time:4424281ms step_avg:43.72ms -step:101400/500000 train_loss:2.0684 train_time:4433025ms step_avg:43.72ms -step:101600/500000 train_loss:2.0624 train_time:4441766ms step_avg:43.72ms -step:101800/500000 train_loss:1.9881 train_time:4450520ms step_avg:43.72ms -step:102000/500000 train_loss:2.0370 train_time:4459258ms step_avg:43.72ms -step:102200/500000 train_loss:2.0265 train_time:4468004ms step_avg:43.72ms -step:102400/500000 train_loss:2.0059 train_time:4476745ms step_avg:43.72ms -step:102600/500000 train_loss:2.1766 train_time:4485487ms step_avg:43.72ms -step:102800/500000 train_loss:1.9531 train_time:4494235ms step_avg:43.72ms -step:103000/500000 train_loss:2.0770 train_time:4502984ms step_avg:43.72ms -step:103200/500000 train_loss:1.9543 train_time:4511735ms step_avg:43.72ms -step:103400/500000 train_loss:2.0411 train_time:4520487ms step_avg:43.72ms -step:103600/500000 train_loss:2.0500 train_time:4529227ms step_avg:43.72ms -step:103800/500000 train_loss:1.9965 train_time:4537967ms step_avg:43.72ms -step:104000/500000 train_loss:2.0624 train_time:4546714ms step_avg:43.72ms -step:104200/500000 train_loss:1.9853 train_time:4555461ms step_avg:43.72ms -step:104400/500000 train_loss:1.9747 train_time:4564204ms step_avg:43.72ms -step:104600/500000 train_loss:2.0848 train_time:4572941ms step_avg:43.72ms -step:104800/500000 train_loss:2.0619 train_time:4581686ms step_avg:43.72ms -step:105000/500000 train_loss:2.0278 train_time:4590551ms step_avg:43.72ms -step:105200/500000 train_loss:1.9722 train_time:4599291ms step_avg:43.72ms -step:105400/500000 train_loss:2.1234 train_time:4608037ms step_avg:43.72ms -step:105600/500000 train_loss:1.9276 train_time:4616777ms step_avg:43.72ms -step:105800/500000 train_loss:2.1207 train_time:4625512ms step_avg:43.72ms -step:106000/500000 train_loss:1.9724 train_time:4634245ms step_avg:43.72ms -step:106200/500000 train_loss:1.8480 train_time:4642984ms step_avg:43.72ms -step:106400/500000 train_loss:1.9962 train_time:4651724ms step_avg:43.72ms -step:106600/500000 train_loss:2.0682 train_time:4660463ms step_avg:43.72ms -step:106800/500000 train_loss:1.9596 train_time:4669202ms step_avg:43.72ms -step:107000/500000 train_loss:1.9887 train_time:4677945ms step_avg:43.72ms -step:107200/500000 train_loss:2.0878 train_time:4686687ms step_avg:43.72ms -step:107400/500000 train_loss:1.9820 train_time:4695424ms step_avg:43.72ms -step:107600/500000 train_loss:1.8612 train_time:4704155ms step_avg:43.72ms -step:107800/500000 train_loss:2.1396 train_time:4712892ms step_avg:43.72ms -step:108000/500000 train_loss:2.0850 train_time:4721621ms step_avg:43.72ms -step:108200/500000 train_loss:2.0261 train_time:4730365ms step_avg:43.72ms -step:108400/500000 train_loss:2.0620 train_time:4739111ms step_avg:43.72ms -step:108600/500000 train_loss:1.9826 train_time:4747853ms step_avg:43.72ms -step:108800/500000 train_loss:2.1798 train_time:4756598ms step_avg:43.72ms -step:109000/500000 train_loss:2.0688 train_time:4765478ms step_avg:43.72ms -step:109200/500000 train_loss:2.1944 train_time:4774228ms step_avg:43.72ms -step:109400/500000 train_loss:2.0273 train_time:4782964ms step_avg:43.72ms -step:109600/500000 train_loss:1.9818 train_time:4791698ms step_avg:43.72ms -step:109800/500000 train_loss:1.8197 train_time:4800433ms step_avg:43.72ms -step:110000/500000 train_loss:1.9632 train_time:4809167ms step_avg:43.72ms -step:110200/500000 train_loss:2.0929 train_time:4817911ms step_avg:43.72ms -step:110400/500000 train_loss:1.9924 train_time:4826648ms step_avg:43.72ms -step:110600/500000 train_loss:1.9691 train_time:4835381ms step_avg:43.72ms -step:110800/500000 train_loss:2.0186 train_time:4844110ms step_avg:43.72ms -step:111000/500000 train_loss:2.0252 train_time:4852847ms step_avg:43.72ms -step:111200/500000 train_loss:2.2204 train_time:4861590ms step_avg:43.72ms -step:111400/500000 train_loss:2.0328 train_time:4870299ms step_avg:43.72ms -step:111600/500000 train_loss:2.0633 train_time:4879034ms step_avg:43.72ms -step:111800/500000 train_loss:2.0296 train_time:4887776ms step_avg:43.72ms -step:112000/500000 train_loss:2.0067 train_time:4896638ms step_avg:43.72ms -step:112200/500000 train_loss:2.0360 train_time:4905378ms step_avg:43.72ms -step:112400/500000 train_loss:1.8986 train_time:4914115ms step_avg:43.72ms -step:112600/500000 train_loss:1.8944 train_time:4922857ms step_avg:43.72ms -step:112800/500000 train_loss:1.9784 train_time:4931604ms step_avg:43.72ms -step:113000/500000 train_loss:2.0355 train_time:4940342ms step_avg:43.72ms -step:113200/500000 train_loss:1.8432 train_time:4949080ms step_avg:43.72ms -step:113400/500000 train_loss:2.2008 train_time:4957827ms step_avg:43.72ms -step:113600/500000 train_loss:1.9759 train_time:4966563ms step_avg:43.72ms -step:113800/500000 train_loss:1.8877 train_time:4975306ms step_avg:43.72ms -step:114000/500000 train_loss:2.0328 train_time:4984057ms step_avg:43.72ms -step:114200/500000 train_loss:1.9748 train_time:4992790ms step_avg:43.72ms -step:114400/500000 train_loss:1.9246 train_time:5001527ms step_avg:43.72ms -step:114600/500000 train_loss:2.1953 train_time:5010259ms step_avg:43.72ms -step:114800/500000 train_loss:1.9761 train_time:5018992ms step_avg:43.72ms -step:115000/500000 train_loss:1.9908 train_time:5027726ms step_avg:43.72ms -step:115200/500000 train_loss:2.0108 train_time:5036461ms step_avg:43.72ms -step:115400/500000 train_loss:2.0292 train_time:5045199ms step_avg:43.72ms -step:115600/500000 train_loss:2.0888 train_time:5053924ms step_avg:43.72ms -step:115800/500000 train_loss:2.2096 train_time:5062667ms step_avg:43.72ms -step:116000/500000 train_loss:1.9721 train_time:5071397ms step_avg:43.72ms -step:116200/500000 train_loss:2.0199 train_time:5080251ms step_avg:43.72ms -step:116400/500000 train_loss:2.1796 train_time:5088980ms step_avg:43.72ms -step:116600/500000 train_loss:1.9663 train_time:5097712ms step_avg:43.72ms -step:116800/500000 train_loss:2.0318 train_time:5106443ms step_avg:43.72ms -step:117000/500000 train_loss:1.9711 train_time:5115180ms step_avg:43.72ms -step:117200/500000 train_loss:1.7761 train_time:5123922ms step_avg:43.72ms -step:117400/500000 train_loss:2.0077 train_time:5132656ms step_avg:43.72ms -step:117600/500000 train_loss:1.9814 train_time:5141384ms step_avg:43.72ms -step:117800/500000 train_loss:2.1267 train_time:5150118ms step_avg:43.72ms -step:118000/500000 train_loss:1.8956 train_time:5158844ms step_avg:43.72ms -step:118200/500000 train_loss:1.8866 train_time:5167573ms step_avg:43.72ms -step:118400/500000 train_loss:2.0781 train_time:5176309ms step_avg:43.72ms -step:118600/500000 train_loss:2.3821 train_time:5185055ms step_avg:43.72ms -step:118800/500000 train_loss:2.0305 train_time:5193784ms step_avg:43.72ms -step:119000/500000 train_loss:1.9897 train_time:5202517ms step_avg:43.72ms -step:119200/500000 train_loss:2.0434 train_time:5211253ms step_avg:43.72ms -step:119400/500000 train_loss:2.0924 train_time:5220001ms step_avg:43.72ms -step:119600/500000 train_loss:2.0761 train_time:5228736ms step_avg:43.72ms -step:119800/500000 train_loss:2.2817 train_time:5237464ms step_avg:43.72ms -step:120000/500000 train_loss:2.1493 train_time:5246197ms step_avg:43.72ms -step:120000/500000 val_loss:2.0210 val_bpb:1.1969 train_time:5246213ms step_avg:43.72ms -step:120200/500000 train_loss:1.9144 train_time:5255050ms step_avg:43.72ms -step:120400/500000 train_loss:2.0053 train_time:5263788ms step_avg:43.72ms -step:120600/500000 train_loss:1.9489 train_time:5272526ms step_avg:43.72ms -step:120800/500000 train_loss:1.9367 train_time:5281265ms step_avg:43.72ms -step:121000/500000 train_loss:1.9876 train_time:5289998ms step_avg:43.72ms -step:121200/500000 train_loss:2.0388 train_time:5298737ms step_avg:43.72ms -step:121400/500000 train_loss:1.9117 train_time:5307477ms step_avg:43.72ms -step:121600/500000 train_loss:2.1033 train_time:5316209ms step_avg:43.72ms -step:121800/500000 train_loss:1.9801 train_time:5324944ms step_avg:43.72ms -step:122000/500000 train_loss:1.9211 train_time:5333677ms step_avg:43.72ms -step:122200/500000 train_loss:1.9922 train_time:5342421ms step_avg:43.72ms -step:122400/500000 train_loss:2.0145 train_time:5351162ms step_avg:43.72ms -step:122600/500000 train_loss:1.9822 train_time:5359901ms step_avg:43.72ms -step:122800/500000 train_loss:1.9503 train_time:5368632ms step_avg:43.72ms -step:123000/500000 train_loss:2.0314 train_time:5377369ms step_avg:43.72ms -step:123200/500000 train_loss:2.0501 train_time:5386094ms step_avg:43.72ms -step:123400/500000 train_loss:2.0650 train_time:5394818ms step_avg:43.72ms -step:123600/500000 train_loss:2.0346 train_time:5403547ms step_avg:43.72ms -step:123800/500000 train_loss:1.9980 train_time:5412273ms step_avg:43.72ms -step:124000/500000 train_loss:1.8607 train_time:5421006ms step_avg:43.72ms -step:124200/500000 train_loss:2.0196 train_time:5429737ms step_avg:43.72ms -step:124400/500000 train_loss:1.9967 train_time:5438575ms step_avg:43.72ms -step:124600/500000 train_loss:2.0685 train_time:5447313ms step_avg:43.72ms -step:124800/500000 train_loss:2.0213 train_time:5456041ms step_avg:43.72ms -step:125000/500000 train_loss:2.0980 train_time:5464778ms step_avg:43.72ms -step:125200/500000 train_loss:1.9944 train_time:5473524ms step_avg:43.72ms -step:125400/500000 train_loss:2.0300 train_time:5482273ms step_avg:43.72ms -step:125600/500000 train_loss:2.0471 train_time:5491013ms step_avg:43.72ms -step:125800/500000 train_loss:2.0354 train_time:5499761ms step_avg:43.72ms -step:126000/500000 train_loss:2.0454 train_time:5508491ms step_avg:43.72ms -step:126200/500000 train_loss:1.9727 train_time:5517233ms step_avg:43.72ms -step:126400/500000 train_loss:2.1303 train_time:5525963ms step_avg:43.72ms -step:126600/500000 train_loss:2.1323 train_time:5534702ms step_avg:43.72ms -step:126800/500000 train_loss:2.0353 train_time:5543447ms step_avg:43.72ms -step:127000/500000 train_loss:2.0481 train_time:5552191ms step_avg:43.72ms -step:127200/500000 train_loss:1.9966 train_time:5560928ms step_avg:43.72ms -step:127400/500000 train_loss:2.0024 train_time:5569671ms step_avg:43.72ms -step:127600/500000 train_loss:2.1352 train_time:5578409ms step_avg:43.72ms -step:127800/500000 train_loss:1.9181 train_time:5587143ms step_avg:43.72ms -step:128000/500000 train_loss:2.0442 train_time:5595877ms step_avg:43.72ms -step:128200/500000 train_loss:1.9795 train_time:5604623ms step_avg:43.72ms -step:128400/500000 train_loss:2.0676 train_time:5613487ms step_avg:43.72ms -step:128600/500000 train_loss:1.9502 train_time:5622229ms step_avg:43.72ms -step:128800/500000 train_loss:1.9356 train_time:5630963ms step_avg:43.72ms -step:129000/500000 train_loss:1.8350 train_time:5639707ms step_avg:43.72ms -step:129200/500000 train_loss:1.9958 train_time:5648439ms step_avg:43.72ms -step:129400/500000 train_loss:1.7850 train_time:5657189ms step_avg:43.72ms -step:129600/500000 train_loss:2.1434 train_time:5665932ms step_avg:43.72ms -step:129800/500000 train_loss:1.9649 train_time:5674682ms step_avg:43.72ms -step:130000/500000 train_loss:2.0598 train_time:5683416ms step_avg:43.72ms -step:130200/500000 train_loss:2.0008 train_time:5692156ms step_avg:43.72ms -step:130400/500000 train_loss:2.1846 train_time:5700896ms step_avg:43.72ms -step:130600/500000 train_loss:2.1258 train_time:5709644ms step_avg:43.72ms -step:130800/500000 train_loss:2.0358 train_time:5718386ms step_avg:43.72ms -step:131000/500000 train_loss:1.9617 train_time:5727113ms step_avg:43.72ms -step:131200/500000 train_loss:1.9527 train_time:5735854ms step_avg:43.72ms -step:131400/500000 train_loss:1.9986 train_time:5744605ms step_avg:43.72ms -step:131600/500000 train_loss:1.8493 train_time:5753346ms step_avg:43.72ms -step:131800/500000 train_loss:2.0235 train_time:5762090ms step_avg:43.72ms -step:132000/500000 train_loss:1.9611 train_time:5770826ms step_avg:43.72ms -step:132200/500000 train_loss:2.0716 train_time:5779577ms step_avg:43.72ms -step:132400/500000 train_loss:2.0732 train_time:5788321ms step_avg:43.72ms -step:132600/500000 train_loss:1.9754 train_time:5797192ms step_avg:43.72ms -step:132800/500000 train_loss:2.0154 train_time:5805936ms step_avg:43.72ms -step:133000/500000 train_loss:1.9978 train_time:5814679ms step_avg:43.72ms -step:133200/500000 train_loss:2.1977 train_time:5823425ms step_avg:43.72ms -step:133400/500000 train_loss:2.0363 train_time:5832175ms step_avg:43.72ms -step:133600/500000 train_loss:1.8937 train_time:5840923ms step_avg:43.72ms -step:133800/500000 train_loss:1.9849 train_time:5849669ms step_avg:43.72ms -step:134000/500000 train_loss:2.1667 train_time:5858407ms step_avg:43.72ms -step:134200/500000 train_loss:2.1546 train_time:5867169ms step_avg:43.72ms -step:134400/500000 train_loss:2.1969 train_time:5875915ms step_avg:43.72ms -step:134600/500000 train_loss:1.9611 train_time:5884654ms step_avg:43.72ms -step:134800/500000 train_loss:1.9517 train_time:5893390ms step_avg:43.72ms -step:135000/500000 train_loss:2.0010 train_time:5902136ms step_avg:43.72ms -step:135200/500000 train_loss:1.9772 train_time:5910876ms step_avg:43.72ms -step:135400/500000 train_loss:2.1341 train_time:5919617ms step_avg:43.72ms -step:135600/500000 train_loss:2.1205 train_time:5928363ms step_avg:43.72ms -step:135800/500000 train_loss:2.0861 train_time:5937104ms step_avg:43.72ms -step:136000/500000 train_loss:1.9799 train_time:5945849ms step_avg:43.72ms -step:136200/500000 train_loss:1.9827 train_time:5954598ms step_avg:43.72ms -step:136400/500000 train_loss:2.0416 train_time:5963337ms step_avg:43.72ms -step:136600/500000 train_loss:1.9804 train_time:5972081ms step_avg:43.72ms -step:136800/500000 train_loss:2.0562 train_time:5980950ms step_avg:43.72ms -step:137000/500000 train_loss:1.9513 train_time:5989697ms step_avg:43.72ms -step:137200/500000 train_loss:2.0339 train_time:5998445ms step_avg:43.72ms -step:137400/500000 train_loss:2.0287 train_time:6007174ms step_avg:43.72ms -step:137600/500000 train_loss:1.9576 train_time:6015910ms step_avg:43.72ms -step:137800/500000 train_loss:2.0302 train_time:6024660ms step_avg:43.72ms -step:138000/500000 train_loss:2.0371 train_time:6033400ms step_avg:43.72ms -step:138200/500000 train_loss:1.8320 train_time:6042150ms step_avg:43.72ms -step:138400/500000 train_loss:2.0225 train_time:6050887ms step_avg:43.72ms -step:138600/500000 train_loss:1.8439 train_time:6059629ms step_avg:43.72ms -step:138800/500000 train_loss:2.1438 train_time:6068362ms step_avg:43.72ms -step:139000/500000 train_loss:1.9947 train_time:6077102ms step_avg:43.72ms -step:139200/500000 train_loss:2.1038 train_time:6085847ms step_avg:43.72ms -step:139400/500000 train_loss:2.0055 train_time:6094592ms step_avg:43.72ms -step:139600/500000 train_loss:1.9556 train_time:6103326ms step_avg:43.72ms -step:139800/500000 train_loss:1.9420 train_time:6112064ms step_avg:43.72ms -step:140000/500000 train_loss:1.9340 train_time:6120801ms step_avg:43.72ms -step:140000/500000 val_loss:2.0168 val_bpb:1.1945 train_time:6120818ms step_avg:43.72ms -step:140200/500000 train_loss:2.0384 train_time:6129542ms step_avg:43.72ms -step:140400/500000 train_loss:1.9519 train_time:6138277ms step_avg:43.72ms -step:140600/500000 train_loss:1.9898 train_time:6147017ms step_avg:43.72ms -step:140800/500000 train_loss:2.2550 train_time:6155876ms step_avg:43.72ms -step:141000/500000 train_loss:1.9867 train_time:6164617ms step_avg:43.72ms -step:141200/500000 train_loss:2.0205 train_time:6173359ms step_avg:43.72ms -step:141400/500000 train_loss:1.8807 train_time:6182099ms step_avg:43.72ms -step:141600/500000 train_loss:1.9898 train_time:6190842ms step_avg:43.72ms -step:141800/500000 train_loss:1.8427 train_time:6199583ms step_avg:43.72ms -step:142000/500000 train_loss:2.1160 train_time:6208323ms step_avg:43.72ms -step:142200/500000 train_loss:1.8754 train_time:6217067ms step_avg:43.72ms -step:142400/500000 train_loss:2.0096 train_time:6225813ms step_avg:43.72ms -step:142600/500000 train_loss:1.9781 train_time:6234549ms step_avg:43.72ms -step:142800/500000 train_loss:2.0586 train_time:6243289ms step_avg:43.72ms -step:143000/500000 train_loss:2.1796 train_time:6252012ms step_avg:43.72ms -step:143200/500000 train_loss:2.0443 train_time:6260744ms step_avg:43.72ms -step:143400/500000 train_loss:2.0865 train_time:6269501ms step_avg:43.72ms -step:143600/500000 train_loss:2.0975 train_time:6278235ms step_avg:43.72ms -step:143800/500000 train_loss:2.0190 train_time:6286972ms step_avg:43.72ms -step:144000/500000 train_loss:2.0697 train_time:6295708ms step_avg:43.72ms -step:144200/500000 train_loss:1.9661 train_time:6304436ms step_avg:43.72ms -step:144400/500000 train_loss:2.0613 train_time:6313163ms step_avg:43.72ms -step:144600/500000 train_loss:2.2161 train_time:6321898ms step_avg:43.72ms -step:144800/500000 train_loss:2.0155 train_time:6330626ms step_avg:43.72ms -step:145000/500000 train_loss:2.0206 train_time:6339481ms step_avg:43.72ms -step:145200/500000 train_loss:2.0652 train_time:6348218ms step_avg:43.72ms -step:145400/500000 train_loss:2.0265 train_time:6356943ms step_avg:43.72ms -step:145600/500000 train_loss:2.0159 train_time:6365669ms step_avg:43.72ms -step:145800/500000 train_loss:1.8650 train_time:6374392ms step_avg:43.72ms -step:146000/500000 train_loss:2.0364 train_time:6383124ms step_avg:43.72ms -step:146200/500000 train_loss:2.0379 train_time:6391852ms step_avg:43.72ms -step:146400/500000 train_loss:1.9428 train_time:6400597ms step_avg:43.72ms -step:146600/500000 train_loss:1.9223 train_time:6409330ms step_avg:43.72ms -step:146800/500000 train_loss:2.0064 train_time:6418075ms step_avg:43.72ms -step:147000/500000 train_loss:2.0725 train_time:6426812ms step_avg:43.72ms -step:147200/500000 train_loss:1.9570 train_time:6435542ms step_avg:43.72ms -step:147400/500000 train_loss:2.0533 train_time:6444273ms step_avg:43.72ms -step:147600/500000 train_loss:2.0055 train_time:6453012ms step_avg:43.72ms -step:147800/500000 train_loss:2.0093 train_time:6461744ms step_avg:43.72ms -step:148000/500000 train_loss:1.9966 train_time:6470481ms step_avg:43.72ms -step:148200/500000 train_loss:1.9797 train_time:6479223ms step_avg:43.72ms -step:148400/500000 train_loss:2.0464 train_time:6487966ms step_avg:43.72ms -step:148600/500000 train_loss:1.9383 train_time:6496787ms step_avg:43.72ms -step:148800/500000 train_loss:2.1146 train_time:6505531ms step_avg:43.72ms -step:149000/500000 train_loss:2.0930 train_time:6514273ms step_avg:43.72ms -step:149200/500000 train_loss:2.0785 train_time:6523008ms step_avg:43.72ms -step:149400/500000 train_loss:2.0231 train_time:6531742ms step_avg:43.72ms -step:149600/500000 train_loss:2.2172 train_time:6540477ms step_avg:43.72ms -step:149800/500000 train_loss:2.0192 train_time:6549214ms step_avg:43.72ms -step:150000/500000 train_loss:2.0543 train_time:6557949ms step_avg:43.72ms -step:150200/500000 train_loss:1.9361 train_time:6566691ms step_avg:43.72ms -step:150400/500000 train_loss:2.0412 train_time:6575420ms step_avg:43.72ms -step:150600/500000 train_loss:1.9852 train_time:6584160ms step_avg:43.72ms -step:150800/500000 train_loss:2.0049 train_time:6592900ms step_avg:43.72ms -step:151000/500000 train_loss:2.0040 train_time:6601638ms step_avg:43.72ms -step:151200/500000 train_loss:2.1101 train_time:6610372ms step_avg:43.72ms -step:151400/500000 train_loss:2.0910 train_time:6619105ms step_avg:43.72ms -step:151600/500000 train_loss:2.0578 train_time:6627844ms step_avg:43.72ms -step:151800/500000 train_loss:2.0207 train_time:6636581ms step_avg:43.72ms -step:152000/500000 train_loss:2.0865 train_time:6645440ms step_avg:43.72ms -step:152200/500000 train_loss:2.0833 train_time:6654166ms step_avg:43.72ms -step:152400/500000 train_loss:2.1092 train_time:6662895ms step_avg:43.72ms -step:152600/500000 train_loss:1.9655 train_time:6671624ms step_avg:43.72ms -step:152800/500000 train_loss:1.9758 train_time:6680366ms step_avg:43.72ms -step:153000/500000 train_loss:2.2251 train_time:6689098ms step_avg:43.72ms -step:153200/500000 train_loss:2.0474 train_time:6697837ms step_avg:43.72ms -step:153400/500000 train_loss:1.9763 train_time:6706572ms step_avg:43.72ms -step:153600/500000 train_loss:1.8641 train_time:6715312ms step_avg:43.72ms -step:153800/500000 train_loss:2.0611 train_time:6724045ms step_avg:43.72ms -step:154000/500000 train_loss:2.0526 train_time:6732775ms step_avg:43.72ms -step:154200/500000 train_loss:2.1270 train_time:6741515ms step_avg:43.72ms -step:154400/500000 train_loss:1.9840 train_time:6750256ms step_avg:43.72ms -step:154600/500000 train_loss:2.0056 train_time:6758997ms step_avg:43.72ms -step:154800/500000 train_loss:2.0772 train_time:6767752ms step_avg:43.72ms -step:155000/500000 train_loss:2.1605 train_time:6776484ms step_avg:43.72ms -step:155200/500000 train_loss:1.9979 train_time:6785227ms step_avg:43.72ms -step:155400/500000 train_loss:2.0597 train_time:6793955ms step_avg:43.72ms -step:155600/500000 train_loss:1.9873 train_time:6802697ms step_avg:43.72ms -step:155800/500000 train_loss:1.9067 train_time:6811437ms step_avg:43.72ms -step:156000/500000 train_loss:2.0082 train_time:6820174ms step_avg:43.72ms -step:156200/500000 train_loss:1.8543 train_time:6829025ms step_avg:43.72ms -step:156400/500000 train_loss:2.0477 train_time:6837757ms step_avg:43.72ms -step:156600/500000 train_loss:2.0584 train_time:6846488ms step_avg:43.72ms -step:156800/500000 train_loss:2.0802 train_time:6855230ms step_avg:43.72ms -step:157000/500000 train_loss:2.0046 train_time:6863969ms step_avg:43.72ms -step:157200/500000 train_loss:2.0375 train_time:6872709ms step_avg:43.72ms -step:157400/500000 train_loss:2.1092 train_time:6881444ms step_avg:43.72ms -step:157600/500000 train_loss:1.9823 train_time:6890182ms step_avg:43.72ms -step:157800/500000 train_loss:2.0278 train_time:6898919ms step_avg:43.72ms -step:158000/500000 train_loss:2.0222 train_time:6907648ms step_avg:43.72ms -step:158200/500000 train_loss:1.9798 train_time:6916376ms step_avg:43.72ms -step:158400/500000 train_loss:1.9678 train_time:6925111ms step_avg:43.72ms -step:158600/500000 train_loss:2.0966 train_time:6933843ms step_avg:43.72ms -step:158800/500000 train_loss:1.9167 train_time:6942577ms step_avg:43.72ms -step:159000/500000 train_loss:1.9557 train_time:6951310ms step_avg:43.72ms -step:159200/500000 train_loss:1.9954 train_time:6960046ms step_avg:43.72ms -step:159400/500000 train_loss:1.9841 train_time:6968776ms step_avg:43.72ms -step:159600/500000 train_loss:2.0950 train_time:6977509ms step_avg:43.72ms -step:159800/500000 train_loss:2.0244 train_time:6986242ms step_avg:43.72ms -step:160000/500000 train_loss:1.9207 train_time:6994982ms step_avg:43.72ms -step:160000/500000 val_loss:2.0151 val_bpb:1.1935 train_time:6994998ms step_avg:43.72ms -step:160200/500000 train_loss:2.0018 train_time:7003716ms step_avg:43.72ms -step:160400/500000 train_loss:2.0475 train_time:7012579ms step_avg:43.72ms -step:160600/500000 train_loss:1.9469 train_time:7021324ms step_avg:43.72ms -step:160800/500000 train_loss:1.9277 train_time:7030059ms step_avg:43.72ms -step:161000/500000 train_loss:1.9321 train_time:7038802ms step_avg:43.72ms -step:161200/500000 train_loss:1.9661 train_time:7047535ms step_avg:43.72ms -step:161400/500000 train_loss:2.0205 train_time:7056270ms step_avg:43.72ms -step:161600/500000 train_loss:1.9691 train_time:7065010ms step_avg:43.72ms -step:161800/500000 train_loss:1.9580 train_time:7073749ms step_avg:43.72ms -step:162000/500000 train_loss:2.1239 train_time:7082482ms step_avg:43.72ms -step:162200/500000 train_loss:2.1069 train_time:7091216ms step_avg:43.72ms -step:162400/500000 train_loss:1.9907 train_time:7099955ms step_avg:43.72ms -step:162600/500000 train_loss:2.0223 train_time:7108693ms step_avg:43.72ms -step:162800/500000 train_loss:1.9429 train_time:7117432ms step_avg:43.72ms -step:163000/500000 train_loss:2.0619 train_time:7126171ms step_avg:43.72ms -step:163200/500000 train_loss:1.9096 train_time:7134910ms step_avg:43.72ms -step:163400/500000 train_loss:1.8932 train_time:7143647ms step_avg:43.72ms -step:163600/500000 train_loss:2.0096 train_time:7152383ms step_avg:43.72ms -step:163800/500000 train_loss:1.9822 train_time:7161125ms step_avg:43.72ms -step:164000/500000 train_loss:1.8641 train_time:7169864ms step_avg:43.72ms -step:164200/500000 train_loss:2.0342 train_time:7178599ms step_avg:43.72ms -step:164400/500000 train_loss:1.9387 train_time:7187462ms step_avg:43.72ms -step:164600/500000 train_loss:2.0189 train_time:7196201ms step_avg:43.72ms -step:164800/500000 train_loss:1.9743 train_time:7204932ms step_avg:43.72ms -step:165000/500000 train_loss:2.0910 train_time:7213663ms step_avg:43.72ms -step:165200/500000 train_loss:2.0131 train_time:7222396ms step_avg:43.72ms -step:165400/500000 train_loss:1.8053 train_time:7231130ms step_avg:43.72ms -step:165600/500000 train_loss:2.1171 train_time:7239864ms step_avg:43.72ms -step:165800/500000 train_loss:2.0570 train_time:7248591ms step_avg:43.72ms -step:166000/500000 train_loss:2.1153 train_time:7257330ms step_avg:43.72ms -step:166200/500000 train_loss:2.0100 train_time:7266063ms step_avg:43.72ms -step:166400/500000 train_loss:1.9842 train_time:7274794ms step_avg:43.72ms -step:166600/500000 train_loss:2.0044 train_time:7283530ms step_avg:43.72ms -step:166800/500000 train_loss:1.9264 train_time:7292273ms step_avg:43.72ms -step:167000/500000 train_loss:2.0284 train_time:7301000ms step_avg:43.72ms -step:167200/500000 train_loss:2.0714 train_time:7309735ms step_avg:43.72ms -step:167400/500000 train_loss:1.8683 train_time:7318470ms step_avg:43.72ms -step:167600/500000 train_loss:2.0481 train_time:7327205ms step_avg:43.72ms -step:167800/500000 train_loss:2.0047 train_time:7335945ms step_avg:43.72ms -step:168000/500000 train_loss:2.0068 train_time:7344682ms step_avg:43.72ms -step:168200/500000 train_loss:1.9203 train_time:7353418ms step_avg:43.72ms -step:168400/500000 train_loss:1.9711 train_time:7362156ms step_avg:43.72ms -step:168600/500000 train_loss:2.0499 train_time:7371013ms step_avg:43.72ms -step:168800/500000 train_loss:2.0193 train_time:7379753ms step_avg:43.72ms -step:169000/500000 train_loss:1.9500 train_time:7388495ms step_avg:43.72ms -step:169200/500000 train_loss:1.9424 train_time:7397235ms step_avg:43.72ms -step:169400/500000 train_loss:2.0937 train_time:7405974ms step_avg:43.72ms -step:169600/500000 train_loss:2.0015 train_time:7414717ms step_avg:43.72ms -step:169800/500000 train_loss:1.9852 train_time:7423454ms step_avg:43.72ms -step:170000/500000 train_loss:1.8737 train_time:7432193ms step_avg:43.72ms -step:170200/500000 train_loss:2.0264 train_time:7440931ms step_avg:43.72ms -step:170400/500000 train_loss:2.1094 train_time:7449668ms step_avg:43.72ms -step:170600/500000 train_loss:1.9665 train_time:7458399ms step_avg:43.72ms -step:170800/500000 train_loss:1.9582 train_time:7467133ms step_avg:43.72ms -step:171000/500000 train_loss:1.9925 train_time:7475896ms step_avg:43.72ms -step:171200/500000 train_loss:2.1872 train_time:7484633ms step_avg:43.72ms -step:171400/500000 train_loss:2.0731 train_time:7493380ms step_avg:43.72ms -step:171600/500000 train_loss:1.9710 train_time:7502128ms step_avg:43.72ms -step:171800/500000 train_loss:1.9976 train_time:7510868ms step_avg:43.72ms -step:172000/500000 train_loss:1.9109 train_time:7519611ms step_avg:43.72ms -step:172200/500000 train_loss:1.9541 train_time:7528366ms step_avg:43.72ms -step:172400/500000 train_loss:2.0850 train_time:7537121ms step_avg:43.72ms -step:172600/500000 train_loss:1.8652 train_time:7545995ms step_avg:43.72ms -step:172800/500000 train_loss:1.9614 train_time:7554747ms step_avg:43.72ms -step:173000/500000 train_loss:2.1891 train_time:7563490ms step_avg:43.72ms -step:173200/500000 train_loss:1.9574 train_time:7572235ms step_avg:43.72ms -step:173400/500000 train_loss:2.0368 train_time:7580959ms step_avg:43.72ms -step:173600/500000 train_loss:2.0617 train_time:7589684ms step_avg:43.72ms -step:173800/500000 train_loss:2.1104 train_time:7598407ms step_avg:43.72ms -step:174000/500000 train_loss:2.0506 train_time:7607136ms step_avg:43.72ms -step:174200/500000 train_loss:1.9299 train_time:7615871ms step_avg:43.72ms -step:174400/500000 train_loss:2.1104 train_time:7624607ms step_avg:43.72ms -step:174600/500000 train_loss:2.1805 train_time:7633333ms step_avg:43.72ms -step:174800/500000 train_loss:2.0631 train_time:7642073ms step_avg:43.72ms -step:175000/500000 train_loss:2.1022 train_time:7650819ms step_avg:43.72ms -step:175200/500000 train_loss:1.8115 train_time:7659557ms step_avg:43.72ms -step:175400/500000 train_loss:2.0692 train_time:7668291ms step_avg:43.72ms -step:175600/500000 train_loss:2.1445 train_time:7677029ms step_avg:43.72ms -step:175800/500000 train_loss:1.9551 train_time:7685758ms step_avg:43.72ms -step:176000/500000 train_loss:1.9635 train_time:7694493ms step_avg:43.72ms -step:176200/500000 train_loss:1.9418 train_time:7703234ms step_avg:43.72ms -step:176400/500000 train_loss:2.1112 train_time:7711975ms step_avg:43.72ms -step:176600/500000 train_loss:2.1513 train_time:7720710ms step_avg:43.72ms -step:176800/500000 train_loss:2.0037 train_time:7729571ms step_avg:43.72ms -step:177000/500000 train_loss:1.9129 train_time:7738307ms step_avg:43.72ms -step:177200/500000 train_loss:2.1441 train_time:7747046ms step_avg:43.72ms -step:177400/500000 train_loss:1.9549 train_time:7755780ms step_avg:43.72ms -step:177600/500000 train_loss:1.9830 train_time:7764513ms step_avg:43.72ms -step:177800/500000 train_loss:1.9886 train_time:7773254ms step_avg:43.72ms -step:178000/500000 train_loss:2.1259 train_time:7781991ms step_avg:43.72ms -step:178200/500000 train_loss:1.9354 train_time:7790727ms step_avg:43.72ms -step:178400/500000 train_loss:2.1140 train_time:7799471ms step_avg:43.72ms -step:178600/500000 train_loss:1.8590 train_time:7808207ms step_avg:43.72ms -step:178800/500000 train_loss:2.0454 train_time:7816945ms step_avg:43.72ms -step:179000/500000 train_loss:1.9407 train_time:7825676ms step_avg:43.72ms -step:179200/500000 train_loss:1.9950 train_time:7834409ms step_avg:43.72ms -step:179400/500000 train_loss:1.9914 train_time:7843141ms step_avg:43.72ms -step:179600/500000 train_loss:2.0901 train_time:7851877ms step_avg:43.72ms -step:179800/500000 train_loss:1.9138 train_time:7860608ms step_avg:43.72ms -step:180000/500000 train_loss:2.0217 train_time:7869340ms step_avg:43.72ms -step:180000/500000 val_loss:2.0117 val_bpb:1.1914 train_time:7869356ms step_avg:43.72ms -step:180200/500000 train_loss:2.0164 train_time:7878077ms step_avg:43.72ms -step:180400/500000 train_loss:2.0996 train_time:7886806ms step_avg:43.72ms -step:180600/500000 train_loss:2.1237 train_time:7895545ms step_avg:43.72ms -step:180800/500000 train_loss:2.0568 train_time:7904280ms step_avg:43.72ms -step:181000/500000 train_loss:2.0839 train_time:7913142ms step_avg:43.72ms -step:181200/500000 train_loss:1.9792 train_time:7921880ms step_avg:43.72ms -step:181400/500000 train_loss:1.8169 train_time:7930628ms step_avg:43.72ms -step:181600/500000 train_loss:2.0093 train_time:7939374ms step_avg:43.72ms -step:181800/500000 train_loss:1.8723 train_time:7948111ms step_avg:43.72ms -step:182000/500000 train_loss:2.0589 train_time:7956844ms step_avg:43.72ms -step:182200/500000 train_loss:2.2871 train_time:7965596ms step_avg:43.72ms -step:182400/500000 train_loss:1.9727 train_time:7974326ms step_avg:43.72ms -step:182600/500000 train_loss:2.1237 train_time:7983068ms step_avg:43.72ms -step:182800/500000 train_loss:1.9695 train_time:7991805ms step_avg:43.72ms -step:183000/500000 train_loss:1.9939 train_time:8000541ms step_avg:43.72ms -step:183200/500000 train_loss:2.1586 train_time:8009273ms step_avg:43.72ms -step:183400/500000 train_loss:1.9795 train_time:8018006ms step_avg:43.72ms -step:183600/500000 train_loss:2.0676 train_time:8026744ms step_avg:43.72ms -step:183800/500000 train_loss:1.9571 train_time:8035482ms step_avg:43.72ms -step:184000/500000 train_loss:2.0105 train_time:8044212ms step_avg:43.72ms -step:184200/500000 train_loss:2.0779 train_time:8052949ms step_avg:43.72ms -step:184400/500000 train_loss:2.0069 train_time:8061687ms step_avg:43.72ms -step:184600/500000 train_loss:2.0897 train_time:8070420ms step_avg:43.72ms -step:184800/500000 train_loss:1.9466 train_time:8079148ms step_avg:43.72ms -step:185000/500000 train_loss:2.0431 train_time:8088000ms step_avg:43.72ms -step:185200/500000 train_loss:2.0700 train_time:8096735ms step_avg:43.72ms -step:185400/500000 train_loss:1.8829 train_time:8105470ms step_avg:43.72ms -step:185600/500000 train_loss:2.0099 train_time:8114181ms step_avg:43.72ms -step:185800/500000 train_loss:1.9734 train_time:8122913ms step_avg:43.72ms -step:186000/500000 train_loss:2.0107 train_time:8131657ms step_avg:43.72ms -step:186200/500000 train_loss:1.8863 train_time:8140408ms step_avg:43.72ms -step:186400/500000 train_loss:1.8992 train_time:8149145ms step_avg:43.72ms -step:186600/500000 train_loss:2.0194 train_time:8157881ms step_avg:43.72ms -step:186800/500000 train_loss:2.0629 train_time:8166619ms step_avg:43.72ms -step:187000/500000 train_loss:2.0338 train_time:8175361ms step_avg:43.72ms -step:187200/500000 train_loss:2.0185 train_time:8184101ms step_avg:43.72ms -step:187400/500000 train_loss:2.1287 train_time:8192842ms step_avg:43.72ms -step:187600/500000 train_loss:1.9641 train_time:8201582ms step_avg:43.72ms -step:187800/500000 train_loss:1.8271 train_time:8210327ms step_avg:43.72ms -step:188000/500000 train_loss:1.8761 train_time:8219195ms step_avg:43.72ms -step:188200/500000 train_loss:2.0407 train_time:8227939ms step_avg:43.72ms -step:188400/500000 train_loss:1.9929 train_time:8236687ms step_avg:43.72ms -step:188600/500000 train_loss:2.1855 train_time:8245431ms step_avg:43.72ms -step:188800/500000 train_loss:1.9462 train_time:8254187ms step_avg:43.72ms -step:189000/500000 train_loss:2.0223 train_time:8262930ms step_avg:43.72ms -step:189200/500000 train_loss:2.0223 train_time:8271678ms step_avg:43.72ms -step:189400/500000 train_loss:1.9818 train_time:8280415ms step_avg:43.72ms -step:189600/500000 train_loss:1.9641 train_time:8289171ms step_avg:43.72ms -step:189800/500000 train_loss:2.0795 train_time:8297919ms step_avg:43.72ms -step:190000/500000 train_loss:1.9859 train_time:8306658ms step_avg:43.72ms -step:190200/500000 train_loss:2.0013 train_time:8315407ms step_avg:43.72ms -step:190400/500000 train_loss:1.8687 train_time:8324156ms step_avg:43.72ms -step:190600/500000 train_loss:2.0704 train_time:8332901ms step_avg:43.72ms -step:190800/500000 train_loss:2.0072 train_time:8341640ms step_avg:43.72ms -step:191000/500000 train_loss:1.9459 train_time:8350385ms step_avg:43.72ms -step:191200/500000 train_loss:2.0498 train_time:8359126ms step_avg:43.72ms -step:191400/500000 train_loss:2.0077 train_time:8367862ms step_avg:43.72ms -step:191600/500000 train_loss:2.0151 train_time:8376604ms step_avg:43.72ms -step:191800/500000 train_loss:2.0355 train_time:8385354ms step_avg:43.72ms -step:192000/500000 train_loss:1.9724 train_time:8394096ms step_avg:43.72ms -step:192200/500000 train_loss:2.0984 train_time:8402951ms step_avg:43.72ms -step:192400/500000 train_loss:1.8789 train_time:8411689ms step_avg:43.72ms -step:192600/500000 train_loss:1.8561 train_time:8420432ms step_avg:43.72ms -step:192800/500000 train_loss:2.0320 train_time:8429163ms step_avg:43.72ms -step:193000/500000 train_loss:2.1582 train_time:8437899ms step_avg:43.72ms -step:193200/500000 train_loss:1.8950 train_time:8446638ms step_avg:43.72ms -step:193400/500000 train_loss:1.9427 train_time:8455369ms step_avg:43.72ms -step:193600/500000 train_loss:2.0281 train_time:8464117ms step_avg:43.72ms -step:193800/500000 train_loss:2.1803 train_time:8472865ms step_avg:43.72ms -step:194000/500000 train_loss:2.0341 train_time:8481600ms step_avg:43.72ms -step:194200/500000 train_loss:2.0818 train_time:8490324ms step_avg:43.72ms -step:194400/500000 train_loss:2.0853 train_time:8499050ms step_avg:43.72ms -step:194600/500000 train_loss:2.1322 train_time:8507769ms step_avg:43.72ms -step:194800/500000 train_loss:1.9417 train_time:8516499ms step_avg:43.72ms -step:195000/500000 train_loss:1.9583 train_time:8525228ms step_avg:43.72ms -step:195200/500000 train_loss:2.1288 train_time:8533957ms step_avg:43.72ms -step:195400/500000 train_loss:1.8735 train_time:8542684ms step_avg:43.72ms -step:195600/500000 train_loss:1.9380 train_time:8551417ms step_avg:43.72ms -step:195800/500000 train_loss:2.0153 train_time:8560153ms step_avg:43.72ms -step:196000/500000 train_loss:2.0793 train_time:8568904ms step_avg:43.72ms -step:196200/500000 train_loss:1.9494 train_time:8577766ms step_avg:43.72ms -step:196400/500000 train_loss:2.0027 train_time:8586500ms step_avg:43.72ms -step:196600/500000 train_loss:1.8933 train_time:8595250ms step_avg:43.72ms -step:196800/500000 train_loss:2.1100 train_time:8603995ms step_avg:43.72ms -step:197000/500000 train_loss:1.9367 train_time:8612745ms step_avg:43.72ms -step:197200/500000 train_loss:2.0373 train_time:8621490ms step_avg:43.72ms -step:197400/500000 train_loss:1.9918 train_time:8630234ms step_avg:43.72ms -step:197600/500000 train_loss:1.9524 train_time:8638976ms step_avg:43.72ms -step:197800/500000 train_loss:2.0744 train_time:8647712ms step_avg:43.72ms -step:198000/500000 train_loss:2.0879 train_time:8656445ms step_avg:43.72ms -step:198200/500000 train_loss:1.8853 train_time:8665173ms step_avg:43.72ms -step:198400/500000 train_loss:1.9999 train_time:8673907ms step_avg:43.72ms -step:198600/500000 train_loss:2.0703 train_time:8682644ms step_avg:43.72ms -step:198800/500000 train_loss:1.9407 train_time:8691376ms step_avg:43.72ms -step:199000/500000 train_loss:1.9825 train_time:8700102ms step_avg:43.72ms -step:199200/500000 train_loss:2.2138 train_time:8708839ms step_avg:43.72ms -step:199400/500000 train_loss:2.0556 train_time:8717569ms step_avg:43.72ms -step:199600/500000 train_loss:1.8311 train_time:8726291ms step_avg:43.72ms -step:199800/500000 train_loss:1.8928 train_time:8735019ms step_avg:43.72ms -step:200000/500000 train_loss:1.9818 train_time:8743744ms step_avg:43.72ms -step:200000/500000 val_loss:2.0105 val_bpb:1.1907 train_time:8743760ms step_avg:43.72ms -step:200200/500000 train_loss:1.9868 train_time:8752474ms step_avg:43.72ms -step:200400/500000 train_loss:2.1167 train_time:8761320ms step_avg:43.72ms -step:200600/500000 train_loss:1.9449 train_time:8770043ms step_avg:43.72ms -step:200800/500000 train_loss:2.0322 train_time:8778769ms step_avg:43.72ms -step:201000/500000 train_loss:2.0473 train_time:8787500ms step_avg:43.72ms -step:201200/500000 train_loss:1.9667 train_time:8796227ms step_avg:43.72ms -step:201400/500000 train_loss:2.1328 train_time:8804954ms step_avg:43.72ms -step:201600/500000 train_loss:2.0573 train_time:8813695ms step_avg:43.72ms -step:201800/500000 train_loss:1.9931 train_time:8822429ms step_avg:43.72ms -step:202000/500000 train_loss:2.1286 train_time:8831155ms step_avg:43.72ms -step:202200/500000 train_loss:2.0619 train_time:8839888ms step_avg:43.72ms -step:202400/500000 train_loss:2.1035 train_time:8848616ms step_avg:43.72ms -step:202600/500000 train_loss:2.0375 train_time:8857346ms step_avg:43.72ms -step:202800/500000 train_loss:1.9938 train_time:8866075ms step_avg:43.72ms -step:203000/500000 train_loss:1.8688 train_time:8874808ms step_avg:43.72ms -step:203200/500000 train_loss:1.9930 train_time:8883537ms step_avg:43.72ms -step:203400/500000 train_loss:1.9899 train_time:8892264ms step_avg:43.72ms -step:203600/500000 train_loss:1.7814 train_time:8901000ms step_avg:43.72ms -step:203800/500000 train_loss:1.9964 train_time:8909742ms step_avg:43.72ms -step:204000/500000 train_loss:2.0461 train_time:8918478ms step_avg:43.72ms -step:204200/500000 train_loss:1.9670 train_time:8927209ms step_avg:43.72ms -step:204400/500000 train_loss:2.0828 train_time:8936045ms step_avg:43.72ms -step:204600/500000 train_loss:2.0625 train_time:8944774ms step_avg:43.72ms -step:204800/500000 train_loss:1.9721 train_time:8953505ms step_avg:43.72ms -step:205000/500000 train_loss:1.8662 train_time:8962232ms step_avg:43.72ms -step:205200/500000 train_loss:1.9807 train_time:8970959ms step_avg:43.72ms -step:205400/500000 train_loss:2.0544 train_time:8979687ms step_avg:43.72ms -step:205600/500000 train_loss:1.9625 train_time:8988423ms step_avg:43.72ms -step:205800/500000 train_loss:2.0482 train_time:8997154ms step_avg:43.72ms -step:206000/500000 train_loss:1.9984 train_time:9005880ms step_avg:43.72ms -step:206200/500000 train_loss:1.9716 train_time:9014613ms step_avg:43.72ms -step:206400/500000 train_loss:2.0728 train_time:9023343ms step_avg:43.72ms -step:206600/500000 train_loss:1.9063 train_time:9032073ms step_avg:43.72ms -step:206800/500000 train_loss:1.9478 train_time:9040806ms step_avg:43.72ms -step:207000/500000 train_loss:2.0179 train_time:9049527ms step_avg:43.72ms -step:207200/500000 train_loss:2.0109 train_time:9058254ms step_avg:43.72ms -step:207400/500000 train_loss:2.0661 train_time:9066986ms step_avg:43.72ms -step:207600/500000 train_loss:1.8557 train_time:9075723ms step_avg:43.72ms -step:207800/500000 train_loss:1.8529 train_time:9084456ms step_avg:43.72ms -step:208000/500000 train_loss:2.0855 train_time:9093188ms step_avg:43.72ms -step:208200/500000 train_loss:2.0335 train_time:9101921ms step_avg:43.72ms -step:208400/500000 train_loss:1.8988 train_time:9110652ms step_avg:43.72ms -step:208600/500000 train_loss:1.9468 train_time:9119502ms step_avg:43.72ms -step:208800/500000 train_loss:1.9999 train_time:9128232ms step_avg:43.72ms -step:209000/500000 train_loss:2.0396 train_time:9136961ms step_avg:43.72ms -step:209200/500000 train_loss:2.0371 train_time:9145693ms step_avg:43.72ms -step:209400/500000 train_loss:2.1630 train_time:9154422ms step_avg:43.72ms -step:209600/500000 train_loss:1.9490 train_time:9163147ms step_avg:43.72ms -step:209800/500000 train_loss:1.9805 train_time:9171877ms step_avg:43.72ms -step:210000/500000 train_loss:2.0601 train_time:9180605ms step_avg:43.72ms -step:210200/500000 train_loss:2.0570 train_time:9189333ms step_avg:43.72ms -step:210400/500000 train_loss:2.0318 train_time:9198063ms step_avg:43.72ms -step:210600/500000 train_loss:2.1689 train_time:9206795ms step_avg:43.72ms -step:210800/500000 train_loss:1.9921 train_time:9215524ms step_avg:43.72ms -step:211000/500000 train_loss:2.0106 train_time:9224247ms step_avg:43.72ms -step:211200/500000 train_loss:2.0720 train_time:9232976ms step_avg:43.72ms -step:211400/500000 train_loss:2.1134 train_time:9241710ms step_avg:43.72ms -step:211600/500000 train_loss:2.1559 train_time:9250441ms step_avg:43.72ms -step:211800/500000 train_loss:2.0516 train_time:9259174ms step_avg:43.72ms -step:212000/500000 train_loss:2.0935 train_time:9267899ms step_avg:43.72ms -step:212200/500000 train_loss:2.0999 train_time:9276632ms step_avg:43.72ms -step:212400/500000 train_loss:2.0021 train_time:9285357ms step_avg:43.72ms -step:212600/500000 train_loss:2.0492 train_time:9294090ms step_avg:43.72ms -step:212800/500000 train_loss:2.0105 train_time:9302943ms step_avg:43.72ms -step:213000/500000 train_loss:1.9174 train_time:9311671ms step_avg:43.72ms -step:213200/500000 train_loss:2.1274 train_time:9320401ms step_avg:43.72ms -step:213400/500000 train_loss:1.9693 train_time:9329132ms step_avg:43.72ms -step:213600/500000 train_loss:1.8956 train_time:9337862ms step_avg:43.72ms -step:213800/500000 train_loss:2.0412 train_time:9346591ms step_avg:43.72ms -step:214000/500000 train_loss:2.0581 train_time:9355323ms step_avg:43.72ms -step:214200/500000 train_loss:2.0161 train_time:9364052ms step_avg:43.72ms -step:214400/500000 train_loss:2.0053 train_time:9372785ms step_avg:43.72ms -step:214600/500000 train_loss:2.0402 train_time:9381520ms step_avg:43.72ms -step:214800/500000 train_loss:1.9926 train_time:9390259ms step_avg:43.72ms -step:215000/500000 train_loss:2.0061 train_time:9398986ms step_avg:43.72ms -step:215200/500000 train_loss:1.9774 train_time:9407719ms step_avg:43.72ms -step:215400/500000 train_loss:2.0146 train_time:9416449ms step_avg:43.72ms -step:215600/500000 train_loss:2.0522 train_time:9425178ms step_avg:43.72ms -step:215800/500000 train_loss:1.9747 train_time:9433914ms step_avg:43.72ms -step:216000/500000 train_loss:2.0128 train_time:9442642ms step_avg:43.72ms -step:216200/500000 train_loss:2.0113 train_time:9451366ms step_avg:43.72ms -step:216400/500000 train_loss:1.8583 train_time:9460100ms step_avg:43.72ms -step:216600/500000 train_loss:2.0709 train_time:9468838ms step_avg:43.72ms -step:216800/500000 train_loss:2.0007 train_time:9477686ms step_avg:43.72ms -step:217000/500000 train_loss:2.0354 train_time:9486424ms step_avg:43.72ms -step:217200/500000 train_loss:2.0861 train_time:9495155ms step_avg:43.72ms -step:217400/500000 train_loss:1.9562 train_time:9503889ms step_avg:43.72ms -step:217600/500000 train_loss:2.0061 train_time:9512625ms step_avg:43.72ms -step:217800/500000 train_loss:1.9055 train_time:9521360ms step_avg:43.72ms -step:218000/500000 train_loss:2.0533 train_time:9530097ms step_avg:43.72ms -step:218200/500000 train_loss:2.0140 train_time:9538834ms step_avg:43.72ms -step:218400/500000 train_loss:2.0301 train_time:9547568ms step_avg:43.72ms -step:218600/500000 train_loss:2.1542 train_time:9556315ms step_avg:43.72ms -step:218800/500000 train_loss:2.0486 train_time:9565045ms step_avg:43.72ms -step:219000/500000 train_loss:1.9843 train_time:9573785ms step_avg:43.72ms -step:219200/500000 train_loss:1.9709 train_time:9582517ms step_avg:43.72ms -step:219400/500000 train_loss:2.0735 train_time:9591258ms step_avg:43.72ms -step:219600/500000 train_loss:1.9672 train_time:9600003ms step_avg:43.72ms -step:219800/500000 train_loss:1.9823 train_time:9608744ms step_avg:43.72ms -step:220000/500000 train_loss:1.9377 train_time:9617487ms step_avg:43.72ms -step:220000/500000 val_loss:2.0076 val_bpb:1.1890 train_time:9617503ms step_avg:43.72ms -step:220200/500000 train_loss:2.0217 train_time:9626225ms step_avg:43.72ms -step:220400/500000 train_loss:2.0827 train_time:9634959ms step_avg:43.72ms -step:220600/500000 train_loss:2.0294 train_time:9643691ms step_avg:43.72ms -step:220800/500000 train_loss:2.1254 train_time:9652438ms step_avg:43.72ms -step:221000/500000 train_loss:2.1577 train_time:9661294ms step_avg:43.72ms -step:221200/500000 train_loss:2.0463 train_time:9670034ms step_avg:43.72ms -step:221400/500000 train_loss:2.1585 train_time:9678765ms step_avg:43.72ms -step:221600/500000 train_loss:1.9449 train_time:9687499ms step_avg:43.72ms -step:221800/500000 train_loss:2.1088 train_time:9696229ms step_avg:43.72ms -step:222000/500000 train_loss:1.9759 train_time:9704959ms step_avg:43.72ms -step:222200/500000 train_loss:1.9192 train_time:9713682ms step_avg:43.72ms -step:222400/500000 train_loss:1.9609 train_time:9722414ms step_avg:43.72ms -step:222600/500000 train_loss:1.9626 train_time:9731139ms step_avg:43.72ms -step:222800/500000 train_loss:1.9320 train_time:9739836ms step_avg:43.72ms -step:223000/500000 train_loss:2.0679 train_time:9748562ms step_avg:43.72ms -step:223200/500000 train_loss:1.9997 train_time:9757297ms step_avg:43.72ms -step:223400/500000 train_loss:1.8773 train_time:9766028ms step_avg:43.72ms -step:223600/500000 train_loss:2.0217 train_time:9774760ms step_avg:43.72ms -step:223800/500000 train_loss:2.1673 train_time:9783485ms step_avg:43.72ms -step:224000/500000 train_loss:1.9545 train_time:9792341ms step_avg:43.72ms -step:224200/500000 train_loss:1.9287 train_time:9801075ms step_avg:43.72ms -step:224400/500000 train_loss:1.9226 train_time:9809802ms step_avg:43.72ms -step:224600/500000 train_loss:2.0093 train_time:9818539ms step_avg:43.72ms -step:224800/500000 train_loss:2.0917 train_time:9827270ms step_avg:43.72ms -step:225000/500000 train_loss:1.8178 train_time:9836011ms step_avg:43.72ms -step:225200/500000 train_loss:1.9130 train_time:9844760ms step_avg:43.72ms -step:225400/500000 train_loss:2.0751 train_time:9853506ms step_avg:43.72ms -step:225600/500000 train_loss:1.9437 train_time:9862235ms step_avg:43.72ms -step:225800/500000 train_loss:2.0723 train_time:9870970ms step_avg:43.72ms -step:226000/500000 train_loss:1.9647 train_time:9879707ms step_avg:43.72ms -step:226200/500000 train_loss:2.0678 train_time:9888446ms step_avg:43.72ms -step:226400/500000 train_loss:2.0256 train_time:9897181ms step_avg:43.72ms -step:226600/500000 train_loss:2.0307 train_time:9905921ms step_avg:43.72ms -step:226800/500000 train_loss:2.1227 train_time:9914660ms step_avg:43.72ms -step:227000/500000 train_loss:2.0353 train_time:9923391ms step_avg:43.72ms -step:227200/500000 train_loss:1.9588 train_time:9932120ms step_avg:43.72ms -step:227400/500000 train_loss:2.0887 train_time:9940856ms step_avg:43.72ms -step:227600/500000 train_loss:1.9178 train_time:9949585ms step_avg:43.72ms -step:227800/500000 train_loss:2.0666 train_time:9958313ms step_avg:43.72ms -step:228000/500000 train_loss:1.9225 train_time:9967159ms step_avg:43.72ms -step:228200/500000 train_loss:1.9397 train_time:9975883ms step_avg:43.72ms -step:228400/500000 train_loss:1.9910 train_time:9984618ms step_avg:43.72ms -step:228600/500000 train_loss:1.9247 train_time:9993362ms step_avg:43.72ms -step:228800/500000 train_loss:1.8871 train_time:10002097ms step_avg:43.72ms -step:229000/500000 train_loss:2.0247 train_time:10010830ms step_avg:43.72ms -step:229200/500000 train_loss:1.9405 train_time:10019566ms step_avg:43.72ms -step:229400/500000 train_loss:2.0672 train_time:10028294ms step_avg:43.72ms -step:229600/500000 train_loss:2.0012 train_time:10037022ms step_avg:43.72ms -step:229800/500000 train_loss:2.0390 train_time:10045754ms step_avg:43.72ms -step:230000/500000 train_loss:1.9441 train_time:10054490ms step_avg:43.72ms -step:230200/500000 train_loss:1.9243 train_time:10063234ms step_avg:43.72ms -step:230400/500000 train_loss:2.0568 train_time:10071966ms step_avg:43.72ms -step:230600/500000 train_loss:1.7721 train_time:10080697ms step_avg:43.72ms -step:230800/500000 train_loss:1.9805 train_time:10089428ms step_avg:43.72ms -step:231000/500000 train_loss:1.8422 train_time:10098161ms step_avg:43.71ms -step:231200/500000 train_loss:2.0484 train_time:10106900ms step_avg:43.71ms -step:231400/500000 train_loss:2.0169 train_time:10115633ms step_avg:43.71ms -step:231600/500000 train_loss:2.0759 train_time:10124365ms step_avg:43.71ms -step:231800/500000 train_loss:1.9531 train_time:10133092ms step_avg:43.71ms -step:232000/500000 train_loss:2.0485 train_time:10141817ms step_avg:43.71ms -step:232200/500000 train_loss:1.9279 train_time:10150679ms step_avg:43.72ms -step:232400/500000 train_loss:2.1539 train_time:10159418ms step_avg:43.72ms -step:232600/500000 train_loss:2.0334 train_time:10168164ms step_avg:43.72ms -step:232800/500000 train_loss:1.7609 train_time:10176898ms step_avg:43.72ms -step:233000/500000 train_loss:2.0500 train_time:10185626ms step_avg:43.72ms -step:233200/500000 train_loss:1.9637 train_time:10194360ms step_avg:43.72ms -step:233400/500000 train_loss:2.0705 train_time:10203093ms step_avg:43.72ms -step:233600/500000 train_loss:2.0482 train_time:10211826ms step_avg:43.72ms -step:233800/500000 train_loss:2.1037 train_time:10220562ms step_avg:43.71ms -step:234000/500000 train_loss:1.9615 train_time:10229297ms step_avg:43.71ms -step:234200/500000 train_loss:1.9095 train_time:10238033ms step_avg:43.71ms -step:234400/500000 train_loss:1.9257 train_time:10246759ms step_avg:43.71ms -step:234600/500000 train_loss:2.2710 train_time:10255487ms step_avg:43.71ms -step:234800/500000 train_loss:1.9472 train_time:10264212ms step_avg:43.71ms -step:235000/500000 train_loss:2.0240 train_time:10272947ms step_avg:43.71ms -step:235200/500000 train_loss:2.0640 train_time:10281676ms step_avg:43.71ms -step:235400/500000 train_loss:1.8869 train_time:10290411ms step_avg:43.71ms -step:235600/500000 train_loss:1.9445 train_time:10299144ms step_avg:43.71ms -step:235800/500000 train_loss:2.0251 train_time:10307873ms step_avg:43.71ms -step:236000/500000 train_loss:1.9722 train_time:10316597ms step_avg:43.71ms -step:236200/500000 train_loss:2.0008 train_time:10325324ms step_avg:43.71ms -step:236400/500000 train_loss:2.1860 train_time:10334172ms step_avg:43.71ms -step:236600/500000 train_loss:1.9035 train_time:10342897ms step_avg:43.71ms -step:236800/500000 train_loss:1.9173 train_time:10351628ms step_avg:43.71ms -step:237000/500000 train_loss:2.0712 train_time:10360356ms step_avg:43.71ms -step:237200/500000 train_loss:1.9455 train_time:10369078ms step_avg:43.71ms -step:237400/500000 train_loss:2.0799 train_time:10377812ms step_avg:43.71ms -step:237600/500000 train_loss:1.9462 train_time:10386543ms step_avg:43.71ms -step:237800/500000 train_loss:2.0383 train_time:10395286ms step_avg:43.71ms -step:238000/500000 train_loss:1.9322 train_time:10404027ms step_avg:43.71ms -step:238200/500000 train_loss:1.9640 train_time:10412765ms step_avg:43.71ms -step:238400/500000 train_loss:2.0001 train_time:10421496ms step_avg:43.71ms -step:238600/500000 train_loss:2.1391 train_time:10430231ms step_avg:43.71ms -step:238800/500000 train_loss:1.9923 train_time:10438966ms step_avg:43.71ms -step:239000/500000 train_loss:2.0646 train_time:10447711ms step_avg:43.71ms -step:239200/500000 train_loss:1.8668 train_time:10456446ms step_avg:43.71ms -step:239400/500000 train_loss:2.4861 train_time:10465183ms step_avg:43.71ms -step:239600/500000 train_loss:1.8213 train_time:10473921ms step_avg:43.71ms -step:239800/500000 train_loss:2.0816 train_time:10482658ms step_avg:43.71ms -step:240000/500000 train_loss:1.8781 train_time:10491396ms step_avg:43.71ms -step:240000/500000 val_loss:2.0063 val_bpb:1.1883 train_time:10491412ms step_avg:43.71ms -step:240200/500000 train_loss:1.9510 train_time:10500136ms step_avg:43.71ms -step:240400/500000 train_loss:1.9606 train_time:10508991ms step_avg:43.71ms -step:240600/500000 train_loss:1.9815 train_time:10517721ms step_avg:43.71ms -step:240800/500000 train_loss:1.9837 train_time:10526454ms step_avg:43.71ms -step:241000/500000 train_loss:2.0349 train_time:10535193ms step_avg:43.71ms -step:241200/500000 train_loss:1.8801 train_time:10543928ms step_avg:43.71ms -step:241400/500000 train_loss:1.9354 train_time:10552661ms step_avg:43.71ms -step:241600/500000 train_loss:1.8986 train_time:10561395ms step_avg:43.71ms -step:241800/500000 train_loss:2.1174 train_time:10570128ms step_avg:43.71ms -step:242000/500000 train_loss:1.9755 train_time:10578858ms step_avg:43.71ms -step:242200/500000 train_loss:1.9471 train_time:10587594ms step_avg:43.71ms -step:242400/500000 train_loss:2.0002 train_time:10596330ms step_avg:43.71ms -step:242600/500000 train_loss:1.9670 train_time:10605064ms step_avg:43.71ms -step:242800/500000 train_loss:1.8907 train_time:10613816ms step_avg:43.71ms -step:243000/500000 train_loss:2.0368 train_time:10622560ms step_avg:43.71ms -step:243200/500000 train_loss:2.0311 train_time:10631301ms step_avg:43.71ms -step:243400/500000 train_loss:1.8762 train_time:10640036ms step_avg:43.71ms -step:243600/500000 train_loss:1.9528 train_time:10648777ms step_avg:43.71ms -step:243800/500000 train_loss:2.0887 train_time:10657512ms step_avg:43.71ms -step:244000/500000 train_loss:2.0231 train_time:10666246ms step_avg:43.71ms -step:244200/500000 train_loss:2.0592 train_time:10674982ms step_avg:43.71ms -step:244400/500000 train_loss:2.0455 train_time:10683715ms step_avg:43.71ms -step:244600/500000 train_loss:1.9829 train_time:10692567ms step_avg:43.71ms -step:244800/500000 train_loss:1.8873 train_time:10701300ms step_avg:43.71ms -step:245000/500000 train_loss:1.9246 train_time:10710035ms step_avg:43.71ms -step:245200/500000 train_loss:1.9532 train_time:10718767ms step_avg:43.71ms -step:245400/500000 train_loss:2.0863 train_time:10727502ms step_avg:43.71ms -step:245600/500000 train_loss:2.0128 train_time:10736234ms step_avg:43.71ms -step:245800/500000 train_loss:2.0420 train_time:10744968ms step_avg:43.71ms -step:246000/500000 train_loss:1.9984 train_time:10753707ms step_avg:43.71ms -step:246200/500000 train_loss:1.9309 train_time:10762446ms step_avg:43.71ms -step:246400/500000 train_loss:2.0014 train_time:10771184ms step_avg:43.71ms -step:246600/500000 train_loss:1.9898 train_time:10779923ms step_avg:43.71ms -step:246800/500000 train_loss:1.9892 train_time:10788658ms step_avg:43.71ms -step:247000/500000 train_loss:2.0507 train_time:10797397ms step_avg:43.71ms -step:247200/500000 train_loss:1.8842 train_time:10806137ms step_avg:43.71ms -step:247400/500000 train_loss:1.9936 train_time:10814874ms step_avg:43.71ms -step:247600/500000 train_loss:2.0653 train_time:10823607ms step_avg:43.71ms -step:247800/500000 train_loss:2.0729 train_time:10832353ms step_avg:43.71ms -step:248000/500000 train_loss:1.9083 train_time:10841094ms step_avg:43.71ms -step:248200/500000 train_loss:1.9371 train_time:10849830ms step_avg:43.71ms -step:248400/500000 train_loss:1.9112 train_time:10858566ms step_avg:43.71ms -step:248600/500000 train_loss:2.0996 train_time:10867428ms step_avg:43.71ms -step:248800/500000 train_loss:1.9620 train_time:10876162ms step_avg:43.71ms -step:249000/500000 train_loss:2.0224 train_time:10884895ms step_avg:43.71ms -step:249200/500000 train_loss:1.9578 train_time:10893631ms step_avg:43.71ms -step:249400/500000 train_loss:1.8937 train_time:10902369ms step_avg:43.71ms -step:249600/500000 train_loss:1.8381 train_time:10911107ms step_avg:43.71ms -step:249800/500000 train_loss:2.0133 train_time:10919845ms step_avg:43.71ms -step:250000/500000 train_loss:2.0581 train_time:10928578ms step_avg:43.71ms -step:250200/500000 train_loss:1.9867 train_time:10937312ms step_avg:43.71ms -step:250400/500000 train_loss:1.9796 train_time:10946038ms step_avg:43.71ms -step:250600/500000 train_loss:2.0323 train_time:10954770ms step_avg:43.71ms -step:250800/500000 train_loss:1.8523 train_time:10963508ms step_avg:43.71ms -step:251000/500000 train_loss:1.9193 train_time:10972245ms step_avg:43.71ms -step:251200/500000 train_loss:1.9896 train_time:10980982ms step_avg:43.71ms -step:251400/500000 train_loss:1.9091 train_time:10989715ms step_avg:43.71ms -step:251600/500000 train_loss:1.9960 train_time:10998454ms step_avg:43.71ms -step:251800/500000 train_loss:2.0287 train_time:11007188ms step_avg:43.71ms -step:252000/500000 train_loss:1.9612 train_time:11015923ms step_avg:43.71ms -step:252200/500000 train_loss:2.0508 train_time:11024659ms step_avg:43.71ms -step:252400/500000 train_loss:1.9624 train_time:11033389ms step_avg:43.71ms -step:252600/500000 train_loss:1.9766 train_time:11042125ms step_avg:43.71ms -step:252800/500000 train_loss:2.0657 train_time:11050997ms step_avg:43.71ms -step:253000/500000 train_loss:1.9535 train_time:11059732ms step_avg:43.71ms -step:253200/500000 train_loss:2.0059 train_time:11068474ms step_avg:43.71ms -step:253400/500000 train_loss:1.9705 train_time:11077211ms step_avg:43.71ms -step:253600/500000 train_loss:2.0051 train_time:11085946ms step_avg:43.71ms -step:253800/500000 train_loss:1.8818 train_time:11094681ms step_avg:43.71ms -step:254000/500000 train_loss:2.0705 train_time:11103415ms step_avg:43.71ms -step:254200/500000 train_loss:1.9941 train_time:11112156ms step_avg:43.71ms -step:254400/500000 train_loss:1.9379 train_time:11120890ms step_avg:43.71ms -step:254600/500000 train_loss:1.9741 train_time:11129625ms step_avg:43.71ms -step:254800/500000 train_loss:1.8901 train_time:11138356ms step_avg:43.71ms -step:255000/500000 train_loss:1.9534 train_time:11147089ms step_avg:43.71ms -step:255200/500000 train_loss:2.0204 train_time:11155825ms step_avg:43.71ms -step:255400/500000 train_loss:1.8525 train_time:11164563ms step_avg:43.71ms -step:255600/500000 train_loss:2.0028 train_time:11173305ms step_avg:43.71ms -step:255800/500000 train_loss:1.9450 train_time:11182037ms step_avg:43.71ms -step:256000/500000 train_loss:2.0375 train_time:11190773ms step_avg:43.71ms -step:256200/500000 train_loss:1.9353 train_time:11199508ms step_avg:43.71ms -step:256400/500000 train_loss:2.1361 train_time:11208247ms step_avg:43.71ms -step:256600/500000 train_loss:2.0129 train_time:11216977ms step_avg:43.71ms -step:256800/500000 train_loss:2.0725 train_time:11225838ms step_avg:43.71ms -step:257000/500000 train_loss:2.1332 train_time:11234569ms step_avg:43.71ms -step:257200/500000 train_loss:1.7644 train_time:11243314ms step_avg:43.71ms -step:257400/500000 train_loss:2.0439 train_time:11252048ms step_avg:43.71ms -step:257600/500000 train_loss:1.9964 train_time:11260788ms step_avg:43.71ms -step:257800/500000 train_loss:2.0681 train_time:11269525ms step_avg:43.71ms -step:258000/500000 train_loss:1.9432 train_time:11278255ms step_avg:43.71ms -step:258200/500000 train_loss:1.9413 train_time:11286980ms step_avg:43.71ms -step:258400/500000 train_loss:2.1118 train_time:11295708ms step_avg:43.71ms -step:258600/500000 train_loss:1.9258 train_time:11304449ms step_avg:43.71ms -step:258800/500000 train_loss:2.0245 train_time:11313180ms step_avg:43.71ms -step:259000/500000 train_loss:2.0141 train_time:11321910ms step_avg:43.71ms -step:259200/500000 train_loss:1.9900 train_time:11330641ms step_avg:43.71ms -step:259400/500000 train_loss:2.4748 train_time:11339371ms step_avg:43.71ms -step:259600/500000 train_loss:2.0246 train_time:11348100ms step_avg:43.71ms -step:259800/500000 train_loss:1.9782 train_time:11356835ms step_avg:43.71ms -step:260000/500000 train_loss:1.9953 train_time:11365662ms step_avg:43.71ms -step:260000/500000 val_loss:2.0041 val_bpb:1.1870 train_time:11365677ms step_avg:43.71ms -step:260200/500000 train_loss:1.9975 train_time:11374405ms step_avg:43.71ms -step:260400/500000 train_loss:2.0585 train_time:11383133ms step_avg:43.71ms -step:260600/500000 train_loss:2.0042 train_time:11391866ms step_avg:43.71ms -step:260800/500000 train_loss:1.9907 train_time:11400605ms step_avg:43.71ms -step:261000/500000 train_loss:1.8036 train_time:11409345ms step_avg:43.71ms -step:261200/500000 train_loss:1.9688 train_time:11418070ms step_avg:43.71ms -step:261400/500000 train_loss:2.0001 train_time:11426801ms step_avg:43.71ms -step:261600/500000 train_loss:1.7854 train_time:11435534ms step_avg:43.71ms -step:261800/500000 train_loss:1.9908 train_time:11444271ms step_avg:43.71ms -step:262000/500000 train_loss:2.0764 train_time:11453013ms step_avg:43.71ms -step:262200/500000 train_loss:1.8727 train_time:11461753ms step_avg:43.71ms -step:262400/500000 train_loss:2.6147 train_time:11470498ms step_avg:43.71ms -step:262600/500000 train_loss:2.0029 train_time:11479235ms step_avg:43.71ms -step:262800/500000 train_loss:2.2071 train_time:11487964ms step_avg:43.71ms -step:263000/500000 train_loss:2.0746 train_time:11496701ms step_avg:43.71ms -step:263200/500000 train_loss:2.1073 train_time:11505436ms step_avg:43.71ms -step:263400/500000 train_loss:1.9842 train_time:11514169ms step_avg:43.71ms -step:263600/500000 train_loss:2.0114 train_time:11522905ms step_avg:43.71ms -step:263800/500000 train_loss:2.0245 train_time:11531645ms step_avg:43.71ms -step:264000/500000 train_loss:2.0311 train_time:11540511ms step_avg:43.71ms -step:264200/500000 train_loss:1.9754 train_time:11549243ms step_avg:43.71ms -step:264400/500000 train_loss:2.0189 train_time:11557976ms step_avg:43.71ms -step:264600/500000 train_loss:2.0092 train_time:11566712ms step_avg:43.71ms -step:264800/500000 train_loss:1.9114 train_time:11575445ms step_avg:43.71ms -step:265000/500000 train_loss:1.8924 train_time:11584189ms step_avg:43.71ms -step:265200/500000 train_loss:2.0347 train_time:11592928ms step_avg:43.71ms -step:265400/500000 train_loss:2.0220 train_time:11601670ms step_avg:43.71ms -step:265600/500000 train_loss:2.0149 train_time:11610396ms step_avg:43.71ms -step:265800/500000 train_loss:2.0498 train_time:11619137ms step_avg:43.71ms -step:266000/500000 train_loss:2.0556 train_time:11627871ms step_avg:43.71ms -step:266200/500000 train_loss:1.9562 train_time:11636604ms step_avg:43.71ms -step:266400/500000 train_loss:2.0086 train_time:11645338ms step_avg:43.71ms -step:266600/500000 train_loss:2.0249 train_time:11654089ms step_avg:43.71ms -step:266800/500000 train_loss:1.9725 train_time:11662827ms step_avg:43.71ms -step:267000/500000 train_loss:2.0500 train_time:11671566ms step_avg:43.71ms -step:267200/500000 train_loss:2.0668 train_time:11680299ms step_avg:43.71ms -step:267400/500000 train_loss:1.9282 train_time:11689039ms step_avg:43.71ms -step:267600/500000 train_loss:1.9192 train_time:11697772ms step_avg:43.71ms -step:267800/500000 train_loss:1.9734 train_time:11706515ms step_avg:43.71ms -step:268000/500000 train_loss:2.0209 train_time:11715254ms step_avg:43.71ms -step:268200/500000 train_loss:1.9353 train_time:11724109ms step_avg:43.71ms -step:268400/500000 train_loss:2.0996 train_time:11732847ms step_avg:43.71ms -step:268600/500000 train_loss:1.9886 train_time:11741584ms step_avg:43.71ms -step:268800/500000 train_loss:2.0374 train_time:11750317ms step_avg:43.71ms -step:269000/500000 train_loss:2.0219 train_time:11759057ms step_avg:43.71ms -step:269200/500000 train_loss:1.8802 train_time:11767793ms step_avg:43.71ms -step:269400/500000 train_loss:2.0656 train_time:11776526ms step_avg:43.71ms -step:269600/500000 train_loss:1.9443 train_time:11785264ms step_avg:43.71ms -step:269800/500000 train_loss:1.9677 train_time:11793999ms step_avg:43.71ms -step:270000/500000 train_loss:2.0042 train_time:11802735ms step_avg:43.71ms -step:270200/500000 train_loss:1.9749 train_time:11811467ms step_avg:43.71ms -step:270400/500000 train_loss:2.0969 train_time:11820205ms step_avg:43.71ms -step:270600/500000 train_loss:2.1183 train_time:11828939ms step_avg:43.71ms -step:270800/500000 train_loss:1.8985 train_time:11837662ms step_avg:43.71ms -step:271000/500000 train_loss:2.0729 train_time:11846398ms step_avg:43.71ms -step:271200/500000 train_loss:2.0153 train_time:11855135ms step_avg:43.71ms -step:271400/500000 train_loss:2.0576 train_time:11863864ms step_avg:43.71ms -step:271600/500000 train_loss:1.9227 train_time:11872593ms step_avg:43.71ms -step:271800/500000 train_loss:1.9194 train_time:11881327ms step_avg:43.71ms -step:272000/500000 train_loss:2.0508 train_time:11890061ms step_avg:43.71ms -step:272200/500000 train_loss:1.9958 train_time:11898917ms step_avg:43.71ms -step:272400/500000 train_loss:2.0412 train_time:11907652ms step_avg:43.71ms -step:272600/500000 train_loss:2.1157 train_time:11916387ms step_avg:43.71ms -step:272800/500000 train_loss:1.9742 train_time:11925122ms step_avg:43.71ms -step:273000/500000 train_loss:1.9292 train_time:11933857ms step_avg:43.71ms -step:273200/500000 train_loss:2.2138 train_time:11942593ms step_avg:43.71ms -step:273400/500000 train_loss:2.0001 train_time:11951335ms step_avg:43.71ms -step:273600/500000 train_loss:2.0742 train_time:11960070ms step_avg:43.71ms -step:273800/500000 train_loss:2.0358 train_time:11968805ms step_avg:43.71ms -step:274000/500000 train_loss:1.9466 train_time:11977538ms step_avg:43.71ms -step:274200/500000 train_loss:2.1008 train_time:11986276ms step_avg:43.71ms -step:274400/500000 train_loss:1.9161 train_time:11995003ms step_avg:43.71ms -step:274600/500000 train_loss:1.9558 train_time:12003738ms step_avg:43.71ms -step:274800/500000 train_loss:2.2432 train_time:12012473ms step_avg:43.71ms -step:275000/500000 train_loss:2.0272 train_time:12021207ms step_avg:43.71ms -step:275200/500000 train_loss:2.0214 train_time:12029946ms step_avg:43.71ms -step:275400/500000 train_loss:1.9085 train_time:12038674ms step_avg:43.71ms -step:275600/500000 train_loss:2.0615 train_time:12047412ms step_avg:43.71ms -step:275800/500000 train_loss:2.1550 train_time:12056142ms step_avg:43.71ms -step:276000/500000 train_loss:2.1168 train_time:12064879ms step_avg:43.71ms -step:276200/500000 train_loss:1.9004 train_time:12073619ms step_avg:43.71ms -step:276400/500000 train_loss:1.8931 train_time:12082474ms step_avg:43.71ms -step:276600/500000 train_loss:2.0315 train_time:12091206ms step_avg:43.71ms -step:276800/500000 train_loss:2.0176 train_time:12099933ms step_avg:43.71ms -step:277000/500000 train_loss:2.0328 train_time:12108668ms step_avg:43.71ms -step:277200/500000 train_loss:1.9919 train_time:12117401ms step_avg:43.71ms -step:277400/500000 train_loss:2.0349 train_time:12126135ms step_avg:43.71ms -step:277600/500000 train_loss:1.9732 train_time:12134862ms step_avg:43.71ms -step:277800/500000 train_loss:2.0866 train_time:12143592ms step_avg:43.71ms -step:278000/500000 train_loss:1.9639 train_time:12152324ms step_avg:43.71ms -step:278200/500000 train_loss:1.8576 train_time:12161051ms step_avg:43.71ms -step:278400/500000 train_loss:1.9091 train_time:12169780ms step_avg:43.71ms -step:278600/500000 train_loss:1.9311 train_time:12178513ms step_avg:43.71ms -step:278800/500000 train_loss:2.0320 train_time:12187246ms step_avg:43.71ms -step:279000/500000 train_loss:2.2007 train_time:12195978ms step_avg:43.71ms -step:279200/500000 train_loss:1.9399 train_time:12204704ms step_avg:43.71ms -step:279400/500000 train_loss:2.1437 train_time:12213438ms step_avg:43.71ms -step:279600/500000 train_loss:1.8597 train_time:12222170ms step_avg:43.71ms -step:279800/500000 train_loss:1.8687 train_time:12230902ms step_avg:43.71ms -step:280000/500000 train_loss:1.9473 train_time:12239631ms step_avg:43.71ms -step:280000/500000 val_loss:2.0059 val_bpb:1.1880 train_time:12239647ms step_avg:43.71ms -step:280200/500000 train_loss:1.8587 train_time:12248371ms step_avg:43.71ms -step:280400/500000 train_loss:1.9275 train_time:12257226ms step_avg:43.71ms -step:280600/500000 train_loss:2.0865 train_time:12265963ms step_avg:43.71ms -step:280800/500000 train_loss:2.0359 train_time:12274697ms step_avg:43.71ms -step:281000/500000 train_loss:2.0256 train_time:12283428ms step_avg:43.71ms -step:281200/500000 train_loss:2.0721 train_time:12292168ms step_avg:43.71ms -step:281400/500000 train_loss:1.9527 train_time:12300901ms step_avg:43.71ms -step:281600/500000 train_loss:2.0435 train_time:12309642ms step_avg:43.71ms -step:281800/500000 train_loss:1.9409 train_time:12318374ms step_avg:43.71ms -step:282000/500000 train_loss:1.8839 train_time:12327110ms step_avg:43.71ms -step:282200/500000 train_loss:1.9873 train_time:12335849ms step_avg:43.71ms -step:282400/500000 train_loss:2.0421 train_time:12344583ms step_avg:43.71ms -step:282600/500000 train_loss:2.1513 train_time:12353319ms step_avg:43.71ms -step:282800/500000 train_loss:1.9599 train_time:12362052ms step_avg:43.71ms -step:283000/500000 train_loss:2.1058 train_time:12370792ms step_avg:43.71ms -step:283200/500000 train_loss:2.0484 train_time:12379527ms step_avg:43.71ms -step:283400/500000 train_loss:2.0516 train_time:12388269ms step_avg:43.71ms -step:283600/500000 train_loss:1.9050 train_time:12397005ms step_avg:43.71ms -step:283800/500000 train_loss:1.9629 train_time:12405746ms step_avg:43.71ms -step:284000/500000 train_loss:1.9291 train_time:12414495ms step_avg:43.71ms -step:284200/500000 train_loss:2.0472 train_time:12423237ms step_avg:43.71ms -step:284400/500000 train_loss:2.0072 train_time:12431976ms step_avg:43.71ms -step:284600/500000 train_loss:1.9916 train_time:12440841ms step_avg:43.71ms -step:284800/500000 train_loss:1.9631 train_time:12449579ms step_avg:43.71ms -step:285000/500000 train_loss:2.0340 train_time:12458321ms step_avg:43.71ms -step:285200/500000 train_loss:1.9090 train_time:12467053ms step_avg:43.71ms -step:285400/500000 train_loss:2.1183 train_time:12475792ms step_avg:43.71ms -step:285600/500000 train_loss:2.0143 train_time:12484524ms step_avg:43.71ms -step:285800/500000 train_loss:2.0114 train_time:12493265ms step_avg:43.71ms -step:286000/500000 train_loss:2.1389 train_time:12502004ms step_avg:43.71ms -step:286200/500000 train_loss:1.9397 train_time:12510739ms step_avg:43.71ms -step:286400/500000 train_loss:1.9368 train_time:12519480ms step_avg:43.71ms -step:286600/500000 train_loss:2.0340 train_time:12528218ms step_avg:43.71ms -step:286800/500000 train_loss:2.0173 train_time:12536952ms step_avg:43.71ms -step:287000/500000 train_loss:2.0685 train_time:12545692ms step_avg:43.71ms -step:287200/500000 train_loss:2.0053 train_time:12554435ms step_avg:43.71ms -step:287400/500000 train_loss:1.9727 train_time:12563172ms step_avg:43.71ms -step:287600/500000 train_loss:1.9363 train_time:12571910ms step_avg:43.71ms -step:287800/500000 train_loss:2.4742 train_time:12580644ms step_avg:43.71ms -step:288000/500000 train_loss:2.0207 train_time:12589381ms step_avg:43.71ms -step:288200/500000 train_loss:1.9881 train_time:12598121ms step_avg:43.71ms -step:288400/500000 train_loss:2.0589 train_time:12606861ms step_avg:43.71ms -step:288600/500000 train_loss:1.9856 train_time:12615613ms step_avg:43.71ms -step:288800/500000 train_loss:1.9577 train_time:12624470ms step_avg:43.71ms -step:289000/500000 train_loss:2.0183 train_time:12633200ms step_avg:43.71ms -step:289200/500000 train_loss:1.8817 train_time:12641936ms step_avg:43.71ms -step:289400/500000 train_loss:2.1841 train_time:12650673ms step_avg:43.71ms -step:289600/500000 train_loss:2.0147 train_time:12659407ms step_avg:43.71ms -step:289800/500000 train_loss:2.0655 train_time:12668147ms step_avg:43.71ms -step:290000/500000 train_loss:2.1097 train_time:12676876ms step_avg:43.71ms -step:290200/500000 train_loss:1.9592 train_time:12685620ms step_avg:43.71ms -step:290400/500000 train_loss:2.0262 train_time:12694358ms step_avg:43.71ms -step:290600/500000 train_loss:1.9529 train_time:12703095ms step_avg:43.71ms -step:290800/500000 train_loss:2.0480 train_time:12711840ms step_avg:43.71ms -step:291000/500000 train_loss:1.8834 train_time:12720580ms step_avg:43.71ms -step:291200/500000 train_loss:2.3031 train_time:12729321ms step_avg:43.71ms -step:291400/500000 train_loss:1.9826 train_time:12738057ms step_avg:43.71ms -step:291600/500000 train_loss:2.1083 train_time:12746796ms step_avg:43.71ms -step:291800/500000 train_loss:2.0582 train_time:12755531ms step_avg:43.71ms -step:292000/500000 train_loss:1.9632 train_time:12764274ms step_avg:43.71ms -step:292200/500000 train_loss:2.3024 train_time:12773015ms step_avg:43.71ms -step:292400/500000 train_loss:2.0066 train_time:12781752ms step_avg:43.71ms -step:292600/500000 train_loss:1.9131 train_time:12790492ms step_avg:43.71ms -step:292800/500000 train_loss:1.7997 train_time:12799355ms step_avg:43.71ms -step:293000/500000 train_loss:1.9488 train_time:12808094ms step_avg:43.71ms -step:293200/500000 train_loss:1.9176 train_time:12816830ms step_avg:43.71ms -step:293400/500000 train_loss:1.9941 train_time:12825563ms step_avg:43.71ms -step:293600/500000 train_loss:1.9393 train_time:12834298ms step_avg:43.71ms -step:293800/500000 train_loss:2.0138 train_time:12843033ms step_avg:43.71ms -step:294000/500000 train_loss:1.9956 train_time:12851779ms step_avg:43.71ms -step:294200/500000 train_loss:2.2086 train_time:12860523ms step_avg:43.71ms -step:294400/500000 train_loss:2.0443 train_time:12869263ms step_avg:43.71ms -step:294600/500000 train_loss:2.0434 train_time:12878001ms step_avg:43.71ms -step:294800/500000 train_loss:1.9583 train_time:12886742ms step_avg:43.71ms -step:295000/500000 train_loss:2.0998 train_time:12895484ms step_avg:43.71ms -step:295200/500000 train_loss:2.0722 train_time:12904220ms step_avg:43.71ms -step:295400/500000 train_loss:1.9461 train_time:12912961ms step_avg:43.71ms -step:295600/500000 train_loss:2.1570 train_time:12921702ms step_avg:43.71ms -step:295800/500000 train_loss:2.0004 train_time:12930439ms step_avg:43.71ms -step:296000/500000 train_loss:2.0067 train_time:12939179ms step_avg:43.71ms -step:296200/500000 train_loss:2.0775 train_time:12947917ms step_avg:43.71ms -step:296400/500000 train_loss:1.9923 train_time:12956652ms step_avg:43.71ms -step:296600/500000 train_loss:1.9149 train_time:12965382ms step_avg:43.71ms -step:296800/500000 train_loss:2.0540 train_time:12974118ms step_avg:43.71ms -step:297000/500000 train_loss:2.0463 train_time:12982942ms step_avg:43.71ms -step:297200/500000 train_loss:2.0057 train_time:12991673ms step_avg:43.71ms -step:297400/500000 train_loss:2.0463 train_time:13000411ms step_avg:43.71ms -step:297600/500000 train_loss:1.9426 train_time:13009145ms step_avg:43.71ms -step:297800/500000 train_loss:2.1075 train_time:13017877ms step_avg:43.71ms -step:298000/500000 train_loss:2.0239 train_time:13026611ms step_avg:43.71ms -step:298200/500000 train_loss:2.0610 train_time:13035344ms step_avg:43.71ms -step:298400/500000 train_loss:1.9703 train_time:13044080ms step_avg:43.71ms -step:298600/500000 train_loss:2.0391 train_time:13052807ms step_avg:43.71ms -step:298800/500000 train_loss:2.0636 train_time:13061543ms step_avg:43.71ms -step:299000/500000 train_loss:2.0894 train_time:13070275ms step_avg:43.71ms -step:299200/500000 train_loss:2.0629 train_time:13079003ms step_avg:43.71ms -step:299400/500000 train_loss:1.9541 train_time:13087740ms step_avg:43.71ms -step:299600/500000 train_loss:1.9493 train_time:13096471ms step_avg:43.71ms -step:299800/500000 train_loss:1.9702 train_time:13105202ms step_avg:43.71ms -step:300000/500000 train_loss:2.1034 train_time:13114057ms step_avg:43.71ms -step:300000/500000 val_loss:2.0019 val_bpb:1.1856 train_time:13114073ms step_avg:43.71ms -step:300200/500000 train_loss:2.0228 train_time:13122789ms step_avg:43.71ms -step:300400/500000 train_loss:1.9918 train_time:13131517ms step_avg:43.71ms -step:300600/500000 train_loss:1.9048 train_time:13140246ms step_avg:43.71ms -step:300800/500000 train_loss:1.9942 train_time:13148989ms step_avg:43.71ms -step:301000/500000 train_loss:2.1526 train_time:13157717ms step_avg:43.71ms -step:301200/500000 train_loss:1.9766 train_time:13166455ms step_avg:43.71ms -step:301400/500000 train_loss:2.0515 train_time:13175188ms step_avg:43.71ms -step:301600/500000 train_loss:1.8824 train_time:13183920ms step_avg:43.71ms -step:301800/500000 train_loss:1.8421 train_time:13192653ms step_avg:43.71ms -step:302000/500000 train_loss:2.0031 train_time:13201380ms step_avg:43.71ms -step:302200/500000 train_loss:1.9957 train_time:13210119ms step_avg:43.71ms -step:302400/500000 train_loss:2.1431 train_time:13218856ms step_avg:43.71ms -step:302600/500000 train_loss:2.0113 train_time:13227593ms step_avg:43.71ms -step:302800/500000 train_loss:2.0139 train_time:13236331ms step_avg:43.71ms -step:303000/500000 train_loss:1.9975 train_time:13245066ms step_avg:43.71ms -step:303200/500000 train_loss:2.1686 train_time:13253809ms step_avg:43.71ms -step:303400/500000 train_loss:1.9958 train_time:13262553ms step_avg:43.71ms -step:303600/500000 train_loss:2.0634 train_time:13271294ms step_avg:43.71ms -step:303800/500000 train_loss:1.8341 train_time:13280026ms step_avg:43.71ms -step:304000/500000 train_loss:2.0196 train_time:13288891ms step_avg:43.71ms -step:304200/500000 train_loss:2.0038 train_time:13297621ms step_avg:43.71ms -step:304400/500000 train_loss:2.1533 train_time:13306354ms step_avg:43.71ms -step:304600/500000 train_loss:1.9385 train_time:13315087ms step_avg:43.71ms -step:304800/500000 train_loss:1.9258 train_time:13323823ms step_avg:43.71ms -step:305000/500000 train_loss:2.0313 train_time:13332561ms step_avg:43.71ms -step:305200/500000 train_loss:2.0963 train_time:13341296ms step_avg:43.71ms -step:305400/500000 train_loss:1.9291 train_time:13350029ms step_avg:43.71ms -step:305600/500000 train_loss:2.1894 train_time:13358764ms step_avg:43.71ms -step:305800/500000 train_loss:1.9851 train_time:13367498ms step_avg:43.71ms -step:306000/500000 train_loss:2.0476 train_time:13376229ms step_avg:43.71ms -step:306200/500000 train_loss:1.9089 train_time:13384965ms step_avg:43.71ms -step:306400/500000 train_loss:1.9578 train_time:13393701ms step_avg:43.71ms -step:306600/500000 train_loss:1.9322 train_time:13402431ms step_avg:43.71ms -step:306800/500000 train_loss:1.9681 train_time:13411166ms step_avg:43.71ms -step:307000/500000 train_loss:1.9248 train_time:13419899ms step_avg:43.71ms -step:307200/500000 train_loss:2.0327 train_time:13428633ms step_avg:43.71ms -step:307400/500000 train_loss:2.1124 train_time:13437369ms step_avg:43.71ms -step:307600/500000 train_loss:1.8934 train_time:13446105ms step_avg:43.71ms -step:307800/500000 train_loss:1.9565 train_time:13454841ms step_avg:43.71ms -step:308000/500000 train_loss:1.9663 train_time:13463575ms step_avg:43.71ms -step:308200/500000 train_loss:2.0848 train_time:13472431ms step_avg:43.71ms -step:308400/500000 train_loss:1.9285 train_time:13481167ms step_avg:43.71ms -step:308600/500000 train_loss:2.0670 train_time:13489905ms step_avg:43.71ms -step:308800/500000 train_loss:1.8746 train_time:13498638ms step_avg:43.71ms -step:309000/500000 train_loss:1.9112 train_time:13507371ms step_avg:43.71ms -step:309200/500000 train_loss:1.9494 train_time:13516103ms step_avg:43.71ms -step:309400/500000 train_loss:1.9585 train_time:13524834ms step_avg:43.71ms -step:309600/500000 train_loss:2.0978 train_time:13533563ms step_avg:43.71ms -step:309800/500000 train_loss:1.7424 train_time:13542293ms step_avg:43.71ms -step:310000/500000 train_loss:2.0507 train_time:13551024ms step_avg:43.71ms -step:310200/500000 train_loss:1.9832 train_time:13559759ms step_avg:43.71ms -step:310400/500000 train_loss:2.0606 train_time:13568497ms step_avg:43.71ms -step:310600/500000 train_loss:1.9982 train_time:13577225ms step_avg:43.71ms -step:310800/500000 train_loss:1.9846 train_time:13585959ms step_avg:43.71ms -step:311000/500000 train_loss:2.1402 train_time:13594684ms step_avg:43.71ms -step:311200/500000 train_loss:1.9878 train_time:13603418ms step_avg:43.71ms -step:311400/500000 train_loss:2.0720 train_time:13612147ms step_avg:43.71ms -step:311600/500000 train_loss:1.9734 train_time:13620878ms step_avg:43.71ms -step:311800/500000 train_loss:2.0040 train_time:13629609ms step_avg:43.71ms -step:312000/500000 train_loss:1.9738 train_time:13638336ms step_avg:43.71ms -step:312200/500000 train_loss:2.0241 train_time:13647067ms step_avg:43.71ms -step:312400/500000 train_loss:1.9978 train_time:13655930ms step_avg:43.71ms -step:312600/500000 train_loss:2.0128 train_time:13664662ms step_avg:43.71ms -step:312800/500000 train_loss:2.0693 train_time:13673392ms step_avg:43.71ms -step:313000/500000 train_loss:2.0319 train_time:13682118ms step_avg:43.71ms -step:313200/500000 train_loss:1.9284 train_time:13690852ms step_avg:43.71ms -step:313400/500000 train_loss:1.8951 train_time:13699586ms step_avg:43.71ms -step:313600/500000 train_loss:2.0394 train_time:13708322ms step_avg:43.71ms -step:313800/500000 train_loss:2.0955 train_time:13717054ms step_avg:43.71ms -step:314000/500000 train_loss:1.9987 train_time:13725790ms step_avg:43.71ms -step:314200/500000 train_loss:2.0925 train_time:13734528ms step_avg:43.71ms -step:314400/500000 train_loss:1.9578 train_time:13743269ms step_avg:43.71ms -step:314600/500000 train_loss:2.0680 train_time:13752011ms step_avg:43.71ms -step:314800/500000 train_loss:2.0073 train_time:13760745ms step_avg:43.71ms -step:315000/500000 train_loss:2.0866 train_time:13769479ms step_avg:43.71ms -step:315200/500000 train_loss:1.9914 train_time:13778210ms step_avg:43.71ms -step:315400/500000 train_loss:1.9911 train_time:13786945ms step_avg:43.71ms -step:315600/500000 train_loss:2.1094 train_time:13795674ms step_avg:43.71ms -step:315800/500000 train_loss:1.9628 train_time:13804407ms step_avg:43.71ms -step:316000/500000 train_loss:1.9922 train_time:13813140ms step_avg:43.71ms -step:316200/500000 train_loss:1.9326 train_time:13821872ms step_avg:43.71ms -step:316400/500000 train_loss:1.9757 train_time:13830721ms step_avg:43.71ms -step:316600/500000 train_loss:1.9799 train_time:13839462ms step_avg:43.71ms -step:316800/500000 train_loss:2.0344 train_time:13848182ms step_avg:43.71ms -step:317000/500000 train_loss:1.9756 train_time:13856912ms step_avg:43.71ms -step:317200/500000 train_loss:1.8813 train_time:13865639ms step_avg:43.71ms -step:317400/500000 train_loss:2.0688 train_time:13874365ms step_avg:43.71ms -step:317600/500000 train_loss:2.1282 train_time:13883090ms step_avg:43.71ms -step:317800/500000 train_loss:2.0416 train_time:13891818ms step_avg:43.71ms -step:318000/500000 train_loss:1.9459 train_time:13900547ms step_avg:43.71ms -step:318200/500000 train_loss:2.1129 train_time:13909275ms step_avg:43.71ms -step:318400/500000 train_loss:1.9522 train_time:13918006ms step_avg:43.71ms -step:318600/500000 train_loss:2.0662 train_time:13926738ms step_avg:43.71ms -step:318800/500000 train_loss:2.0459 train_time:13935471ms step_avg:43.71ms -step:319000/500000 train_loss:2.0817 train_time:13944202ms step_avg:43.71ms -step:319200/500000 train_loss:2.0718 train_time:13952933ms step_avg:43.71ms -step:319400/500000 train_loss:2.0248 train_time:13961668ms step_avg:43.71ms -step:319600/500000 train_loss:2.0153 train_time:13970395ms step_avg:43.71ms -step:319800/500000 train_loss:1.9507 train_time:13979128ms step_avg:43.71ms -step:320000/500000 train_loss:1.9200 train_time:13987858ms step_avg:43.71ms -step:320000/500000 val_loss:2.0010 val_bpb:1.1851 train_time:13987873ms step_avg:43.71ms -step:320200/500000 train_loss:2.0915 train_time:13996588ms step_avg:43.71ms -step:320400/500000 train_loss:1.9669 train_time:14005321ms step_avg:43.71ms -step:320600/500000 train_loss:2.0496 train_time:14014180ms step_avg:43.71ms -step:320800/500000 train_loss:2.1033 train_time:14022909ms step_avg:43.71ms -step:321000/500000 train_loss:1.8803 train_time:14031642ms step_avg:43.71ms -step:321200/500000 train_loss:1.9655 train_time:14040378ms step_avg:43.71ms -step:321400/500000 train_loss:2.0748 train_time:14049121ms step_avg:43.71ms -step:321600/500000 train_loss:2.0919 train_time:14057851ms step_avg:43.71ms -step:321800/500000 train_loss:2.0200 train_time:14066585ms step_avg:43.71ms -step:322000/500000 train_loss:2.0471 train_time:14075312ms step_avg:43.71ms -step:322200/500000 train_loss:1.9971 train_time:14084041ms step_avg:43.71ms -step:322400/500000 train_loss:1.8674 train_time:14092769ms step_avg:43.71ms -step:322600/500000 train_loss:1.9619 train_time:14101496ms step_avg:43.71ms -step:322800/500000 train_loss:2.0699 train_time:14110228ms step_avg:43.71ms -step:323000/500000 train_loss:2.4914 train_time:14118957ms step_avg:43.71ms -step:323200/500000 train_loss:1.9907 train_time:14127689ms step_avg:43.71ms -step:323400/500000 train_loss:2.0156 train_time:14136416ms step_avg:43.71ms -step:323600/500000 train_loss:2.0626 train_time:14145143ms step_avg:43.71ms -step:323800/500000 train_loss:1.9674 train_time:14153871ms step_avg:43.71ms -step:324000/500000 train_loss:2.0524 train_time:14162599ms step_avg:43.71ms -step:324200/500000 train_loss:2.0582 train_time:14171336ms step_avg:43.71ms -step:324400/500000 train_loss:1.8915 train_time:14180068ms step_avg:43.71ms -step:324600/500000 train_loss:2.0369 train_time:14188927ms step_avg:43.71ms -step:324800/500000 train_loss:2.0923 train_time:14197654ms step_avg:43.71ms -step:325000/500000 train_loss:1.9624 train_time:14206393ms step_avg:43.71ms -step:325200/500000 train_loss:2.0010 train_time:14215134ms step_avg:43.71ms -step:325400/500000 train_loss:2.0460 train_time:14223861ms step_avg:43.71ms -step:325600/500000 train_loss:2.0737 train_time:14232589ms step_avg:43.71ms -step:325800/500000 train_loss:1.8715 train_time:14241323ms step_avg:43.71ms -step:326000/500000 train_loss:1.9732 train_time:14250055ms step_avg:43.71ms -step:326200/500000 train_loss:1.8768 train_time:14258799ms step_avg:43.71ms -step:326400/500000 train_loss:1.9460 train_time:14267540ms step_avg:43.71ms -step:326600/500000 train_loss:1.9567 train_time:14276274ms step_avg:43.71ms -step:326800/500000 train_loss:1.9932 train_time:14285014ms step_avg:43.71ms -step:327000/500000 train_loss:2.0671 train_time:14293752ms step_avg:43.71ms -step:327200/500000 train_loss:1.9528 train_time:14302485ms step_avg:43.71ms -step:327400/500000 train_loss:2.0451 train_time:14311217ms step_avg:43.71ms -step:327600/500000 train_loss:2.0360 train_time:14319952ms step_avg:43.71ms -step:327800/500000 train_loss:2.0273 train_time:14328692ms step_avg:43.71ms -step:328000/500000 train_loss:2.1318 train_time:14337428ms step_avg:43.71ms -step:328200/500000 train_loss:2.0219 train_time:14346167ms step_avg:43.71ms -step:328400/500000 train_loss:1.8702 train_time:14354904ms step_avg:43.71ms -step:328600/500000 train_loss:2.2112 train_time:14363648ms step_avg:43.71ms -step:328800/500000 train_loss:2.0007 train_time:14372507ms step_avg:43.71ms -step:329000/500000 train_loss:1.9333 train_time:14381249ms step_avg:43.71ms -step:329200/500000 train_loss:1.8833 train_time:14389989ms step_avg:43.71ms -step:329400/500000 train_loss:1.9639 train_time:14398731ms step_avg:43.71ms -step:329430/500000 val_loss:1.9837 val_bpb:1.1749 train_time:14400039ms step_avg:43.71ms -stopping_early: wallclock_cap train_time:14400039ms step:329430/500000 -peak memory allocated: 10184 MiB reserved: 10588 MiB -Serialized model: 67224983 bytes -Code size: 47642 bytes -Total submission size: 67272625 bytes -Serialized model int8+zlib: 15762519 bytes (payload:17178912 raw_torch:17224025 payload_ratio:3.91x) -Total submission size int8+zlib: 15810161 bytes -final_int8_zlib_roundtrip val_loss:2.0386 val_bpb:1.2074 eval_time:1356ms -final_int8_zlib_roundtrip_exact val_loss:2.03860961 val_bpb:1.20737944 diff --git a/records/track_non_record_16mb/2026-03-18_Quasi10Bfrom50B_SP1024_9x512_KV4_4h_pgut3/train_gpt.py b/records/track_non_record_16mb/2026-03-18_Quasi10Bfrom50B_SP1024_9x512_KV4_4h_pgut3/train_gpt.py deleted file mode 100644 index 0deb0565f5..0000000000 --- a/records/track_non_record_16mb/2026-03-18_Quasi10Bfrom50B_SP1024_9x512_KV4_4h_pgut3/train_gpt.py +++ /dev/null @@ -1,1126 +0,0 @@ -""" -The `train_gpt.py` and `train_gpt_mlx.py` scripts are intended as good launching-off points for new participants, not SOTA configs. We'll accept PRs that tune, improve, or simplify these scripts without significantly increasing complexity, but competitive submissions should stay in the `/records` folder. - -Hard stop: `train_gpt.py` and `train_gpt_mlx.py` must never be longer than 1500 lines. -""" - -from __future__ import annotations - -import copy -import glob -import io -import math -import os -import random -import subprocess -import sys -import time -import uuid -import zlib -from pathlib import Path - -import numpy as np -import sentencepiece as spm -import torch -import torch.distributed as dist -import torch.nn.functional as F -from torch import Tensor, nn -from torch.nn.parallel import DistributedDataParallel as DDP - -# ----------------------------- -# HYPERPARAMETERS -# ----------------------------- -# Default Simple Baseline run: -# - 9 transformer blocks at width 512 -# - 8 attention heads with 4 KV heads (GQA) and 2x MLP expansion -# - vocab size 1024, sequence length 1024, tied embeddings -# - 524,288 train tokens per step for 20,000 iterations with a ~10 minute cap - -class Hyperparameters: - # Data paths are shard globs produced by the existing preprocessing pipeline. - data_path = os.environ.get("DATA_PATH", "./data/datasets/fineweb10B_sp1024") - train_files = os.path.join(data_path, "fineweb_train_*.bin") - val_files = os.path.join(data_path, "fineweb_val_*.bin") - tokenizer_path = os.environ.get("TOKENIZER_PATH", "./data/tokenizers/fineweb_1024_bpe.model") - run_id = os.environ.get("RUN_ID", str(uuid.uuid4())) - seed = int(os.environ.get("SEED", 1337)) - - # Validation cadence and batch size. Validation always uses the full fineweb_val split. - val_batch_size = int(os.environ.get("VAL_BATCH_SIZE", 524_288)) - val_loss_every = int(os.environ.get("VAL_LOSS_EVERY", 1000)) - train_log_every = int(os.environ.get("TRAIN_LOG_EVERY", 200)) - - # Training length. - iterations = int(os.environ.get("ITERATIONS", 20000)) - warmdown_iters = int(os.environ.get("WARMDOWN_ITERS", 1200)) - warmup_steps = int(os.environ.get("WARMUP_STEPS", 20)) - train_batch_tokens = int(os.environ.get("TRAIN_BATCH_TOKENS", 524_288)) - train_seq_len = int(os.environ.get("TRAIN_SEQ_LEN", 1024)) - max_wallclock_seconds = float(os.environ.get("MAX_WALLCLOCK_SECONDS", 600.0)) - qk_gain_init = float(os.environ.get("QK_GAIN_INIT", 1.5)) - - # Model shape. - vocab_size = int(os.environ.get("VOCAB_SIZE", 1024)) - num_layers = int(os.environ.get("NUM_LAYERS", 9)) - num_kv_heads = int(os.environ.get("NUM_KV_HEADS", 4)) - model_dim = int(os.environ.get("MODEL_DIM", 512)) - num_heads = int(os.environ.get("NUM_HEADS", 8)) - mlp_mult = int(os.environ.get("MLP_MULT", 2)) - tie_embeddings = bool(int(os.environ.get("TIE_EMBEDDINGS", "1"))) - rope_base = float(os.environ.get("ROPE_BASE", 10000.0)) - logit_softcap = float(os.environ.get("LOGIT_SOFTCAP", 30.0)) - - # Optimizer hyperparameters. - embed_lr = float(os.environ.get("EMBED_LR", 0.6)) - head_lr = float(os.environ.get("HEAD_LR", 0.008)) - tied_embed_lr = float(os.environ.get("TIED_EMBED_LR", 0.05)) - tied_embed_init_std = float(os.environ.get("TIED_EMBED_INIT_STD", 0.005)) - matrix_lr = float(os.environ.get("MATRIX_LR", 0.04)) - scalar_lr = float(os.environ.get("SCALAR_LR", 0.04)) - muon_momentum = float(os.environ.get("MUON_MOMENTUM", 0.95)) - muon_backend_steps = int(os.environ.get("MUON_BACKEND_STEPS", 5)) - muon_momentum_warmup_start = float(os.environ.get("MUON_MOMENTUM_WARMUP_START", 0.85)) - muon_momentum_warmup_steps = int(os.environ.get("MUON_MOMENTUM_WARMUP_STEPS", 500)) - beta1 = float(os.environ.get("BETA1", 0.9)) - beta2 = float(os.environ.get("BETA2", 0.95)) - adam_eps = float(os.environ.get("ADAM_EPS", 1e-8)) - grad_clip_norm = float(os.environ.get("GRAD_CLIP_NORM", 0.0)) - -# ----------------------------- -# MUON OPTIMIZER -# ----------------------------- -# -# As borrowed from modded-nanogpt -# Background on Muon: https://kellerjordan.github.io/posts/muon/ - -def zeropower_via_newtonschulz5(G: Tensor, steps: int = 10, eps: float = 1e-7) -> Tensor: - # Orthogonalize a 2D update matrix with a fast Newton-Schulz iteration. - # Muon uses this to normalize matrix-shaped gradients before applying them. - a, b, c = (3.4445, -4.7750, 2.0315) - X = G.bfloat16() - X /= X.norm() + eps - transposed = G.size(0) > G.size(1) - if transposed: - X = X.T - for _ in range(steps): - A = X @ X.T - B = b * A + c * A @ A - X = a * X + B @ X - return X.T if transposed else X - - -class Muon(torch.optim.Optimizer): - def __init__(self, params, lr: float, momentum: float, backend_steps: int, nesterov: bool = True): - super().__init__( - params, - dict(lr=lr, momentum=momentum, backend_steps=backend_steps, nesterov=nesterov), - ) - - @torch.no_grad() - def step(self, closure=None): - loss = None - if closure is not None: - with torch.enable_grad(): - loss = closure() - - distributed = dist.is_available() and dist.is_initialized() - world_size = dist.get_world_size() if distributed else 1 - rank = dist.get_rank() if distributed else 0 - - for group in self.param_groups: - params = group["params"] - if not params: - continue - lr = group["lr"] - momentum = group["momentum"] - backend_steps = group["backend_steps"] - nesterov = group["nesterov"] - - total_params = sum(int(p.numel()) for p in params) - updates_flat = torch.zeros(total_params, device=params[0].device, dtype=torch.bfloat16) - - curr = 0 - for i, p in enumerate(params): - if i % world_size == rank and p.grad is not None: - g = p.grad - state = self.state[p] - if "momentum_buffer" not in state: - state["momentum_buffer"] = torch.zeros_like(g) - buf = state["momentum_buffer"] - buf.mul_(momentum).add_(g) - if nesterov: - g = g.add(buf, alpha=momentum) - g = zeropower_via_newtonschulz5(g, steps=backend_steps) - # Scale correction from Muon reference implementations. - g *= max(1, g.size(0) / g.size(1)) ** 0.5 - updates_flat[curr : curr + p.numel()] = g.reshape(-1) - curr += p.numel() - - if distributed: - dist.all_reduce(updates_flat, op=dist.ReduceOp.SUM) - - curr = 0 - for p in params: - g = updates_flat[curr : curr + p.numel()].view_as(p).to(dtype=p.dtype) - p.add_(g, alpha=-lr) - curr += p.numel() - - return loss - - -# ----------------------------- -# TOKENIZER-AGNOSTIC EVALUATION SETUP -# ----------------------------- -# -# It's common for small models have a large fraction of their parameters be embeddings, since the 2 * d_model * d_vocab vectors can be gigantic. -# Instead of locking the tokenizer, we let you bring your own and calculate our validation metrics on the average compression of the validation set. -# We calculate BPB (bits-per-byte) instead of validation loss, so we need methods to count the number of bits per token in the tokenizer. -# Note: Submissions that edit the tokenizer will be examined more carefully, since screwing this up might unjustly improve your score. - -def build_sentencepiece_luts( - sp: spm.SentencePieceProcessor, vocab_size: int, device: torch.device -) -> tuple[Tensor, Tensor, Tensor]: - sp_vocab_size = int(sp.vocab_size()) - table_size = max(sp_vocab_size, vocab_size) - base_bytes_np = np.zeros((table_size,), dtype=np.int16) - has_leading_space_np = np.zeros((table_size,), dtype=np.bool_) - is_boundary_token_np = np.ones((table_size,), dtype=np.bool_) - for token_id in range(sp_vocab_size): - if sp.is_control(token_id) or sp.is_unknown(token_id) or sp.is_unused(token_id): - continue - is_boundary_token_np[token_id] = False - if sp.is_byte(token_id): - base_bytes_np[token_id] = 1 - continue - piece = sp.id_to_piece(token_id) - if piece.startswith("▁"): - has_leading_space_np[token_id] = True - piece = piece[1:] - base_bytes_np[token_id] = len(piece.encode("utf-8")) - return ( - torch.tensor(base_bytes_np, dtype=torch.int16, device=device), - torch.tensor(has_leading_space_np, dtype=torch.bool, device=device), - torch.tensor(is_boundary_token_np, dtype=torch.bool, device=device), - ) - - -def load_validation_tokens(pattern: str, seq_len: int) -> Tensor: - files = [Path(p) for p in sorted(glob.glob(pattern))] - if not files: - raise FileNotFoundError(f"No files found for pattern: {pattern}") - # The export pipeline writes the fixed first-50k-doc validation set to fineweb_val_*. - tokens = torch.cat([load_data_shard(file) for file in files]).contiguous() - usable = ((tokens.numel() - 1) // seq_len) * seq_len - if usable <= 0: - raise ValueError(f"Validation split is too short for TRAIN_SEQ_LEN={seq_len}") - return tokens[: usable + 1] - - -def eval_val( - args: Hyperparameters, - model: nn.Module, - rank: int, - world_size: int, - device: torch.device, - grad_accum_steps: int, - val_tokens: Tensor, - base_bytes_lut: Tensor, - has_leading_space_lut: Tensor, - is_boundary_token_lut: Tensor, -) -> tuple[float, float]: - # Validation computes two metrics: - # - val_loss: token cross-entropy (natural log) - # - val_bpb: tokenizer-agnostic compression metric used by the challenge - local_batch_tokens = args.val_batch_size // (world_size * grad_accum_steps) - if local_batch_tokens < args.train_seq_len: - raise ValueError( - "VAL_BATCH_SIZE must provide at least one sequence per rank; " - f"got VAL_BATCH_SIZE={args.val_batch_size}, WORLD_SIZE={world_size}, " - f"GRAD_ACCUM_STEPS={grad_accum_steps}, TRAIN_SEQ_LEN={args.train_seq_len}" - ) - local_batch_seqs = local_batch_tokens // args.train_seq_len - total_seqs = (val_tokens.numel() - 1) // args.train_seq_len - seq_start = (total_seqs * rank) // world_size - seq_end = (total_seqs * (rank + 1)) // world_size - val_loss_sum = torch.zeros((), device=device, dtype=torch.float64) - val_token_count = torch.zeros((), device=device, dtype=torch.float64) - val_byte_count = torch.zeros((), device=device, dtype=torch.float64) - - model.eval() - with torch.inference_mode(): - for batch_seq_start in range(seq_start, seq_end, local_batch_seqs): - batch_seq_end = min(batch_seq_start + local_batch_seqs, seq_end) - raw_start = batch_seq_start * args.train_seq_len - raw_end = batch_seq_end * args.train_seq_len + 1 - local = val_tokens[raw_start:raw_end].to(device=device, dtype=torch.int64, non_blocking=True) - x = local[:-1].reshape(-1, args.train_seq_len) - y = local[1:].reshape(-1, args.train_seq_len) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - batch_loss = model(x, y).detach() - batch_token_count = float(y.numel()) - val_loss_sum += batch_loss.to(torch.float64) * batch_token_count - val_token_count += batch_token_count - prev_ids = x.reshape(-1) - tgt_ids = y.reshape(-1) - token_bytes = base_bytes_lut[tgt_ids].to(dtype=torch.int16) - token_bytes += (has_leading_space_lut[tgt_ids] & ~is_boundary_token_lut[prev_ids]).to(dtype=torch.int16) - val_byte_count += token_bytes.to(torch.float64).sum() - - if dist.is_available() and dist.is_initialized(): - dist.all_reduce(val_loss_sum, op=dist.ReduceOp.SUM) - dist.all_reduce(val_token_count, op=dist.ReduceOp.SUM) - dist.all_reduce(val_byte_count, op=dist.ReduceOp.SUM) - - val_loss = val_loss_sum / val_token_count - bits_per_token = val_loss.item() / math.log(2.0) - tokens_per_byte = val_token_count.item() / val_byte_count.item() - model.train() - return float(val_loss.item()), float(bits_per_token * tokens_per_byte) - -# ----------------------------- -# POST-TRAINING QUANTIZATION -# ----------------------------- -# -# It's silly to export our model, which is trained in bf16 and fp32, at that same precision. -# Instead, we get approximately the same model (with a small hit) by quantizing the model to int8 & zlib compressing. -# We can then decompress the model and run in higher precision for evaluation, after closing in under the size limit. - -CONTROL_TENSOR_NAME_PATTERNS = tuple( - pattern - for pattern in os.environ.get( - "CONTROL_TENSOR_NAME_PATTERNS", - "attn_scale,attn_scales,mlp_scale,mlp_scales,resid_mix,resid_mixes,q_gain,skip_weight,skip_weights", - ).split(",") - if pattern -) -INT8_KEEP_FLOAT_FP32_NAME_PATTERNS = tuple( - pattern - for pattern in os.environ.get( - "INT8_KEEP_FLOAT_FP32_NAME_PATTERNS", - ",".join(CONTROL_TENSOR_NAME_PATTERNS), - ).split(",") - if pattern -) -INT8_KEEP_FLOAT_MAX_NUMEL = 65_536 -INT8_KEEP_FLOAT_STORE_DTYPE = torch.float16 -INT8_PER_ROW_SCALE_DTYPE = torch.float16 -INT8_CLIP_PERCENTILE = 99.99984 -INT8_CLIP_Q = INT8_CLIP_PERCENTILE / 100.0 - -def tensor_nbytes(t: Tensor) -> int: - return int(t.numel()) * int(t.element_size()) - -def keep_float_tensor(name: str, t: Tensor, passthrough_orig_dtypes: dict[str, str]) -> Tensor: - if any(pattern in name for pattern in INT8_KEEP_FLOAT_FP32_NAME_PATTERNS): - return t.float().contiguous() - if t.dtype in {torch.float32, torch.bfloat16}: - passthrough_orig_dtypes[name] = str(t.dtype).removeprefix("torch.") - return t.to(dtype=INT8_KEEP_FLOAT_STORE_DTYPE).contiguous() - return t - -def quantize_float_tensor(t: Tensor) -> tuple[Tensor, Tensor]: - t32 = t.float() - if t32.ndim == 2: - # Matrices get one scale per row, which usually tracks output-channel - # ranges much better than a single tensor-wide scale. - clip_abs = ( - torch.quantile(t32.abs(), INT8_CLIP_Q, dim=1) - if t32.numel() - else torch.empty((t32.shape[0],), dtype=torch.float32) - ) - clipped = torch.maximum(torch.minimum(t32, clip_abs[:, None]), -clip_abs[:, None]) - scale = (clip_abs / 127.0).clamp_min(1.0 / 127.0) - q = torch.clamp(torch.round(clipped / scale[:, None]), -127, 127).to(torch.int8).contiguous() - return q, scale.to(dtype=INT8_PER_ROW_SCALE_DTYPE).contiguous() - - # Vectors / scalars use a simpler per-tensor scale. - clip_abs = float(torch.quantile(t32.abs().flatten(), INT8_CLIP_Q).item()) if t32.numel() else 0.0 - scale = torch.tensor(clip_abs / 127.0 if clip_abs > 0 else 1.0, dtype=torch.float32) - q = torch.clamp(torch.round(torch.clamp(t32, -clip_abs, clip_abs) / scale), -127, 127).to(torch.int8).contiguous() - return q, scale - -def quantize_state_dict_int8(state_dict: dict[str, Tensor]): - # Single supported clean-script export format: - # - per-row int8 for 2D float tensors - # - per-tensor int8 for other float tensors - # - exact passthrough for non-floats - # - passthrough for small float tensors, stored as fp16 to save bytes - quantized: dict[str, Tensor] = {} - scales: dict[str, Tensor] = {} - dtypes: dict[str, str] = {} - passthrough: dict[str, Tensor] = {} - passthrough_orig_dtypes: dict[str, str] = {} - qmeta: dict[str, dict[str, object]] = {} - stats = dict.fromkeys( - ("param_count", "num_tensors", "num_float_tensors", "num_nonfloat_tensors", "baseline_tensor_bytes", "int8_payload_bytes"), - 0, - ) - - for name, tensor in state_dict.items(): - t = tensor.detach().to("cpu").contiguous() - stats["param_count"] += int(t.numel()) - stats["num_tensors"] += 1 - stats["baseline_tensor_bytes"] += tensor_nbytes(t) - - if not t.is_floating_point(): - stats["num_nonfloat_tensors"] += 1 - passthrough[name] = t - stats["int8_payload_bytes"] += tensor_nbytes(t) - continue - - # Small float tensors are cheap enough to keep directly. We still downcast - # fp32/bf16 passthrough tensors to fp16 so metadata does not dominate size. - if t.numel() <= INT8_KEEP_FLOAT_MAX_NUMEL: - kept = keep_float_tensor(name, t, passthrough_orig_dtypes) - passthrough[name] = kept - stats["int8_payload_bytes"] += tensor_nbytes(kept) - continue - - stats["num_float_tensors"] += 1 - q, s = quantize_float_tensor(t) - if s.ndim > 0: - qmeta[name] = {"scheme": "per_row", "axis": 0} - quantized[name] = q - scales[name] = s - dtypes[name] = str(t.dtype).removeprefix("torch.") - stats["int8_payload_bytes"] += tensor_nbytes(q) + tensor_nbytes(s) - - obj: dict[str, object] = { - "__quant_format__": "int8_clean_per_row_v1", - "quantized": quantized, - "scales": scales, - "dtypes": dtypes, - "passthrough": passthrough, - } - if qmeta: - obj["qmeta"] = qmeta - if passthrough_orig_dtypes: - obj["passthrough_orig_dtypes"] = passthrough_orig_dtypes - return obj, stats - -def dequantize_state_dict_int8(obj: dict[str, object]) -> dict[str, Tensor]: - out: dict[str, Tensor] = {} - qmeta = obj.get("qmeta", {}) - passthrough_orig_dtypes = obj.get("passthrough_orig_dtypes", {}) - for name, q in obj["quantized"].items(): - dtype = getattr(torch, obj["dtypes"][name]) - s = obj["scales"][name] - if qmeta.get(name, {}).get("scheme") == "per_row" or s.ndim > 0: - s = s.to(dtype=torch.float32) - # Broadcast the saved row scale back across trailing dimensions. - out[name] = (q.float() * s.view(q.shape[0], *([1] * (q.ndim - 1)))).to(dtype=dtype).contiguous() - else: - scale = float(s.item()) - out[name] = (q.float() * scale).to(dtype=dtype).contiguous() - for name, t in obj["passthrough"].items(): - # Restore small tensors, undoing the temporary fp16 storage cast if needed. - out_t = t.detach().to("cpu").contiguous() - orig_dtype = passthrough_orig_dtypes.get(name) - if isinstance(orig_dtype, str): - out_t = out_t.to(dtype=getattr(torch, orig_dtype)).contiguous() - out[name] = out_t - return out - - -# ----------------------------- -# DATA LOADING -# ----------------------------- - -def load_data_shard(file: Path) -> Tensor: - header_bytes = 256 * np.dtype(" None: - self.file_idx = (self.file_idx + 1) % len(self.files) - self.tokens = load_data_shard(self.files[self.file_idx]) - self.pos = 0 - - def take(self, n: int) -> Tensor: - chunks: list[Tensor] = [] - remaining = n - while remaining > 0: - avail = self.tokens.numel() - self.pos - if avail <= 0: - self._advance_file() - continue - k = min(remaining, avail) - chunks.append(self.tokens[self.pos : self.pos + k]) - self.pos += k - remaining -= k - return chunks[0] if len(chunks) == 1 else torch.cat(chunks) - - -class DistributedTokenLoader: - # Each call consumes a contiguous chunk from the shared token stream, then slices out - # one disjoint span per rank. The extra "+1" token lets us build (x, y) by shifting. - def __init__(self, pattern: str, rank: int, world_size: int, device: torch.device): - self.rank = rank - self.world_size = world_size - self.device = device - self.stream = TokenStream(pattern) - - def next_batch(self, global_tokens: int, seq_len: int, grad_accum_steps: int) -> tuple[Tensor, Tensor]: - local_tokens = global_tokens // (self.world_size * grad_accum_steps) - per_rank_span = local_tokens + 1 - chunk = self.stream.take(per_rank_span * self.world_size) - start = self.rank * per_rank_span - local = chunk[start : start + per_rank_span].to(dtype=torch.int64) - x = local[:-1].reshape(-1, seq_len) - y = local[1:].reshape(-1, seq_len) - return x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True) - -# ----------------------------- -# TRANSFORMER MODULES -# ----------------------------- - -class RMSNorm(nn.Module): - def __init__(self, eps: float | None = None): - super().__init__() - self.eps = eps - - def forward(self, x: Tensor) -> Tensor: - return F.rms_norm(x, (x.size(-1),), eps=self.eps) - - -class CastedLinear(nn.Linear): - # Keep weights in fp32 for optimizer/state quality, cast at matmul time for bf16 compute. - def forward(self, x: Tensor) -> Tensor: - bias = self.bias.to(x.dtype) if self.bias is not None else None - return F.linear(x, self.weight.to(x.dtype), bias) - - -def restore_low_dim_params_to_fp32(module: nn.Module) -> None: - # Keep small/control parameters in fp32 even when the model body runs in bf16. - with torch.no_grad(): - for name, param in module.named_parameters(): - if (param.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS)) and param.dtype != torch.float32: - param.data = param.data.float() - - -class Rotary(nn.Module): - # Caches cos/sin tables per sequence length on the current device. - def __init__(self, dim: int, base: float = 10000.0): - super().__init__() - inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self._seq_len_cached = 0 - self._cos_cached: Tensor | None = None - self._sin_cached: Tensor | None = None - - def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> tuple[Tensor, Tensor]: - if ( - self._cos_cached is None - or self._sin_cached is None - or self._seq_len_cached != seq_len - or self._cos_cached.device != device - ): - t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) - freqs = torch.outer(t, self.inv_freq.to(device)) - self._cos_cached = freqs.cos()[None, None, :, :] - self._sin_cached = freqs.sin()[None, None, :, :] - self._seq_len_cached = seq_len - return self._cos_cached.to(dtype=dtype), self._sin_cached.to(dtype=dtype) - - -def apply_rotary_emb(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: - half = x.size(-1) // 2 - x1, x2 = x[..., :half], x[..., half:] - return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1) - - -class CausalSelfAttention(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - num_kv_heads: int, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - if dim % num_heads != 0: - raise ValueError("model_dim must be divisible by num_heads") - if num_heads % num_kv_heads != 0: - raise ValueError("num_heads must be divisible by num_kv_heads") - self.num_heads = num_heads - self.num_kv_heads = num_kv_heads - self.head_dim = dim // num_heads - if self.head_dim % 2 != 0: - raise ValueError("head_dim must be even for RoPE") - kv_dim = self.num_kv_heads * self.head_dim - self.c_q = CastedLinear(dim, dim, bias=False) - self.c_k = CastedLinear(dim, kv_dim, bias=False) - self.c_v = CastedLinear(dim, kv_dim, bias=False) - self.proj = CastedLinear(dim, dim, bias=False) - self.proj._zero_init = True - self.q_gain = nn.Parameter(torch.full((num_heads,), qk_gain_init, dtype=torch.float32)) - self.rotary = Rotary(self.head_dim, base=rope_base) - - def forward(self, x: Tensor) -> Tensor: - bsz, seqlen, dim = x.shape - q = self.c_q(x).reshape(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2) - k = self.c_k(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - v = self.c_v(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - q = F.rms_norm(q, (q.size(-1),)) - k = F.rms_norm(k, (k.size(-1),)) - cos, sin = self.rotary(seqlen, x.device, q.dtype) - q = apply_rotary_emb(q, cos, sin) - k = apply_rotary_emb(k, cos, sin) - q = q * self.q_gain.to(dtype=q.dtype)[None, :, None, None] - y = F.scaled_dot_product_attention( - q, - k, - v, - attn_mask=None, - is_causal=True, - enable_gqa=(self.num_kv_heads != self.num_heads), - ) - y = y.transpose(1, 2).contiguous().reshape(bsz, seqlen, dim) - return self.proj(y) - - -class MLP(nn.Module): - # relu^2 MLP from the original modded-nanogpt setup - def __init__(self, dim: int, mlp_mult: int): - super().__init__() - hidden = mlp_mult * dim - self.fc = CastedLinear(dim, hidden, bias=False) - self.proj = CastedLinear(hidden, dim, bias=False) - self.proj._zero_init = True - - def forward(self, x: Tensor) -> Tensor: - x = torch.relu(self.fc(x)) - return self.proj(x.square()) - - -class Block(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - num_kv_heads: int, - mlp_mult: int, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - self.attn_norm = RMSNorm() - self.mlp_norm = RMSNorm() - self.attn = CausalSelfAttention(dim, num_heads, num_kv_heads, rope_base, qk_gain_init) - self.mlp = MLP(dim, mlp_mult) - self.attn_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) - self.mlp_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) - self.resid_mix = nn.Parameter(torch.stack((torch.ones(dim), torch.zeros(dim))).float()) - - def forward(self, x: Tensor, x0: Tensor) -> Tensor: - mix = self.resid_mix.to(dtype=x.dtype) - x = mix[0][None, None, :] * x + mix[1][None, None, :] * x0 - attn_out = self.attn(self.attn_norm(x)) - x = x + self.attn_scale.to(dtype=x.dtype)[None, None, :] * attn_out - x = x + self.mlp_scale.to(dtype=x.dtype)[None, None, :] * self.mlp(self.mlp_norm(x)) - return x - - -class GPT(nn.Module): - def __init__( - self, - vocab_size: int, - num_layers: int, - model_dim: int, - num_heads: int, - num_kv_heads: int, - mlp_mult: int, - tie_embeddings: bool, - tied_embed_init_std: float, - logit_softcap: float, - rope_base: float, - qk_gain_init: float, - ): - super().__init__() - if logit_softcap <= 0.0: - raise ValueError(f"logit_softcap must be positive, got {logit_softcap}") - self.tie_embeddings = tie_embeddings - self.tied_embed_init_std = tied_embed_init_std - self.logit_softcap = logit_softcap - self.tok_emb = nn.Embedding(vocab_size, model_dim) - self.num_encoder_layers = num_layers // 2 - self.num_decoder_layers = num_layers - self.num_encoder_layers - self.num_skip_weights = min(self.num_encoder_layers, self.num_decoder_layers) - self.skip_weights = nn.Parameter(torch.ones(self.num_skip_weights, model_dim, dtype=torch.float32)) - self.blocks = nn.ModuleList( - [ - Block( - model_dim, - num_heads, - num_kv_heads, - mlp_mult, - rope_base, - qk_gain_init, - ) - for i in range(num_layers) - ] - ) - self.final_norm = RMSNorm() - self.lm_head = None if tie_embeddings else CastedLinear(model_dim, vocab_size, bias=False) - if self.lm_head is not None: - self.lm_head._zero_init = True - self._init_weights() - - def _init_weights(self) -> None: - if self.tie_embeddings: - nn.init.normal_(self.tok_emb.weight, mean=0.0, std=self.tied_embed_init_std) - for module in self.modules(): - if isinstance(module, nn.Linear) and getattr(module, "_zero_init", False): - nn.init.zeros_(module.weight) - - def forward(self, input_ids: Tensor, target_ids: Tensor) -> Tensor: - x = self.tok_emb(input_ids) - x = F.rms_norm(x, (x.size(-1),)) - x0 = x - skips: list[Tensor] = [] - - # First half stores skips; second half reuses them in reverse order. - for i in range(self.num_encoder_layers): - x = self.blocks[i](x, x0) - skips.append(x) - for i in range(self.num_decoder_layers): - if skips: - x = x + self.skip_weights[i].to(dtype=x.dtype)[None, None, :] * skips.pop() - x = self.blocks[self.num_encoder_layers + i](x, x0) - - x = self.final_norm(x).reshape(-1, x.size(-1)) - targets = target_ids.reshape(-1) - if self.tie_embeddings: - logits_proj = F.linear(x, self.tok_emb.weight) - else: - if self.lm_head is None: - raise RuntimeError("lm_head is required when tie_embeddings=False") - logits_proj = self.lm_head(x) - logits = self.logit_softcap * torch.tanh(logits_proj / self.logit_softcap) - return F.cross_entropy(logits.float(), targets, reduction="mean") - - -# ----------------------------- -# TRAINING -# ----------------------------- - -def main() -> None: - global zeropower_via_newtonschulz5 - - code = Path(__file__).read_text(encoding="utf-8") - args = Hyperparameters() - zeropower_via_newtonschulz5 = torch.compile(zeropower_via_newtonschulz5) - - # ----------------------------- - # DISTRIBUTED + CUDA SETUP - # ----------------------------- - - distributed = "RANK" in os.environ and "WORLD_SIZE" in os.environ - rank = int(os.environ.get("RANK", "0")) - world_size = int(os.environ.get("WORLD_SIZE", "1")) - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - if world_size <= 0: - raise ValueError(f"WORLD_SIZE must be positive, got {world_size}") - if 8 % world_size != 0: - raise ValueError(f"WORLD_SIZE={world_size} must divide 8 so grad_accum_steps stays integral") - grad_accum_steps = 8 // world_size - grad_scale = 1.0 / grad_accum_steps - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required") - device = torch.device("cuda", local_rank) - torch.cuda.set_device(device) - if distributed: - dist.init_process_group(backend="nccl", device_id=device) - dist.barrier() - master_process = rank == 0 - - # Fast math knobs - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True - from torch.backends.cuda import enable_cudnn_sdp, enable_flash_sdp, enable_math_sdp, enable_mem_efficient_sdp - - enable_cudnn_sdp(False) - enable_flash_sdp(True) - enable_mem_efficient_sdp(False) - enable_math_sdp(False) - - logfile = None - if master_process: - os.makedirs("logs", exist_ok=True) - logfile = f"logs/{args.run_id}.txt" - print(logfile) - - def log0(msg: str, console: bool = True) -> None: - if not master_process: - return - if console: - print(msg) - if logfile is not None: - with open(logfile, "a", encoding="utf-8") as f: - print(msg, file=f) - - log0(code, console=False) - log0("=" * 100, console=False) - log0(f"Running Python {sys.version}", console=False) - log0(f"Running PyTorch {torch.__version__}", console=False) - log0( - subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False).stdout, - console=False, - ) - log0("=" * 100, console=False) - - # ----------------------------- - # TOKENIZER + VALIDATION METRIC SETUP - # ----------------------------- - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed_all(args.seed) - - if not args.tokenizer_path.endswith(".model"): - raise ValueError(f"Script only setup for SentencePiece .model file: {args.tokenizer_path}") - sp = spm.SentencePieceProcessor(model_file=args.tokenizer_path) - if int(sp.vocab_size()) != args.vocab_size: - raise ValueError( - f"VOCAB_SIZE={args.vocab_size} does not match tokenizer vocab_size={int(sp.vocab_size())}" - ) - dataset_dir = Path(args.data_path).resolve() - actual_train_files = len(list(dataset_dir.glob("fineweb_train_*.bin"))) - val_tokens = load_validation_tokens(args.val_files, args.train_seq_len) - base_bytes_lut, has_leading_space_lut, is_boundary_token_lut = build_sentencepiece_luts( - sp, args.vocab_size, device - ) - log0(f"val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path={args.tokenizer_path}") - log0(f"train_loader:dataset:{dataset_dir.name} train_shards:{actual_train_files}") - log0(f"val_loader:shards pattern={args.val_files} tokens:{val_tokens.numel() - 1}") - - # ----------------------------- - # MODEL + OPTIMIZER SETUP - # ----------------------------- - - base_model = GPT( - vocab_size=args.vocab_size, - num_layers=args.num_layers, - model_dim=args.model_dim, - num_heads=args.num_heads, - num_kv_heads=args.num_kv_heads, - mlp_mult=args.mlp_mult, - tie_embeddings=args.tie_embeddings, - tied_embed_init_std=args.tied_embed_init_std, - logit_softcap=args.logit_softcap, - rope_base=args.rope_base, - qk_gain_init=args.qk_gain_init, - ).to(device).bfloat16() - for module in base_model.modules(): - if isinstance(module, CastedLinear): - module.float() - restore_low_dim_params_to_fp32(base_model) - compiled_model = torch.compile(base_model, dynamic=False, fullgraph=True) - model: nn.Module = DDP(compiled_model, device_ids=[local_rank], broadcast_buffers=False) if distributed else compiled_model - - # Optimizer split: - # - token embedding (Adam) uses EMBED_LR - # - untied lm_head (Adam) uses HEAD_LR - # - matrix params in transformer blocks use MATRIX_LR via Muon - # - vectors/scalars use SCALAR_LR via Adam - block_named_params = list(base_model.blocks.named_parameters()) - matrix_params = [ - p - for name, p in block_named_params - if p.ndim == 2 and not any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - scalar_params = [ - p - for name, p in block_named_params - if p.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - if base_model.skip_weights.numel() > 0: - scalar_params.append(base_model.skip_weights) - token_lr = args.tied_embed_lr if args.tie_embeddings else args.embed_lr - optimizer_tok = torch.optim.Adam( - [{"params": [base_model.tok_emb.weight], "lr": token_lr, "base_lr": token_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizer_muon = Muon( - matrix_params, - lr=args.matrix_lr, - momentum=args.muon_momentum, - backend_steps=args.muon_backend_steps, - ) - for group in optimizer_muon.param_groups: - group["base_lr"] = args.matrix_lr - optimizer_scalar = torch.optim.Adam( - [{"params": scalar_params, "lr": args.scalar_lr, "base_lr": args.scalar_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers: list[torch.optim.Optimizer] = [optimizer_tok, optimizer_muon, optimizer_scalar] - if base_model.lm_head is not None: - optimizer_head = torch.optim.Adam( - [{"params": [base_model.lm_head.weight], "lr": args.head_lr, "base_lr": args.head_lr}], - betas=(args.beta1, args.beta2), - eps=args.adam_eps, - fused=True, - ) - optimizers.insert(1, optimizer_head) - - n_params = sum(p.numel() for p in base_model.parameters()) - log0(f"model_params:{n_params}") - log0(f"world_size:{world_size} grad_accum_steps:{grad_accum_steps}") - log0("sdp_backends:cudnn=False flash=True mem_efficient=False math=False") - log0(f"attention_mode:gqa num_heads:{args.num_heads} num_kv_heads:{args.num_kv_heads}") - log0( - f"tie_embeddings:{args.tie_embeddings} embed_lr:{token_lr} " - f"head_lr:{args.head_lr if base_model.lm_head is not None else 0.0} " - f"matrix_lr:{args.matrix_lr} scalar_lr:{args.scalar_lr}" - ) - log0( - f"train_batch_tokens:{args.train_batch_tokens} train_seq_len:{args.train_seq_len} " - f"iterations:{args.iterations} warmup_steps:{args.warmup_steps} " - f"max_wallclock_seconds:{args.max_wallclock_seconds:.3f}" - ) - log0(f"seed:{args.seed}") - - # ----------------------------- - # DATA LOADER & MODEL WARMUP - # ----------------------------- - - train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device) - - def zero_grad_all() -> None: - for opt in optimizers: - opt.zero_grad(set_to_none=True) - - max_wallclock_ms = 1000.0 * args.max_wallclock_seconds if args.max_wallclock_seconds > 0 else None - - def lr_mul(step: int, elapsed_ms: float) -> float: - if args.warmdown_iters <= 0: - return 1.0 - if max_wallclock_ms is None: - warmdown_start = max(args.iterations - args.warmdown_iters, 0) - return max((args.iterations - step) / max(args.warmdown_iters, 1), 0.0) if warmdown_start <= step < args.iterations else 1.0 - step_ms = elapsed_ms / max(step, 1) - warmdown_ms = args.warmdown_iters * step_ms - remaining_ms = max(max_wallclock_ms - elapsed_ms, 0.0) - return remaining_ms / max(warmdown_ms, 1e-9) if remaining_ms <= warmdown_ms else 1.0 - - # Warmup primes the compiled forward/backward/optimizer paths, then we restore the - # initial weights/optimizer state so measured training starts from the true init. - if args.warmup_steps > 0: - initial_model_state = {name: tensor.detach().cpu().clone() for name, tensor in base_model.state_dict().items()} - initial_optimizer_states = [copy.deepcopy(opt.state_dict()) for opt in optimizers] - model.train() - for warmup_step in range(args.warmup_steps): - zero_grad_all() - for micro_step in range(grad_accum_steps): - if distributed: - model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 - x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - warmup_loss = model(x, y) - (warmup_loss * grad_scale).backward() - for opt in optimizers: - opt.step() - zero_grad_all() - if args.warmup_steps <= 20 or (warmup_step + 1) % 10 == 0 or warmup_step + 1 == args.warmup_steps: - log0(f"warmup_step:{warmup_step + 1}/{args.warmup_steps}") - base_model.load_state_dict(initial_model_state, strict=True) - for opt, state in zip(optimizers, initial_optimizer_states, strict=True): - opt.load_state_dict(state) - zero_grad_all() - if distributed: - model.require_backward_grad_sync = True - train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device) - - # ----------------------------- - # MAIN TRAINING LOOP - # ----------------------------- - - training_time_ms = 0.0 - stop_after_step: int | None = None - torch.cuda.synchronize() - t0 = time.perf_counter() - - step = 0 - while True: - last_step = step == args.iterations or (stop_after_step is not None and step >= stop_after_step) - - should_validate = last_step or (args.val_loss_every > 0 and step % args.val_loss_every == 0) - if should_validate: - torch.cuda.synchronize() - training_time_ms += 1000.0 * (time.perf_counter() - t0) - val_loss, val_bpb = eval_val( - args, - model, - rank, - world_size, - device, - grad_accum_steps, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - ) - log0( - f"step:{step}/{args.iterations} val_loss:{val_loss:.4f} val_bpb:{val_bpb:.4f} " - f"train_time:{training_time_ms:.0f}ms step_avg:{training_time_ms / max(step, 1):.2f}ms" - ) - torch.cuda.synchronize() - t0 = time.perf_counter() - - if last_step: - if stop_after_step is not None and step < args.iterations: - log0( - f"stopping_early: wallclock_cap train_time:{training_time_ms:.0f}ms " - f"step:{step}/{args.iterations}" - ) - break - - elapsed_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) - scale = lr_mul(step, elapsed_ms) - zero_grad_all() - train_loss = torch.zeros((), device=device) - for micro_step in range(grad_accum_steps): - if distributed: - model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 - x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): - loss = model(x, y) - train_loss += loss.detach() - (loss * grad_scale).backward() - train_loss /= grad_accum_steps - - frac = min(step / args.muon_momentum_warmup_steps, 1.0) if args.muon_momentum_warmup_steps > 0 else 1.0 - muon_momentum = (1 - frac) * args.muon_momentum_warmup_start + frac * args.muon_momentum - for group in optimizer_muon.param_groups: - group["momentum"] = muon_momentum - - for opt in optimizers: - for group in opt.param_groups: - group["lr"] = group["base_lr"] * scale - - if args.grad_clip_norm > 0: - torch.nn.utils.clip_grad_norm_(base_model.parameters(), args.grad_clip_norm) - for opt in optimizers: - opt.step() - zero_grad_all() - - step += 1 - approx_training_time_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) - should_log_train = ( - args.train_log_every > 0 - and (step <= 10 or step % args.train_log_every == 0 or stop_after_step is not None) - ) - if should_log_train: - log0( - f"step:{step}/{args.iterations} train_loss:{train_loss.item():.4f} " - f"train_time:{approx_training_time_ms:.0f}ms step_avg:{approx_training_time_ms / step:.2f}ms" - ) - - # Needed to sync whether we've reached the wallclock cap. - reached_cap = max_wallclock_ms is not None and approx_training_time_ms >= max_wallclock_ms - if distributed and max_wallclock_ms is not None: - reached_cap_tensor = torch.tensor(int(reached_cap), device=device) - dist.all_reduce(reached_cap_tensor, op=dist.ReduceOp.MAX) - reached_cap = bool(reached_cap_tensor.item()) - if stop_after_step is None and reached_cap: - stop_after_step = step - - log0( - f"peak memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB " - f"reserved: {torch.cuda.max_memory_reserved() // 1024 // 1024} MiB" - ) - - # ----------------------------- - # SERIALIZATION + ROUNDTRIP VALIDATION - # ----------------------------- - # Save the raw state (useful for debugging/loading in PyTorch directly), then always produce - # the compressed int8+zlib artifact and validate the round-tripped weights. - - if master_process: - torch.save(base_model.state_dict(), "final_model.pt") - model_bytes = os.path.getsize("final_model.pt") - code_bytes = len(code.encode("utf-8")) - log0(f"Serialized model: {model_bytes} bytes") - log0(f"Code size: {code_bytes} bytes") - log0(f"Total submission size: {model_bytes + code_bytes} bytes") - - quant_obj, quant_stats = quantize_state_dict_int8(base_model.state_dict()) - quant_buf = io.BytesIO() - torch.save(quant_obj, quant_buf) - quant_raw = quant_buf.getvalue() - quant_blob = zlib.compress(quant_raw, level=9) - quant_raw_bytes = len(quant_raw) - if master_process: - with open("final_model.int8.ptz", "wb") as f: - f.write(quant_blob) - quant_file_bytes = os.path.getsize("final_model.int8.ptz") - code_bytes = len(code.encode("utf-8")) - ratio = quant_stats["baseline_tensor_bytes"] / max(quant_stats["int8_payload_bytes"], 1) - log0( - f"Serialized model int8+zlib: {quant_file_bytes} bytes " - f"(payload:{quant_stats['int8_payload_bytes']} raw_torch:{quant_raw_bytes} payload_ratio:{ratio:.2f}x)" - ) - log0(f"Total submission size int8+zlib: {quant_file_bytes + code_bytes} bytes") - - if distributed: - dist.barrier() - with open("final_model.int8.ptz", "rb") as f: - quant_blob_disk = f.read() - quant_state = torch.load(io.BytesIO(zlib.decompress(quant_blob_disk)), map_location="cpu") - base_model.load_state_dict(dequantize_state_dict_int8(quant_state), strict=True) - torch.cuda.synchronize() - t_qeval = time.perf_counter() - q_val_loss, q_val_bpb = eval_val( - args, - model, - rank, - world_size, - device, - grad_accum_steps, - val_tokens, - base_bytes_lut, - has_leading_space_lut, - is_boundary_token_lut, - ) - torch.cuda.synchronize() - log0( - f"final_int8_zlib_roundtrip val_loss:{q_val_loss:.4f} val_bpb:{q_val_bpb:.4f} " - f"eval_time:{1000.0 * (time.perf_counter() - t_qeval):.0f}ms" - ) - log0(f"final_int8_zlib_roundtrip_exact val_loss:{q_val_loss:.8f} val_bpb:{q_val_bpb:.8f}") - - if distributed: - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/train_gpt.py b/train_gpt.py index 651beb2b89..438d0bac5e 100644 --- a/train_gpt.py +++ b/train_gpt.py @@ -4,6 +4,12 @@ Hard stop: To keep readable for newcomers, let's make sure `train_gpt.py` and `train_gpt_mlx.py` never are longer than 1500 lines. """ +### nbefore, install zstandard and causal_conv1d +# pip install zstandard +# pip install --no-cache-dir "https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.6.1.post4/causal_conv1d-1.6.1%2Bcu12torch2.9cxx11abiTRUE-cp312-cp312-linux_x86_64.whl" +# pip install --no-cache-dir "https://download.pytorch.org/whl/cu128/flash_attn_3-3.0.0-cp39-abi3-manylinux_2_28_x86_64.whl" + + from __future__ import annotations import copy @@ -19,6 +25,31 @@ import zlib from pathlib import Path +try: + import zstandard as zstd +except Exception: + zstd = None + print("zstandard not found, please install it with `pip install zstandard`") + + +try: + from causal_conv1d import causal_conv1d_fn +except Exception: + causal_conv1d_fn = None + print("causal_conv1d not found, please install it with `pip install causal_conv1d`") + +try: + from flash_attn.cute import flash_attn_func as flash_attn_4_func +except Exception: + flash_attn_4_func = None + print("flash_attn.cute not found, please install it with `pip install flash-attn-4`") + +try: + from flash_attn_interface import flash_attn_func as flash_attn_3_func +except Exception: + flash_attn_3_func = None + print("flash_attn_interface not found, please install it with `pip install flash-attn`") + import numpy as np import sentencepiece as spm import torch @@ -49,6 +80,10 @@ class Hyperparameters: val_batch_size = int(os.environ.get("VAL_BATCH_SIZE", 524_288)) val_loss_every = int(os.environ.get("VAL_LOSS_EVERY", 1000)) train_log_every = int(os.environ.get("TRAIN_LOG_EVERY", 200)) + timing_log_every = int(os.environ.get("TIMING_LOG_EVERY", 100)) + eval_seq_len = int(os.environ.get("EVAL_SEQ_LEN", 2048)) + eval_stride = int(os.environ.get("EVAL_STRIDE", 64)) + eval_batch_seqs = int(os.environ.get("EVAL_BATCH_SEQS", 32)) # Training length. iterations = int(os.environ.get("ITERATIONS", 20000)) @@ -65,10 +100,30 @@ class Hyperparameters: num_kv_heads = int(os.environ.get("NUM_KV_HEADS", 4)) model_dim = int(os.environ.get("MODEL_DIM", 512)) num_heads = int(os.environ.get("NUM_HEADS", 8)) - mlp_mult = int(os.environ.get("MLP_MULT", 2)) + # MLP_MULT=3.0 means hidden size is 3x model width (e.g. 512 -> 1536). + mlp_mult = float(os.environ.get("MLP_MULT", 3.0)) tie_embeddings = bool(int(os.environ.get("TIE_EMBEDDINGS", "1"))) rope_base = float(os.environ.get("ROPE_BASE", 10000.0)) logit_softcap = float(os.environ.get("LOGIT_SOFTCAP", 30.0)) + bigram_vocab_size = int(os.environ.get("BIGRAM_VOCAB_SIZE", 2048)) + bigram_dim = int(os.environ.get("BIGRAM_DIM", 128)) + canon_set = os.environ.get("CANON_SET", "ABCD") + canon_kernel = int(os.environ.get("CANON_KERNEL", 4)) + canon_first_n = int(os.environ.get("CANON_FIRST_N", 0)) + canon_last_n = int(os.environ.get("CANON_LAST_N", 0)) + canon_layers = tuple( + int(part.strip()) + for part in os.environ.get("CANON_LAYERS", os.environ.get("CANON_LAYER_LIST", "")) + .replace("[", "") + .replace("]", "") + .split(",") + if part.strip() + ) + canon_residual = bool(int(os.environ.get("CANON_RESIDUAL", "1"))) + canon_activation = bool(int(os.environ.get("CANON_ACTIVATION", "0"))) + canon_bias = bool(int(os.environ.get("CANON_BIAS", "0"))) + canon_delta_gate = bool(int(os.environ.get("CANON_DELTA_GATE", "0"))) + canon_delta_gate_init = float(os.environ.get("CANON_DELTA_GATE_INIT", -4.0)) # Optimizer hyperparameters. embed_lr = float(os.environ.get("EMBED_LR", 0.6)) @@ -85,6 +140,34 @@ class Hyperparameters: beta2 = float(os.environ.get("BETA2", 0.95)) adam_eps = float(os.environ.get("ADAM_EPS", 1e-8)) grad_clip_norm = float(os.environ.get("GRAD_CLIP_NORM", 0.0)) + muon_weight_decay = float(os.environ.get("MUON_WEIGHT_DECAY", os.environ.get("MUON_WD", 0.04))) + adam_weight_decay = float(os.environ.get("ADAM_WEIGHT_DECAY", os.environ.get("ADAM_WD", 0.04))) + swa_enabled = bool(int(os.environ.get("SWA_ENABLED", "1"))) + swa_every = int(os.environ.get("SWA_EVERY", 200)) + swa_start_lrmul = float(os.environ.get("SWA_START_LRMUL", 0.5)) + tight_swa = bool(int(os.environ.get("TIGHT_SWA", "0"))) + tight_swa_every = int(os.environ.get("TIGHT_SWA_EVERY", 50)) + tight_swa_start_lrmul = float(os.environ.get("TIGHT_SWA_START_LRMUL", 0.2)) + tight_swa_max_checkpoints = int(os.environ.get("TIGHT_SWA_MAX_CHECKPOINTS", 0)) + xsa_last_n = int(os.environ.get("XSA_LAST_N", 0)) + xsa_learnable_gate = bool(int(os.environ.get("XSA_LEARNABLE_GATE", "0"))) + xsa_gate_init = float(os.environ.get("XSA_GATE_INIT", 2.0)) + ema_enabled = bool(int(os.environ.get("EMA_ENABLED", "0"))) + ema_decay = float(os.environ.get("EMA_DECAY", 0.997)) + rope_dims = int(os.environ.get("ROPE_DIMS", 0)) + ln_scale = bool(int(os.environ.get("LN_SCALE", "0"))) + boundary_delta_enabled = bool(int(os.environ.get("BOUNDARY_DELTA_ENABLED", "0"))) + boundary_delta_first_n = int(os.environ.get("BOUNDARY_DELTA_FIRST_N", 4)) + boundary_delta_gate_vector = bool(int(os.environ.get("BOUNDARY_DELTA_GATE_VECTOR", "0"))) + boundary_delta_gate_init = float(os.environ.get("BOUNDARY_DELTA_GATE_INIT", -4.0)) + qat_enabled = bool(int(os.environ.get("QAT_ENABLED", "0"))) + late_qat = bool(int(os.environ.get("LATE_QAT", "0"))) + qat_threshold = float(os.environ.get("QAT_THRESHOLD", 0.1)) + int6_categories = tuple( + part.strip().lower() + for part in os.environ.get("INT6_CATEGORIES", "mlp,attn").split(",") + if part.strip() + ) # ----------------------------- # MUON OPTIMIZER @@ -110,10 +193,24 @@ def zeropower_via_newtonschulz5(G: Tensor, steps: int = 10, eps: float = 1e-7) - class Muon(torch.optim.Optimizer): - def __init__(self, params, lr: float, momentum: float, backend_steps: int, nesterov: bool = True): + def __init__( + self, + params, + lr: float, + momentum: float, + backend_steps: int, + nesterov: bool = True, + weight_decay: float = 0.0, + ): super().__init__( params, - dict(lr=lr, momentum=momentum, backend_steps=backend_steps, nesterov=nesterov), + dict( + lr=lr, + momentum=momentum, + backend_steps=backend_steps, + nesterov=nesterov, + weight_decay=weight_decay, + ), ) @torch.no_grad() @@ -159,8 +256,11 @@ def step(self, closure=None): if distributed: dist.all_reduce(updates_flat, op=dist.ReduceOp.SUM) + wd = group.get("weight_decay", 0.0) curr = 0 for p in params: + if wd > 0.0: + p.data.mul_(1.0 - lr * wd) g = updates_flat[curr : curr + p.numel()].view_as(p).to(dtype=p.dtype) p.add_(g, alpha=-lr) curr += p.numel() @@ -227,19 +327,21 @@ def eval_val( base_bytes_lut: Tensor, has_leading_space_lut: Tensor, is_boundary_token_lut: Tensor, + eval_seq_len: int | None = None, ) -> tuple[float, float]: # Validation computes two metrics: # - val_loss: token cross-entropy (natural log) # - val_bpb: tokenizer-agnostic compression metric used by the challenge + seq_len = eval_seq_len or args.train_seq_len local_batch_tokens = args.val_batch_size // (world_size * grad_accum_steps) - if local_batch_tokens < args.train_seq_len: + if local_batch_tokens < seq_len: raise ValueError( "VAL_BATCH_SIZE must provide at least one sequence per rank; " f"got VAL_BATCH_SIZE={args.val_batch_size}, WORLD_SIZE={world_size}, " - f"GRAD_ACCUM_STEPS={grad_accum_steps}, TRAIN_SEQ_LEN={args.train_seq_len}" + f"GRAD_ACCUM_STEPS={grad_accum_steps}, seq_len={seq_len}" ) - local_batch_seqs = local_batch_tokens // args.train_seq_len - total_seqs = (val_tokens.numel() - 1) // args.train_seq_len + local_batch_seqs = local_batch_tokens // seq_len + total_seqs = (val_tokens.numel() - 1) // seq_len seq_start = (total_seqs * rank) // world_size seq_end = (total_seqs * (rank + 1)) // world_size val_loss_sum = torch.zeros((), device=device, dtype=torch.float64) @@ -250,11 +352,11 @@ def eval_val( with torch.inference_mode(): for batch_seq_start in range(seq_start, seq_end, local_batch_seqs): batch_seq_end = min(batch_seq_start + local_batch_seqs, seq_end) - raw_start = batch_seq_start * args.train_seq_len - raw_end = batch_seq_end * args.train_seq_len + 1 + raw_start = batch_seq_start * seq_len + raw_end = batch_seq_end * seq_len + 1 local = val_tokens[raw_start:raw_end].to(device=device, dtype=torch.int64, non_blocking=True) - x = local[:-1].reshape(-1, args.train_seq_len) - y = local[1:].reshape(-1, args.train_seq_len) + x = local[:-1].reshape(-1, seq_len) + y = local[1:].reshape(-1, seq_len) with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): batch_loss = model(x, y).detach() batch_token_count = float(y.numel()) @@ -277,6 +379,83 @@ def eval_val( model.train() return float(val_loss.item()), float(bits_per_token * tokens_per_byte) + +def eval_val_sliding( + args: Hyperparameters, + base_model: nn.Module, + rank: int, + world_size: int, + device: torch.device, + val_tokens: Tensor, + base_bytes_lut: Tensor, + has_leading_space_lut: Tensor, + is_boundary_token_lut: Tensor, + stride: int, + eval_seq_len: int, +) -> tuple[float, float]: + # Sliding eval: score only the tail slice of each window so most tokens see + # long context. This matches current leaderboard-style setup. + if stride <= 0 or stride >= eval_seq_len: + raise ValueError(f"stride must satisfy 0 < stride < eval_seq_len, got {stride} and {eval_seq_len}") + total_tokens = val_tokens.numel() - 1 + window_starts = [start for start in range(0, total_tokens, stride) if min(start + eval_seq_len, total_tokens) - start >= 1] + total_windows = len(window_starts) + my_start = (total_windows * rank) // world_size + my_end = (total_windows * (rank + 1)) // world_size + my_windows = window_starts[my_start:my_end] + + loss_sum = torch.zeros((), device=device, dtype=torch.float64) + token_count = torch.zeros((), device=device, dtype=torch.float64) + byte_count = torch.zeros((), device=device, dtype=torch.float64) + batch_seqs = max(args.eval_batch_seqs, 1) + + base_model.eval() + with torch.inference_mode(): + for batch_start in range(0, len(my_windows), batch_seqs): + batch_ws = my_windows[batch_start : batch_start + batch_seqs] + bsz = len(batch_ws) + x_batch = torch.zeros((bsz, eval_seq_len), dtype=torch.int64, device=device) + y_batch = torch.zeros((bsz, eval_seq_len), dtype=torch.int64, device=device) + window_lengths: list[int] = [] + for i, ws in enumerate(batch_ws): + end = min(ws + eval_seq_len, total_tokens) + wlen = end - ws + window_lengths.append(wlen) + chunk = val_tokens[ws : end + 1].to(device=device, dtype=torch.int64, non_blocking=True) + x_batch[i, :wlen] = chunk[:-1] + y_batch[i, :wlen] = chunk[1:] + + with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): + logits = base_model.forward_logits(x_batch) + nll = F.cross_entropy( + logits.reshape(-1, logits.size(-1)).float(), + y_batch.reshape(-1), + reduction="none", + ).reshape(bsz, eval_seq_len) + + for i, ws in enumerate(batch_ws): + wlen = window_lengths[i] + score_start = 0 if ws == 0 else max(wlen - stride, 0) + scored_nll = nll[i, score_start:wlen].to(torch.float64) + loss_sum += scored_nll.sum() + token_count += float(wlen - score_start) + tgt = y_batch[i, score_start:wlen] + prev = x_batch[i, score_start:wlen] + token_bytes = base_bytes_lut[tgt].to(torch.float64) + token_bytes += (has_leading_space_lut[tgt] & ~is_boundary_token_lut[prev]).to(torch.float64) + byte_count += token_bytes.sum() + + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM) + dist.all_reduce(token_count, op=dist.ReduceOp.SUM) + dist.all_reduce(byte_count, op=dist.ReduceOp.SUM) + + val_loss = (loss_sum / token_count).item() + bits_per_token = val_loss / math.log(2.0) + tokens_per_byte = token_count.item() / byte_count.item() + base_model.train() + return float(val_loss), float(bits_per_token * tokens_per_byte) + # ----------------------------- # POST-TRAINING QUANTIZATION # ----------------------------- @@ -289,7 +468,7 @@ def eval_val( pattern for pattern in os.environ.get( "CONTROL_TENSOR_NAME_PATTERNS", - "attn_scale,attn_scales,mlp_scale,mlp_scales,resid_mix,resid_mixes,q_gain,skip_weight,skip_weights", + "attn_scale,attn_scales,mlp_scale,mlp_scales,resid_mix,resid_mixes,q_gain,xsa_gate,boundary_delta_gate,skip_weight,skip_weights,smear,bigram.scale", ).split(",") if pattern ) @@ -422,6 +601,121 @@ def dequantize_state_dict_int8(obj: dict[str, object]) -> dict[str, Tensor]: return out +def _classify_quant_category(name: str) -> str: + if "tok_emb" in name or "lm_head" in name: + return "embed" + if ".mlp." in name: + return "mlp" + if ".attn." in name: + return "attn" + return "other" + + +def quantize_float_tensor_bits(t: Tensor, bits: int) -> tuple[Tensor, Tensor]: + maxq = (1 << (bits - 1)) - 1 + minq = -(1 << (bits - 1)) + t32 = t.float() + if t32.ndim == 2: + clip_abs = ( + torch.quantile(t32.abs(), INT8_CLIP_Q, dim=1) + if t32.numel() + else torch.empty((t32.shape[0],), dtype=torch.float32) + ) + clipped = torch.maximum(torch.minimum(t32, clip_abs[:, None]), -clip_abs[:, None]) + scale = (clip_abs / float(maxq)).clamp_min(1.0 / float(maxq)) + q = torch.clamp(torch.round(clipped / scale[:, None]), minq, maxq).to(torch.int8).contiguous() + return q, scale.to(dtype=INT8_PER_ROW_SCALE_DTYPE).contiguous() + clip_abs = float(torch.quantile(t32.abs().flatten(), INT8_CLIP_Q).item()) if t32.numel() else 0.0 + scale = torch.tensor(clip_abs / float(maxq) if clip_abs > 0 else 1.0, dtype=torch.float32) + q = torch.clamp(torch.round(torch.clamp(t32, -clip_abs, clip_abs) / scale), minq, maxq).to(torch.int8).contiguous() + return q, scale + + +def quantize_state_dict_mixed_int6(state_dict: dict[str, Tensor], int6_categories: tuple[str, ...]): + int6_set = set(int6_categories) + quantized: dict[str, Tensor] = {} + scales: dict[str, Tensor] = {} + dtypes: dict[str, str] = {} + passthrough: dict[str, Tensor] = {} + passthrough_orig_dtypes: dict[str, str] = {} + qmeta: dict[str, dict[str, object]] = {} + stats = dict.fromkeys( + ( + "param_count", + "num_tensors", + "num_float_tensors", + "num_nonfloat_tensors", + "baseline_tensor_bytes", + "quant_payload_bytes", + ), + 0, + ) + + for name, tensor in state_dict.items(): + t = tensor.detach().to("cpu").contiguous() + stats["param_count"] += int(t.numel()) + stats["num_tensors"] += 1 + stats["baseline_tensor_bytes"] += tensor_nbytes(t) + + if not t.is_floating_point(): + stats["num_nonfloat_tensors"] += 1 + passthrough[name] = t + stats["quant_payload_bytes"] += tensor_nbytes(t) + continue + + if t.numel() <= INT8_KEEP_FLOAT_MAX_NUMEL: + kept = keep_float_tensor(name, t, passthrough_orig_dtypes) + passthrough[name] = kept + stats["quant_payload_bytes"] += tensor_nbytes(kept) + continue + + stats["num_float_tensors"] += 1 + category = _classify_quant_category(name) + bits = 6 if category in int6_set else 8 + q, s = quantize_float_tensor_bits(t, bits=bits) + if s.ndim > 0: + qmeta[name] = {"scheme": "per_row", "axis": 0, "bits": bits} + else: + qmeta[name] = {"scheme": "per_tensor", "bits": bits} + quantized[name] = q + scales[name] = s + dtypes[name] = str(t.dtype).removeprefix("torch.") + stats["quant_payload_bytes"] += tensor_nbytes(q) + tensor_nbytes(s) + + obj: dict[str, object] = { + "__quant_format__": "mixed_int6_int8_per_row_v1", + "quantized": quantized, + "scales": scales, + "dtypes": dtypes, + "passthrough": passthrough, + "qmeta": qmeta, + } + if passthrough_orig_dtypes: + obj["passthrough_orig_dtypes"] = passthrough_orig_dtypes + return obj, stats + + +def dequantize_state_dict_mixed_int6(obj: dict[str, object]) -> dict[str, Tensor]: + out: dict[str, Tensor] = {} + qmeta = obj.get("qmeta", {}) + passthrough_orig_dtypes = obj.get("passthrough_orig_dtypes", {}) + for name, q in obj["quantized"].items(): + dtype = getattr(torch, obj["dtypes"][name]) + s = obj["scales"][name] + if qmeta.get(name, {}).get("scheme") == "per_row" or s.ndim > 0: + s = s.to(dtype=torch.float32) + out[name] = (q.float() * s.view(q.shape[0], *([1] * (q.ndim - 1)))).to(dtype=dtype).contiguous() + else: + out[name] = (q.float() * float(s.item())).to(dtype=dtype).contiguous() + for name, t in obj["passthrough"].items(): + out_t = t.detach().to("cpu").contiguous() + orig_dtype = passthrough_orig_dtypes.get(name) + if isinstance(orig_dtype, str): + out_t = out_t.to(dtype=getattr(torch, orig_dtype)).contiguous() + out[name] = out_t + return out + + # ----------------------------- # DATA LOADING # ----------------------------- @@ -508,9 +802,20 @@ def forward(self, x: Tensor) -> Tensor: class CastedLinear(nn.Linear): # Keep weights in fp32 for optimizer/state quality, cast at matmul time for bf16 compute. + _qat_enabled: bool = False + def forward(self, x: Tensor) -> Tensor: + w = self.weight.to(x.dtype) + if CastedLinear._qat_enabled and self.training and self.weight.ndim == 2: + with torch.no_grad(): + w32 = self.weight.float() + row_max = w32.abs().amax(dim=1) + scale = (row_max / 31.0).clamp_min(1.0 / 31.0) + w_q = (torch.clamp(torch.round(w32 / scale[:, None]), -32, 31) * scale[:, None]).to(x.dtype) + # STE: forward on fake-quantized weights, gradients through full-precision weights. + w = w + (w_q - w).detach() bias = self.bias.to(x.dtype) if self.bias is not None else None - return F.linear(x, self.weight.to(x.dtype), bias) + return F.linear(x, w, bias) def restore_low_dim_params_to_fp32(module: nn.Module) -> None: @@ -523,9 +828,16 @@ def restore_low_dim_params_to_fp32(module: nn.Module) -> None: class Rotary(nn.Module): # Caches cos/sin tables per sequence length on the current device. - def __init__(self, dim: int, base: float = 10000.0): + # Supports optional partial RoPE and NTK-aware base scaling above train_seq_len. + def __init__(self, dim: int, base: float = 10000.0, train_seq_len: int = 1024, rope_dims: int = 0): super().__init__() - inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + self.dim = dim + self.base = base + self.train_seq_len = train_seq_len + self.rope_dims = rope_dims if rope_dims > 0 else dim + if self.rope_dims > dim or self.rope_dims <= 0 or self.rope_dims % 2 != 0: + raise ValueError(f"rope_dims must be even and in [2, {dim}], got {self.rope_dims}") + inv_freq = 1.0 / (base ** (torch.arange(0, self.rope_dims, 2, dtype=torch.float32) / self.rope_dims)) self.register_buffer("inv_freq", inv_freq, persistent=False) self._seq_len_cached = 0 self._cos_cached: Tensor | None = None @@ -538,8 +850,15 @@ def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> tup or self._seq_len_cached != seq_len or self._cos_cached.device != device ): - t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) - freqs = torch.outer(t, self.inv_freq.to(device)) + rd = self.rope_dims + if seq_len > self.train_seq_len: + scale = seq_len / self.train_seq_len + new_base = self.base * (scale ** (rd / (rd - 2))) + inv_freq = 1.0 / (new_base ** (torch.arange(0, rd, 2, dtype=torch.float32, device=device) / rd)) + else: + inv_freq = self.inv_freq.to(device) + t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype) + freqs = torch.outer(t, inv_freq) self._cos_cached = freqs.cos()[None, None, :, :] self._sin_cached = freqs.sin()[None, None, :, :] self._seq_len_cached = seq_len @@ -547,11 +866,73 @@ def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> tup def apply_rotary_emb(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: + rd = cos.size(-1) * 2 + if rd < x.size(-1): + x_rope = x[..., :rd] + x_pass = x[..., rd:] + half = rd // 2 + x1, x2 = x_rope[..., :half], x_rope[..., half:] + x_rot = torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1) + return torch.cat((x_rot, x_pass), dim=-1) half = x.size(-1) // 2 x1, x2 = x[..., :half], x[..., half:] return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1) +class CanonLayer(nn.Module): + # Lightweight Canon layer: depthwise causal conv over sequence. + # This mirrors PhysicsLM4's A/B/C/D placement semantics without extra dependencies. + def __init__( + self, + dim: int, + kernel: int = 4, + bias: bool = False, + activation: bool = False, + residual: bool = True, + delta_gate: bool = False, + delta_gate_init: float = -4.0, + ): + super().__init__() + if kernel <= 0: + raise ValueError(f"canon_kernel must be positive, got {kernel}") + self.kernel = kernel + self.residual = residual + self.activation = activation + self.delta_gate = delta_gate + if self.delta_gate: + self.delta_gate_logit = nn.Parameter(torch.tensor(float(delta_gate_init), dtype=torch.float32)) + else: + self.register_parameter("delta_gate_logit", None) + # Use built-in Conv1d padding and crop to keep strict causality while avoiding + # an explicit F.pad allocation every forward. + self.conv = nn.Conv1d(dim, dim, kernel_size=kernel, groups=dim, bias=bias, padding=kernel - 1) + self.use_fast_conv1d = ( + bool(int(os.environ.get("CANON_FAST_CONV1D", "1"))) + and causal_conv1d_fn is not None + and kernel in (2, 3, 4) + ) + + def forward(self, x: Tensor) -> Tensor: + if self.use_fast_conv1d: + # causal_conv1d's custom op can fail under torch.compile(fullgraph=True) + # when fed non-contiguous channel-first views. Make it explicitly contiguous. + y = causal_conv1d_fn( + x=x.transpose(1, 2).contiguous(), + weight=self.conv.weight.squeeze(1), + bias=self.conv.bias, + activation="silu" if self.activation else None, + ).transpose(1, 2) + else: + y = self.conv(x.transpose(1, 2)) + y = y[..., : x.size(1)].transpose(1, 2) + if self.activation: + y = F.silu(y) + if self.delta_gate: + gate = torch.sigmoid(self.delta_gate_logit.to(dtype=x.dtype)) + return x + gate * y if self.residual else gate * y + return x + y if self.residual else y + + class CausalSelfAttention(nn.Module): def __init__( self, @@ -560,6 +941,17 @@ def __init__( num_kv_heads: int, rope_base: float, qk_gain_init: float, + rope_dims: int, + canon_set: str, + canon_kernel: int, + canon_bias: bool, + canon_activation: bool, + canon_residual: bool, + canon_delta_gate: bool, + canon_delta_gate_init: float, + use_xsa: bool, + xsa_learnable_gate: bool, + xsa_gate_init: float, ): super().__init__() if dim % num_heads != 0: @@ -578,42 +970,168 @@ def __init__( self.proj = CastedLinear(dim, dim, bias=False) self.proj._zero_init = True self.q_gain = nn.Parameter(torch.full((num_heads,), qk_gain_init, dtype=torch.float32)) - self.rotary = Rotary(self.head_dim, base=rope_base) + self.rotary = Rotary(self.head_dim, base=rope_base, train_seq_len=1024, rope_dims=rope_dims) + self.use_xsa = use_xsa + self.xsa_gate = ( + nn.Parameter(torch.tensor(float(xsa_gate_init), dtype=torch.float32)) + if xsa_learnable_gate and use_xsa + else None + ) + self._fa4_failed = False + self._fa3_failed = False + total_qkv = dim + kv_dim + kv_dim + self.canon_b = ( + CanonLayer( + total_qkv, + kernel=canon_kernel, + bias=canon_bias, + activation=canon_activation, + residual=canon_residual, + delta_gate=canon_delta_gate, + delta_gate_init=canon_delta_gate_init, + ) + if "B" in canon_set + else None + ) + + def _xsa_efficient(self, y: Tensor, v: Tensor) -> Tensor: + # y: [B, H, T, D], v: [B, Hkv, T, D] + bsz, heads, seqlen, head_dim = y.shape + kv_heads = v.size(1) + group = heads // kv_heads + y_grouped = y.reshape(bsz, kv_heads, group, seqlen, head_dim) + v_norm = F.normalize(v, dim=-1).unsqueeze(2) + proj = (y_grouped * v_norm).sum(dim=-1, keepdim=True) * v_norm + return (y_grouped - proj).reshape(bsz, heads, seqlen, head_dim) def forward(self, x: Tensor) -> Tensor: bsz, seqlen, dim = x.shape - q = self.c_q(x).reshape(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2) - k = self.c_k(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) - v = self.c_v(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) + q = self.c_q(x) + k = self.c_k(x) + v = self.c_v(x) + if self.canon_b is not None: + qkv = self.canon_b(torch.cat([q, k, v], dim=-1)) + q, k, v = qkv.split((self.num_heads * self.head_dim, self.num_kv_heads * self.head_dim, self.num_kv_heads * self.head_dim), dim=-1) + q = q.reshape(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2) + k = k.reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) + v = v.reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) q = F.rms_norm(q, (q.size(-1),)) k = F.rms_norm(k, (k.size(-1),)) cos, sin = self.rotary(seqlen, x.device, q.dtype) q = apply_rotary_emb(q, cos, sin) k = apply_rotary_emb(k, cos, sin) q = q * self.q_gain.to(dtype=q.dtype)[None, :, None, None] - y = F.scaled_dot_product_attention( - q, - k, - v, - attn_mask=None, - is_causal=True, - enable_gqa=(self.num_kv_heads != self.num_heads), - ) + y: Tensor | None = None + if q.is_cuda and k.is_cuda and v.is_cuda: + q_t = q.transpose(1, 2) + k_t = k.transpose(1, 2) + v_t = v.transpose(1, 2) + if flash_attn_4_func is not None and not self._fa4_failed: + try: + y = flash_attn_4_func(q_t, k_t, v_t, causal=True).transpose(1, 2) + except Exception as e: + self._fa4_failed = True + print(f"flash_attn_4 failed, falling back to flash_attn_interface/SDPA: {e}") + if y is None and flash_attn_3_func is not None and not self._fa3_failed: + try: + y = flash_attn_3_func(q_t, k_t, v_t, causal=True).transpose(1, 2) + except Exception as e: + self._fa3_failed = True + print(f"flash_attn_interface failed, falling back to SDPA: {e}") + if y is None: + y = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=None, + is_causal=True, + enable_gqa=(self.num_kv_heads != self.num_heads), + ) + if self.use_xsa: + y_xsa = self._xsa_efficient(y, v) + if self.xsa_gate is None: + y = y_xsa + else: + alpha = torch.sigmoid(self.xsa_gate.to(dtype=y.dtype)) + y = y + alpha * (y_xsa - y) y = y.transpose(1, 2).contiguous().reshape(bsz, seqlen, dim) return self.proj(y) +class SmearGate(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.gate = nn.Parameter(torch.zeros(dim, dtype=torch.float32)) + + def forward(self, x: Tensor) -> Tensor: + g = torch.sigmoid(self.gate.to(dtype=x.dtype))[None, None, :] + x_prev = torch.cat([torch.zeros_like(x[:, :1]), x[:, :-1]], dim=1) + return (1 - g) * x + g * x_prev + + +class BigramHashEmbedding(nn.Module): + def __init__(self, bigram_vocab_size: int, bigram_dim: int, model_dim: int): + super().__init__() + self.bigram_vocab_size = bigram_vocab_size + self.embed = nn.Embedding(bigram_vocab_size, bigram_dim) + nn.init.zeros_(self.embed.weight) + self.proj = CastedLinear(bigram_dim, model_dim, bias=False) if bigram_dim != model_dim else None + if self.proj is not None: + nn.init.zeros_(self.proj.weight) + self.scale = nn.Parameter(torch.tensor(0.05, dtype=torch.float32)) + + def bigram_hash(self, tokens: Tensor) -> Tensor: + t = tokens.to(torch.int32) + mod = self.bigram_vocab_size - 1 + out = torch.empty_like(t) + out[..., 0] = mod + out[..., 1:] = torch.bitwise_xor(36313 * t[..., 1:], 27191 * t[..., :-1]) % mod + return out.long() + + def forward(self, token_ids: Tensor) -> Tensor: + h = self.embed(self.bigram_hash(token_ids)) + if self.proj is not None: + h = self.proj(h) + return h * self.scale.to(dtype=h.dtype) + + class MLP(nn.Module): # relu^2 MLP from the original modded-nanogpt setup - def __init__(self, dim: int, mlp_mult: int): + def __init__( + self, + dim: int, + mlp_mult: float, + canon_set: str, + canon_kernel: int, + canon_bias: bool, + canon_activation: bool, + canon_residual: bool, + canon_delta_gate: bool, + canon_delta_gate_init: float, + ): super().__init__() - hidden = mlp_mult * dim + hidden = int(mlp_mult * dim) self.fc = CastedLinear(dim, hidden, bias=False) self.proj = CastedLinear(hidden, dim, bias=False) self.proj._zero_init = True + self.canon_d = ( + CanonLayer( + hidden, + kernel=canon_kernel, + bias=canon_bias, + activation=canon_activation, + residual=canon_residual, + delta_gate=canon_delta_gate, + delta_gate_init=canon_delta_gate_init, + ) + if "D" in canon_set + else None + ) def forward(self, x: Tensor) -> Tensor: x = torch.relu(self.fc(x)) + if self.canon_d is not None: + x = self.canon_d(x) return self.proj(x.square()) @@ -623,25 +1141,113 @@ def __init__( dim: int, num_heads: int, num_kv_heads: int, - mlp_mult: int, + mlp_mult: float, rope_base: float, qk_gain_init: float, + rope_dims: int, + layer_idx: int, + ln_scale: bool, + canon_set: str, + canon_kernel: int, + canon_bias: bool, + canon_activation: bool, + canon_residual: bool, + canon_delta_gate: bool, + canon_delta_gate_init: float, + use_xsa: bool, + xsa_learnable_gate: bool, + xsa_gate_init: float, + boundary_delta_enabled: bool, + boundary_delta_gate_vector: bool, + boundary_delta_gate_init: float, ): super().__init__() self.attn_norm = RMSNorm() self.mlp_norm = RMSNorm() - self.attn = CausalSelfAttention(dim, num_heads, num_kv_heads, rope_base, qk_gain_init) - self.mlp = MLP(dim, mlp_mult) + self.attn = CausalSelfAttention( + dim, + num_heads, + num_kv_heads, + rope_base, + qk_gain_init, + rope_dims, + canon_set, + canon_kernel, + canon_bias, + canon_activation, + canon_residual, + canon_delta_gate, + canon_delta_gate_init, + use_xsa, + xsa_learnable_gate, + xsa_gate_init, + ) + self.mlp = MLP( + dim, + mlp_mult, + canon_set, + canon_kernel, + canon_bias, + canon_activation, + canon_residual, + canon_delta_gate, + canon_delta_gate_init, + ) self.attn_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) self.mlp_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32)) self.resid_mix = nn.Parameter(torch.stack((torch.ones(dim), torch.zeros(dim))).float()) + self.ln_scale_factor = 1.0 / math.sqrt(layer_idx + 1) if ln_scale else 1.0 + self.boundary_delta_gate = None + if boundary_delta_enabled: + gate_shape = (dim,) if boundary_delta_gate_vector else (1,) + self.boundary_delta_gate = nn.Parameter( + torch.full(gate_shape, float(boundary_delta_gate_init), dtype=torch.float32) + ) + self.canon_a = ( + CanonLayer( + dim, + kernel=canon_kernel, + bias=canon_bias, + activation=canon_activation, + residual=canon_residual, + delta_gate=canon_delta_gate, + delta_gate_init=canon_delta_gate_init, + ) + if "A" in canon_set + else None + ) + self.canon_c = ( + CanonLayer( + dim, + kernel=canon_kernel, + bias=canon_bias, + activation=canon_activation, + residual=canon_residual, + delta_gate=canon_delta_gate, + delta_gate_init=canon_delta_gate_init, + ) + if "C" in canon_set + else None + ) def forward(self, x: Tensor, x0: Tensor) -> Tensor: mix = self.resid_mix.to(dtype=x.dtype) x = mix[0][None, None, :] * x + mix[1][None, None, :] * x0 - attn_out = self.attn(self.attn_norm(x)) + s = self.ln_scale_factor + attn_in = self.attn_norm(x) * s + if self.boundary_delta_gate is not None: + x_prev = torch.cat([torch.zeros_like(x[:, :1]), x[:, :-1]], dim=1) + dx = x - x_prev + g = torch.sigmoid(self.boundary_delta_gate.to(dtype=x.dtype)) + attn_in = attn_in + dx * g[None, None, :] + if self.canon_a is not None: + attn_in = self.canon_a(attn_in) + attn_out = self.attn(attn_in) x = x + self.attn_scale.to(dtype=x.dtype)[None, None, :] * attn_out - x = x + self.mlp_scale.to(dtype=x.dtype)[None, None, :] * self.mlp(self.mlp_norm(x)) + mlp_in = self.mlp_norm(x) * s + if self.canon_c is not None: + mlp_in = self.canon_c(mlp_in) + x = x + self.mlp_scale.to(dtype=x.dtype)[None, None, :] * self.mlp(mlp_in) return x @@ -653,12 +1259,33 @@ def __init__( model_dim: int, num_heads: int, num_kv_heads: int, - mlp_mult: int, + mlp_mult: float, tie_embeddings: bool, tied_embed_init_std: float, logit_softcap: float, rope_base: float, qk_gain_init: float, + bigram_vocab_size: int, + bigram_dim: int, + xsa_last_n: int, + xsa_learnable_gate: bool, + xsa_gate_init: float, + rope_dims: int, + ln_scale: bool, + canon_set: str, + canon_kernel: int, + canon_first_n: int, + canon_last_n: int, + canon_layers: tuple[int, ...], + canon_bias: bool, + canon_activation: bool, + canon_residual: bool, + canon_delta_gate: bool, + canon_delta_gate_init: float, + boundary_delta_enabled: bool, + boundary_delta_first_n: int, + boundary_delta_gate_vector: bool, + boundary_delta_gate_init: float, ): super().__init__() if logit_softcap <= 0.0: @@ -667,12 +1294,45 @@ def __init__( self.tied_embed_init_std = tied_embed_init_std self.logit_softcap = logit_softcap self.tok_emb = nn.Embedding(vocab_size, model_dim) + self.bigram = BigramHashEmbedding(bigram_vocab_size, bigram_dim, model_dim) if bigram_vocab_size > 0 else None + self.smear = SmearGate(model_dim) self.num_encoder_layers = num_layers // 2 self.num_decoder_layers = num_layers - self.num_encoder_layers self.num_skip_weights = min(self.num_encoder_layers, self.num_decoder_layers) self.skip_weights = nn.Parameter(torch.ones(self.num_skip_weights, model_dim, dtype=torch.float32)) - self.blocks = nn.ModuleList( - [ + if canon_first_n < 0 or canon_last_n < 0: + raise ValueError( + f"canon layer scopes must be non-negative, got CANON_FIRST_N={canon_first_n}, CANON_LAST_N={canon_last_n}" + ) + if canon_first_n > num_layers or canon_last_n > num_layers: + raise ValueError( + f"canon layer scopes exceed num_layers={num_layers}: CANON_FIRST_N={canon_first_n}, CANON_LAST_N={canon_last_n}" + ) + if boundary_delta_first_n < 0 or boundary_delta_first_n > num_layers: + raise ValueError( + f"BOUNDARY_DELTA_FIRST_N must be in [0, {num_layers}], got {boundary_delta_first_n}" + ) + canon_layers_uniq = tuple(dict.fromkeys(canon_layers)) + if any(layer < 1 or layer > num_layers for layer in canon_layers_uniq): + raise ValueError( + f"CANON_LAYERS values must be in [1, {num_layers}], got {canon_layers_uniq}" + ) + explicit_canon_layers = {layer - 1 for layer in canon_layers_uniq} if canon_layers_uniq else None + use_scoped_canon = explicit_canon_layers is not None or canon_first_n > 0 or canon_last_n > 0 + blocks: list[Block] = [] + for i in range(num_layers): + if explicit_canon_layers is not None: + use_layer_canon = i in explicit_canon_layers + else: + use_layer_canon = ( + not use_scoped_canon + or i < canon_first_n + or i >= num_layers - canon_last_n + ) + layer_canon_set = canon_set if use_layer_canon else "" + use_boundary_delta = boundary_delta_enabled and i < boundary_delta_first_n + use_xsa = i >= max(0, num_layers - xsa_last_n) + blocks.append( Block( model_dim, num_heads, @@ -680,10 +1340,25 @@ def __init__( mlp_mult, rope_base, qk_gain_init, + rope_dims, + i, + ln_scale, + layer_canon_set, + canon_kernel, + canon_bias, + canon_activation, + canon_residual, + canon_delta_gate, + canon_delta_gate_init, + use_xsa, + xsa_learnable_gate, + xsa_gate_init, + use_boundary_delta, + boundary_delta_gate_vector, + boundary_delta_gate_init, ) - for i in range(num_layers) - ] - ) + ) + self.blocks = nn.ModuleList(blocks) self.final_norm = RMSNorm() self.lm_head = None if tie_embeddings else CastedLinear(model_dim, vocab_size, bias=False) if self.lm_head is not None: @@ -693,13 +1368,23 @@ def __init__( def _init_weights(self) -> None: if self.tie_embeddings: nn.init.normal_(self.tok_emb.weight, mean=0.0, std=self.tied_embed_init_std) - for module in self.modules(): - if isinstance(module, nn.Linear) and getattr(module, "_zero_init", False): - nn.init.zeros_(module.weight) + num_layers = len(self.blocks) + for name, module in self.named_modules(): + if isinstance(module, nn.Linear): + if getattr(module, "_zero_init", False): + nn.init.zeros_(module.weight) + elif module.weight.ndim == 2 and module.weight.shape[0] >= 64 and module.weight.shape[1] >= 64: + nn.init.orthogonal_(module.weight, gain=1.0) + if ".proj." in name or name.endswith(".proj"): + with torch.no_grad(): + module.weight.mul_(1.0 / math.sqrt(2 * num_layers)) def forward(self, input_ids: Tensor, target_ids: Tensor) -> Tensor: x = self.tok_emb(input_ids) + if self.bigram is not None: + x = x + self.bigram(input_ids) x = F.rms_norm(x, (x.size(-1),)) + x = self.smear(x) x0 = x skips: list[Tensor] = [] @@ -723,6 +1408,30 @@ def forward(self, input_ids: Tensor, target_ids: Tensor) -> Tensor: logits = self.logit_softcap * torch.tanh(logits_proj / self.logit_softcap) return F.cross_entropy(logits.float(), targets, reduction="mean") + def forward_logits(self, input_ids: Tensor) -> Tensor: + x = self.tok_emb(input_ids) + if self.bigram is not None: + x = x + self.bigram(input_ids) + x = F.rms_norm(x, (x.size(-1),)) + x = self.smear(x) + x0 = x + skips: list[Tensor] = [] + for i in range(self.num_encoder_layers): + x = self.blocks[i](x, x0) + skips.append(x) + for i in range(self.num_decoder_layers): + if skips: + x = x + self.skip_weights[i].to(dtype=x.dtype)[None, None, :] * skips.pop() + x = self.blocks[self.num_encoder_layers + i](x, x0) + x = self.final_norm(x) + if self.tie_embeddings: + logits_proj = F.linear(x, self.tok_emb.weight) + else: + if self.lm_head is None: + raise RuntimeError("lm_head is required when tie_embeddings=False") + logits_proj = self.lm_head(x) + return self.logit_softcap * torch.tanh(logits_proj / self.logit_softcap) + # ----------------------------- # TRAINING @@ -733,7 +1442,12 @@ def main() -> None: code = Path(__file__).read_text(encoding="utf-8") args = Hyperparameters() - zeropower_via_newtonschulz5 = torch.compile(zeropower_via_newtonschulz5) + CastedLinear._qat_enabled = args.qat_enabled + compile_newton = bool(int(os.environ.get("COMPILE_NEWTONSCHULZ", "1"))) + compile_model = bool(int(os.environ.get("COMPILE_MODEL", "1"))) + compile_fullgraph = bool(int(os.environ.get("COMPILE_FULLGRAPH", "1"))) + if compile_newton: + zeropower_via_newtonschulz5 = torch.compile(zeropower_via_newtonschulz5) # ----------------------------- # DISTRIBUTED + CUDA SETUP @@ -811,7 +1525,9 @@ def log0(msg: str, console: bool = True) -> None: ) dataset_dir = Path(args.data_path).resolve() actual_train_files = len(list(dataset_dir.glob("fineweb_train_*.bin"))) - val_tokens = load_validation_tokens(args.val_files, args.train_seq_len) + effective_eval_seq_len = args.eval_seq_len if args.eval_seq_len > 0 else args.train_seq_len + val_seq_len = max(args.train_seq_len, effective_eval_seq_len) + val_tokens = load_validation_tokens(args.val_files, val_seq_len) base_bytes_lut, has_leading_space_lut, is_boundary_token_lut = build_sentencepiece_luts( sp, args.vocab_size, device ) @@ -835,12 +1551,33 @@ def log0(msg: str, console: bool = True) -> None: logit_softcap=args.logit_softcap, rope_base=args.rope_base, qk_gain_init=args.qk_gain_init, + bigram_vocab_size=args.bigram_vocab_size, + bigram_dim=args.bigram_dim, + xsa_last_n=args.xsa_last_n, + xsa_learnable_gate=args.xsa_learnable_gate, + xsa_gate_init=args.xsa_gate_init, + rope_dims=args.rope_dims, + ln_scale=args.ln_scale, + canon_set=args.canon_set, + canon_kernel=args.canon_kernel, + canon_first_n=args.canon_first_n, + canon_last_n=args.canon_last_n, + canon_layers=args.canon_layers, + canon_bias=args.canon_bias, + canon_activation=args.canon_activation, + canon_residual=args.canon_residual, + canon_delta_gate=args.canon_delta_gate, + canon_delta_gate_init=args.canon_delta_gate_init, + boundary_delta_enabled=args.boundary_delta_enabled, + boundary_delta_first_n=args.boundary_delta_first_n, + boundary_delta_gate_vector=args.boundary_delta_gate_vector, + boundary_delta_gate_init=args.boundary_delta_gate_init, ).to(device).bfloat16() for module in base_model.modules(): if isinstance(module, CastedLinear): module.float() restore_low_dim_params_to_fp32(base_model) - compiled_model = torch.compile(base_model, dynamic=False, fullgraph=True) + compiled_model = torch.compile(base_model, dynamic=False, fullgraph=compile_fullgraph) if compile_model else base_model model: nn.Module = DDP(compiled_model, device_ids=[local_rank], broadcast_buffers=False) if distributed else compiled_model # Optimizer split: @@ -848,24 +1585,28 @@ def log0(msg: str, console: bool = True) -> None: # - untied lm_head (Adam) uses HEAD_LR # - matrix params in transformer blocks use MATRIX_LR via Muon # - vectors/scalars use SCALAR_LR via Adam - block_named_params = list(base_model.blocks.named_parameters()) - matrix_params = [ - p - for name, p in block_named_params - if p.ndim == 2 and not any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - scalar_params = [ - p - for name, p in block_named_params - if p.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) - ] - if base_model.skip_weights.numel() > 0: - scalar_params.append(base_model.skip_weights) + named_params = dict(base_model.named_parameters()) token_lr = args.tied_embed_lr if args.tie_embeddings else args.embed_lr - optimizer_tok = torch.optim.Adam( - [{"params": [base_model.tok_emb.weight], "lr": token_lr, "base_lr": token_lr}], + token_param_names = ["tok_emb.weight"] + if base_model.bigram is not None: + token_param_names.append("bigram.embed.weight") + head_param_names = ["lm_head.weight"] if base_model.lm_head is not None else [] + matrix_params: list[Tensor] = [] + scalar_params: list[Tensor] = [] + for name, p in named_params.items(): + if name in token_param_names or name in head_param_names: + continue + if p.ndim == 2 and not any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS): + matrix_params.append(p) + else: + scalar_params.append(p) + + tok_groups = [{"params": [named_params[name] for name in token_param_names], "lr": token_lr, "base_lr": token_lr}] + optimizer_tok = torch.optim.AdamW( + tok_groups, betas=(args.beta1, args.beta2), eps=args.adam_eps, + weight_decay=args.adam_weight_decay, fused=True, ) optimizer_muon = Muon( @@ -873,13 +1614,15 @@ def log0(msg: str, console: bool = True) -> None: lr=args.matrix_lr, momentum=args.muon_momentum, backend_steps=args.muon_backend_steps, + weight_decay=args.muon_weight_decay, ) for group in optimizer_muon.param_groups: group["base_lr"] = args.matrix_lr - optimizer_scalar = torch.optim.Adam( + optimizer_scalar = torch.optim.AdamW( [{"params": scalar_params, "lr": args.scalar_lr, "base_lr": args.scalar_lr}], betas=(args.beta1, args.beta2), eps=args.adam_eps, + weight_decay=args.adam_weight_decay, fused=True, ) optimizers: list[torch.optim.Optimizer] = [optimizer_tok, optimizer_muon, optimizer_scalar] @@ -892,16 +1635,51 @@ def log0(msg: str, console: bool = True) -> None: ) optimizers.insert(1, optimizer_head) + active_swa_every = args.tight_swa_every if args.tight_swa else args.swa_every + active_swa_start_lrmul = args.tight_swa_start_lrmul if args.tight_swa else args.swa_start_lrmul + active_swa_max_checkpoints = args.tight_swa_max_checkpoints if args.tight_swa else 0 + if active_swa_every < 0: + raise ValueError(f"swa_every must be >= 0, got {active_swa_every}") + if active_swa_max_checkpoints < 0: + raise ValueError(f"swa_max_checkpoints must be >= 0, got {active_swa_max_checkpoints}") + n_params = sum(p.numel() for p in base_model.parameters()) log0(f"model_params:{n_params}") log0(f"world_size:{world_size} grad_accum_steps:{grad_accum_steps}") log0("sdp_backends:cudnn=False flash=True mem_efficient=False math=False") log0(f"attention_mode:gqa num_heads:{args.num_heads} num_kv_heads:{args.num_kv_heads}") + log0( + f"arch:num_layers:{args.num_layers} model_dim:{args.model_dim} mlp_mult:{args.mlp_mult} " + f"bigram_vocab_size:{args.bigram_vocab_size} bigram_dim:{args.bigram_dim}" + ) + canon_layers_str = ",".join(str(layer) for layer in args.canon_layers) if args.canon_layers else "-" + log0( + f"canon:set:{args.canon_set} kernel:{args.canon_kernel} residual:{args.canon_residual} " + f"activation:{args.canon_activation} bias:{args.canon_bias} " + f"delta_gate:{args.canon_delta_gate} delta_gate_init:{args.canon_delta_gate_init} " + f"first_n:{args.canon_first_n} last_n:{args.canon_last_n} layers:{canon_layers_str} " + f"fast_conv1d:{bool(int(os.environ.get('CANON_FAST_CONV1D', '1'))) and causal_conv1d_fn is not None and args.canon_kernel in (2, 3, 4)}" + ) + log0(f"compile:model:{compile_model} fullgraph:{compile_fullgraph} newton:{compile_newton}") log0( f"tie_embeddings:{args.tie_embeddings} embed_lr:{token_lr} " f"head_lr:{args.head_lr if base_model.lm_head is not None else 0.0} " f"matrix_lr:{args.matrix_lr} scalar_lr:{args.scalar_lr}" ) + log0( + f"weight_decay:muon:{args.muon_weight_decay} adam:{args.adam_weight_decay} " + f"swa_enabled:{args.swa_enabled} swa_every:{active_swa_every} swa_start_lrmul:{active_swa_start_lrmul} " + f"tight_swa:{args.tight_swa} swa_max_checkpoints:{active_swa_max_checkpoints}" + ) + log0( + f"frontier:xsa_last_n:{args.xsa_last_n} xsa_learnable_gate:{args.xsa_learnable_gate} xsa_gate_init:{args.xsa_gate_init} " + f"boundary_delta_enabled:{args.boundary_delta_enabled} boundary_delta_first_n:{args.boundary_delta_first_n} " + f"boundary_delta_gate_vector:{args.boundary_delta_gate_vector} boundary_delta_gate_init:{args.boundary_delta_gate_init} " + f"ema_enabled:{args.ema_enabled} ema_decay:{args.ema_decay} " + f"rope_dims:{args.rope_dims} ln_scale:{args.ln_scale} " + f"qat_enabled:{args.qat_enabled} late_qat:{args.late_qat} qat_threshold:{args.qat_threshold} " + f"fa4_available:{flash_attn_4_func is not None} fa3_available:{flash_attn_3_func is not None}" + ) log0( f"train_batch_tokens:{args.train_batch_tokens} train_seq_len:{args.train_seq_len} " f"iterations:{args.iterations} warmup_steps:{args.warmup_steps} " @@ -964,8 +1742,17 @@ def lr_mul(step: int, elapsed_ms: float) -> float: # MAIN TRAINING LOOP # ----------------------------- + swa_state: dict[str, Tensor] | None = None + swa_count = 0 + ema_state: dict[str, Tensor] | None = None + if args.ema_enabled: + ema_state = {name: t.detach().float().clone() for name, t in base_model.state_dict().items()} training_time_ms = 0.0 + data_loading_time_ms = 0.0 stop_after_step: int | None = None + last_timing_log_step = 0 + last_timing_log_train_time_ms = 0.0 + last_timing_log_data_time_ms = 0.0 torch.cuda.synchronize() t0 = time.perf_counter() @@ -991,7 +1778,8 @@ def lr_mul(step: int, elapsed_ms: float) -> float: ) log0( f"step:{step}/{args.iterations} val_loss:{val_loss:.4f} val_bpb:{val_bpb:.4f} " - f"train_time:{training_time_ms:.0f}ms step_avg:{training_time_ms / max(step, 1):.2f}ms" + f"train_time:{training_time_ms:.0f}ms step_avg:{training_time_ms / max(step, 1):.2f}ms " + f"data_time:{data_loading_time_ms:.0f}ms data_avg:{data_loading_time_ms / max(step, 1):.2f}ms" ) torch.cuda.synchronize() t0 = time.perf_counter() @@ -1006,12 +1794,17 @@ def lr_mul(step: int, elapsed_ms: float) -> float: elapsed_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) scale = lr_mul(step, elapsed_ms) + if args.late_qat and scale < args.qat_threshold and not CastedLinear._qat_enabled: + CastedLinear._qat_enabled = True + log0(f"late_qat:enabled step:{step} scale:{scale:.4f}") zero_grad_all() train_loss = torch.zeros((), device=device) for micro_step in range(grad_accum_steps): if distributed: model.require_backward_grad_sync = micro_step == grad_accum_steps - 1 + load_t0 = time.perf_counter() x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps) + data_loading_time_ms += 1000.0 * (time.perf_counter() - load_t0) with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): loss = model(x, y) train_loss += loss.detach() @@ -1034,6 +1827,27 @@ def lr_mul(step: int, elapsed_ms: float) -> float: zero_grad_all() step += 1 + if ema_state is not None: + decay = args.ema_decay + with torch.no_grad(): + for name, t in base_model.state_dict().items(): + ema_state[name].mul_(decay).add_(t.detach().float(), alpha=1.0 - decay) + if ( + args.swa_enabled + and ema_state is None + and active_swa_every > 0 + and scale < active_swa_start_lrmul + and step % active_swa_every == 0 + and (active_swa_max_checkpoints <= 0 or swa_count < active_swa_max_checkpoints) + ): + if swa_state is None: + swa_state = {name: t.detach().cpu().clone() for name, t in base_model.state_dict().items()} + swa_count = 1 + log0(f"swa:start step:{step}") + else: + for name, t in base_model.state_dict().items(): + swa_state[name].add_(t.detach().cpu()) + swa_count += 1 approx_training_time_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0) should_log_train = ( args.train_log_every > 0 @@ -1042,8 +1856,28 @@ def lr_mul(step: int, elapsed_ms: float) -> float: if should_log_train: log0( f"step:{step}/{args.iterations} train_loss:{train_loss.item():.4f} " - f"train_time:{approx_training_time_ms:.0f}ms step_avg:{approx_training_time_ms / step:.2f}ms" + f"train_time:{approx_training_time_ms:.0f}ms step_avg:{approx_training_time_ms / step:.2f}ms " + f"data_time:{data_loading_time_ms:.0f}ms data_avg:{data_loading_time_ms / step:.2f}ms" + ) + should_log_timing_window = ( + args.timing_log_every > 0 + and step >= 1 + and (step % args.timing_log_every == 0 or stop_after_step is not None) + and step > last_timing_log_step + ) + if should_log_timing_window: + window_steps = step - last_timing_log_step + window_train_ms = approx_training_time_ms - last_timing_log_train_time_ms + window_data_ms = data_loading_time_ms - last_timing_log_data_time_ms + log0( + f"timing_window:steps:{last_timing_log_step + 1}-{step} " + f"step_time:{window_train_ms / window_steps:.2f}ms " + f"data_time:{window_data_ms / window_steps:.2f}ms " + f"data_frac:{(100.0 * window_data_ms / max(window_train_ms, 1e-9)):.2f}%" ) + last_timing_log_step = step + last_timing_log_train_time_ms = approx_training_time_ms + last_timing_log_data_time_ms = data_loading_time_ms # Needed to sync whether we've reached the wallclock cap. reached_cap = max_wallclock_ms is not None and approx_training_time_ms >= max_wallclock_ms @@ -1058,12 +1892,27 @@ def lr_mul(step: int, elapsed_ms: float) -> float: f"peak memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB " f"reserved: {torch.cuda.max_memory_reserved() // 1024 // 1024} MiB" ) + log0( + f"data_loading_total:{data_loading_time_ms:.0f}ms " + f"data_loading_step_avg:{data_loading_time_ms / max(step, 1):.2f}ms" + ) + if ema_state is not None: + log0("ema:applying EMA weights") + ema_weights = { + name: t.to(dtype=base_model.state_dict()[name].dtype) + for name, t in ema_state.items() + } + base_model.load_state_dict(ema_weights, strict=True) + elif args.swa_enabled and swa_state is not None and swa_count > 1: + log0(f"swa:applying averaged {swa_count} checkpoints") + avg_state = {name: (t / swa_count).to(dtype=base_model.state_dict()[name].dtype) for name, t in swa_state.items()} + base_model.load_state_dict(avg_state, strict=True) # ----------------------------- # SERIALIZATION + ROUNDTRIP VALIDATION # ----------------------------- - # Save the raw state (useful for debugging/loading in PyTorch directly), then always produce - # the compressed int8+zlib artifact and validate the round-tripped weights. + # Save the raw state (useful for debugging/loading in PyTorch directly), then produce + # a mixed int6/int8 artifact and validate the round-tripped weights. if master_process: torch.save(base_model.state_dict(), "final_model.pt") @@ -1073,30 +1922,32 @@ def lr_mul(step: int, elapsed_ms: float) -> float: log0(f"Code size: {code_bytes} bytes") log0(f"Total submission size: {model_bytes + code_bytes} bytes") - quant_obj, quant_stats = quantize_state_dict_int8(base_model.state_dict()) + quant_obj, quant_stats = quantize_state_dict_mixed_int6(base_model.state_dict(), args.int6_categories) quant_buf = io.BytesIO() torch.save(quant_obj, quant_buf) quant_raw = quant_buf.getvalue() - quant_blob = zlib.compress(quant_raw, level=9) + compress_label = "int6+zstd" if zstd is not None else "int6+zlib" + quant_blob = zstd.ZstdCompressor(level=22).compress(quant_raw) if zstd is not None else zlib.compress(quant_raw, level=9) quant_raw_bytes = len(quant_raw) if master_process: - with open("final_model.int8.ptz", "wb") as f: + with open("final_model.int6.ptz", "wb") as f: f.write(quant_blob) - quant_file_bytes = os.path.getsize("final_model.int8.ptz") + quant_file_bytes = os.path.getsize("final_model.int6.ptz") code_bytes = len(code.encode("utf-8")) - ratio = quant_stats["baseline_tensor_bytes"] / max(quant_stats["int8_payload_bytes"], 1) + ratio = quant_stats["baseline_tensor_bytes"] / max(quant_stats["quant_payload_bytes"], 1) log0( - f"Serialized model int8+zlib: {quant_file_bytes} bytes " - f"(payload:{quant_stats['int8_payload_bytes']} raw_torch:{quant_raw_bytes} payload_ratio:{ratio:.2f}x)" + f"Serialized model {compress_label}: {quant_file_bytes} bytes " + f"(payload:{quant_stats['quant_payload_bytes']} raw_torch:{quant_raw_bytes} payload_ratio:{ratio:.2f}x)" ) - log0(f"Total submission size int8+zlib: {quant_file_bytes + code_bytes} bytes") + log0(f"Total submission size {compress_label}: {quant_file_bytes + code_bytes} bytes") if distributed: dist.barrier() - with open("final_model.int8.ptz", "rb") as f: + with open("final_model.int6.ptz", "rb") as f: quant_blob_disk = f.read() - quant_state = torch.load(io.BytesIO(zlib.decompress(quant_blob_disk)), map_location="cpu") - base_model.load_state_dict(dequantize_state_dict_int8(quant_state), strict=True) + quant_buf_in = zstd.ZstdDecompressor().decompress(quant_blob_disk) if zstd is not None else zlib.decompress(quant_blob_disk) + quant_state = torch.load(io.BytesIO(quant_buf_in), map_location="cpu") + base_model.load_state_dict(dequantize_state_dict_mixed_int6(quant_state), strict=True) torch.cuda.synchronize() t_qeval = time.perf_counter() q_val_loss, q_val_bpb = eval_val( @@ -1110,13 +1961,37 @@ def lr_mul(step: int, elapsed_ms: float) -> float: base_bytes_lut, has_leading_space_lut, is_boundary_token_lut, + eval_seq_len=effective_eval_seq_len, ) torch.cuda.synchronize() log0( - f"final_int8_zlib_roundtrip val_loss:{q_val_loss:.4f} val_bpb:{q_val_bpb:.4f} " + f"final_int6_roundtrip val_loss:{q_val_loss:.4f} val_bpb:{q_val_bpb:.4f} " f"eval_time:{1000.0 * (time.perf_counter() - t_qeval):.0f}ms" ) - log0(f"final_int8_zlib_roundtrip_exact val_loss:{q_val_loss:.8f} val_bpb:{q_val_bpb:.8f}") + log0(f"final_int6_roundtrip_exact val_loss:{q_val_loss:.8f} val_bpb:{q_val_bpb:.8f}") + + if 0 < args.eval_stride < effective_eval_seq_len: + torch.cuda.synchronize() + t_slide = time.perf_counter() + sw_val_loss, sw_val_bpb = eval_val_sliding( + args, + base_model, + rank, + world_size, + device, + val_tokens, + base_bytes_lut, + has_leading_space_lut, + is_boundary_token_lut, + stride=args.eval_stride, + eval_seq_len=effective_eval_seq_len, + ) + torch.cuda.synchronize() + log0( + f"final_int6_sliding_window val_loss:{sw_val_loss:.4f} val_bpb:{sw_val_bpb:.4f} " + f"stride:{args.eval_stride} eval_time:{1000.0 * (time.perf_counter() - t_slide):.0f}ms" + ) + log0(f"final_int6_sliding_window_exact val_loss:{sw_val_loss:.8f} val_bpb:{sw_val_bpb:.8f}") if distributed: dist.destroy_process_group()