first main run setup ready to begin

layer_0.py

@@ -0,0 +1,133 @@
+"""Layer 0: the binary input.
+
+No parameters. Owns the u8 id streams (train/val) and turns them into
+model-ready batches: x (B,T) and y (B,T) int64 tensors, where y is the
+same stream shifted by one id (next-token target).
+
+Independence: imports only numpy/torch. Knows nothing about layer 1+.
+"""
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import torch
+
+
+@dataclass
+class Layer0Config:
+    block_size: int = 1024       # T: context length handed to layer 1
+    batch_size: int = 32        # B
+    chunk_size: int = 65536    # shuffle granularity, in ids
+    seed: Optional[int] = 1337
+
+
+class Layer0:
+    """Read-only view of the binary; chunk-shuffle sampling."""
+
+    def __init__(self, cfg, train_path="data/train.bin", val_path="data/val.bin"):
+        self.cfg = cfg
+        # memmaps: fork-inheritable later, no copies in this process now
+        self.train = np.memmap(train_path, dtype=np.uint8, mode="r")
+        self.val = np.memmap(val_path, dtype=np.uint8, mode="r")
+        self.n_train = int(self.train.size)
+        self.n_val = int(self.val.size)
+        self.n_chunks = self.n_train // cfg.chunk_size   # tail partial chunk dropped
+        self.rng = np.random.default_rng(cfg.seed)
+
+    def sample_batch(self):
+        """One training batch: random chunk, random offset, x and y shifted by 1."""
+        B, T, C = self.cfg.batch_size, self.cfg.block_size, self.cfg.chunk_size
+        xs = np.empty((B, T + 1), dtype=np.uint8)
+        for b in range(B):
+            c = int(self.rng.integers(0, self.n_chunks))
+            o = int(self.rng.integers(0, C - T))          # keeps y inside the chunk
+            base = c * C + o
+            xs[b] = self.train[base: base + T + 1]
+        x = torch.tensor(xs[:, :T], dtype=torch.long)
+        y = torch.tensor(xs[:, 1:], dtype=torch.long)
+        return x, y
+
+    def val_block(self, i):
+        """The i-th contiguous block of the val stream (sequential, no shuffle)."""
+        T = self.cfg.block_size
+        base = i * T
+        assert base + T < self.n_val, f"val block {i} out of range"
+        seg = self.val[base: base + T + 1]
+        return (torch.tensor(seg[:T], dtype=torch.long),
+                torch.tensor(seg[1:], dtype=torch.long))
+
+    @property
+    def n_val_blocks(self):
+        # blocks i with base+T+1 <= n_val  ->  i <= (n_val - T - 1) // T
+        return (self.n_val - self.cfg.block_size - 1) // self.cfg.block_size + 1
+
+    @property
+    def steps_per_epoch(self):
+        return self.n_train // (self.cfg.batch_size * self.cfg.block_size)
+
+    # ---- stateful truncated BPTT mode: sequential windows over contiguous tapes ----
+
+    @property
+    def tape_len(self):
+        """Contiguous ids per row, truncated to whole windows."""
+        return (self.n_train // self.cfg.batch_size) // self.cfg.block_size * self.cfg.block_size
+
+    @property
+    def n_seq_windows(self):
+        return self.tape_len // self.cfg.block_size
+
+    def seq_windows(self):
+        """Yield (x, y) for each T-window, all B rows advancing in lockstep.
+
+        Row b walks tape b = train[b*tape_len : (b+1)*tape_len] in order.
+        Tapes are contiguous in the original stream, so the caller may carry
+        RNN state across windows and reset only at sweep end (the one real
+        discontinuity). Deterministic: no RNG.
+        """
+        B, T = self.cfg.batch_size, self.cfg.block_size
+        rows = np.arange(B, dtype=np.int64) * self.tape_len
+        cols = np.arange(T + 1, dtype=np.int64)
+        idx = rows[:, None] + cols[None, :]     # (B, T+1) base indices, built once
+        for w in range(self.n_seq_windows):
+            off = (idx + w * T).ravel()         # per-window offsets
+            xs = self.train[off].reshape(B, T + 1)   # one C-level gather, no python loop
+            yield (torch.tensor(xs[:, :T], dtype=torch.long),
+                   torch.tensor(xs[:, 1:], dtype=torch.long))
+
+
+def _self_test():
+    PROJ = Path(__file__).resolve().parents[1]
+    import sys
+    sys.path.insert(0, str(PROJ / "scripts"))
+    import tokenizer as tk
+    tk.load_vocab()
+
+    cfg = Layer0Config(batch_size=8, block_size=256)
+    l0 = Layer0(cfg)
+    x, y = l0.sample_batch()
+
+    assert x.shape == (8, 256) and y.shape == (8, 256), "batch shape"
+    assert int(x.min()) >= 0 and int(x.max()) <= 108, "ids out of vocab"
+    assert torch.equal(y[:, :255], x[:, 1:]), "shift law violated"
+
+    v0, _ = l0.val_block(0)
+    assert torch.equal(v0, torch.tensor(l0.val[:256].astype(np.int64))), "val walk"
+
+    def render(row):
+        out = bytes(row.numpy().astype(np.uint8)).translate(tk.TABLE_DEC)
+        for s, t in tk.SENT_DEC:
+            out = out.replace(s, t)
+        return tk.display(out).decode()
+
+    print("x[0] (model input):")
+    print(render(x[0]))
+    print("y[0] (model target):")
+    print(render(y[0]))
+    print(f"n_train={l0.n_train:,}  n_chunks={l0.n_chunks:,}  "
+          f"steps/epoch={l0.steps_per_epoch:,}  n_val_blocks={l0.n_val_blocks:,}")
+    print("layer0 self-test: PASS")
+
+
+if __name__ == "__main__":
+    _self_test()

layer_1.py

@@ -0,0 +1,105 @@
+"""Layer 1: the embedding - id -> vector.
+
+wte: token table    (vocab_size, d_model)   "what"
+wpe: position table (block_size, d_model)   "where"
+x = wte(ids) + wpe(pos), then optional dropout.
+
+Independence: imports only torch. Knows nothing about layer 0 (data)
+or anything above it. Exposes tie_lm_head for the (later) output layer.
+"""
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+
+
+@dataclass
+class Layer1Config:
+    vocab_size: int = 109
+    block_size: int = 1024
+    d_model: int = 512
+    init_std: float = 0.02
+    dropout: float = 0.0
+
+
+class Layer1(nn.Module):
+    def __init__(self, cfg: Layer1Config):
+        super().__init__()
+        self.cfg = cfg
+        self.wte = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.wpe = nn.Embedding(cfg.block_size, cfg.d_model)
+        self.drop = nn.Dropout(cfg.dropout)
+        self.init_weights()
+
+    def init_weights(self):
+        nn.init.normal_(self.wte.weight, mean=0.0, std=self.cfg.init_std)
+        nn.init.normal_(self.wpe.weight, mean=0.0, std=self.cfg.init_std)
+
+    def forward(self, ids, pos_offset=0):
+        """ids: (B, T) long -> (B, T, d_model). pos_offset: position embedding
+        start, used by stateful generation so wpe advances with the stream."""
+        B, T = ids.shape
+        assert pos_offset + T <= self.cfg.block_size, \
+            f"pos_offset+T={pos_offset + T} exceeds block_size {self.cfg.block_size}"
+        pos = torch.arange(T, device=ids.device) + pos_offset
+        return self.drop(self.wte(ids) + self.wpe(pos))
+
+    def tie_lm_head(self, head):
+        """Share wte with the (later) output head: head.weight IS wte.weight."""
+        head.weight = self.wte.weight
+        return head
+
+    @property
+    def n_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+
+def _self_test():
+    torch.manual_seed(0)
+    cfg = Layer1Config()
+    l1 = Layer1(cfg)
+    B, T = 4, 16
+    ids = torch.randint(0, cfg.vocab_size, (B, T))
+    x = l1(ids)
+
+    assert x.shape == (B, T, cfg.d_model), "embed shape"
+
+    # same id everywhere -> the wte part is constant across positions
+    same = torch.full((1, T), 42, dtype=torch.long)
+    xs = l1(same)
+    for p in range(T - 1):
+        d = (xs[0, p] - l1.wpe.weight[p]) - (xs[0, p + 1] - l1.wpe.weight[p + 1])
+        assert d.abs().max().item() < 1e-6, "wte part varies with position"
+
+    # same position, different ids -> the wpe part is constant
+    for b in range(B - 1):
+        d = (x[b, 0] - l1.wte(ids[b, 0])) - (x[b + 1, 0] - l1.wte(ids[b + 1, 0]))
+        assert d.abs().max().item() < 1e-6, "wpe part varies with token"
+
+    # gradients flow through both tables
+    x.sum().backward()
+    assert l1.wte.weight.grad is not None and l1.wte.weight.grad.abs().sum().item() > 0
+    assert l1.wpe.weight.grad is not None and l1.wpe.weight.grad.abs().sum().item() > 0
+
+    # param count: 109*256 + 256*256
+    assert l1.n_params == cfg.vocab_size * cfg.d_model + cfg.block_size * cfg.d_model == 93440
+
+    # weight tying is a true share, not a copy
+    head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+    l1.tie_lm_head(head)
+    assert head.weight is l1.wte.weight, "tie is not a shared tensor"
+
+    # block_size guard
+    try:
+        l1(torch.zeros((1, cfg.block_size + 1), dtype=torch.long))
+        raise SystemExit("expected assertion for T > block_size")
+    except AssertionError:
+        pass
+
+    print(f"wte {tuple(l1.wte.weight.shape)}  wpe {tuple(l1.wpe.weight.shape)}  "
+          f"params {l1.n_params:,}")
+    print("layer1 self-test: PASS")
+
+
+if __name__ == "__main__":
+    _self_test()

layer_11.py

@@ -0,0 +1,29 @@
+"""Layer 11: the output head.
+
+ln_f + lm_head over the vocab. Supports weight tying with layer 1's wte.
+The 12-layer stack: layer 0 + layer 1 + 9x layer 2 + this.
+"""
+from dataclasses import dataclass
+
+import torch.nn as nn
+
+
+@dataclass
+class Layer11Config:
+    vocab_size: int = 109
+    d_model: int = 512
+
+
+class Layer11(nn.Module):
+    def __init__(self, cfg: Layer11Config):
+        super().__init__()
+        self.cfg = cfg
+        self.ln_f = nn.LayerNorm(cfg.d_model)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+
+    def tie_to(self, layer1):
+        """Share layer 1's wte: head.weight IS wte.weight (0 extra params)."""
+        self.lm_head.weight = layer1.wte.weight
+
+    def forward(self, x):
+        return self.lm_head(self.ln_f(x))

layer_2.py

@@ -0,0 +1,40 @@
+"""Layer 2: one recurrent block (pre-LN + residual, LSTM or GRU).
+
+The 12-layer stack instantiates this 9 times: layers 2..10.
+Causal by construction - no mask needed; memory is O(T), not O(T^2).
+Hidden state resets each forward (windowed training); stateful sampling
+can pass an initial state later.
+
+Independence: imports only torch. Knows nothing about other layers.
+"""
+from dataclasses import dataclass
+
+import torch.nn as nn
+
+
+@dataclass
+class Layer2Config:
+    d_model: int = 512
+    rnn_type: str = "gru"    # "lstm" | "gru"
+    dropout: float = 0.0
+
+
+class Layer2(nn.Module):
+    def __init__(self, cfg: Layer2Config):
+        super().__init__()
+        self.cfg = cfg
+        self.ln = nn.LayerNorm(cfg.d_model)
+        if cfg.rnn_type == "lstm":
+            self.rnn = nn.LSTM(cfg.d_model, cfg.d_model, batch_first=True)
+        elif cfg.rnn_type == "gru":
+            self.rnn = nn.GRU(cfg.d_model, cfg.d_model, batch_first=True)
+        else:
+            raise ValueError(f"unknown rnn_type: {cfg.rnn_type}")
+
+    def forward(self, x, state=None):
+        out, state = self.rnn(self.ln(x), state)
+        return x + out, state
+
+    @property
+    def n_params(self):
+        return sum(p.numel() for p in self.parameters())

notio.py

@@ -0,0 +1,75 @@
+"""The assembler: layer 0 + layer 1 + 9x layer 2 + layer 11 = 12 layers.
+
+The only file that knows the whole stack. forward(ids) -> logits (B, T, vocab).
+"""
+from dataclasses import dataclass, field
+
+import torch.nn as nn
+
+from .layer_0 import Layer0, Layer0Config
+from .layer_1 import Layer1, Layer1Config
+from .layer_2 import Layer2, Layer2Config
+from .layer_11 import Layer11, Layer11Config
+
+
+@dataclass
+class NotioConfig:
+    layer0: Layer0Config = field(default_factory=Layer0Config)
+    layer1: Layer1Config = field(default_factory=Layer1Config)
+    block: Layer2Config = field(default_factory=Layer2Config)
+    head: Layer11Config = field(default_factory=Layer11Config)
+    n_blocks: int = 2   # layers 2..10; + head = layer 11 -> 12 layers total
+
+
+class Notio(nn.Module):
+    def __init__(self, cfg: NotioConfig):
+        super().__init__()
+        self.cfg = cfg
+        # cross-layer contracts (the one place they can be checked)
+        assert cfg.layer1.block_size >= cfg.layer0.block_size,             "layer 1 capacity must be >= layer 0 sequence length"
+        assert cfg.layer1.d_model == cfg.block.d_model == cfg.head.d_model,             "d_model must match everywhere"
+        assert cfg.layer1.vocab_size == cfg.head.vocab_size,             "vocab_size must match layer 1 and head"
+
+        self.layer0 = Layer0(cfg.layer0)      # data (not an nn.Module)
+        self.layer1 = Layer1(cfg.layer1)
+        self.blocks = nn.ModuleList([Layer2(cfg.block) for _ in range(cfg.n_blocks)])
+        self.head = Layer11(cfg.head)
+        self.tie_head()
+        self.init_weights()
+
+    def init_weights(self):
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.normal_(module.weight, mean=0.0, std=self.cfg.layer1.init_std)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+        # residual-path projections get the GPT-2 depth scaling
+        for name, p in self.named_parameters():
+            if name.endswith("c_proj.weight"):
+                nn.init.normal_(p, mean=0.0,
+                                std=self.cfg.layer1.init_std / (2 * self.cfg.n_blocks) ** 0.5)
+
+    def tie_head(self):
+        self.head.tie_to(self.layer1)
+
+    def forward(self, ids, states=None, pos_offset=0):
+        """ids: (B, T) long -> (logits (B, T, vocab), states list).
+
+        states: per-block recurrent state from the previous window, or None
+        for zeros. The caller decides when to detach/reset (truncated BPTT).
+        pos_offset: position embedding offset for stateful generation.
+        """
+        x = self.layer1(ids, pos_offset)
+        new_states = []
+        for i, block in enumerate(self.blocks):
+            st = None if states is None else states[i]
+            x, st = block(x, st)
+            new_states.append(st)
+        return self.head(x), new_states
+
+    @property
+    def n_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+    def sample_batch(self):
+        return self.layer0.sample_batch()

train_seq.py

@@ -0,0 +1,253 @@
+#!/usr/bin/env python3
+"""Stateful truncated BPTT training loop for the RNN stack.
+
+- layer_0.seq_windows(): B contiguous tapes walked in order
+- RNN state carried across windows but detached at every boundary:
+  gradients are truncated to one window, while information flows forward
+  through the carried state (the k1 == k2 == T form of truncated BPTT)
+- state reset to zeros only at sweep boundaries (the one real discontinuity)
+
+Usage:
+    python3 src/train_seq.py                # train (defaults: 1 epoch, cpu/cuda auto)
+    python3 src/train_seq.py --check        # prove truncation + statefulness
+    python3 src/train_seq.py --epochs 3 --lr 3e-4 --save out.pt
+"""
+import argparse
+import math
+import sys
+import time
+from itertools import islice
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+PROJ = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJ))
+sys.path.insert(0, str(PROJ / "scripts"))
+import tokenizer as tk
+from model import Notio, NotioConfig
+
+tk.load_vocab()   # also asserts vocab.txt == vocab.bin at train start
+
+
+def detach_states(states):
+    """Detach every hidden/cell tensor: no gradient flows across windows.
+    Handles LSTM (h, c) tuples and GRU bare tensors."""
+    out = []
+    for st in states:
+        out.append(tuple(s.detach() for s in st) if isinstance(st, tuple) else st.detach())
+    return out
+
+
+def ids_to_display(ids):
+    """id tensor -> human text via the runtime decoder (tags dropped/mapped)."""
+    out = bytes(ids.cpu().numpy().astype(np.uint8)).translate(tk.TABLE_DEC)
+    for sent, tok in tk.SENT_DEC:
+        out = out.replace(sent, tok)
+    return tk.display(out).decode()
+
+
+@torch.no_grad()
+def generate(m, device, max_tokens=256, temperature=0.9, top_k=0):
+    """Sample one story: <bos> prompt, carry RNN state, stop at <eos>."""
+    m.eval()
+    ids = torch.tensor([[1]], dtype=torch.long, device=device)  # <bos>
+    states = None
+    pos = 0
+    for _ in range(max_tokens):
+        logits, states = m(ids[:, -1:], states, pos_offset=pos)
+        states = detach_states(states)
+        pos += 1
+        logits = logits[:, -1, :] / max(temperature, 1e-6)
+        if top_k > 0:
+            v, _ = torch.topk(logits, top_k)
+            logits = torch.where(logits < v[:, -1:],
+                                 torch.full_like(logits, -float("inf")), logits)
+        probs = F.softmax(logits, dim=-1)
+        nxt = torch.multinomial(probs, 1)
+        ids = torch.cat([ids, nxt], dim=1)
+        if nxt.item() == 2:  # <eos>
+            break
+    m.train()
+    return ids
+
+
+def sample_and_print(m, device, n=2):
+    for i in range(n):
+        ids = generate(m, device)
+        print(f"--- sample {i} ({ids.numel()} ids) ---")
+        print(ids_to_display(ids))
+
+
+@torch.no_grad()
+def val_loss(m, device, n_blocks=20):
+    """Mean next-token loss over the first n_blocks contiguous val blocks,
+    carrying state across them (mirrors training)."""
+    m.eval()
+    states = None
+    total = 0.0
+    n = min(n_blocks, m.layer0.n_val_blocks)
+    for i in range(n):
+        x, y = m.layer0.val_block(i)
+        x, y = x.unsqueeze(0).to(device), y.unsqueeze(0).to(device)
+        logits, states = m(x, states)
+        total += F.cross_entropy(logits.view(-1, m.cfg.head.vocab_size), y.view(-1)).item()
+        states = detach_states(states)
+    m.train()
+    return total / n
+
+
+def lr_at(step, lr, warmup, total, min_lr):
+    """Linear warmup, then cosine decay to min_lr over [warmup, total)."""
+    if step < warmup:
+        return lr * step / max(warmup, 1)
+    if step >= total:
+        return min_lr
+    prog = (step - warmup) / max(total - warmup, 1)
+    return min_lr + 0.5 * (lr - min_lr) * (1 + math.cos(math.pi * prog))
+
+
+def check():
+    torch.manual_seed(0)
+    m = Notio(NotioConfig())
+    m.train()
+    it = m.layer0.seq_windows()
+    x1, y1 = next(it)
+    x2, y2 = next(it)
+
+    # 1) statefulness: carried state changes the second window's output
+    _, st = m(x1)
+    logits_carry, _ = m(x2, st)
+    logits_zero, _ = m(x2)
+    assert not torch.equal(logits_carry, logits_zero), "state is not used"
+
+    # 2) truncation: detach cuts the graph back to window 1.
+    #    (values still flow - that is the point of carrying state)
+    V = m.cfg.head.vocab_size
+    assert st[0][0].grad_fn is not None, "state should carry history to window 1"
+    st_d = detach_states(st)
+    assert st_d[0][0].grad_fn is None, "detached state should be a leaf"
+    logits_d, _ = m(x2, st_d)
+    assert torch.equal(logits_carry, logits_d), "detach changed the forward pass"
+
+    def pgrad(states_):
+        m.zero_grad()
+        logits, _ = m(x2, states_)
+        F.cross_entropy(logits.view(-1, V), y2.view(-1)).backward()
+        return sum(p.grad.detach().abs().sum() for p in m.parameters() if p.grad is not None)
+
+    gd = pgrad(st_d)   # detached: graph stops at the carried state
+    gu = pgrad(st)     # undetached: gradient also flows back through window 1
+    assert torch.equal(gd, gu) is False, "detach did not truncate the gradient path"
+
+    # 3) tape continuity: window w+1 continues window w inside a tape
+    assert torch.equal(x2[:, 0], y1[:, -1]), "tapes are not contiguous"
+
+    # 4) a few real optimizer steps with carried+detached state
+    opt = torch.optim.AdamW(m.parameters(), lr=1e-3)
+    states = None
+    it2 = m.layer0.seq_windows()
+    t0 = time.time()
+    for step in range(3):
+        x, y = next(it2)
+        logits, states = m(x, states)
+        loss = F.cross_entropy(logits.view(-1, m.cfg.head.vocab_size), y.view(-1))
+        opt.zero_grad()
+        loss.backward()
+        opt.step()
+        states = detach_states(states)
+        print(f"step {step}: loss {loss.item():.3f}")
+    toks = 3 * x.numel()
+    print(f"truncation: True | statefulness: True | tape continuity: True")
+    print(f"{toks / (time.time() - t0):,.0f} tokens/s (CPU, {m.n_params:,} params, "
+          f"T={m.cfg.layer0.block_size}, B={m.cfg.layer0.batch_size}, d={m.cfg.layer1.d_model}, "
+          f"n_blocks={m.cfg.n_blocks})")
+
+
+def main(device, epochs, lr, log_every, save, max_steps, save_every,
+         warmup, min_lr, resume, val_blocks):
+    torch.backends.cudnn.benchmark = True
+    if resume:
+        ck = torch.load(resume, map_location=device, weights_only=False)  # our own file, contains NotioConfig
+        m = Notio(ck["cfg"])
+        m.load_state_dict(ck["state_dict"])
+        m.to(device)
+        start = ck.get("step", 0)
+        print(f"resumed {resume} @ step {start} "
+              f"(loss {ck.get('loss', float('nan')):.3f}, "
+              f"val loss {ck.get('val_loss', float('nan')):.3f})")
+    else:
+        m = Notio(NotioConfig())
+        m.to(device)
+        start = 0
+    m.train()
+    opt = torch.optim.AdamW(m.parameters(), lr=lr)
+    total = max_steps if max_steps else epochs * m.layer0.n_seq_windows
+    states = None
+    step = start
+    t0 = time.time()
+    it = islice(m.layer0.seq_windows(), start, None)   # continue the sweep position
+    for epoch in range(epochs if max_steps is None else 1):
+        src = islice(it, max_steps) if max_steps is not None else it
+        for x, y in src:
+            new_lr = lr_at(step, lr, warmup, total, min_lr)
+            for g in opt.param_groups:
+                g["lr"] = new_lr
+            x, y = x.to(device), y.to(device)
+            logits, states = m(x, states)
+            loss = F.cross_entropy(logits.view(-1, m.cfg.head.vocab_size), y.view(-1))
+            opt.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0)
+            opt.step()
+            states = detach_states(states)   # truncate gradients at window length T
+            step += 1
+            if step % log_every == 0:
+                tok = step * x.numel()
+                print(f"epoch {epoch} step {step}: loss {loss.item():.3f} | "
+                      f"{tok / (time.time() - t0):,.0f} tok/s | lr {new_lr:.2e}")
+            if step % save_every == 0:
+                vl = val_loss(m, device, val_blocks)
+                print(f"val loss @ step {step}: {vl:.3f}")
+                sample_and_print(m, device)
+                if save:
+                    torch.save({"step": step, "loss": loss.item(), "val_loss": vl,
+                                "lr": new_lr, "window": step, "cfg": m.cfg,
+                                "state_dict": m.state_dict()}, save)
+                    print(f"checkpoint saved: {save} @ step {step}")
+        states = None   # sweep ended: reset at the real discontinuity
+        if max_steps is not None:
+            break
+        it = m.layer0.seq_windows()
+    total_tok = (step - start) * m.cfg.layer0.batch_size * m.cfg.layer0.block_size
+    vl = val_loss(m, device, val_blocks)
+    print(f"done: {step} steps | {total_tok / (time.time() - t0):,.0f} tok/s on {device} "
+          f"| final loss {loss.item():.3f} | val loss {vl:.3f}")
+    if save:
+        torch.save({"step": step, "loss": loss.item(), "val_loss": vl, "lr": new_lr,
+                    "window": step, "cfg": m.cfg, "state_dict": m.state_dict()}, save)
+        print(f"saved {save}")
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--check", action="store_true", help="prove truncation + statefulness, then exit")
+    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    ap.add_argument("--epochs", type=int, default=1)
+    ap.add_argument("--lr", type=float, default=4e-4)
+    ap.add_argument("--log-every", type=int, default=20)
+    ap.add_argument("--max-steps", type=int, default=None, help="benchmark: train this many steps then stop")
+    ap.add_argument("--save", default=None, help="checkpoint path (saved every --save-every steps + at end)")
+    ap.add_argument("--save-every", type=int, default=500)
+    ap.add_argument("--warmup", type=int, default=500)
+    ap.add_argument("--min-lr", type=float, default=4e-5)
+    ap.add_argument("--resume", default=None, help="checkpoint to resume from (continues sweep position)")
+    ap.add_argument("--val-blocks", type=int, default=20)
+    a = ap.parse_args()
+    if a.check:
+        check()
+    else:
+        main(a.device, a.epochs, a.lr, a.log_every, a.save, a.max_steps, a.save_every,
+             a.warmup, a.min_lr, a.resume, a.val_blocks)