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