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