| --- |
| license: apache-2.0 |
| library_name: pytorch |
| datasets: |
| - Skylion007/openwebtext |
| tags: |
| - discrete-diffusion |
| - uniform-state-diffusion |
| - duo |
| - scaling-laws |
| - checkpoints |
| --- |
| |
| # Duo data-fraction scaling study: `tiny` models |
|
|
| 10 uniform-state discrete diffusion language models (Duo, `duo_base`), identical |
| except for the size of their training set: nested subsets of OpenWebText holding |
| 10%, 20%, ..., 100% of the training blocks. Each model was trained for 1,000,000 steps, |
| and each folder holds **53 full checkpoints** on a fixed step grid, so the whole |
| training trajectory can be analysed (e.g. test loss against weight norm, in the |
| spirit of D'Amico, Bocchi & Negri, *Implicit bias produces neural scaling laws in |
| learning curves*). |
|
|
| ## Models |
|
|
| | fraction | training blocks (1,024 tokens) | epochs at 1M steps | val NELBO at 1M (nats/token) | best val NELBO | training hours | |
| |---|---|---|---|---|---| |
| | 0.1 | 873,094 | 586 | 3.9718 | 3.9717 (step 995,000) | 153 | |
| | 0.2 | 1,746,188 | 293 | 3.9657 | 3.9656 (step 995,000) | 146 | |
| | 0.3 | 2,619,283 | 195 | 3.9998 | 3.9996 (step 990,000) | 151 | |
| | 0.4 | 3,492,377 | 147 | 3.9671 | 3.9669 (step 990,000) | 144 | |
| | 0.5 | 4,365,471 | 117 | 3.9682 | 3.9682 (step 1,000,000) | 150 | |
| | 0.6 | 5,238,565 | 98 | 3.9831 | 3.9826 (step 975,000) | 139 | |
| | 0.7 | 6,111,659 | 84 | 3.9633 | 3.9633 (step 1,000,000) | 148 | |
| | 0.8 | 6,984,754 | 73 | 3.9898 | 3.9885 (step 990,000) | 140 | |
| | 0.9 | 7,857,848 | 65 | 3.9842 | 3.9835 (step 980,000) | 145 | |
| | 1.0 | 8,730,942 | 59 | 3.9902 | 3.9902 (step 1,000,000) | 151 | |
|
|
| Validation NELBO is computed with the EMA weights on the full OpenWebText validation split |
| (110,464 blocks), with fixed noise, every 5,000 steps (`val/nll` in the metrics files). |
|
|
| ## Training setup (same for every model) |
|
|
| - **Model:** DiT backbone, config `tiny` (33,789,010 parameters in the backbone), |
| context length 1,024, GPT-2 tokenizer (vocabulary 50,258), untied embeddings. |
| - **Algorithm:** `duo_base`: uniform-state discrete diffusion trained on the exact NELBO, |
| log-linear noise schedule (no curriculum or distillation). |
| - **Optimisation:** AdamW, learning rate 3e-4, constant after 2,500 warm-up steps, |
| **weight decay 0**, gradient clipping 1.0, EMA decay 0.9999, bf16 mixed precision, |
| global batch 512 sequences x 1,024 tokens, 1,000,000 steps, seed 1. |
| - **Data:** OpenWebText (`Skylion007/openwebtext`, revision `79d93d7`), GPT-2 tokenized and |
| wrapped into 1,024-token blocks: 8,730,942 training blocks; the last 100,000 |
| documents form the validation split. Fraction *f* uses the first floor(f x N + 0.5) |
| blocks of one fixed permutation (seed 20250927, sha256 in `run_info.json`), so the |
| subsets are nested (0.1 is contained in 0.2, and so on). The data order depends only on |
| the optimizer step, not on the GPU layout. |
| - **Code:** a modified version of [s-sahoo/duo](https://github.com/s-sahoo/duo) |
| (scaling-study branch, commit `7c9e7d04d4a4`, recorded in `code_sha`). |
|
|
| ## Files |
|
|
| ``` |
| frac_0.10/ ... frac_1.00/ |
| checkpoints/step_0000000.ckpt ... step_1000000.ckpt 53 PyTorch Lightning checkpoints |
| checkpoints/index.csv step, wall-clock training hours, tokens and epochs per checkpoint |
| metrics/version_*/metrics.csv training/validation metrics (one file per job segment) |
| norms/norms_raw.jsonl, norms_ema.jsonl weight norms every 500 steps, raw and EMA weights |
| run_info.json, config_resolved.yaml, code_sha, segments.jsonl |
| ``` |
|
|
| Checkpoint steps: 0, 100, 200, 400, 700, 1k, 1.5k, 2.5k, 4k, 6k, 10k, 15k, 20k, then every 25k to 1M. |
|
|
| Each checkpoint has the raw weights, the EMA weights, the optimizer state and the |
| loop/sampler state, so training can be resumed exactly. Step 0 is about half the size |
| (no optimizer state yet). |
|
|
| ```python |
| import torch |
| |
| ck = torch.load("frac_0.10/checkpoints/step_1000000.ckpt", map_location="cpu", weights_only=False) |
| raw = {k: v for k, v in ck["state_dict"].items() if k.startswith("backbone.")} |
| # EMA tensors follow the backbone's trainable parameters in order (all of the |
| # state_dict except the rotary-embedding buffer). |
| names = [k for k in raw if not k.endswith("rotary_emb.inv_freq")] |
| ema = dict(zip(names, ck["ema"]["shadow_params"])) |
| print(ck["global_step"], ck["scaling_meta"]["clock"]["train_s"] / 3600, "training hours") |
| ``` |
|
|
| **Norms** (`norms/*.jsonl`): per record the step, aggregate norms of the token-path matrices |
| (embedding, per-block Q/K/V, attention output, MLP, output layer) and, separately, of the |
| timestep-conditioning path, plus per-matrix statistics (spectral norm, Frobenius, L1, |
| row and column (2,1)-norms) in float64. `log_R` is the Bartlett spectral complexity: |
| `sum log||W||_sigma + 1.5 log sum (||W^T||_2,1^(2/3) / ||W||_sigma^(2/3))`. It is null at |
| step 0 because the output layer starts at zero. |
|
|
| ## Caveats |
|
|
| - **Wall-clock time** (`train_hours`) mixes H100 and H200 segments (the GPU of every |
| segment is in `segments.jsonl`). For time-based comparisons, use steps or normalise |
| per GPU type. |
| - `trainer/loss` in the metrics is one micro-batch on one GPU and depends on the GPU |
| layout; compare runs with `val/nll`. |
| - **frac_0.80** has a transient rise in EMA validation NELBO at steps 590k-625k |
| (peak 4.75 at 600k against about 4.00). A replay from the 575k checkpoint on different |
| hardware reproduced it, so it is a genuine event in the training trajectory, not a |
| hardware fault. |
| |
