Baby LLM

A developer-first PyTorch engine for explicit, modular LLM training.

Baby LLM trains language models with a loop you can read. No Trainer class, no callback graph, no hidden hook order — a recipe is one Python file, the loop inside it is ~60 lines of raw PyTorch, and every config key maps to a line of code you can point at.

The package is baby-llm on PyPI and imports as baby_llm (Python modules can't contain hyphens). The CLI is baby-llm, or the alias babyllm.

Why this exists

Most training libraries are configuration surfaces: you describe what you want and a framework assembles the loop. That's efficient right up until you need to change the mechanics — and with LLMs, changing the mechanics is the job.

Say you want to drop in a fused Triton attention kernel, or clip gradients differently per parameter group, or find out why memory spikes on step 2. In a high-level trainer you're now reverse-engineering someone else's abstraction: which of eleven callbacks fires first, whether accelerator.backward() already unscaled your gradients, what the config key actually toggles three layers down. The code you need to understand is real; it's just not in front of you.

Baby LLM inverts that. The things you're most likely to modify — the loop, the attention kernel, gradient accumulation, the memory trade-offs — are the things written out in plain sight:

• The loop is the trainer. baby_llm/recipes/minimal_sft.py is a file you read top to bottom, then edit. Forking a recipe is the supported way to change training.
• Kernels are a registry, not a fork. One Attention module with interchangeable execution paths — SDPA, FlashAttention-2/3, or a Triton kernel you register in five lines — switchable by a string, including on a live module for A/B tests.
• Memory is a budget you set, not a mystery. Gradient accumulation, activation checkpointing, and activation offloading are each one visible construct in the loop.
• Mistakes fail at parse time, not step 40 000. Configs are Pydantic V2 models with cross-field validators, so a bad GQA head count or a flash backend in float32 is rejected before the GPU warms up.

This is aimed at people modifying training, not just running it: kernel work, custom optimizers and backward passes, teaching, and debugging. See Status and limitations for what it deliberately isn't.

Install

Requires Python 3.11+.

git clone https://github.com/700799/baby-llm baby-llm && cd baby-llm
pip install -e ".[dev,yaml]"

torch>=2.3 resolves to the default CUDA wheel, which is several GB. For CPU-only work, install torch from PyTorch's CPU index first:

pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install -e .

Tutorial

Seven steps, each one runnable. No dataset, no GPU, and no tokenizer setup needed.

1. Train something

baby-llm run --config configs/minimal_sft.toml \
  --set training.max_steps=4 --set training.checkpoint_interval=2 \
  --set training.log_interval=2 --set model.dtype=float32 \
  --set training.seq_len=64 --set training.micro_batch_size=2 \
  --set training.output_dir=runs/tutorial

recipe: minimal_sft
data: 3,988 tokens from built-in demo corpus
model: training from random init
model: 3,082,496 params | device: cpu | attention backend: sdpa
step      2/4 | loss  5.6940 | lr 2.60e-04 | grad_norm  6.118 | 1,962 tok/s
checkpoint: saved step 2 -> runs/tutorial/step_000002.pt
step      4/4 | loss  5.2823 | lr 6.95e-05 | grad_norm  4.962 | 1,734 tok/s
checkpoint: saved step 4 -> runs/tutorial/step_000004.pt
checkpoint: saved step 4 -> runs/tutorial/final.pt

(Throughput varies by machine; the losses are from a real CPU run.)

That's a real LLaMA-style decoder — RMSNorm, rotary embeddings, SwiGLU, grouped-query attention — training on a built-in demo corpus with a byte-level tokenizer. Loss starts near ln(vocab_size) (random guessing) and falls from there.

Two commands worth knowing before you go further:

baby-llm recipes                            # what recipes are registered
baby-llm validate configs/minimal_sft.toml  # parse + validate, without importing torch

validate is genuinely torch-free, which means you can check configs on a laptop or in a CI job with no ML stack installed. CI has a dedicated job that proves it stays that way.

2. Read the loop you just ran

This is the part that matters. Open baby_llm/recipes/minimal_sft.py and you'll find the whole thing — here it is, unabridged:

for step in range(start_step, train.max_steps + 1):
    lr = lr_at_step(step - 1, cfg)
    for group in optimizer.param_groups:
        group["lr"] = lr

    optimizer.zero_grad(set_to_none=True)

    # ── gradient accumulation: N micro-batches, one optimizer step ──
    micro_losses = []
    for _ in range(cfg.memory.grad_accum_steps):
        x, y = sample_batch(stream, train.micro_batch_size, train.seq_len, data_rng)
        x, y = x.to(device), y.to(device)
        with offload_ctx():
            with torch.autocast(device.type, dtype=autocast_dtype, enabled=amp_enabled):
                logits = model(x)
            loss = F.cross_entropy(logits.float().view(-1, cfg.model.vocab_size), y.view(-1))
        scaler.scale(loss / cfg.memory.grad_accum_steps).backward()
        micro_losses.append(loss.detach())

    # ── clip in true (unscaled) gradient space, then step ──
    if cfg.optimizer.grad_clip_norm is not None:
        scaler.unscale_(optimizer)
        grad_norm = torch.nn.utils.clip_grad_norm_(
            raw_model.parameters(), cfg.optimizer.grad_clip_norm
        ).item()
    scaler.step(optimizer)
    scaler.update()

Everything is in the open, including the parts other frameworks bury:

• The LR schedule is stateless. lr_at_step(step - 1, cfg) is a pure function of the step, so there's no scheduler object whose internal counter can drift out of sync with your checkpoint.
• The loss is divided by grad_accum_steps before .backward(), so accumulated gradients average rather than sum. That single / is the difference between a correct effective batch and a silently 8× learning rate.
• scaler.unscale_() runs before clipping. Clip a scaled gradient and your "max norm 1.0" is really 1.0 / scale_factor — a classic AMP bug, visible here as ordering you can check.
• logits.float() before cross-entropy keeps the softmax in fp32 under autocast.

You don't have to trust that these are right. You can see them, and change them.

3. Swap the attention kernel

The same module runs any registered backend, so you can compare them on identical weights:

from baby_llm.modules.attention import Attention, force_attention_backend

attn = Attention(hidden_size=256, num_heads=8, num_kv_heads=4, backend="sdpa")
out_fast = attn(x)

with force_attention_backend("eager"):  # explicit matmul→softmax reference
    out_reference = attn(x)

torch.testing.assert_close(out_fast, out_reference, atol=1e-5, rtol=1e-5)

Or switch the whole model from a config file, no code change:

baby-llm run --config configs/minimal_sft.toml --set model.attention.backend=eager

Registering your own kernel is the same mechanism the built-ins use — there's no private fast path:

from baby_llm.modules.attention import register_attention_backend


@register_attention_backend("my_triton_attn")
def my_triton_attn(q, k, v, *, attn_mask, dropout_p, is_causal, scale):
    # q, k, v are (batch, heads, seq, head_dim). k/v may have fewer heads than q (GQA);
    # call repeat_kv() first if your kernel doesn't handle that natively.
    return my_fused_kernel(q, k, v, causal=is_causal, sm_scale=scale)

Then --set model.attention.backend=my_triton_attn. The project's own test suite checks eager against sdpa for both forward and backward agreement — copy that test for your kernel and you'll know immediately whether your gradients are right.

4. Make it fit in less memory

Three independent knobs, each one line in the loop:

# Same effective batch (2 × 8 = 16 sequences), a quarter of the activation memory
baby-llm run --config configs/minimal_sft.toml \
  --set training.micro_batch_size=2 --set memory.grad_accum_steps=8

# Recompute each block in backward instead of storing activations
baby-llm run --config configs/minimal_sft.toml --set memory.activation_checkpointing=true

Offloading and checkpointing together are rejected at parse time — checkpointing discards the very activations offloading would move, so asking for both is a mistake worth an error rather than silent waste.

5. Interrupt it, then resume exactly

Checkpoints are written atomically (temp file → fsync → rename), so a crash mid-save can't leave you a truncated file. And they carry everything a run is: weights, optimizer moments, GradScaler state, and the data/torch RNG streams.

baby-llm run --config configs/minimal_sft.toml --set training.max_steps=6 \
  --set training.log_interval=1 \
  --set training.resume_from=runs/tutorial/step_000002.pt

resume: continuing from runs/tutorial/step_000002.pt at step 3
step      3/6 | loss  5.3751 | lr 2.32e-04 | grad_norm  5.073 | 2,000 tok/s

It picks up at step 3 — not step 1, and not step 2 again. The strength of the guarantee is worth stating plainly: 3 steps plus a resume for 3 more produces bitwise-identical weights to 6 straight steps. That's not a claim in a docstring, it's test_resume_matches_uninterrupted in tests/test_checkpointing.py, asserted with rtol=0, atol=0. No incomplete checkpoint can pass it.

Note the deliberate distinction between two things frameworks often conflate:

[training]
resume_from = "runs/tutorial/step_000002.pt"  # continue THIS run where it stopped
[model]
init_from = "runs/tutorial/final.pt"          # start a NEW run (fresh optimizer,
                                              # schedule, data order) from these weights

Setting both is a validation error, because you have to mean one of them.

6. Train on your own text

Point data.path at any UTF-8 file:

baby-llm run --config configs/minimal_sft.toml --set data.path=my_corpus.txt

The bundled tokenizer is byte-level — 256 byte values plus an EOS token. That's a deliberate demo choice: zero dependencies, zero out-of-vocabulary tokens, and trivially reversible, so nothing about tokenization is hiding. It is not what you'd train a real model with; see the limitations below.

7. Write your own recipe

There's no base class and no hook to implement. Copy the file, edit the loop, register it:

# baby_llm/recipes/my_recipe.py
def run(cfg: TrainerConfig) -> dict[str, float]:
    # Your loop, your rules. Copy the body of minimal_sft.run() and edit it.
    ...

# baby_llm/recipes/__init__.py
RECIPE_REGISTRY = {
    "minimal_sft": "baby_llm.recipes.minimal_sft",
    "my_recipe": "baby_llm.recipes.my_recipe",  # <- add this line
}

baby-llm run --recipe my_recipe --config your.toml

The registry imports lazily, so adding recipes never slows down --help. And because run() is an ordinary function taking a config, you can skip the CLI entirely:

from baby_llm.configs import TrainerConfig
from baby_llm.recipes.minimal_sft import run

cfg = TrainerConfig.from_file("configs/minimal_sft.toml")
metrics = run(cfg.apply_overrides({"training.max_steps": "50"}))

Configuration reference

One TOML (or YAML) file, five sections, each a Pydantic model in baby_llm/configs/base.py. Unknown keys are rejected, so typos fail loudly instead of being ignored.

Any field is overridable with dotted keys; values are coerced by Pydantic, so everything can be passed as a string:

baby-llm run --config configs/minimal_sft.toml \
  --set training.max_steps=500 --set model.dtype=bfloat16 --set optimizer.scheduler=linear

Caught at parse time rather than mid-run: GQA head divisibility, warmup_steps vs warmup_ratio, a flash backend with float32, seq_len exceeding the RoPE cache, resume_from with init_from, and offloading with checkpointing.

Status and limitations

Alpha (0.1.0) — the API will change. What it deliberately does not do yet:

• Not on PyPI yet. Install from source; .github/workflows/publish.yml publishes via Trusted Publishing on a v* tag.
• Single device. No FSDP, DDP, or tensor parallelism. If you need to train across 64 GPUs today, use a framework built for that.
• Demo-scale model and tokenizer. The bundled decoder is small and the byte tokenizer is a stand-in, not a real BPE.
• No Hugging Face checkpoint interop. init_from reads this project's own format, so you can't yet fine-tune a downloaded Llama/Qwen checkpoint. This is the top priority.
• generate() has no KV cache — written for clarity, so it re-prefills per token.

Planned next: HF state-dict conversion, modules/loss.py (chunked cross-entropy, z-loss), first-party Triton kernels, and a LoRA recipe.

Project layout

baby-llm/
├── pyproject.toml                # packaging, deps, `baby-llm` CLI entry point
├── configs/minimal_sft.toml      # example config, one per recipe
├── baby_llm/
│   ├── cli.py                    # typer CLI: run / recipes / validate
│   ├── configs/base.py           # Pydantic V2 schema + TOML/YAML loaders + overrides
│   ├── modules/attention.py      # backend-switchable Attention + kernel registry
│   ├── kernels/triton/           # first-party Triton kernels (same registry user code uses)
│   ├── recipes/minimal_sft.py    # explicit single-device SFT loop
│   └── py.typed                  # ships inline type hints
├── tests/                        # 76 tests + 1 network-marked
└── .github/workflows/            # ci.yml (lint, torch-free proof, test matrix)
                                  # publish.yml (TestPyPI + PyPI via Trusted Publishing)

Development

pytest -m "not network"            # full suite minus tests that download models
ruff check . && ruff format --check .
python -m mypy

CI runs exactly those, plus a job that installs the package without torch and runs the config tests — keeping the torch-free claim honest. See CONTRIBUTING.md for the philosophy guardrails changes are reviewed against.

License

MIT — see LICENSE.