lemoncmd/duo-scaling-medium

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 (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).

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.