Lavender-2

Appearance

v0.2.2 — Notebook decomposed into the lavender_2 package ​

Status: current baseline. Supersedes v0.2.1 as the default main / change_tokenizer state.

Version source: pyproject.toml (version = "0.2.2").

What changed ​

v0.2.1 kept every class, function, and training loop inline in full_llm_pipeline.ipynb. Notebook was ~2,000 lines and any tweak to (say) the checkpoint manager's save path was a notebook edit.

v0.2.2 moves every class/function into a new lavender_2/ Python package and rewrites the notebook as thin orchestration. Notebook is now ~800 lines; each cell is a handful of imports + a call. The modules' behavior is byte-identical to v0.2.1 — this round is a refactor, not a behavior change.

Incidental v0.2.2 additions the refactor made easy:

  • Model-init heartbeat (lavender_2.model.build_model): wraps MoETransformerLM(**cfg).to(device) in a :class:~lavender_2.utils.Heartbeat, so the 32B path cannot go silent for more than 60 s. Closes the gap flagged by the v0.2.1 audit against LLM_CHECK §1.5.
  • Memory-budget probes (lavender_2.utils.log_memory
    • check_memory_budget): the canonical 128 GB ceiling per docs/AGENTS.md §3. Called at every stage boundary in the notebook (post-ingest, post-upsample, post-tokenise, before and after each training epoch).

Package layout ​

lavender_2/
├── __init__.py          — version, public API, no heavy imports
├── utils.py             — Heartbeat, set_seed, get_device,
│                          log_memory, check_memory_budget,
│                          tokenizer_sha, texts_fingerprint,
│                          raw_ds_fingerprint,
│                          MemoryBudgetExceeded
├── cache.py             — CACHE_DIR, DATA_CACHE_INTERVAL_SEC,
│                          WRITE_BATCH_* constants,
│                          read_meta / write_meta
├── tokenizer.py         — TiktokenAdapter, build_tokenizer,
│                          TIKTOKEN_BASE, CHAT_SPECIAL_TOKENS
├── templates.py         — conversation_to_text, upsample_by_weight,
│                          effective_mix
├── ingest.py            — DATASETS registry, NORMALIZERS table,
│                          load_all_datasets, print_preflight
├── data.py              — PretrainDataset, SFTDataset, DPODataset
├── model.py             — MoETransformerLM, KVCache,
│                          demo_cfg, production_32b_cfg, build_model
├── training.py          — pretrain_epoch, sft_train_epoch,
│                          dpo_train_epoch, dpo_loss,
│                          sequence_log_probs
├── checkpoints.py       — CheckpointManager, CKPT_DIR,
│                          CKPT_INTERVAL_SEC, CKPT_COPIES
└── inference.py         — generate, PagedKVCache,
                           build_draft_model, speculative_decode, chat

Per-module documentation: docs/development/modules.md.

Packaging ​

pyproject.toml gains a [build-system] + [tool.hatch.*] block so uv sync now installs the project as an editable Python package:

toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.metadata]
allow-direct-references = true   # for open-mythos @ git+…

[tool.hatch.build.targets.wheel]
packages = ["lavender_2"]

After uv sync, import lavender_2 works from the notebook, from any Python shell, and from a future chat.py rewrite. The [project.name] is still lavender-2 (hyphen); Python's standard hyphen→underscore rule maps it to import name lavender_2.

Host memory budget rule ​

Added as a new section in docs/AGENTS.md §3: host RAM (not VRAM) is hard-capped at 128 GB across every step of the pipeline. The notebook probes RSS at every stage boundary via log_memory(...) and calls check_memory_budget(...) which warns at 90% and raises MemoryBudgetExceeded at 100%.

Existing scripts verified against the rule:

  • Ingest (lavender_2.ingest.load_all_datasets) — keeps one streaming iterator per source; peak RAM ≈ per-source train_data accumulation. With SAMPLE_LIMIT=-1 on the full corpus this may approach 20+ GB of Python objects; set SAMPLE_LIMIT for bounded runs.
  • Tokenisation (PretrainDataset, SFTDataset, DPODataset) — on disk in int32 / int8 shards with mmap reads. Peak RAM per loop ≈ write batch (tens of MB), regardless of corpus size.
  • Training — per-batch activations; bounded by the model / sequence length / batch size config. DEMO_CFG with batch_size=16 and max_seq_len=256 stays comfortably under the ceiling on any modern box.
  • Checkpointing — torch.save serialises the state dict in place (peak ≈ 1× model weights) then shutil.copy2 streams the bytes without full materialisation.

Notebook contract (why this refactor) ​

Per the owner's v0.2.2 request: "include only high level abstracted codes and functions in llm_pipline of ipynb. (so that ipynb code will change little during each incremental updates on the model components for easy debug and understanding)".

The resulting invariants:

  1. Cell IDs are stable. Every cell that existed in v0.2.1 still exists by ID in v0.2.2; only the body shrank.
  2. Every cell is short. The longest code cell is the orchestration of the pretraining loop (~900 bytes). The median is ~400 bytes. No cell exceeds one full screen.
  3. Imports are front-loaded. The imports cell names every symbol the rest of the notebook uses. A reader scans one cell to know what is in scope.
  4. Behavior lives in lavender_2/. Editing a training-loop knob means editing one .py file, not rerunning a cell.

Cache invalidation ​

Storage format is unchanged from v0.1.4 / v0.2.0, so:

  • data/cache/pretrain_tokens.bin from v0.2.1 is still valid if tok_sha + texts_fp + eos_id match (same tiktoken tokenizer + same corpus).
  • Same for sft_* and dpo_* shards — cache keys are preserved.

No forced rebuild. First run after uv sync picks up every existing shard as-is.

What's not in this release ​

  • Not validated end-to-end on real data. Every module imports cleanly, the tokenizer roundtrips (verified: hello world → [24912, 2375] → hello world), and every edited notebook cell passes ast.parse. A full pretrain → SFT → DPO → chat run against the 12-dataset corpus has not been executed in this repo this round.
  • chat.py not updated. Still uses the v0.1.x chat tags and loads a v0.1.x-shape tokenizer.json. Will need a similar refresh to use from lavender_2 import ...; the package API is ready for it.
  • No unit tests yet. Shape and roundtrip checks are manual. Follow-up work: pytest tests that exercise each module in isolation.

Upstream references ​