docs/ARCHITECTURE.md

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

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.

---

Checkpoint format

A single torch.save payload:

{
  "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