--- 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: `medium` 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. **Training is in progress** (target 1,000,000 steps). This repository is a backup of the grid checkpoints reached so far and will be updated as training continues. Each folder holds **full checkpoints** on a fixed step grid (53 at 1M), 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) | step | epochs | val NELBO at that step (nats/token) | training hours | |---|---|---|---|---|---| | 0.1 | 873,094 | 325,000 | 191 | 3.0383 | 155 | | 0.2 | 1,746,188 | 350,000 | 103 | 3.0142 | 173 | | 0.3 | 2,619,283 | 275,000 | 54 | 3.0438 | 138 | | 0.4 | 3,492,377 | 325,000 | 48 | 3.0191 | 159 | | 0.5 | 4,365,471 | 625,000 | 73 | 2.9651 | 307 | | 0.6 | 5,238,565 | 275,000 | 27 | 3.0327 | 138 | | 0.7 | 6,111,659 | 275,000 | 23 | 3.0296 | 133 | | 0.8 | 6,984,754 | 275,000 | 20 | 3.0287 | 139 | | 0.9 | 7,857,848 | 275,000 | 18 | 3.0334 | 135 | | 1.0 | 8,730,942 | 275,000 | 16 | 3.0420 | 133 | 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 `medium` (424,477,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 (target), 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 ... PyTorch Lightning checkpoints (53 at 1M) 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 so far: 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`. - Some medium segments ran on a node with a faulty GPU (alphagpu08, NaN losses). Their outputs were removed and those runs resumed from the last clean checkpoint, so every published checkpoint comes from clean hardware. - **frac_0.30**: steps 0-2,500 ran on a throttled node, so its early `train_hours` are inflated.