lemoncmd/duo-scaling-tiny

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

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.