Metadata-Version: 2.5
Name: baby-llm
Version: 0.1.0
Summary: A developer-first PyTorch engine for explicit, modular LLM training. Zero magic.
Project-URL: Homepage, https://github.com/700799/baby-llm
Project-URL: Repository, https://github.com/700799/baby-llm
Project-URL: Issues, https://github.com/700799/baby-llm/issues
Author: Baby LLM contributors
License-Expression: MIT
License-File: LICENSE
Keywords: fine-tuning,flash-attention,llm,pytorch,training,triton
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: >=3.11
Requires-Dist: pydantic>=2.7
Requires-Dist: torch>=2.3
Requires-Dist: typer>=0.12
Provides-Extra: dev
Requires-Dist: mypy>=1.10; extra == 'dev'
Requires-Dist: pytest>=8.0; extra == 'dev'
Requires-Dist: ruff==0.17.0; extra == 'dev'
Requires-Dist: types-pyyaml>=6.0; extra == 'dev'
Provides-Extra: flash
Requires-Dist: flash-attn>=2.5.8; extra == 'flash'
Provides-Extra: triton
Requires-Dist: triton>=3.0; extra == 'triton'
Provides-Extra: yaml
Requires-Dist: pyyaml>=6.0; extra == 'yaml'
Description-Content-Type: text/markdown

# Baby LLM

[![CI](https://github.com/700799/baby-llm/actions/workflows/ci.yml/badge.svg)](https://github.com/700799/baby-llm/actions/workflows/ci.yml)
![Python](https://img.shields.io/badge/python-3.11%20%7C%203.12-blue)
![License](https://img.shields.io/badge/license-MIT-green)

**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](#status-and-limitations) for what it deliberately isn't.

## Install

Requires **Python 3.11+**.

```bash
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:

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

| Extra | Installs | For |
|---|---|---|
| `yaml` | PyYAML | YAML configs (TOML needs nothing — it's stdlib) |
| `flash` | flash-attn | the `flash_attention_2` backend |
| `triton` | triton | writing your own Triton kernels |
| `dev` | pytest, ruff, mypy | the test suite and checks |

---

# Tutorial

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

## 1. Train something

```bash
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:

```bash
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:

```python
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:

```python
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:

```bash
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:

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

| Backend | Notes |
|---|---|
| `eager` | explicit matmul → mask → softmax reference; materializes the full score matrix |
| `sdpa` | PyTorch picks the fastest fused kernel (default) |
| `sdpa_math`, `sdpa_flash`, `sdpa_mem_efficient` | pinned to one SDPA backend — **errors** instead of silently falling back, which is what you want in a benchmark |
| `flash_attention_2` | needs `flash-attn`; fp16/bf16 only, no dense masks; GQA native |
| `flash_attention_3` | Hopper; no dropout, no dense masks |

## 4. Make it fit in less memory

Three independent knobs, each one line in the loop:

```bash
# 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
```

| Setting | Costs | Buys |
|---|---|---|
| `memory.grad_accum_steps` | wall-clock per step | larger effective batch at fixed memory |
| `memory.activation_checkpointing` | ~30% more compute | large activation-memory savings |
| `memory.activation_offloading` | PCIe traffic; CUDA only | parks activations in pinned CPU RAM |
| `training.compile` | one-off compile latency | `torch.compile` throughput |
| `model.dtype` | precision | autocast speed; params stay fp32 |

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.

```bash
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:

```toml
[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:

```bash
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:

```python
# 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.
    ...
```

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

```bash
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:

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

| Section | Key fields |
|---|---|
| `[model]` | `hidden_size`, `num_layers`, `num_heads`, `num_kv_heads` (GQA), `vocab_size`, `max_seq_len`, `ffn_hidden_size`, `rope_theta`, `norm_eps`, `tie_embeddings`, `dtype`, `init_from` |
| `[model.attention]` | `backend`, `dropout`, `bias` |
| `[optimizer]` | `name` (adamw/adam/sgd), `lr`, `weight_decay`, `betas`, `eps`, `momentum`, `fused`, `scheduler` (cosine/linear/constant), `warmup_steps` \| `warmup_ratio`, `min_lr_ratio`, `grad_clip_norm` |
| `[memory]` | `grad_accum_steps`, `activation_checkpointing`, `activation_offloading` |
| `[training]` | `micro_batch_size`, `seq_len`, `max_steps`, `seed`, `device`, `compile`, `log_interval`, `checkpoint_interval`, `output_dir`, `resume_from` |
| `[data]` | `path` (UTF-8 text; omit for the demo corpus) |

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

```bash
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

```bash
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](https://github.com/700799/baby-llm/blob/main/CONTRIBUTING.md) for the
philosophy guardrails changes are reviewed against.

## License

MIT — see [LICENSE](LICENSE).
