nanoDiff 350M — base
A
masked diffusion language model trained with the LLaDA recipe, from the
nanoDiff project — a minimal, hackable
"nanoGPT for diffusion LLMs".
The next rung in the scaling ladder after
nanodiff-50m-base and
nanodiff-150m-base.
This is the
base model — pretrained on raw web text only, no instruction
tuning. It continues documents; it does not answer questions.
TL;DR
| |
|---|
| Architecture | Masked diffusion LM (LLaDA recipe), bidirectional Transformer |
| Parameters | ~382M total · ~317M non-embedding (tied input/output embeddings) |
| Training data | FineWeb-Edu, ~10B tokens |
| Validation | NLL-bound 3.38 nats/token · perplexity 29.3 |
| LAMBADA accuracy | 36.95% |
| Hardware | NVIDIA DGX-Spark (GB10), ~6 days |
Scaling family
This run completes a 4-point capacity scaling sweep on the same FineWeb-Edu
shard, same masked-diffusion recipe, same WSD schedule structure:
| Model | Tokens | Val NLL | Perplexity | LAMBADA acc |
|---|
| 50M | 2B | 3.92 | 50.6 | 19.83% |
| 50M (matched-token control) | 3B | 3.91 | 50.1 | — |
| 150M | 3B | 3.78 | 43.8 | 21.89% |
| 350M | 10B | 3.38 | 29.3 | 36.95% |
The 150M→350M step is the largest gain in the family — -0.38 nats / ~31% PPL
reduction — consistent with continuing to scale capacity and training tokens
together (10B tokens is Chinchilla-optimal for 350M at ~20 tokens/param,
overshooting slightly).
Provenance
- Code: https://github.com/BY571/nanoDiff
- Trained at commit:
30fa4dd
- Config:
pretrain/configs/350m.py
- Reproduce:
1python scripts/prepare_data.py --out-dir data/fineweb_edu_10b \
2 --num-tokens 10_000_000_000
3python pretrain/train.py --config pretrain/configs/350m.py
Architecture
A decoder-style Transformer with bidirectional self-attention — the one
architectural change from an autoregressive GPT, since a diffusion LM denoises
all positions at once.
| Hyperparameter | Value |
|---|
| Layers | 16 |
| Attention heads | 20 |
Embedding dim (n_embd) | 1280 |
| Head dim | 64 |
Context length (block_size) | 512 |
| Positional encoding | RoPE |
| Vocabulary | 50304 (GPT-2 BPE 50257 + [MASK]@50257 + padding) |
| Embeddings | tied (input = output) |
| Total parameters | ~382M (~317M non-embedding) |
Training objective
Masked (absorbing-state) diffusion, following
LLaDA: each token is independently replaced
by
[MASK] with probability
t ~ U(0, 1); the model predicts the clean tokens
at all masked positions in one pass; the loss is a
1/t-weighted cross-entropy
over masked positions (an upper bound on the negative log-likelihood).
Time-free parameterization —
t is not fed to the model.
Training setup
| |
|---|
| Dataset | FineWeb-Edu, GPT-2 BPE, uint16 memmap |
| Tokens seen | ~10.0B (76,300 iters × batch 64 × 4 grad-accum × 512 ctx) |
| Optimizer | AdamW (β1 0.9, β2 0.95, weight decay 0.1, grad clip 1.0) |
| LR schedule | WSD — warmup 1000, stable, linear decay over the final 25000 iters |
| Peak / min LR | 6e-4 → 1e-5 |
| Batch size | 64 sequences × 4 grad-accum = effective 256 sequences/step |
| Precision | bf16 autocast |
| Compilation | torch.compile(mode="default") |
| Hardware | NVIDIA DGX-Spark (GB10), single device |
| Wall-clock | ~6 days (~6.85 sec/iter steady-state) |
Results
Offline evaluation (eval.py, 500 batches) on the FineWeb-Edu validation split:
- NLL-bound: 3.38 nats/token
- Perplexity: ~29.3
LAMBADA last-word prediction (single-pass diffusion scoring, full 5153-example
test split):
- Accuracy: 36.95% (1904/5153)
- Perplexity (last-word): 55.3
Notable asymmetry vs the 150M: val PPL improved 33% (43.8 → 29.3), but LAMBADA
PPL improved 84% (358 → 55.3) and LAMBADA accuracy jumped +15.06 pp
(21.89 → 36.95). The larger jump on the discourse-dependent benchmark suggests
the added capacity is being spent on long-range representations rather than
sharper local statistics — a capability signal, not just a fitting one.
Intended use
A research / learning artifact, not a product.
- ✅ A base model to supervised-fine-tune, or to study masked-diffusion
training and sampling dynamics at the next-larger-than-toy scale.
- ❌ Not an instruction-following model. Prompt it document-style ("The history
of Rome is"), not question-style.
Sampling — important
Small diffusion LMs collapse into repetition loops under the default
low-confidence remasking sampler. Use the frequency repetition penalty (on by
default in chat.py / sample.py):
1python sample.py --ckpt nanodiff-350m-base.pt \
2 --prompt "The history of Rome is" --rep-penalty 3.0
chat.py and
sample.py enable a number of speed defaults (compile, reduced
steps, persistent inductor cache). See
the repo
README for the full table.
Limitations
- Confabulates. At ~350M params on ~10B tokens, factual recall is
unreliable. It learned English structure, not a dependable world model.
- Coherence degrades on harder prompts.
- English only; FineWeb-Edu domain.
- Base model — no alignment, no safety tuning.
- The training was still improving when it ended — the val curve hadn't
flattened. Extending to ~15B tokens would likely give another 0.05-0.10
nats; this checkpoint is not at its capacity ceiling.
Citation
Built on the LLaDA recipe:
1@article{nie2025llada,
2 title = {Large Language Diffusion Models},
3 author = {Nie, Shen and others},
4 journal = {arXiv preprint arXiv:2502.09992},
5 year = {2025}
6}
