nanoDiff 50M — base
A
masked diffusion language model trained with the LLaDA recipe, from the
nanoDiff project — a minimal, hackable
"nanoGPT for diffusion LLMs" built for learning how diffusion LMs train.
This is the base model — pretrained on raw web text only. It has no
instruction tuning: it continues documents, it does not answer questions.
See Intended use below.
TL;DR
| |
|---|
| Architecture | Masked diffusion LM (LLaDA recipe), bidirectional Transformer |
| Parameters | ~88M total · ~50M non-embedding (tied input/output embeddings) |
| Training data | FineWeb-Edu, ~2.1B tokens (~1 epoch) |
| Validation | NLL-bound 3.92 nats/token · perplexity ~50.6 |
| Hardware | NVIDIA DGX-Spark (GB10), ~12 h wall-clock |
Provenance
- Code: https://github.com/BY571/nanoDiff
- Trained at commit:
de6c9a2
- Config:
pretrain/configs/50m.py
- Reproduce:
python scripts/prepare_data.py --out-dir data/fineweb_edu --num-tokens 2_000_000_000
then python pretrain/train.py --config pretrain/configs/50m.py
Architecture
A decoder-style Transformer with bidirectional self-attention (the one
architectural change from an autoregressive GPT — a diffusion LM denoises all
positions at once, so attention is not causally masked).
| Hyperparameter | Value |
|---|
| Layers | 7 |
| Attention heads | 12 |
Embedding dim (n_embd) | 768 |
Context length (block_size) | 1024 |
| Positional encoding | RoPE |
| Vocabulary | 50304 (GPT-2 BPE 50257 + [MASK]@50257 + padding to a multiple of 64) |
| Embeddings | tied (input = output) |
| Total parameters | ~88M (50M non-embedding) |
Training objective
Masked (absorbing-state) diffusion, following
LLaDA:
- Forward process: each token is independently replaced by
[MASK] with
probability t ~ U(0, 1).
- Model: predicts the clean tokens at all masked positions in one pass.
Time-free parameterization —
t is not fed to the model.
- Loss: cross-entropy over masked positions, weighted by
1/t and
normalized by sequence length (an upper bound on the negative log-likelihood).
Training setup
| |
|---|
| Dataset | FineWeb-Edu, tokenized with GPT-2 BPE, uint16 memmap |
| Tokens seen | ~2.1B (16,000 iters × batch 128 × 1024 ctx) |
| Optimizer | AdamW (β1 0.9, β2 0.95, weight decay 0.1, grad clip 1.0) |
| LR schedule | WSD — warmup 500, stable, linear decay over the final 3000 iters |
| Peak / min LR | 1.2e-3 → 1e-5 |
| Batch size | 128 sequences (no gradient accumulation) |
| Precision | bf16 autocast |
| Compilation | torch.compile (default mode) |
| Hardware | NVIDIA DGX-Spark (GB10), single device |
| Wall-clock | ~12 hours |
Results
Offline evaluation (eval.py, 500 batches) on the FineWeb-Edu validation split:
- NLL-bound: 3.92 nats/token
- Perplexity: ~50.6
Intended use
This is a research / learning artifact, not a product.
- ✅ A base model to supervised-fine-tune (the LLaDA SFT recipe masks only
the response tokens).
- ✅ Studying masked-diffusion training and sampling dynamics.
- ❌ Not an instruction-following model. Prompt it with document-shaped
text ("The capital of France is", "A recipe for bread:"), not questions.
Sampling — important
Small diffusion LMs collapse into repetition loops ("the capital of France is
the capital of France is…") under the default low-confidence remasking sampler.
This is a logit-level bias, not model weakness. Use a frequency
repetition penalty when sampling:
1python sample.py --ckpt nanodiff-50m-base.pt \
2 --prompt "The capital of France is" --rep-penalty 3.0
rep_penalty subtracts penalty × token_count from each token's logits.
chat.py and sample.py default it to 3.0. With it, the same weights
produce varied, fluent English.
Limitations
- Confabulates freely. At 50M non-embedding params on 2.1B tokens, factual
recall is unreliable ("founded by Louis XIV in 1515"). It learned English
syntax, not a world model.
- Coherence degrades on harder prompts.
- English only; FineWeb-Edu domain.
- Base model — no alignment, no safety tuning.
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}
