diff --git a/README.md b/README.md
index 31a23455d..0f1b34b66 100644
--- a/README.md
+++ b/README.md
@@ -1,211 +1,53 @@
-
+# parameter_golf_bj_run
-
-
+This repository contains my local reproduction of the OpenAI Parameter Golf baseline.
-**OpenAI Model Craft Challenge: Parameter Golf** is a challenge to train the best language model that fits in a 16MB artifact and trains in under 10 minutes on 8xH100s, evaluated by compression on the FineWeb validation set (tokenizer-agnostic, bits per byte).
+## 🚀 Baseline Result (Local Run)
-This challenge is heavily inspired by the [NanoGPT Speedrunning](https://github.com/KellerJordan/modded-nanogpt) challenge, where participants compete to train a model that reaches 3.28 FineWeb validation loss as quickly as possible. We're excited to see how optimizing for a parameter-constrained setting pushes people toward unique architectures (test-time compute, aggressive parameter tying, depth recurrence, low-rank training, ...), compression schemes (low precision, QAT, bitnets, novel tokenizers, ...), and other creative submissions (test-time training, long context, megakernels ...).
+- Environment: single GPU (local)
+- Dataset: fineweb10B_sp1024 (train_shards=1)
+- Sequence length: 1024
+- Grad accumulation: 8
-If you're familiar with [neural scaling laws](https://arxiv.org/abs/2001.08361), you can consider this challenge a form of L(N) optimization, where the objective is to optimize the lowest loss given a fixed number of parameters (N) unconstrained by data, compute, steps, or architecture. Challenges like the [NanoGPT Speedrun](https://github.com/KellerJordan/modded-nanogpt), which optimizes for a form of L(T) (~lowest time given constrained loss) or the [NanoGPT Slowrun](https://github.com/qlabs-eng/slowrun), which optimizes for L(D) (lowest loss given constrained dataset size), can be thought of as equivalent challenges in this family.
+### Best Result
-Ideally, we'd allow for submissions to use arbitrary computational resources. But in order to make the challenge not inaccessibly expensive, we're limiting *leaderboard submissions* to 10 minutes on 8xH100s. However, we'd still love to see submissions that don't meet the compute limitation requirements in our 'Non-record Submissions' section: We're excited to see people push the infinite frontier of parameter limited performance as well.
+- **val_bpb = 1.3529 @ step 4200**
-We also know compute is expensive, so **OpenAI is sponsoring $1,000,000 in compute credits** to help people get started training their models. To request a compute grant, use this form: [Request a Compute Grant](https://openai.com/index/parameter-golf/#credit-form).
-When requesting compute, please make sure you choose the appropriate level, write sufficient justification, and **submit with an email tied to a OpenAI / ChatGPT account**.
+### Training Summary
-## Participant Form
+- step 0: val_bpb = 4.1077
+- step 1000: val_bpb = 1.3944
+- step 2000: val_bpb = 1.3651
+- step 2800: val_bpb = 1.3548
+- step 4200: val_bpb = 1.3529 (best)
+- step 5000: val_bpb = 1.3576
-If you enjoy solving very difficult technical problems, please introduce yourself via the [Challenge Participant Form](https://jobs.ashbyhq.com/openai/form/open-ai-challenge-parameter-golf). It helps us attribute challenge submissions and reach out about opportunities with OpenAI. _Completing the form is not required to participate._
+## 📊 Observations
-Many researchers at OpenAI first distinguished themselves through elite mathematics and programming competitions. The Model Craft Challenge is designed in that spirit: testing the ability to tackle unfamiliar problems with creativity and rigor, qualities we believe are essential for frontier AI research.
+- Rapid improvement in early training
+- Convergence around ~1.35 bpb
+- No further improvement after ~4200 steps
+- Additional training leads to fluctuation, not improvement
-In June, we plan to hire a small cohort of early-career researchers, targeting current undergraduate students and recent graduates, including Olympiad medalists and elite competitors. For exceptional participants, the challenge may also serve as a way to stand out to OpenAI researchers and recruiters.
+## 📁 Files
-The challenge runs from March 18th to April 30th.
+- Log file:
+ - `local_logs/baseline_local_smoke_2026_03_21_v01.log`
-Happy training!
+- Experiment notes:
+ - `notes/baseline_run_2026-03-21.md`
-## Leaderboard
+## 🎯 Status
-| Run | Score | Author | Summary | Date | Info |
-|-----|------:|--------|---------|------|------|
-| 10L Int5-MLP + BigramHash(10240) | 1.1428 | thwu1 | 10 layers, mixed int5/int6 quantization, BigramHash(10240), SWA(0.4), WD=0.04 | 2026-03-20 | [info](records/track_10min_16mb/2026-03-20_10L_Int5MLP_MuonWD04_SWA50/README.md) |
-| Int6 MLP3x + SmearGate + BigramHash | 1.1458 | Raahil Shah | 3x MLP + SmearGate + BigramHash + OrthoInit + Muon WD + SWA | 2026-03-20 | [info](records/track_10min_16mb/2026-03-20_Int6_MLP3x_SmearGate_BigramHash_MuonWD_SWA/README.md) |
-| 11L MLP3x + Int6 QAT | 1.1502 | aruniyer | 11 layers, 3x MLP, int6 QAT, zstd-22, WD=0.04, sliding eval | 2026-03-20 | [info](records/track_10min_16mb/2026-03-19_MLP3x_QAT_Int6_SlidingWindow/README.md) |
-| SmearGate + OrthoInit + Muon WD | 1.1556 | aquariouseworkman | SmearGate + BigramHash + 3x MLP + int6 STE QAT + sliding eval | 2026-03-19 | [info](records/track_10min_16mb/2026-03-19_smeargate_orthoinit_muonwd/README.md) |
-| 10L Int6 QAT + Zstd MLP2.6x | 1.1586 | yahya010 | 10 layers, int6 QAT + zstd-22, MLP 1344, Muon 0.99, sliding eval | 2026-03-19 | [info](records/track_10min_16mb/2026-03-19_Seq2048_FP16Emb_TunedLR/README.md) |
-| Mixed Quant + Sliding Window Eval | 1.1630 | aquariouseworkman | Int6 block weights + int8 embeddings + 3x MLP + sliding eval | 2026-03-19 | [info](records/track_10min_16mb/2026-03-19_MixedQuant_Int6Int8_SlidingWindow/README.md) |
-| Muon WD + 10 layer | 1.1748 | notapplica | Includes prev. wins + Spectral embed init + resid mix | 2026-03-19 | [info](records/track_10min_16mb/2026-03-19_SlidingWindow_FP16Emb_10L_MuonWD_OvertoneInit/README.md) |
-| Sliding Window Eval | 1.1925 | Matthew Li | Sliding window evaluation at stride=64, increasing context for eval | 2026-03-19 | [info](records/track_10min_16mb/2026-03-19_SlidingWindowEval/README.md) |
-| Lora TTT | 1.1928 | samacqua | Test-time training with LORAs | 2026-03-19 | [info](records/track_10min_16mb/2026-03-17_LoRA_TTT/README.md) |
-| 4k seq length| 1.2014 | Spokane Way | 4k seq length + better hypers | 2026-03-19 | [info](records/track_10min_16mb/2026-03-19_TrainingOptSeq4096/README.md) |
-| 2048 seq length | 1.206 | Spokane Way | 2048 seq length (train + val) | 2026-03-18 | [info](records/track_10min_16mb/2026-03-18_LongContextSeq2048/README.md) |
-| int6 mixed precision | 1.2147 | Nan Liu | 10 layers, mixed int8/int6 | 2026-03-18 | [info](records/track_10min_16mb/2026-03-19_10L_MixedPrecision/README.md) |
-| fp16 Embed | 1.2197 | Renier Velazco | FP16 Tied Embedding + LR/Warmdown Tuning | 2026-03-18 | [info](records/track_10min_16mb/2026-03-18_FP16Embed_WD3600/README.md) |
-| Naive Baseline | 1.2244 | Baseline | 9layer 512dim 1024vocab TiedEmbeddings 4 KV heads | 2026-03-18 | [info](records/track_10min_16mb/2026-03-17_NaiveBaseline/README.md) |
+- Baseline successfully reproduced on local hardware
+- This result is used as the reference point for further experiments
-#### Notable Non-Record Runs
+## 🔜 Next Steps
-| Run | Score | Author | Summary | Date | Info |
-|-----|------:|--------|---------|------|------|
-| 4-Hour Baseline | 1.2074 | Will DePue | Testing unlimited compute, 4 hours on 8xH100 | 2026-03-18 | [info](records/track_non_record_16mb/2026-03-18_Quasi10Bfrom50B_SP1024_9x512_KV4_4h_pgut3/README.md) |
+- Modify training parameters (e.g. learning rate)
+- Add evaluation improvements (e.g. sliding window)
+- Compare against baseline (1.3529)
-## Getting Started
+---
-### Training Your First Model (Mac with Apple Silicon)
-
-If you have an Apple laptop or desktop with Apple Silicon, we've set up a simple MLX training script to help you start iterating locally.
-
-If you don't have a Mac with Apple Silicon, you can run an adapted version of this script without MLX support. Just ask [Codex](https://openai.com/codex/) to refactor it; the change is straightforward. It may still be fairly slow, so we recommend jumping straight to cloud GPUs with Runpod.
-
-First, clone the repository, create a fresh Python environment, and install the packages needed for the MLX path plus dataset download:
-
-```bash
-git clone https://github.com/openai/parameter-golf.git
-cd parameter-golf
-python3 -m venv .venv
-source .venv/bin/activate
-python -m pip install --upgrade pip
-pip install mlx numpy sentencepiece huggingface-hub datasets tqdm
-```
-
-Download our cached version of FineWeb with the 1024-token vocabulary:
-
-```bash
-python3 data/cached_challenge_fineweb.py --variant sp1024 --train-shards 10
-```
-
-This populates `./data/datasets/fineweb10B_sp1024/` and `./data/tokenizers/`.
-By default this downloads the full validation split plus 80 training shards (8B tokens). For a smaller local smoke subset, pass `--train-shards 1`, for example `python3 data/cached_challenge_fineweb.py --variant sp1024 --train-shards 1`.
-
-Then run a small MLX training job:
-
-```bash
-RUN_ID=mlx_smoke \
-ITERATIONS=200 \
-TRAIN_BATCH_TOKENS=8192 \
-VAL_LOSS_EVERY=0 \
-VAL_BATCH_SIZE=8192 \
-python3 train_gpt_mlx.py
-```
-
-Validation always runs on the full `fineweb_val_*` split, which is the fixed first-50k-document set. The smoke command above skips periodic validation and just prints the final `val_loss` and `val_bpb` once at the end.
-
-### Scaling Up to a Remote Machine
-
-Once you're happy with your local tests, or you want more compute, switch to a remote CUDA machine.
-
-You can rent GPUs from anywhere, but OpenAI is partnering with Runpod to make setup as easy as possible.
-
-#### Launching a 1xH100 Pod
-
-1. First, [create a Runpod account](https://console.runpod.io/deploy). You should also set up an SSH key in the Settings tab on the left so you can connect to your remote machine. If you're new to this, ask Codex to help you set it up.
-
-2. Once you've set up your account, create a new GPU Cloud Pod. You can choose whichever GPU SKU you'd like. Final leaderboard submissions must run in under 10 minutes on 8xH100s (specifically the SXM variant), but we strongly recommend testing and running experiments on cheaper SKUs first, since an 8xH100 box can cost around $20/hour.
-
-3. Let's start with a 1xH100 pod. Deploy using the official Parameter Golf template: [Launch Template](https://console.runpod.io/deploy?template=y5cejece4j&ref=nl2r56th). Enable SSH terminal access, leaving the other settings at their defaults. Deploy your pod and SSH into it once it's up. You should land in `/workspace/`.
-
-On your remote machine, clone the repo onto local disk. All Python dependencies are already pre-installed in the image.
-
-```bash
-cd /workspace
-git clone https://github.com/openai/parameter-golf.git
-cd parameter-golf
-```
-
-Download our cached version of FineWeb. We'll use the 1024-token vocabulary for now.
-
-```bash
-python3 data/cached_challenge_fineweb.py --variant sp1024
-```
-
-This defaults to the full validation split plus 80 training shards (8B tokens). If you only want a smaller subset while iterating, pass `--train-shards N`, for example `--train-shards 1`.
-
-Launch your first training run. Note that we're passing `nproc_per_node=1` because we're running on a single H100 GPU in this case.
-
-```bash
-RUN_ID=baseline_sp1024 \
-DATA_PATH=./data/datasets/fineweb10B_sp1024/ \
-TOKENIZER_PATH=./data/tokenizers/fineweb_1024_bpe.model \
-VOCAB_SIZE=1024 \
-torchrun --standalone --nproc_per_node=1 train_gpt.py
-```
-
-By default, `train_gpt.py` keeps its ~10 minute wallclock cap. If you want a longer run, override it explicitly, for example `MAX_WALLCLOCK_SECONDS=0`.
-
-By default, this command prints `train_loss` step logs during training and prints `val_loss`, `val_bpb`, and compressed model size in the final `final_int8_zlib_roundtrip` lines at the end. If you want periodic validation logs during the run, set `VAL_LOSS_EVERY`, for example `VAL_LOSS_EVERY=200`. For the baseline config, the final `val_bpb` should land around ~1.2 with a compressed model size under 16MB.
-
-For dataset export, tokenizer export, and docs-cache rebuild instructions, see [data/README.md](data/README.md).
-
-Evaluation will be in the RunPod environment with all packages installed. `requirements.txt` is provided as a reference if you want to self-setup.
-
-## FAQ
-
-**What exactly counts toward the 16MB artifact size?**
-
-The submission artifact is computed as code bytes plus compressed model bytes. All counted code should live in the `train_gpt.py` script.
-The cap is decimal 16MB, i.e. 16,000,000 total bytes, not 16 MiB / 16,777,216 bytes.
-No external downloads, training dataset access, or network calls are allowed during evaluation. The artifact must be fully self-contained and reproducible.
-
-**Are scores independently verified by OpenAI?**
-
-We're not automatically verifying every submission, but we will verify the top leaderboard entries over time. Any non-reproducible results can be disqualified, and issues reproducing submissions should be raised on the PR. If you find an issue with a record on the leaderboard or find a record isn't reproducible, please let us know and add an Github Issue describing your findings.
-
-**What counts as 'external compute'? For example, is it fair to tune my hyperparameters offline?**
-
-There's no perfectly clear answer here and it's hard to draw a clean line around what does or does not count as external compute. For now, we're reserving the right to disqualify runs that are not in the spirit of the challenge. Tuning your Adam hyperparameters across a bunch of runs is fine, but if there's evidence that you're sneaking in additional compute unfairly, such as brute-forcing ridiculous seeds, we won't allow it. Use your best judgment and there's no penalty for asking questions.
-
-**What are the restrictions on evaluation?**
-
-We won't accept submissions that take more than 10 minutes on 8xH100 to evaluate (Note: This limit is in addition to the 10 minutes of training time allowed!), but otherwise you're free to evaluate however. As with modded-nanogpt, we allow evaluation at any sequence length. And, obviously, you aren't allowed to access any training data during evaluation, unless you pay for those bits in the <16MB limit. We encourage competitors to push the bounds of evaluation methods as aggressively as with training methods. You CANNOT access validation data during training, e.g. by compressing it into your 16mb with "paid prefix".
-
-If it isn't abundantly obvious: You can't cheat on your test loss. You can't cheat by training on the validation set before you evaluate on the validation set. The validation language around test-time training has been confusing people: you are only allowed to test-time train on validation set tokens _you've already evaluated your model on_, since those tokens have already been graded!
-
-**What is the process for accepting new submissions?**
-
-Since all submissions are public, we're accepting record submissions chronologically depending on their PR creation time. The leaderboard may take time to update due to verification and review of submissions, so pay consideration to what the current SOTA PR is when submitting. As explained below, submissions should exceed the SOTA record with sufficient statistical significance in order to accepted for the leaderboard. Otherwise, submissions may be accepted as 'non-record submissions' given they are sufficiently unique or interesting.
-
-**Can I import XYZ package or library?**
-
-Yes, you're free to import any package or library you want, so long as it does not unjustly violate the rules on evaluation, compute, training time, code size or otherwise. Just include a requirements.txt in your records folder and mention setup instructions in your README.md. Since you don't pay for bits imported in Python libraries, limitations clearly apply: You can't sneak in extra compute, capabilities, or massively increase effective code size with custom libraries, but importing FlashAttention, etc. is completely fine.
-
-
-## Submission Process
-
-New SOTA records must fulfill the following criteria:
-
-1. They must beat the existing SOTA by at least 0.005 nats. As in modded-nanogpt, because of inter-run variance all submissions must provide enough run logs to show at `p < 0.01` that they achieved the required 0.005-nat improvement. For submissions that improve speed through systems optimization without changing the ML, this requirement is waived.
-
-2. If changes are made to the tokenizer or dataset, prove with certainty that the val_bpb is correctly calculated. Submissions that edit the tokenizer will be examined much more carefully, since bugs may unjustly improve your score.
-
-3. Reproducibly run in under 10 minutes on 8xH100s.
-
-All submissions should be made as a pull request that only adds a new folder to the appropriate `/records` subfolder and includes the following files. Submissions without the full set of requirements will not be accepted.
-
-1. A README.md file that explains the submission in reasonable detail.
-
-2. A `submission.json` file (see the example runs) that includes your name, GitHub ID, `val_bpb`, and related metadata.
-
-3. A train log, automatically produced by your script. Please demonstrate a statistically significant win. Most often, submitting an average over 3 training runs is sufficient.
-
-4. A `train_gpt.py` script and any other dependencies. Note: this must successfully compile and run within the records folder. Broken scripts will not be accepted.
-
-### Non-record Submissions
-
-Submissions are also open to unique and interesting approaches that might not beat the existing SOTA, but still satisfy the 16MB artifact limit. We strongly encourage participants to submit implementations for weird or out-of-the-box ideas, in-progress or unoptimized solutions, so long as they run successfully, or even interesting negative results. We're excited to see what you come up with. We'll still maintain a high bar for non-record submissions, so be sure to justify your ideas and results in detail when submitting.
-
-We also accept non-record submissions to an unlimited compute track for runs that are not intended to meet the 10-minute cutoff. Just note as such in your README file.
-
-Non-record submissions should be made in the same fashion as SOTA records, as described above.
-
-#### PRs on Core Code
-
-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 the best models should stay in the `/records` folder.
-
-## Support
-
-
-Join the [OpenAI Discord server](https://discord.com/invite/openai) and visit the Parameter Golf channels (#parameter-golf-discussions, #parameter-golf-announcements) and ask questions.
-
-This repository adapts code from `modded-nanogpt`, see [THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md) for attribution.
+> This repository is based on the OpenAI Parameter Golf baseline, with added local experiment results.
diff --git a/records/track_non_record_16mb/2026-03-21_local_baseline_bj/README.md b/records/track_non_record_16mb/2026-03-21_local_baseline_bj/README.md
new file mode 100644
index 000000000..64577c841
--- /dev/null
+++ b/records/track_non_record_16mb/2026-03-21_local_baseline_bj/README.md
@@ -0,0 +1,22 @@
+# Local Baseline Reproduction
+
+## Summary
+
+This submission records a local single-GPU reproduction of the OpenAI Parameter Golf NaiveBaseline.
+
+## Environment
+
+- local single-GPU run
+- train_shards=1
+- seq_len=1024
+- grad_accum_steps=8
+
+## Result
+
+- best observed val_bpb: **1.3529 @ step 4200**
+
+## Notes
+
+- validation improved from 4.1077 to 1.3529
+- performance plateaued after around step 4200
+- this is a local baseline reproduction and not a leaderboard SOTA submission
diff --git a/records/track_non_record_16mb/2026-03-21_local_baseline_bj/requirements.txt b/records/track_non_record_16mb/2026-03-21_local_baseline_bj/requirements.txt
new file mode 100644
index 000000000..911b0e52f
--- /dev/null
+++ b/records/track_non_record_16mb/2026-03-21_local_baseline_bj/requirements.txt
@@ -0,0 +1,10 @@
+numpy
+tqdm
+torch
+huggingface-hub
+kernels
+setuptools
+typing-extensions==4.15.0
+datasets
+tiktoken
+sentencepiece
\ No newline at end of file
diff --git a/records/track_non_record_16mb/2026-03-21_local_baseline_bj/submission.json b/records/track_non_record_16mb/2026-03-21_local_baseline_bj/submission.json
new file mode 100644
index 000000000..d89da32c0
--- /dev/null
+++ b/records/track_non_record_16mb/2026-03-21_local_baseline_bj/submission.json
@@ -0,0 +1,10 @@
+{
+ "name": "Local Baseline Reproduction",
+ "github_id": "bjbjbjbjbjbj",
+ "author": "bjbjbjbjbjbj",
+ "track": "track_non_record_16mb",
+ "category": "non-record",
+ "val_bpb": 1.3529,
+ "best_step": 4200,
+ "description": "Local single-GPU reproduction of the OpenAI Parameter Golf NaiveBaseline."
+}
diff --git a/records/track_non_record_16mb/2026-03-21_local_baseline_bj/train.log b/records/track_non_record_16mb/2026-03-21_local_baseline_bj/train.log
new file mode 100644
index 000000000..960064667
--- /dev/null
+++ b/records/track_non_record_16mb/2026-03-21_local_baseline_bj/train.log
@@ -0,0 +1,128 @@
+logs/baseline_local_smoke.txt
+val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path=./data/tokenizers/fineweb_1024_bpe.model
+train_loader:dataset:fineweb10B_sp1024 train_shards:1
+val_loader:shards pattern=./data/datasets/fineweb10B_sp1024/fineweb_val_*.bin tokens:62021632
+model_params:17059912
+world_size:1 grad_accum_steps:8
+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:0.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.9357 train_time:877ms step_avg:877.22ms
+step:2/20000 train_loss:16.7399 train_time:1754ms step_avg:876.79ms
+step:3/20000 train_loss:8.7524 train_time:2630ms step_avg:876.72ms
+step:4/20000 train_loss:6.5863 train_time:3509ms step_avg:877.29ms
+step:5/20000 train_loss:6.6533 train_time:4383ms step_avg:876.55ms
+step:6/20000 train_loss:6.6960 train_time:5256ms step_avg:875.99ms
+step:7/20000 train_loss:6.3296 train_time:6130ms step_avg:875.72ms
+step:8/20000 train_loss:6.1789 train_time:7005ms step_avg:875.62ms
+step:9/20000 train_loss:6.0812 train_time:7879ms step_avg:875.40ms
+step:10/20000 train_loss:5.9924 train_time:8756ms step_avg:875.64ms
+step:200/20000 train_loss:2.8115 train_time:174889ms step_avg:874.45ms
+step:200/20000 val_loss:2.8262 val_bpb:1.6738 train_time:174890ms step_avg:874.45ms
+step:400/20000 train_loss:2.5579 train_time:349749ms step_avg:874.37ms
+step:400/20000 val_loss:2.5627 val_bpb:1.5178 train_time:349750ms step_avg:874.37ms
+step:600/20000 train_loss:2.5009 train_time:524619ms step_avg:874.36ms
+step:600/20000 val_loss:2.4512 val_bpb:1.4517 train_time:524619ms step_avg:874.37ms
+step:800/20000 train_loss:2.3082 train_time:699469ms step_avg:874.34ms
+step:800/20000 val_loss:2.3889 val_bpb:1.4149 train_time:699470ms step_avg:874.34ms
+step:1000/20000 train_loss:2.4521 train_time:874398ms step_avg:874.40ms
+step:1000/20000 val_loss:2.3543 val_bpb:1.3944 train_time:874399ms step_avg:874.40ms
+step:1200/20000 train_loss:2.2483 train_time:1049333ms step_avg:874.44ms
+step:1200/20000 val_loss:2.3310 val_bpb:1.3806 train_time:1049334ms step_avg:874.44ms
+step:1400/20000 train_loss:2.2245 train_time:1224297ms step_avg:874.50ms
+step:1400/20000 val_loss:2.3244 val_bpb:1.3767 train_time:1224298ms step_avg:874.50ms
+step:1600/20000 train_loss:2.2640 train_time:1399206ms step_avg:874.50ms
+step:1600/20000 val_loss:2.3126 val_bpb:1.3697 train_time:1399206ms step_avg:874.50ms
+step:1800/20000 train_loss:2.2040 train_time:1574162ms step_avg:874.53ms
+step:1800/20000 val_loss:2.3030 val_bpb:1.3640 train_time:1574162ms step_avg:874.53ms
+step:2000/20000 train_loss:2.2418 train_time:1749086ms step_avg:874.54ms
+step:2000/20000 val_loss:2.3049 val_bpb:1.3651 train_time:1749086ms step_avg:874.54ms
+step:2200/20000 train_loss:2.1885 train_time:1924124ms step_avg:874.60ms
+step:2200/20000 val_loss:2.2987 val_bpb:1.3614 train_time:1924124ms step_avg:874.60ms
+step:2400/20000 train_loss:2.1676 train_time:2098920ms step_avg:874.55ms
+step:2400/20000 val_loss:2.2931 val_bpb:1.3581 train_time:2098920ms step_avg:874.55ms
+step:2600/20000 train_loss:2.1875 train_time:2273690ms step_avg:874.50ms
+step:2600/20000 val_loss:2.2975 val_bpb:1.3607 train_time:2273690ms step_avg:874.50ms
+step:2800/20000 train_loss:2.2359 train_time:2448740ms step_avg:874.55ms
+step:2800/20000 val_loss:2.2875 val_bpb:1.3548 train_time:2448741ms step_avg:874.55ms
+step:3000/20000 train_loss:2.1136 train_time:2623603ms step_avg:874.53ms
+step:3000/20000 val_loss:2.2932 val_bpb:1.3582 train_time:2623603ms step_avg:874.53ms
+step:3200/20000 train_loss:2.1406 train_time:2798419ms step_avg:874.51ms
+step:3200/20000 val_loss:2.3014 val_bpb:1.3630 train_time:2798420ms step_avg:874.51ms
+step:3400/20000 train_loss:2.1058 train_time:2973275ms step_avg:874.49ms
+step:3400/20000 val_loss:2.2962 val_bpb:1.3600 train_time:2973275ms step_avg:874.49ms
+step:3600/20000 train_loss:2.1619 train_time:3146783ms step_avg:874.11ms
+step:3600/20000 val_loss:2.2934 val_bpb:1.3583 train_time:3146784ms step_avg:874.11ms
+step:3800/20000 train_loss:2.0911 train_time:3320199ms step_avg:873.74ms
+step:3800/20000 val_loss:2.2942 val_bpb:1.3588 train_time:3320199ms step_avg:873.74ms
+step:4000/20000 train_loss:2.1452 train_time:3492381ms step_avg:873.10ms
+step:4000/20000 val_loss:2.2926 val_bpb:1.3578 train_time:3492381ms step_avg:873.10ms
+step:4200/20000 train_loss:2.2554 train_time:3664294ms step_avg:872.45ms
+step:4200/20000 val_loss:2.2844 val_bpb:1.3529 train_time:3664294ms step_avg:872.45ms
+step:4400/20000 train_loss:2.1171 train_time:3835777ms step_avg:871.77ms
+step:4400/20000 val_loss:2.2900 val_bpb:1.3563 train_time:3835778ms step_avg:871.77ms
+step:4600/20000 train_loss:2.0984 train_time:4008484ms step_avg:871.41ms
+step:4600/20000 val_loss:2.2988 val_bpb:1.3615 train_time:4008484ms step_avg:871.41ms
+step:4800/20000 train_loss:2.0940 train_time:4180506ms step_avg:870.94ms
+step:4800/20000 val_loss:2.2980 val_bpb:1.3610 train_time:4180507ms step_avg:870.94ms
+step:5000/20000 train_loss:2.0744 train_time:4352219ms step_avg:870.44ms
+step:5000/20000 val_loss:2.2923 val_bpb:1.3576 train_time:4352220ms step_avg:870.44ms
+step:5200/20000 train_loss:2.1048 train_time:4523487ms step_avg:869.90ms
+step:5200/20000 val_loss:2.2970 val_bpb:1.3604 train_time:4523487ms step_avg:869.90ms
+step:5400/20000 train_loss:2.0611 train_time:4694162ms step_avg:869.29ms
+step:5400/20000 val_loss:2.2943 val_bpb:1.3588 train_time:4694162ms step_avg:869.29ms
+step:5600/20000 train_loss:2.1174 train_time:4864944ms step_avg:868.74ms
+step:5600/20000 val_loss:2.2966 val_bpb:1.3602 train_time:4864945ms step_avg:868.74ms
+step:5800/20000 train_loss:2.1285 train_time:5035549ms step_avg:868.20ms
+step:5800/20000 val_loss:2.2986 val_bpb:1.3614 train_time:5035550ms step_avg:868.20ms
+step:6000/20000 train_loss:2.0747 train_time:5206131ms step_avg:867.69ms
+step:6000/20000 val_loss:2.2974 val_bpb:1.3606 train_time:5206131ms step_avg:867.69ms
+step:6200/20000 train_loss:2.0419 train_time:5376432ms step_avg:867.17ms
+step:6200/20000 val_loss:2.3130 val_bpb:1.3699 train_time:5376432ms step_avg:867.17ms
+step:6400/20000 train_loss:2.1215 train_time:5546742ms step_avg:866.68ms
+step:6400/20000 val_loss:2.3097 val_bpb:1.3679 train_time:5546742ms step_avg:866.68ms
+step:6600/20000 train_loss:2.0574 train_time:5716551ms step_avg:866.14ms
+step:6600/20000 val_loss:2.3073 val_bpb:1.3665 train_time:5716551ms step_avg:866.14ms
+step:6800/20000 train_loss:2.0724 train_time:5886183ms step_avg:865.62ms
+step:6800/20000 val_loss:2.3004 val_bpb:1.3624 train_time:5886184ms step_avg:865.62ms
+step:7000/20000 train_loss:2.1208 train_time:6055649ms step_avg:865.09ms
+step:7000/20000 val_loss:2.2989 val_bpb:1.3615 train_time:6055649ms step_avg:865.09ms
+step:7200/20000 train_loss:2.1087 train_time:6225230ms step_avg:864.62ms
+step:7200/20000 val_loss:2.3082 val_bpb:1.3670 train_time:6225230ms step_avg:864.62ms
+step:7400/20000 train_loss:2.1120 train_time:6394607ms step_avg:864.14ms
+step:7400/20000 val_loss:2.3146 val_bpb:1.3708 train_time:6394607ms step_avg:864.14ms
+step:7600/20000 train_loss:2.0722 train_time:6563899ms step_avg:863.67ms
+step:7600/20000 val_loss:2.3132 val_bpb:1.3700 train_time:6563899ms step_avg:863.67ms
+step:7800/20000 train_loss:2.0393 train_time:6733081ms step_avg:863.22ms
+step:7800/20000 val_loss:2.3250 val_bpb:1.3770 train_time:6733082ms step_avg:863.22ms
+step:8000/20000 train_loss:2.0460 train_time:6901636ms step_avg:862.70ms
+step:8000/20000 val_loss:2.3273 val_bpb:1.3784 train_time:6901636ms step_avg:862.70ms
+step:8200/20000 train_loss:2.1143 train_time:7071240ms step_avg:862.35ms
+step:8200/20000 val_loss:2.3048 val_bpb:1.3650 train_time:7071241ms step_avg:862.35ms
+step:8400/20000 train_loss:2.0629 train_time:7241308ms step_avg:862.06ms
+step:8400/20000 val_loss:2.3078 val_bpb:1.3668 train_time:7241309ms step_avg:862.06ms
+step:8600/20000 train_loss:2.0461 train_time:7411531ms step_avg:861.81ms
+step:8600/20000 val_loss:2.3062 val_bpb:1.3659 train_time:7411531ms step_avg:861.81ms
diff --git a/records/track_non_record_16mb/2026-03-21_local_baseline_bj/train_gpt.py b/records/track_non_record_16mb/2026-03-21_local_baseline_bj/train_gpt.py
new file mode 100644
index 000000000..651beb2b8
--- /dev/null
+++ b/records/track_non_record_16mb/2026-03-21_local_baseline_bj/train_gpt.py
@@ -0,0 +1,1126 @@
+"""
+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))
+ 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()