Lavender-2

Appearance

Module reference — lavender_2/ ​

Every class, function, and training loop the training-pipeline notebook uses lives in this package. The notebook itself is orchestration-only (per docs/AGENTS.md §2) so each module evolves independently without churning notebook cells.

Source is under lavender_2/ in the repo.

Dependency graph (import-time only) ​

         ┌──────────┐
         │  utils   │   Heartbeat, seeding, memory probes, fingerprints
         └────┬─────┘
              │
         ┌────┴─────┐
         │  cache   │   CACHE_DIR, read_meta / write_meta
         └────┬─────┘
              │
   ┌──────────┼──────────┐
   ▼          ▼          ▼
tokenizer  templates   ingest
   │                      │
   └──────────┬──────────┘
              ▼
            data
              │
              ▼          ┌──────────────────┐
            model ◄──────┤  checkpoints,    │
              │          │  inference       │
              ▼          └──────────────────┘
           training

Leaf modules (utils, templates) have no internal imports beyond stdlib + numpy/torch/tqdm. Higher modules import from lower ones but never reverse; import cycles are physically impossible.

lavender_2.utils ​

What: building-block primitives used across every other module.

Public API:

  • set_seed(seed=42) — seed Python, NumPy, and PyTorch.
  • get_device() — return torch.device("cuda") if available else "cpu".
  • class Heartbeat — background keep-alive context manager. Parameters label, every_sec=60, status=callable, quiet_exit=False. Prints via tqdm.write. Use around any blocking call that can exceed 60 s.
  • log_memory(label="memory") — print RSS / available / VRAM.
  • check_memory_budget(ceiling_gb=128.0, warn_at=0.9, label) — warn at 90%, raise MemoryBudgetExceeded at 100%. Returns the current RSS in GB.
  • class MemoryBudgetExceeded(RuntimeError) — raised when RSS crosses the hard ceiling.
  • tokenizer_sha(tokenizer) -> str — short sha of the tokenizer's to_str() serialisation; cache-key input.
  • texts_fingerprint(texts) -> dict — content fingerprint of a list of strings (n, total_chars, head/tail sha16).
  • raw_ds_fingerprint(raw_ds) -> dict — same for a list of {conversation, source, weight} dicts.

Why it exists: keeps the 60-second visibility rule and the 128 GB budget rule in one file, so the rest of the package doesn't have to re-implement them.

lavender_2.cache ​

What: cache filesystem layout + atomic JSON meta IO.

Constants:

  • CACHE_DIR = "data/cache" — gitignored shard directory.
  • DATA_CACHE_INTERVAL_SEC = 1800 — 30 min between partial flushes during tokenisation.
  • WRITE_BATCH_PRETRAIN = 2000, WRITE_BATCH_SFT = 500, WRITE_BATCH_DPO = 500 — docs/conversations per flush.

Functions:

  • read_meta(path, expected_key) -> dict | None — returns meta on key match, None on miss / corruption / mismatch. Caller rebuilds on None.
  • write_meta(path, key, **extra) — atomic write via tmp + os.replace. Disk-full errors caught and logged.

lavender_2.tokenizer ​

What: tiktoken o200k_base + 4 chat specials, wrapped.

Public API:

  • TIKTOKEN_BASE = "o200k_base" — GPT-4o encoding.
  • CHAT_SPECIAL_TOKENS = ("<|im_start|>", "<|im_sep|>", "<|im_end|>", "<|pad|>").
  • build_encoding() -> tiktoken.Encoding — base + specials combined; safe to call repeatedly.
  • build_tokenizer() -> TiktokenAdapter — one-shot helper.
  • class TiktokenAdapter — HF-tokenizers-shape API: encode(text).ids, decode(ids), token_to_id(s), get_vocab_size(), save(path), from_file(path), to_str(). Frozen allow-list of special tokens = the six documented above.

Why it exists: the rest of the pipeline calls the HF API; this module is the bridge.

lavender_2.templates ​

What: chat-format rendering + weighted upsampling.

Public API:

  • conversation_to_text(example) -> str — flatten one normalised conversation into a ChatML-rendered string.
  • upsample_by_weight(train_data) -> list[str] — weighted expansion + shuffle (with Heartbeat around the shuffle).
  • effective_mix(train_data) -> list[(name, count, pct)] — post-weighting mix for the printed summary.

Template per turn: <|im_start|>{role}<|im_sep|>{content}<|im_end|>, no whitespace between turns.

lavender_2.ingest ​

What: HuggingFace dataset registry, per-source normalisers, combined loader.

Public API:

  • DATASETS: dict — 12-source registry; each entry carries path / parquet + split + weight + optional gated.
  • NORMALIZERS: dict[str, callable] — dispatch table keyed by registry name.
  • _norm_*(sample) -> list | None — per-source normalisers; emit [{role, content}, ...] or None if the sample fails minimum length / content checks. All 12 are public-by-convention (no leading underscore in the table; leading underscore on the function name follows an older convention and is kept for git-diff continuity).
  • load_all_datasets(sample_limit=-1) -> (list[dict], Counter) — iterate every registry entry, normalise, return combined train_data and per-source counts.
  • print_preflight() — probe metadata + print download / on-disk / RAM / disk totals before the long ingest loop.

lavender_2.data ​

What: three torch.utils.data.Dataset subclasses — pretrain, SFT, DPO — all backed by memory-mapped binary shards.

Public API:

  • class PretrainDataset(texts, tokenizer, max_len, eos_id) — one flat int32 token stream in pretrain_tokens.bin; chunked into (max_len + 1) windows at __getitem__.
  • class SFTDataset(raw_ds, tokenizer, max_len, pad_id, limit=-1) — parallel sft_ids.bin (int32) + sft_mask.bin (int8); loss mask is 1 only on assistant turns.
  • class DPODataset(raw_ds, tokenizer, max_len, pad_id, limit=-1) — parallel dpo_ids.bin (int32) + dpo_mask.bin (int8) with rows laid out as concat(chosen, rejected).

All three: resumable per-conversation, 30-min flush cadence, atomic meta updates. Peak construction RAM is O(write batch).

lavender_2.model ​

What: OpenMythos backbone wrapped to the notebook's legacy interface.

Public API:

  • class MoETransformerLM(**kwargs) — adapter over open_mythos.main.OpenMythos. Legacy kwargs (d_model, d_ff, num_experts, top_k, n_layers) are mapped to MythosConfig fields. Returns (logits, aux_loss); aux_loss is always 0.0 because OpenMythos folds routing loss inside MoEFFN.
  • class KVCache — thin adapter exposing seq_len on top of OpenMythos's dict cache.
  • demo_cfg(vocab_size) -> dict — single-GPU / CPU config.
  • production_32b_cfg(vocab_size) -> dict — 32B target config; requires 8× A100-80GB + FSDP / DeepSpeed ZeRO-3.
  • build_model(cfg, device) -> MoETransformerLM — instantiate under a Heartbeat and move to device.

lavender_2.training ​

What: one epoch runner per training regime + the DPO loss helpers.

Public API:

  • pretrain_epoch(model, dataloader, optimizer, *, device, pad_id, scaler, use_amp, …) -> (ce, aux)
  • sft_train_epoch(model, dataloader, optimizer, *, device, …) -> loss
  • dpo_train_epoch(policy, ref_model, dataloader, optimizer, *, device, …) -> (loss, acc)
  • sequence_log_probs(model, input_ids, labels, loss_mask) -> Tensor
  • dpo_loss(policy, ref_model, …, beta=0.1) -> (loss, margin, accuracy)

Each epoch runner: mininterval=1.0 + miniters=1 tqdm bar + first-batch timing print + slow-batch log (>60 s) + outer Heartbeat with a live batch N/total, loss=… status. Matches LLM_CHECK §1.5's "progress at least once per minute" contract.

lavender_2.checkpoints (rewritten v0.2.4) ​

What: per-run, timestamp-named checkpointing with latest-N retention and Bark milestone pushes on save / load.

Public API:

  • class CheckpointManager(model, tokenizer, model_cfg, optimizer=None, ckpt_dir="checkpoints", interval_sec=1800, retention=3, logger=None, notify=True). Creates checkpoints/run-<YYYY-MM-DD_HH-MM-SS>/ on construction; saves land inside it.
  • .set_optimizer(optimizer) — attach / switch optimizer.
  • .maybe_save(stage, **extra) — save if ≥ interval_sec since last save. Safe to call every batch.
  • .save(stage, **extra) -> str — write <stage>-<HH-MM-SS>-<counter>.pt under run_dir, enforce retention, log SAVE + RETAIN lines via logger, fire notify_milestone(...) with the filename + ETA / loss extras. Wraps torch.save in Heartbeat.
  • .load(path, device="cpu", *, notify=True) -> dict — static; single-file load under Heartbeat; fires notify_milestone on success. No copy-fallback scan (v0.2.3 had one; retired with the _copy suffix).
  • .list_runs(ckpt_dir) -> list[str] — static; every run-<ts> subdir, chronological.
  • .list_checkpoints(where) -> list[str] — static; pass a run directory or the root. Union of nested + flat (so pre-v0.2.4 legacy flat files are still discoverable for migration).

Module-level constants: CKPT_DIR = "checkpoints", CKPT_INTERVAL_SEC = 1800, CKPT_RETENTION = 3.

See docs/versions/0.2.4.md for the layout rewrite rationale and the pre-/post-v0.2.4 API migration table.

lavender_2.logging_setup (v0.2.5) ​

What: per-run stdlib :mod:logging configuration with level-aware retention.

Public API:

  • setup_logging(level=DEBUG, name="lavender_2", log_dir="logs") -> str — installs a single :class:DebugRetentionFileHandler on the root logger, returns the per-run file path. Idempotent: re-running the imports cell strips prior handlers first.
  • get_logger(suffix=None) -> logging.Logger — child logger under the lavender_2.<suffix> namespace; canonical for every module.
  • class DebugRetentionFileHandler — file handler that drops DEBUG records older than max_debug_age_sec (default 3 h). Pruning runs at most once per prune_interval_sec (default 5 min); INFO+ records are kept indefinitely.
  • Constants: LOGS_DIR = "logs", DEBUG_RETENTION_SEC = 3 h, PRUNE_INTERVAL_SEC = 5 min, LOG_FORMAT, DATE_FORMAT.

Why it exists: the v0.2.3 TrainingLogger deduped pings inline to fit a 1 KB cap. v0.2.5 trades the byte budget for level-awareness — INFO breadcrumbs persist forever, DEBUG fades after 3 h, and the rest of the package speaks the same :mod:logging API as every other Python codebase.

lavender_2.logs (deprecated v0.2.5) ​

Forwards every call to the package's stdlib logger so old notebook cells that imported TrainingLogger still parse. log() → INFO; log_heartbeat_start/end() → DEBUG. New code should import from :mod:lavender_2.logging_setup directly.

lavender_2.notify (v0.2.3, extended v0.2.4) ​

What: Bark push-notification client + background swap-memory watchdog + milestone helper.

Public API:

  • send_bark(message, *, root=".", timeout_sec=5.0, silent=True) -> (ok, fail) — URL-encode message, substitute into every *.bark.env template's $0, append ?level=passive&isArchive=0 when silent=True. Returns counts; network errors are printed as one short line and counted as fail.
  • notify_milestone(message, *, silent=True, prefix="[lavender-2] ", root=".") -> (ok, fail) — (v0.2.4) wrapper over send_bark that prepends a fixed [lavender-2] prefix (with a trailing space) for grep-ability. Called from the notebook at every stage boundary and from inside CheckpointManager.save / .load.
  • class SwapWatchdog(interval_sec=600, threshold=0.5, logger=None) — daemon poller. .start() fires one immediate poll then spins up the background thread. .stop() joins with a 1 s timeout. Fires send_bark(...) when swap.used / swap.total >= threshold.
  • BARK_ENV_GLOB = "*.bark.env" — the config-file glob matched at root.

Why it exists: long training runs shouldn't require the owner to be watching the terminal — the moment the host starts paging RAM to disk, every subsequent tokenisation / training step slows to a crawl, so a silent push lets them intervene before hours are wasted.

lavender_2.logs (v0.2.3) ​

What: per-run append-only training log with a hard 1 KB size cap.

Public API:

  • class TrainingLogger(path=None, max_bytes=1024, header=None) — creates logs/run-<ts>.log; seeded with a one-line header + an optional user header.
  • .log(msg) — free-form timestamped line.
  • .log_heartbeat_start(label) / .log_heartbeat_end(label, duration_sec) — persists START / END markers for Heartbeat blocks. Intermediate pings are deliberately not persisted.
  • Module constants: LOGS_DIR = "logs", MAX_LOG_BYTES = 1024.

Overflow behavior: when an append would exceed max_bytes, the logger rewrites the file with head + "# [...compacted N lines, round M]" + tail. Verified to stay under cap under 100-write synthetic load.

Why it exists: the owner's brief: persistent, timestamped trace of each run that survives kernel restarts, that fits in email / ticket quotes without cropping, and that does not interfere with terminal visibility (which stays byte-identical to pre-v0.2.3).

lavender_2.inference ​

What: autoregressive generation, paged KV cache, speculative decoding, interactive chat.

Public API:

  • generate(model, prompt_ids, eos_id, im_end_id, *, max_new_tokens, temperature, top_k, progress=True) -> Tensor — KV-cached sampling; stops on either im_end_id or eos_id.
  • class PagedKVCache(block_size, max_blocks, n_layers, n_kv_heads, head_dim, device) — block-structured KV pool.
  • build_draft_model(cfg, device) -> MoETransformerLM — small dense variant of the target config for speculative decoding.
  • speculative_decode(target, draft, prompt_ids, eos_id, *, …) -> Tensor — draft-and-verify loop with a tqdm bar + outer Heartbeat.
  • chat(model, tokenizer, device, eos_id, im_end_id, max_seq_len, …) — stdin-driven ChatML chat loop.

Editing workflow ​

When a training knob changes (say, you want AdamW → SOAP in pretraining):

  1. Edit lavender_2/training.py.
  2. Re-run the pretrain cell in the notebook. No notebook code change.
  3. Commit with the usual feat(training): ... / fix(training): ... per-commit message style.

When a new component lands (e.g. an evaluation harness):

  1. Add a new module lavender_2/eval.py.
  2. Re-export in lavender_2/__init__.py.
  3. Add a single orchestration cell to the notebook.
  4. Document the module in this file.