Lavender-2

Appearance

Lavender-2 ​

A from-scratch training workbench for a small decoder-only language model with a Mixture-of-Experts feed-forward path. Every component — tokenizer, architecture, data pipeline, training loop — is implemented in plain PyTorch so the full recipe is auditable in one notebook.

Initiatives ​

Lavender-2 exists to be a reproducible, readable reference implementation of the modern chat-LM recipe. Concretely, v0 aims to:

  • Ship a single-notebook pipeline (full_llm_pipeline.ipynb) that takes raw HuggingFace datasets and produces a trained, chat- capable model.
  • Use plain PyTorch for the model code — no Megatron, no Transformer Engine internals — so anyone with a basic PyTorch background can read every line.
  • Cover the full modern stack: byte-level BPE, RMSNorm, RoPE, Grouped-Query Attention, SwiGLU Mixture-of-Experts, AdamW + cosine schedule, bf16 mixed precision, KV-cache inference.
  • Be honest about scale — ship a small "dev" config that runs on a single GPU (or CPU) and a 70B "reference" config documented as a design target, not as a measured result.

Incentives ​

Why build this when open-source frontier models already exist?

  • Pedagogy. Reading the weights of Llama 3 teaches nothing about why those weights exist. A readable recipe, walked through one component at a time, does.
  • Control. A flat PyTorch codebase is easy to instrument, patch, and ablate. Frameworks that hide the training loop make research harder, not easier.
  • Provenance. Every technique in the pipeline ships with a paper and a reference-implementation pointer, so the lineage of each design decision is one click away. See techniques.md.

What's shipped at v0.2.x (current baseline) ​

  • Data ingest — ~12 HuggingFace datasets (dialogue, reasoning, Wikipedia) downloaded via download_data.py with a pinned HF cache at data/.hf_cache/.
  • Normalization — per-dataset _norm_<name> functions emit a uniform [{role, content}, ...] chat-turn schema.
  • Tokenizer — tiktoken o200k_base (the GPT-4o encoding) extended with four chat specials (<|im_start|>, <|im_sep|>, <|im_end|>, <|pad|>) for a final vocab of 200,023. No in-repo BPE training; tiktoken's vocabulary ships with the wheel. v0.1.x trained a byte-level BPE from scratch and was retired in v0.2.0.
  • Model — plain-PyTorch decoder-only Transformer with RMSNorm, RoPE, GQA, and SwiGLU Mixture-of-Experts.
  • Training — AdamW + cosine LR, bf16 AMP on GPU, gradient clipping at 1.0, MoE load-balancing auxiliary loss.
  • Inference — chat.py with KV cache, sharing the same attention code path training builds.
  • Progress visibility (v0.1.1 / v0.1.2 / v0.2.1) — every notebook stage that could run silently for more than a few seconds emits a tqdm progress bar, and every blocking call that could exceed a minute (torch.save, copy.deepcopy, large shutil.copy2, MoETransformerLM(**cfg).to(device)) is wrapped in a background Heartbeat so stdout never goes quiet long enough for an external watchdog to kill the notebook. See versions/0.1.2 and versions/0.2.1 for details.
  • Python package (v0.2.5) — every class, function, and training loop the notebook uses lives in the lavender_2/ Python package. The notebook is pure orchestration; editing a training knob touches one .py file, not a notebook cell. Per-module reference: development/modules.md.
  • Host memory budget (v0.2.2) — docs/AGENTS.md §3 formalises a 128 GB host RSS ceiling across every stage of the pipeline. lavender_2.utils.log_memory and check_memory_budget are called at every stage boundary and raise MemoryBudgetExceeded at 100%.
  • Push notifications + run log (v0.2.3) — a silent Bark push fires when host swap crosses 50% full (polled every 10 min, plus a sync poll on run start) so the owner knows the host has started paging without having to watch the terminal. Every run writes a compact logs/run-<ts>.log capped at 1 KB; the terminal view is unchanged.
  • Milestone pushes + per-run checkpoint layout (v0.2.4) — silent Bark pushes on every stage boundary (run start, data loaded, tokeniser ready, each dataset built, each epoch save, final save, load) so you know training is progressing without watching the terminal. Checkpoints now land under checkpoints/run-<YYYY-MM-DD_HH-MM-SS>/ with timestamp-named files (<stage>-<HH-MM-SS>-<counter>.pt), latest-3 retention per run, and no more _copy byte-identical redundancy. See versions/0.2.4.
  • Stdlib logging + level-aware retention + ETA in pushes (v0.2.5) — per-run logs use stdlib :mod:logging to a single logs/run-<ts>.log with no hard size cap. INFO+ lines persist forever; DEBUG (Heartbeat pings, full checkpoint paths) fade after 3 h via the new DebugRetentionFileHandler. Every periodic CheckpointManager.maybe_save now carries a per-batch ETA in its Bark push body (ETA 128.5h), and a progress: X/Y (P%) loss=L stage-ETA=Hh INFO line lands in the log every 5 minutes so the file is the canonical replay of the run. Resume guide: operations/resume.
  • Docs — this VitePress site, deployed by Cloudflare Pages from docs/.

A full component-by-component walkthrough lives under Training → v0 (the v0 label covers the pre-1.0 era, not a specific MAJOR). Current-baseline release notes are at versions/0.2.1; the v0.1.x baseline is at versions/0.1.0, incremental notes live beside it as 0.1.<digit>.md, and the archived notebook is at versions/snapshots/0.1.x/full_llm_pipeline.ipynb.

Repository invariants ​

These are the project-specific rules the agent should not break. The canonical cross-project rules live in the AGENTS.md submodule; this section lists only the Lavender-2 overrides.

  • Dataset-download streaming mode. For any dataset above ~1M rows, the registry entry in download_data.py sets streaming=False. The download loop uses memory-mapped Arrow + .select(range(limit)) in that mode, avoiding the accumulate-to- list OOM. See download_data.py:173.
  • Notebook ↔ disk schema contract. The column names written by download_data.py are the interface consumed by the _norm_* functions in the notebook. Renaming columns on either side without updating the other silently breaks training.
  • No secrets in commits. hf-token.txt is real; data/ may contain user PII from WildChat / LMSYS / Arena. The .gitignore covers both — keep it covering them.
  • Submodule is upstream-owned. Do not edit files inside the AGENTS.md/ submodule from this repo. Bumping the pointer is fine; editing in place is not.
  • Docs-only tooling stays in docs/. The site is a self- contained VitePress project deployed from docs/. No docs-related files should appear at the repo root.

Known gaps ​

Being explicit so the next agent does not waste a round rediscovering them:

  • No automated test suite. Smoke-testing is manual (re-run affected notebook cells; python download_data.py --limit 10).
  • No CI beyond Cloudflare Pages' docs build. Lint / typecheck / model-smoke jobs are not wired up.
  • The 32B reference config has never been trained end-to-end in this repo; the 32B numbers in the docs are targets, not measurements.
  • Gated datasets (lmsys_chat, arena) require a HF token and manual terms acceptance. download_data.py surfaces this as an error; it is not auto-handled.

The v0.1.0 "hand-mirrored requirements.txt" gap was closed in v0.1.1 — it is now generated by uv export and must be regenerated in the same commit as any pyproject.toml or uv.lock change. See docs/AGENTS.md §1.4.

Where to go next ​