# Architecture Five layers, strictly one-directional. Nothing below a layer imports from above it. ``` ┌─────────────────────────────┐ │ scripts/ │ smoke.py, e2e.py └──────────────┬──────────────┘ │ ┌───────────────────────────────────┼───────────────────────────────────┐ │ │ │ ┌─▼──────────────┐ ┌───────▼────────┐ ┌─────────▼─────────┐ │ rsi.py │ │ train.py │ │ eval.py │ │ accept/revert │◄───────────┤ loop, ckpt │───────────────►│ score, verify, │ └────────────────┘ scores │ resume, CLI │ model+tokenizer│ gate, repair │ └───────┬────────┘ └─────────┬─────────┘ │ │ ┌────────────▼─────────────┐ │ │ model.py │◄──────────────────┘ │ GQA, RoPE, RMSNorm, │ │ SwiGLU, MiniTransformer│ └────────────┬─────────────┘ │ ┌────────────────────────────────────────┼──────────────────────────┐ │ data.py │ config.py │ │ tokenizer, manifest, streams │ frozen config objects │ └────────────────────────────────────────┴──────────────────────────┘ tools.py — orthogonal: called by eval (tool scoring) and by any agent runtime ``` `tools.py` is deliberately a leaf. The model never imports it; only the evaluation loop and external callers do. That keeps the sandbox auditable in isolation. --- ## Module contracts ### `config.py` — validation at the boundary Four frozen dataclasses: `ModelConfig`, `TrainingConfig`, `EvalConfig`, `SafetyConfig`. Every one validates itself in `__post_init__`, so an invalid configuration cannot be constructed — it raises at construction, not deep inside a training step. Cross-field invariants that are enforced (and would otherwise be silent bugs): - `num_query_heads % num_kv_heads == 0` — GQA cannot broadcast otherwise. - `d_model % num_query_heads == 0` — head_dim must be integral. - `head_dim % 2 == 0` — RoPE rotates pairs of dimensions. - `max_steps > warmup_steps` — the LR schedule would divide by zero. - `precision == "fp32"` and `gradient_checkpointing is True` — the CPU contract is not negotiable. - `schedule == "cosine"` — the only implemented schedule. `ModelConfig` also owns the **parameter accounting** (`parameter_count_tied` / `parameter_count_untied`) as a closed-form formula, so the number in the README is derived, not measured. `scripts/smoke.py` asserts it against the constructed model. ### `model.py` — the network `RMSNorm` → `GroupedQueryAttention` → `SwiGLU` → `DecoderBlock` → `MiniTransformer`. `MiniTransformer.forward` returns logits directly: with tied embeddings the output projection is `hidden @ E.T` and `self.output_projection is None`. There is no bias anywhere in the projections. Attention is computed in **query blocks of 128** so the score matrix is never materialised for the full sequence. This is the single most important memory decision in the model and the reason a 8,192-token context fits the inference budget. Two guards matter for numerical safety: - rows where **no** key is allowed (fully padded) would make `softmax` produce `NaN`; they are zeroed before the softmax and re-masked after it. - masked scores are filled with `torch.finfo(dtype).min`, not `-inf`, so a `NaN` in one block cannot propagate. `generate` is a straightforward autoregressive loop with a sliding window of `config.context_length`; it extends the attention mask by one position per step and stops as soon as any sequence in the batch emits `eos_id`. There is no KV cache — see [THEORY.md](THEORY.md#inference-memory) for what that means for the memory estimate. `estimate_inference_memory_bytes()` is an **analytic** upper bound: weights + KV + hidden + one attention block + next-token logits. It is a budget-checking tool, not a measurement. See [OPERATIONS.md](OPERATIONS.md#verifying-the-memory-budget). ### `data.py` — bytes in, token streams out `SimpleTokenizer` is deliberately a byte-level tokenizer: 6 reserved ids (`pad`, `eos`, `unk`, `tool`, `python`, `shell`) and one id per byte. It is not a good tokenizer — it is a *reproducible* one, which matters more at this stage, and it makes the whole pipeline Unicode-safe with no vocab file to ship. `build_train_stream` produces fixed-length token windows from a list of examples, cycling the source, and accepts `start_token` so training can resume at an exact token offset. `build_eval_stream` pads to a fixed length instead. `validate_manifest` enforces the provenance contract: a source may only be `enabled: true` if it has a 64-hex SHA-256 and a positive token budget; a disabled source must have `token_budget: 0`. URLs must be http(s). ### `train.py` — the loop `train_model` owns the whole lifecycle: 1. move the model to the device, enable gradient checkpointing, wrap in `nn.DataParallel` when more than one CUDA device is visible; 2. build the optimizer; 3. optionally restore a checkpoint (weights, optimizer, step, token cursor, RNG states); 4. iterate the token stream, accumulating 8 sequences of 1024 tokens into one optimizer step; 5. checkpoint on the token grid, emit structured progress events, return `{step, tokens_seen, last_loss}`. All telemetry is JSON Lines on **stderr** (`train_start`, `gpu_ready`, `train_progress`, `checkpoint_saved`, `train_done`), leaving stdout for the final result object. That split is what makes the run greppable in a Kaggle/Colab log. ### `eval.py` — score, verify, gate `evaluate_model` moves the model to CPU, switches to eval mode, and per example: - **math / coding** — generate, then run `VerifierRepairLoop(max_repairs=1)`: verify the answer, and on failure feed a feedback continuation back through the model once. A score counts only if the logits were finite *and* the verifier passed. - **tool** — hand the parsed tool call to the executor; the score is the executor's `ok`. The training/eval mode flag of the model is restored in a `finally` block, so an exception mid-eval cannot leave the model stuck in eval mode. `strict_gate` returns an `EvalGateResult` that is truthy, so `if strict_gate(scores):` reads naturally. ### `rsi.py` — the policy `rsi_accept` is a pure function over two `EvalScores`. It is the only place that decides whether a change is allowed to survive. `RSIController` adds apply/accept/revert state; `run_rsi_cycle` wraps a full experiment and **always reverts in a `finally`**, so a crashed experiment cannot leave a rejected candidate installed. ### `tools.py` — protocol and sandbox `build_envelope` / `parse_envelope` define a canonical JSON envelope (`sort_keys`, compact separators) with protected `tool` / `request_id` fields, so a tool can never spoof the envelope's own metadata. `parse_tool_call` normalises a call into `{"kind", "payload"}` with `kind ∈ {python, shell}`. Execution is described in [SECURITY.md](SECURITY.md). --- ## Checkpoint format A single `torch.save` payload: ```python { "model": state_dict, # device-agnostic, always saved from the unwrapped module "optimizer": optimizer.state_dict(), "tokens_seen": int, "step": int, "rng_state": {"python": ..., "torch": ..., "cuda": [...]}, "extra": {"last_loss": float, "next_checkpoint_at": int}, "data_state": {"cursor": int}, # must equal tokens_seen, or loading raises } ``` Properties that matter operationally: - **Atomic.** Written to `.{name}.{pid}.tmp` and moved into place with `os.replace`, so a kill mid-write cannot corrupt the previous checkpoint. - **Device-agnostic.** `load_checkpoint` always maps to CPU and transparently adds or strips the `module.` prefix that `nn.DataParallel` introduces, so a CPU checkpoint resumes on a GPU run and vice versa. - **Verifiable cursor.** A checkpoint whose `data_state.cursor` disagrees with `tokens_seen` is rejected rather than silently resuming at the wrong position in the corpus. ## Data flow ``` manifest.json ──validate_manifest──► allowed sources │ examples ──encode──► token ids ──build_train_stream(start_token=…)──► [1024] windows │ ┌──────────────────────┘ ▼ [1,1023] inputs + [1,1023] labels ──► forward ──► cross_entropy ──► backward │ accumulate until 8192 tokens ──┤ ▼ set LR ─► unscale grads ─► clip(1.0) ─► AdamW.step │ every 5M tokens ─► save_checkpoint (atomic) ``` ## Extension points | To add… | Do this | Not this | | --- | --- | --- | | a corpus | add a pinned entry to `data/manifest.json` with sha256 + budget, set `enabled: true` | dropping files into the repo | | a tokenizer | implement `encode(text, add_eos=True) -> list[int]` and `decode(ids) -> str` | editing the training loop | | a model change | subclass or edit `model.py`; re-run `scripts/smoke.py` so the parameter-count assertion still holds | bypassing `ModelConfig` | | a new domain | add a branch in `evaluate_model` and a field on `EvalScores` + `EvalConfig` | adding an unvalidated score | | a new tool | extend `parse_tool_call` + `execute_tool_call` and add it to the allowlists | calling `subprocess` directly from eval | | an RSI experiment | call `run_rsi_cycle` with apply/evaluate/revert callbacks | mutating the baseline in place |