presentation/src/pretrain.py

"""Phase 1: masked-language-model pretraining of the from-scratch encoder (resumable, bf16, VRAM-logged).

This is the long GPU run Nathan launches in tmux, so every choice below favours robustness:

* **Data.** pretrain.parquet is tokenised ONCE into a cached uint16 token stream
  (data/processed/pretrain_tokens.npy, [SEP] appended after every document exactly like
  smoke_test's packing) plus a JSON sidecar holding the tokenizer sha1, row counts and the
  pretrain-val split. The LAST 1% of documents by parquet order is held out as
  pretrain_val_tokens.npy and is used only for the val MLM loss. The cache is reused when the
  sidecar's tokenizer sha1 and parquet row count match; otherwise it is rebuilt and the reason
  is logged. Training sequences are contiguous 512-token windows over the stream (RoBERTa
  FULL-SENTENCES: text crosses document boundaries, no padding, no attention mask, so SDPA
  takes the Flash path). Each epoch visits every window once in a permutation seeded with
  ``seed + epoch``; the data position is therefore (epoch, sequences consumed) and resume
  reproduces the order exactly.
* **Optimisation = the measured smoke-test path.** ``mlm_mask`` and ``build_optimizer`` are
  imported from smoke_test.py (bf16 autocast, fused AdamW with biases/LayerNorm excluded from
  weight decay, dynamic 15% masking with 80/10/10 replacement, loss on masked positions only
  with the hidden states gathered before the MLM head, grad-norm clipping). Gradient
  accumulation is ``round(effective_batch_sequences / micro_batch)`` and the realised batch
  is logged (a warning when it is not exactly 256, e.g. 24 x 11 = 264).
* **Schedule.** Linear warmup over ``min(config.warmup_steps, 6% of total steps)`` then cosine
  (or linear) decay to 10% of the peak over ``epochs x steps_per_epoch`` optimizer steps. The
  schedule is a pure function of the step and is stored in every checkpoint; a resume must
  reproduce it exactly (``--max-steps`` / ``--max-hours`` cap how far THIS invocation runs,
  they never change the plan).
* **Checkpoints.** ``checkpoints/pretrain/<run>/step-NNNNNNN/{model.safetensors,
  trainer_state.pt, summary.json}`` written atomically (tmp dir + rename) every
  ``checkpoint_every_steps``, at every epoch end, on SIGINT/SIGTERM (saved at the next
  optimizer-step boundary, exit 0), when ``--max-hours`` is reached (exit 0) and when the
  VRAM hard cap is breached (exit 3: a controlled stop beats a mid-step OOM when the desktop
  spikes). ``model.safetensors`` is ``EncoderForMLM.state_dict_for_safetensors()`` (fp32 master
  weights; the tied decoder alias is dropped and restored by ``load_tied_state_dict``).
  ``trainer_state.pt`` carries optimizer, schedule, step/epoch/data cursor, every RNG state and
  a config snapshot; ``--resume`` refuses a different model shape or a different plan.
* **Logging.** Every ``log_every_steps`` optimizer steps: loss, lr, grad norm, tokens/s,
  peak allocated/reserved VRAM for the interval (peak stats reset per interval), device memory
  in use, elapsed and ETA, to the logger and to ``metrics.csv``. Val MLM loss and masked-token
  accuracy every ``eval_every_steps`` on a fixed set of val windows with fixed masks.

Hard rule honoured here: the TREC-06 holdout parquet is never opened by this script; the
data pipeline's asserted-zero holdout-vs-pretrain overlap (data_report.md) is the tripwire.

CLI::

    python src/pretrain.py train [--run-name NAME] [--resume [PATH|latest]] [--epochs N]
                                 [--max-steps N] [--max-hours H] [--micro-batch B]
                                 [--grad-checkpointing] [--limit-docs N] [--checkpoint-every N]
                                 [--log-every N] [--eval-every N] [--keep-checkpoints K] [--fresh]
    python src/pretrain.py export --run-name NAME     # final/ from the latest checkpoint
    python src/pretrain.py plot --run-name NAME       # reports/pretrain_curves.png
"""

from __future__ import annotations

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import argparse  # noqa: E402
import csv  # noqa: E402
import json  # noqa: E402
import logging  # noqa: E402
import math  # noqa: E402
import random  # noqa: E402
import re  # noqa: E402
import shutil  # noqa: E402
import signal  # noqa: E402
import sys  # noqa: E402
import time  # noqa: E402
from dataclasses import asdict, dataclass  # noqa: E402
from datetime import datetime  # noqa: E402
from pathlib import Path  # noqa: E402
from types import SimpleNamespace  # noqa: E402
from typing import Any  # noqa: E402

sys.path.insert(0, str(Path(__file__).parent))

import numpy as np  # noqa: E402
import pyarrow.parquet as pq  # noqa: E402
import torch  # noqa: E402
import torch.nn.functional as F  # noqa: E402
from safetensors.torch import load_file, save_file  # noqa: E402
from torch import Tensor  # noqa: E402

import common  # noqa: E402
from model import EncoderConfig, EncoderForMLM, count_parameters  # noqa: E402
from smoke_test import build_optimizer, mlm_mask  # noqa: E402  (the measured code path)
from tokenizer import load_tokenizer  # noqa: E402

GIB = float(2**30)
LOG = logging.getLogger("pretrain")
VAL_FRACTION = 0.01                 # last 1% of pretraining docs -> val MLM loss only
VAL_WINDOWS = 64                    # fixed windows (and fixed masks) scored at every eval
VAL_MASK_SEED_OFFSET = 7919         # eval masks never share a generator with training
ENCODE_CHUNK_DOCS = 4096
WARMUP_FRACTION = 0.06              # config comment: warmup = min(1000, 6% of total steps)
RUN_LOG_MIN_STEPS = 100             # dev runs shorter than this never touch RUN_LOG.md
STEP_DIR_RE = re.compile(r"^step-(\d{7})$")
METRICS_COLUMNS = [
    "step", "epoch", "epoch_frac", "loss", "loss_avg", "lr", "grad_norm", "tokens_per_s",
    "max_alloc_gb", "max_reserved_gb", "device_used_gb", "elapsed_s", "train_seconds_total",
    "val_mlm_loss", "val_mlm_acc",
]
EXIT_OK, EXIT_CONFIG, EXIT_VRAM, EXIT_NAN = 0, 2, 3, 4


# ================================================================================ data cache
@dataclass
class TokenCache:
    train: np.ndarray          # 1-D uint16 token stream, [SEP] after every document
    val: np.ndarray            # same, for the last 1% of documents
    meta: dict[str, Any]


def _cache_paths(processed: Path, limit_docs: int | None) -> tuple[Path, Path, Path]:
    suffix = f".dev{limit_docs}" if limit_docs else ""
    return (processed / f"pretrain_tokens{suffix}.npy",
            processed / f"pretrain_val_tokens{suffix}.npy",
            processed / f"pretrain_tokens{suffix}.json")


def _encode_ids(tok, texts: list[str]) -> list[list[int]]:
    fast = getattr(tok, "encode_batch_fast", None)
    encs = fast(texts, add_special_tokens=False) if fast else tok.encode_batch(texts, add_special_tokens=False)
    return [e.ids for e in encs]


def load_or_build_token_cache(cfg: dict, limit_docs: int | None) -> TokenCache:
    """Reuse the cached token stream when its sidecar matches the tokenizer and the parquet,
    else tokenise the corpus once (~106M tokens; progress and ETA are logged per chunk)."""
    processed = Path(cfg["paths"]["data_processed"])
    parquet = processed / "pretrain.parquet"
    tok_path = Path(cfg["paths"]["tokenizer"])
    if not parquet.exists() or not tok_path.exists():
        raise FileNotFoundError(f"need {parquet} and {tok_path}; run data_prep.py and tokenizer.py first")
    rows = pq.ParquetFile(parquet).metadata.num_rows
    n_docs = min(rows, limit_docs) if limit_docs else rows
    tok_sha1 = common.sha1_hex(tok_path.read_text(encoding="utf-8"))
    train_p, val_p, meta_p = _cache_paths(processed, limit_docs)

    reason = None
    if not (train_p.exists() and val_p.exists() and meta_p.exists()):
        reason = "cache files missing"
    else:
        meta = common.read_json(meta_p)
        if meta.get("tokenizer_sha1") != tok_sha1:
            reason = f"tokenizer sha1 changed ({meta.get('tokenizer_sha1', '?')[:12]} -> {tok_sha1[:12]})"
        elif meta.get("parquet_rows") != rows:
            reason = f"parquet row count changed ({meta.get('parquet_rows')} -> {rows})"
        elif meta.get("n_docs") != n_docs:
            reason = f"document count changed ({meta.get('n_docs')} -> {n_docs})"
    if reason is None:
        meta = common.read_json(meta_p)
        train = np.load(train_p)
        val = np.load(val_p)
        LOG.info("token cache reused: %s (%s tokens train / %s val, %d docs, tokenizer %s)",
                 train_p.name, f"{len(train):,}", f"{len(val):,}", meta["n_docs"], tok_sha1[:12])
        return TokenCache(train, val, meta)

    LOG.info("rebuilding token cache: %s", reason)
    tok = load_tokenizer(tok_path)
    sep_id, mask_id = tok.token_to_id("[SEP]"), tok.token_to_id("[MASK]")
    if sep_id is None or mask_id is None:
        raise RuntimeError("tokenizer lacks [SEP]/[MASK]")
    if tok.get_vocab_size() > np.iinfo(np.uint16).max + 1:
        raise RuntimeError("vocab exceeds uint16; widen the cache dtype")
    n_val_docs = max(1, int(n_docs * VAL_FRACTION))
    split_at = n_docs - n_val_docs
    LOG.info("tokenising %s docs from %s in chunks of %d (val = last %d docs); ~106M tokens expected for the full corpus",
             f"{n_docs:,}", parquet.name, ENCODE_CHUNK_DOCS, n_val_docs)

    train_chunks: list[np.ndarray] = []
    val_chunks: list[np.ndarray] = []
    done = chars = 0
    n_tok_train = n_tok_val = 0
    t0 = common.timer()
    for batch in pq.ParquetFile(parquet).iter_batches(batch_size=ENCODE_CHUNK_DOCS, columns=["text"]):
        if done >= n_docs:
            break
        texts = [t or "" for t in batch.column(0).to_pylist()][: n_docs - done]
        chars += sum(map(len, texts))
        for j, ids in enumerate(_encode_ids(tok, texts)):
            arr = np.asarray(ids + [sep_id], dtype=np.uint16)
            if done + j >= split_at:
                val_chunks.append(arr)
                n_tok_val += arr.size
            else:
                train_chunks.append(arr)
                n_tok_train += arr.size
        done += len(texts)
        elapsed = common.timer() - t0
        eta = elapsed / done * (n_docs - done) if done else 0.0
        LOG.info("tokenising: %s/%s docs (%.0f%%), %.1fM tokens, %.0fs elapsed, ETA %s",
                 f"{done:,}", f"{n_docs:,}", 100.0 * done / n_docs, (n_tok_train + n_tok_val) / 1e6, elapsed, fmt_dur(eta))
    train = np.concatenate(train_chunks) if train_chunks else np.zeros(0, np.uint16)
    val = np.concatenate(val_chunks) if val_chunks else np.zeros(0, np.uint16)
    encode_s = common.timer() - t0
    meta = {
        "created": datetime.now().isoformat(timespec="seconds"),
        "parquet": str(parquet), "parquet_rows": rows, "n_docs": done, "limit_docs": limit_docs,
        "n_docs_train": min(done, split_at), "n_docs_val": max(0, done - split_at),
        "val_rule": f"last {VAL_FRACTION:.0%} of documents by parquet order",
        "tokenizer": str(tok_path), "tokenizer_sha1": tok_sha1, "vocab_size": tok.get_vocab_size(),
        "sep_id": sep_id, "mask_id": mask_id, "dtype": "uint16",
        "n_tokens_train": int(train.size), "n_tokens_val": int(val.size), "chars": chars,
        "tokens_per_char": round((train.size + val.size) / max(1, chars), 5),
        "encode_seconds": round(encode_s, 1),
        "packing": "documents concatenated with [SEP] after each; windows of max_seq_len cut from the stream",
    }
    np.save(train_p, train)
    np.save(val_p, val)
    common.write_json(meta, meta_p)
    LOG.info("token cache written: %s tokens train, %s val (%.4f tok/char) in %.0fs -> %s",
             f"{train.size:,}", f"{val.size:,}", meta["tokens_per_char"], encode_s, train_p)
    return TokenCache(train, val, meta)


class WindowStore:
    """Contiguous windows over a token stream; a batch is gathered on the CPU
    (16 x 512 x 2 bytes) and moved to the device, which costs microseconds against a
    ~350 ms micro-step and keeps the 212 MB corpus out of VRAM.

    With ``prefix_id`` set (the [CLS] id) every window is ``[CLS]`` + ``seq_len - 1`` stream
    tokens, as RoBERTa prepends ``<s>`` to each packed sequence. MLM never targets [CLS]
    (specials are excluded from masking), but attention to and from position 0 is trained,
    so the classifier's pooled [CLS] state does not start fine-tuning from an untouched row.
    """

    def __init__(self, stream: np.ndarray, seq_len: int, prefix_id: int | None = None) -> None:
        self.prefix_id = prefix_id
        self.body_len = seq_len - 1 if prefix_id is not None else seq_len
        self.n = int(stream.size // self.body_len)
        self.windows = stream[: self.n * self.body_len].reshape(self.n, self.body_len)

    def batch(self, idx: np.ndarray, device: torch.device) -> Tensor:
        w = torch.from_numpy(self.windows[idx].astype(np.int64))
        if self.prefix_id is not None:
            w = torch.cat([torch.full((w.shape[0], 1), self.prefix_id, dtype=torch.int64), w], dim=1)
        return w.to(device, non_blocking=True)


def epoch_permutation(seed: int, epoch: int, n: int) -> np.ndarray:
    return np.random.default_rng(seed + epoch).permutation(n)


# ================================================================================= schedule
@dataclass
class Schedule:
    kind: str                  # "cosine" | "linear"
    peak_lr: float
    warmup_steps: int
    total_steps: int
    floor_ratio: float = 0.1   # decay target as a fraction of peak

    @classmethod
    def from_config(cls, pcfg: dict, total_steps: int) -> "Schedule":
        name = str(pcfg.get("schedule", "cosine_to_10pct")).lower()
        kind = "linear" if name.startswith("linear") else "cosine"
        m = re.search(r"to_(\d+)pct", name)
        floor = int(m.group(1)) / 100.0 if m else 0.1
        frac = float(pcfg.get("warmup_fraction", WARMUP_FRACTION))
        warmup = min(int(pcfg["warmup_steps"]), max(1, round(frac * total_steps)))
        return cls(kind, float(pcfg["lr"]), warmup, int(total_steps), floor)

    def lr_at(self, step: int) -> float:
        """LR for optimizer step ``step`` (1-based). Past ``total_steps`` the floor holds."""
        if step <= self.warmup_steps:
            return self.peak_lr * step / self.warmup_steps
        progress = min(1.0, (step - self.warmup_steps) / max(1, self.total_steps - self.warmup_steps))
        decay = 0.5 * (1.0 + math.cos(math.pi * progress)) if self.kind == "cosine" else 1.0 - progress
        return self.peak_lr * (self.floor_ratio + (1.0 - self.floor_ratio) * decay)


# ============================================================================ run bookkeeping
@dataclass
class TrainState:
    step: int = 0                    # optimizer steps completed
    epoch: int = 0                   # 0-based epoch currently being trained
    position: int = 0                # sequences consumed in the current epoch
    tokens_seen: int = 0
    train_seconds_total: float = 0.0 # wall time inside the training loop, summed over resumes
    peak_reserved_gb: float = 0.0
    peak_allocated_gb: float = 0.0
    last_loss: float | None = None
    last_val_loss: float | None = None
    last_val_acc: float | None = None
    invocations: int = 0


class StopSignal:
    """SIGINT/SIGTERM set a flag; the loop saves at the next optimizer-step boundary."""

    def __init__(self) -> None:
        self.name: str | None = None

    def install(self) -> None:
        for s in (signal.SIGINT, signal.SIGTERM):
            signal.signal(s, self._handle)

    def _handle(self, signum: int, frame: Any) -> None:
        self.name = signal.Signals(signum).name
        sys.stderr.write(f"\n[pretrain] {self.name} received: saving a checkpoint at the next optimizer step, "
                         "then exiting 0 (kill -9 aborts without saving)\n")


def fmt_dur(seconds: float) -> str:
    seconds = max(0.0, float(seconds))
    h, rem = divmod(int(seconds), 3600)
    m, s = divmod(rem, 60)
    return f"{h}h {m:02d}m {s:02d}s" if h else f"{m}m {s:02d}s"


def list_step_dirs(run_dir: Path) -> list[Path]:
    if not run_dir.exists():
        return []
    dirs = [p for p in run_dir.iterdir() if p.is_dir() and STEP_DIR_RE.match(p.name)]
    return sorted(dirs, key=lambda p: int(STEP_DIR_RE.match(p.name).group(1)))


def latest_step_dir(run_dir: Path) -> Path | None:
    dirs = list_step_dirs(run_dir)
    return dirs[-1] if dirs else None


def rng_states(mask_gen: torch.Generator) -> dict[str, Any]:
    return {
        "python": random.getstate(),
        "numpy": np.random.get_state(),
        "torch_cpu": torch.get_rng_state(),
        "torch_cuda": torch.cuda.get_rng_state_all(),
        "mask_gen": mask_gen.get_state(),
    }


def restore_rng(states: dict[str, Any], mask_gen: torch.Generator) -> None:
    random.setstate(states["python"])
    np.random.set_state(states["numpy"])
    torch.set_rng_state(states["torch_cpu"])
    torch.cuda.set_rng_state_all(states["torch_cuda"])
    mask_gen.set_state(states["mask_gen"])


def config_snapshot(mcfg: EncoderConfig, pcfg: dict, seed: int, micro_batch: int, accum: int,
                    sched: Schedule, steps_per_epoch: int, epochs: int, cache: TokenCache,
                    n_windows: int, n_val_windows: int) -> dict[str, Any]:
    model = mcfg.to_dict()
    grad_ckpt = model.pop("grad_checkpointing")
    return {
        "model": model,
        "grad_checkpointing": grad_ckpt,
        "pretrain": {k: pcfg[k] for k in ("lr", "betas", "eps", "weight_decay", "grad_clip",
                                           "mlm_probability", "mask_replace_probs", "schedule",
                                           "effective_batch_sequences")},
        "seed": seed,
        "micro_batch": micro_batch,
        "grad_accumulation": accum,
        "effective_batch": micro_batch * accum,
        "schedule": asdict(sched),
        "steps_per_epoch": steps_per_epoch,
        "epochs": epochs,
        "total_steps": sched.total_steps,
        "data": {"n_windows": n_windows, "n_val_windows": n_val_windows,
                 "tokenizer_sha1": cache.meta["tokenizer_sha1"], "n_docs": cache.meta["n_docs"],
                 "limit_docs": cache.meta.get("limit_docs"),
                 "window": f"[CLS] + {mcfg.max_seq_len - 1} stream tokens"},
    }


def check_resume_compatible(saved: dict[str, Any], current: dict[str, Any]) -> None:
    """Refuse to resume into a different model shape or a different plan; micro-batch and
    gradient checkpointing may change (the effective batch may not, so the data cursor and
    the schedule stay valid)."""
    if saved["model"] != current["model"]:
        diff = {k: (saved["model"].get(k), current["model"].get(k))
                for k in set(saved["model"]) | set(current["model"]) if saved["model"].get(k) != current["model"].get(k)}
        raise RuntimeError(f"refusing to resume: model config differs (saved, current) = {diff}")
    for key in ("pretrain", "seed", "effective_batch", "schedule", "steps_per_epoch", "epochs", "total_steps", "data"):
        if saved.get(key) != current.get(key):
            raise RuntimeError(f"refusing to resume: '{key}' differs from the checkpoint "
                               f"(saved={saved.get(key)!r}, current={current.get(key)!r}); "
                               "a resumed run must continue the same schedule and data order")
    for key in ("micro_batch", "grad_accumulation", "grad_checkpointing"):
        if saved.get(key) != current.get(key):
            LOG.warning("resume: %s changed %r -> %r (allowed; effective batch unchanged)", key, saved.get(key), current.get(key))


# ============================================================================== checkpoints
def save_checkpoint(run_dir: Path, model: EncoderForMLM, opt: torch.optim.Optimizer, sched: Schedule,
                    state: TrainState, snapshot: dict[str, Any], mask_gen: torch.Generator,
                    metrics_tail: list[dict[str, Any]], keep: int, reason: str) -> Path:
    """Atomic: everything lands in ``step-NNNNNNN.tmp`` and is renamed at the end, so a kill
    during the 4 GB write never leaves a half-written directory that ``--resume latest`` would
    pick up. Older step dirs beyond ``keep`` are pruned afterwards."""
    t0 = common.timer()
    final = run_dir / f"step-{state.step:07d}"
    tmp = run_dir / f"{final.name}.tmp"
    if tmp.exists():
        shutil.rmtree(tmp)
    tmp.mkdir(parents=True)
    sd = {k: v.detach().to("cpu", copy=True).contiguous() for k, v in model.state_dict_for_safetensors().items()}
    save_file(sd, str(tmp / "model.safetensors"),
              metadata={"format": "pt", "step": str(state.step), "config": json.dumps(model.config.to_dict())})
    del sd
    summary = {
        "step": state.step, "epoch": state.epoch, "position": state.position, "tokens_seen": state.tokens_seen,
        "train_seconds_total": round(state.train_seconds_total, 1),
        "peak_reserved_gb": round(state.peak_reserved_gb, 3), "peak_allocated_gb": round(state.peak_allocated_gb, 3),
        "last_loss": state.last_loss, "last_val_loss": state.last_val_loss, "last_val_acc": state.last_val_acc,
        "lr": sched.lr_at(max(1, state.step)), "reason": reason, "saved_at": datetime.now().isoformat(timespec="seconds"),
        "config": snapshot,
    }
    torch.save({
        "optimizer": opt.state_dict(),
        "scheduler": asdict(sched),
        "step": state.step,
        "epoch": state.epoch,
        "micro_step_in_epoch": (state.step % max(1, snapshot["steps_per_epoch"])) * snapshot["grad_accumulation"],
        "data_cursor": {"seed": snapshot["seed"], "epoch": state.epoch, "position": state.position,
                        "n_windows": snapshot["data"]["n_windows"], "rule": "numpy default_rng(seed + epoch).permutation(n_windows)"},
        "rng": rng_states(mask_gen),
        "config": snapshot,
        "state": asdict(state),
        "metrics_tail": metrics_tail[-50:],
    }, tmp / "trainer_state.pt")
    common.write_json(summary, tmp / "summary.json")
    if final.exists():
        shutil.rmtree(final)
    tmp.rename(final)
    pruned = []
    for old in list_step_dirs(run_dir)[:-keep] if keep > 0 else []:
        shutil.rmtree(old)
        pruned.append(old.name)
    LOG.info("checkpoint saved: %s (%s) in %.1fs%s", final, reason, common.timer() - t0,
             f"; pruned {pruned}" if pruned else "")
    return final


def load_checkpoint_into(step_dir: Path, model: EncoderForMLM, opt: torch.optim.Optimizer,
                         mask_gen: torch.Generator, device: torch.device) -> tuple[dict[str, Any], TrainState]:
    t0 = common.timer()
    weights = load_file(str(step_dir / "model.safetensors"), device=str(device))
    model.load_tied_state_dict(weights, strict=True)
    del weights
    saved = torch.load(step_dir / "trainer_state.pt", map_location="cpu", weights_only=False)
    opt.load_state_dict(saved["optimizer"])
    restore_rng(saved["rng"], mask_gen)
    state = TrainState(**saved["state"])
    LOG.info("resumed from %s: step %d, epoch %d, position %d sequences, %.1fs of training so far (load took %.1fs)",
             step_dir, state.step, state.epoch, state.position, state.train_seconds_total, common.timer() - t0)
    return saved, state


# ================================================================================= metrics
class MetricsCSV:
    def __init__(self, path: Path) -> None:
        self.path = path

    def truncate_after(self, step: int) -> None:
        """Drop rows past the resumed step so a run killed between checkpoints leaves no
        duplicate or orphan steps in the record."""
        if not self.path.exists():
            return
        with self.path.open("r", encoding="utf-8", newline="") as fh:
            rows = list(csv.DictReader(fh))
        keep = [r for r in rows if r.get("step", "").isdigit() and int(r["step"]) <= step]
        if len(keep) != len(rows):
            LOG.warning("metrics.csv: dropping %d rows past step %d from an earlier, unsaved stretch", len(rows) - len(keep), step)
            self._write_all(keep)

    def _write_all(self, rows: list[dict[str, Any]]) -> None:
        with self.path.open("w", encoding="utf-8", newline="") as fh:
            w = csv.DictWriter(fh, fieldnames=METRICS_COLUMNS, extrasaction="ignore")
            w.writeheader()
            w.writerows(rows)

    def append(self, row: dict[str, Any]) -> None:
        new = not self.path.exists() or self.path.stat().st_size == 0
        with self.path.open("a", encoding="utf-8", newline="") as fh:
            w = csv.DictWriter(fh, fieldnames=METRICS_COLUMNS, extrasaction="ignore")
            if new:
                w.writeheader()
            w.writerow({k: ("" if row.get(k) is None else row[k]) for k in METRICS_COLUMNS})

    def read(self) -> list[dict[str, str]]:
        if not self.path.exists():
            return []
        with self.path.open("r", encoding="utf-8", newline="") as fh:
            return list(csv.DictReader(fh))


# ================================================================================== eval
@torch.no_grad()
def evaluate_mlm(model: EncoderForMLM, val_store: WindowStore, val_idx: np.ndarray, micro_batch: int,
                 device: torch.device, mask_seed: int, pool_like: Any, vocab: int, p_mask: float,
                 replace_probs: list[float]) -> tuple[float, float]:
    """Mean MLM loss and masked-token accuracy on fixed val windows with fixed masks (the
    generator is re-seeded every call, so successive evals score identical inputs)."""
    was_training = model.training
    model.eval()
    gen = torch.Generator(device=device).manual_seed(mask_seed)
    loss_sum = torch.zeros((), device=device, dtype=torch.float64)
    correct = torch.zeros((), device=device, dtype=torch.int64)
    n = 0
    for i in range(0, len(val_idx), micro_batch):
        ids = val_store.batch(val_idx[i:i + micro_batch], device)
        inputs, labels, pos = mlm_mask(ids, pool_like, vocab, p_mask, replace_probs, gen)
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            logits = model(inputs, attention_mask=None, masked_positions=pos)
        loss_sum += F.cross_entropy(logits.float(), labels, reduction="sum").double()
        correct += (logits.argmax(dim=-1) == labels).sum()
        n += int(labels.numel())
    if was_training:
        model.train()
    return float(loss_sum) / max(1, n), float(correct) / max(1, n)


# ============================================================================== final / log
def write_final(run_dir: Path, step_dir: Path, cfg: dict, status: str) -> Path:
    """``final/`` per the artifact contract: model.safetensors (copied from the chosen step
    dir), config.json with the model config + tokenizer pointer + training summary, and a
    copy of tokenizer.json so the directory is self-contained for finetune.py."""
    final = run_dir / "final"
    final.mkdir(parents=True, exist_ok=True)
    shutil.copy2(step_dir / "model.safetensors", final / "model.safetensors")
    tok_path = Path(cfg["paths"]["tokenizer"])
    if tok_path.exists():
        shutil.copy2(tok_path, final / "tokenizer.json")
    summary_p = step_dir / "summary.json"
    if summary_p.exists():
        s = common.read_json(summary_p)
    else:
        ts = torch.load(step_dir / "trainer_state.pt", map_location="cpu", weights_only=False)
        s = {**ts["state"], "config": ts["config"], "lr": None, "reason": "trainer_state.pt"}
    snap = s["config"]
    run_summary_p = run_dir / "run_summary.json"
    run_summary = common.read_json(run_summary_p) if run_summary_p.exists() else {}
    training = {
        "run_name": run_dir.name, "status": status, "source_checkpoint": step_dir.name,
        "steps": s["step"], "total_steps_planned": snap["total_steps"], "steps_per_epoch": snap["steps_per_epoch"],
        "epochs_planned": snap["epochs"], "epochs_completed": round(s["step"] / max(1, snap["steps_per_epoch"]), 3),
        "tokens_seen": s["tokens_seen"], "total_training_seconds": s["train_seconds_total"],
        "total_training_time": fmt_dur(s["train_seconds_total"]),
        "final_loss": s.get("last_loss"), "val_loss": s.get("last_val_loss"), "val_masked_acc": s.get("last_val_acc"),
        "peak_vram_reserved_gb": s.get("peak_reserved_gb"), "peak_vram_allocated_gb": s.get("peak_allocated_gb"),
        "micro_batch": snap["micro_batch"], "grad_accumulation": snap["grad_accumulation"], "effective_batch": snap["effective_batch"],
        "precision": "bf16_autocast", "optimizer": "AdamW(fused) " + json.dumps(snap["pretrain"]),
        "schedule": snap["schedule"], "mean_tokens_per_s": run_summary.get("mean_tokens_per_s"),
        "data": snap["data"], "weights_dtype": "float32", "exported_at": datetime.now().isoformat(timespec="seconds"),
    }
    common.write_json({"model": EncoderConfig.from_dict({**snap["model"], "grad_checkpointing": False}).to_dict(),
                       "tokenizer": "tokenizer.json", "training": training}, final / "config.json")
    LOG.info("final/ written from %s -> %s", step_dir.name, final)
    return final


def next_run_log_number(path: Path) -> int:
    nums = [int(m.group(1)) for m in re.finditer(r"^## (\d+)\.", path.read_text(encoding="utf-8"), re.M)] if path.exists() else []
    return (max(nums) + 1) if nums else 0


def run_log_block(run_dir: Path, snap: dict[str, Any], state: TrainState, status: str, mean_tok_s: float,
                  logfile: str | None, counts: dict[str, Any]) -> str:
    n = next_run_log_number(common.project_root() / "RUN_LOG.md")
    sch = snap["schedule"]
    lines = [
        f"## {n}. {datetime.now():%Y-%m-%d %H:%M} — MLM pretraining run `{run_dir.name}` ({status})",
        "",
        f"**Run.** `python src/pretrain.py train --run-name {run_dir.name}`; {counts['total']:,} parameters (EncoderForMLM, tied head). "
        f"micro-batch {snap['micro_batch']} x accumulation {snap['grad_accumulation']} = {snap['effective_batch']} sequences x "
        f"{snap['model']['max_seq_len']} tokens per step; bf16 autocast, fused AdamW {snap['pretrain']['betas']} eps {snap['pretrain']['eps']} "
        f"wd {snap['pretrain']['weight_decay']} (no decay on biases/LayerNorm), clip {snap['pretrain']['grad_clip']}; "
        f"{sch['kind']} schedule peak {sch['peak_lr']:.1e} -> {sch['floor_ratio']:.0%} over {sch['total_steps']} planned steps "
        f"({snap['epochs']} epochs x {snap['steps_per_epoch']} steps) with {sch['warmup_steps']} warmup steps; "
        f"dynamic 15% masking 80/10/10, loss on masked positions only; {snap['data']['n_windows']:,} packed 512-token windows "
        f"({snap['data']['n_docs']:,} docs), {snap['data']['n_val_windows']:,} val windows.",
        "",
        f"**Result.** {state.step} optimizer steps ({state.step / max(1, snap['steps_per_epoch']):.2f} epochs), "
        f"{state.tokens_seen / 1e6:.1f}M tokens seen, total training time {fmt_dur(state.train_seconds_total)} "
        f"({state.train_seconds_total:,.0f} s) over {state.invocations} invocation(s), mean {mean_tok_s:,.0f} tokens/s; "
        f"peak VRAM reserved {state.peak_reserved_gb:.2f} GiB / allocated {state.peak_allocated_gb:.2f} GiB; "
        f"final train MLM loss {state.last_loss if state.last_loss is None else f'{state.last_loss:.4f}'}, "
        f"val MLM loss {state.last_val_loss if state.last_val_loss is None else f'{state.last_val_loss:.4f}'} "
        f"(masked-token accuracy {state.last_val_acc if state.last_val_acc is None else f'{100 * state.last_val_acc:.1f}%'}).",
        "",
        f"**Artifacts.** {run_dir.relative_to(common.project_root()) if run_dir.is_relative_to(common.project_root()) else run_dir}/ "
        f"(metrics.csv, run_summary.json, step-*/ and final/); log {logfile or 'logs/pretrain-*.log'}.",
    ]
    return "\n".join(lines)


def current_logfile() -> str | None:
    for h in LOG.handlers:
        if isinstance(h, logging.FileHandler):
            return h.baseFilename
    return None


def smoke_test_note(cfg: dict) -> str:
    rep = Path(cfg["paths"]["reports"]) / "smoke_test.json"
    try:
        rows = {r["micro_batch"]: r for r in common.read_json(rep)["finder"]["no_ckpt"] if r.get("ok")}
        parts = [f"{b}: {rows[b]['max_reserved_gib']:.2f} GiB reserved, {rows[b]['tokens_per_s']:,} tok/s" for b in sorted(rows)]
        return f"smoke_test measured (no checkpointing) micro-batch {'; '.join(parts)} ({rep.name})"
    except Exception:  # noqa: BLE001 - informational only
        return "smoke_test measured 13.65 GiB reserved at micro-batch 16 and 17.65 GiB at 24 (reports/smoke_test.json)"


# ================================================================================== train
def resolve_run_and_resume(args: argparse.Namespace, cfg: dict) -> tuple[Path, Path | None]:
    base = Path(cfg["paths"]["checkpoints"]) / "pretrain"
    if args.resume is None:
        return base / args.run_name, None
    if args.resume == "latest":
        run_dir = base / args.run_name
        step_dir = latest_step_dir(run_dir)
        if step_dir is None:
            raise FileNotFoundError(f"--resume: no step-* checkpoint under {run_dir}")
        return run_dir, step_dir
    p = Path(args.resume).expanduser().resolve()
    if (p / "trainer_state.pt").exists():
        return p.parent, p
    step_dir = latest_step_dir(p)
    if step_dir is None:
        raise FileNotFoundError(f"--resume {p}: neither a step dir nor a run dir with step-* checkpoints")
    return p, step_dir


def cmd_train(args: argparse.Namespace, cfg: dict) -> int:
    t_invocation = common.timer()
    global LOG
    LOG = common.setup_logging("pretrain", cfg["paths"]["logs"])
    device_info = common.print_device_info()
    if not torch.cuda.is_available():
        LOG.error("CUDA is required for pretraining")
        return EXIT_CONFIG
    LOG.info("PYTORCH_CUDA_ALLOC_CONF=%s | torch %s | %s", os.environ.get("PYTORCH_CUDA_ALLOC_CONF"), torch.__version__, device_info["device"])
    seed = int(cfg["seed"])
    common.set_seed(seed)
    device = torch.device("cuda")
    pcfg, vr = cfg["pretrain"], cfg["vram"]
    hard_cap = float(vr["hard_cap_gb"])
    mcfg = EncoderConfig.from_dict(cfg["model"])
    if args.grad_checkpointing:
        mcfg.grad_checkpointing = True
    L = mcfg.max_seq_len

    run_dir, resume_dir = resolve_run_and_resume(args, cfg)
    if resume_dir is None and list_step_dirs(run_dir):
        if args.fresh:
            LOG.warning("--fresh: removing existing checkpoints and metrics in %s", run_dir)
            shutil.rmtree(run_dir)
        else:
            LOG.error("%s already holds checkpoints %s; pass --resume to continue it, --fresh to start over, or a new --run-name",
                      run_dir, [p.name for p in list_step_dirs(run_dir)])
            return EXIT_CONFIG
    run_dir.mkdir(parents=True, exist_ok=True)
    saved_summary = common.read_json(resume_dir / "summary.json") if resume_dir and (resume_dir / "summary.json").exists() else None

    # -------------------------------------------------------------- batch geometry
    if args.micro_batch:
        micro_batch = int(args.micro_batch)
    elif saved_summary:
        micro_batch = int(saved_summary["config"]["micro_batch"])
    else:
        micro_batch = int(pcfg["micro_batch_start"])
    eff_target = int(pcfg["effective_batch_sequences"])
    accum = max(1, round(eff_target / micro_batch))
    eff_batch = accum * micro_batch
    LOG.info("micro-batch %d x accumulation %d = effective %d sequences (%s tokens/step); %s",
             micro_batch, accum, eff_batch, f"{eff_batch * L:,}", smoke_test_note(cfg))
    if eff_batch != eff_target:
        LOG.warning("realised effective batch %d != configured %d because %d is not a divisor of %d",
                    eff_batch, eff_target, micro_batch, eff_target)

    # ------------------------------------------------------------------------ data
    # Refuse a mismatched --limit-docs BEFORE tokenising anything (resuming a dev run without
    # --limit-docs would otherwise build the full 106M-token cache and then be refused).
    if saved_summary is not None:
        saved_limit = saved_summary["config"]["data"].get("limit_docs")
        if args.limit_docs is None and saved_limit is not None:
            LOG.warning("--resume: using the checkpoint's --limit-docs %s", saved_limit)
            args.limit_docs = saved_limit
        elif args.limit_docs != saved_limit:
            LOG.error("--resume: --limit-docs %s differs from the checkpoint's %s; a resumed run must use the same data",
                      args.limit_docs, saved_limit)
            return EXIT_CONFIG
    cache = load_or_build_token_cache(cfg, args.limit_docs)
    cls_id = int(cfg["tokenizer"]["special_tokens"].index("[CLS]"))
    train_store, val_store = WindowStore(cache.train, L, prefix_id=cls_id), WindowStore(cache.val, L, prefix_id=cls_id)
    steps_per_epoch = train_store.n // eff_batch
    if steps_per_epoch < 1:
        LOG.error("only %d windows of %d tokens; need at least %d for one optimizer step", train_store.n, L, eff_batch)
        return EXIT_CONFIG
    epochs = int(args.epochs) if args.epochs else (int(saved_summary["config"]["epochs"]) if saved_summary else int(pcfg["epochs"]))
    total_steps = epochs * steps_per_epoch
    sched = Schedule.from_config(pcfg, total_steps)
    LOG.info("data: %s train windows (%s tokens) -> %d steps/epoch x %d epochs = %d planned steps; %s val windows "
             "(%d scored per eval); %d leftover windows per epoch are skipped",
             f"{train_store.n:,}", f"{train_store.n * L:,}", steps_per_epoch, epochs, total_steps,
             f"{val_store.n:,}", min(VAL_WINDOWS, val_store.n), train_store.n - steps_per_epoch * eff_batch)
    LOG.info("schedule: %s warmup %d steps (= min(config %d, %.0f%% of %d)) -> peak %.2e -> floor %.0f%% at step %d",
             sched.kind, sched.warmup_steps, int(pcfg["warmup_steps"]), 100 * float(pcfg.get("warmup_fraction", WARMUP_FRACTION)), total_steps,
             sched.peak_lr, 100 * sched.floor_ratio, sched.total_steps)
    val_idx = np.sort(np.random.default_rng(seed + VAL_MASK_SEED_OFFSET).choice(val_store.n, size=min(VAL_WINDOWS, val_store.n), replace=False)) if val_store.n else np.zeros(0, np.int64)

    tok = load_tokenizer(cfg["paths"]["tokenizer"])
    special_ids = [i for i in (tok.token_to_id(s) for s in cfg["tokenizer"]["special_tokens"]) if i is not None]
    pool_like = SimpleNamespace(special_ids=special_ids, mask_id=tok.token_to_id("[MASK]"))
    assert pool_like.mask_id == cache.meta["mask_id"], "tokenizer [MASK] id differs from the token cache"
    p_mask, replace_probs = float(pcfg["mlm_probability"]), [float(x) for x in pcfg["mask_replace_probs"]]

    # ----------------------------------------------------------------------- model
    common.set_seed(seed)
    with torch.device(device):
        model = EncoderForMLM(mcfg)
    model.train()
    counts = count_parameters(model)
    print(f"EncoderForMLM parameters: {counts['total']:,} total | embeddings {counts['embeddings']:,} | "
          f"encoder layers {counts['encoder_layers']:,} | final norm {counts['final_norm']:,} | MLM head {counts['head']:,}"
          f" | grad checkpointing {'ON' if mcfg.grad_checkpointing else 'off'}")
    LOG.info("model: %s params %s", f"{counts['total']:,}", {k: v for k, v in mcfg.to_dict().items()})
    opt = build_optimizer(model, sched.peak_lr, pcfg)
    LOG.info("optimizer: %s fused=%s groups(decay=%d tensors, no_decay=%d tensors) betas=%s eps=%s wd=%s clip=%s",
             type(opt).__name__, opt.defaults.get("fused"), len(opt.param_groups[0]["params"]), len(opt.param_groups[1]["params"]),
             pcfg["betas"], pcfg["eps"], pcfg["weight_decay"], pcfg["grad_clip"])
    mask_gen = torch.Generator(device=device).manual_seed(seed)
    snapshot = config_snapshot(mcfg, pcfg, seed, micro_batch, accum, sched, steps_per_epoch, epochs, cache, train_store.n, val_store.n)
    metrics = MetricsCSV(run_dir / "metrics.csv")

    state = TrainState()
    if resume_dir is not None:
        saved, state = load_checkpoint_into(resume_dir, model, opt, mask_gen, device)
        check_resume_compatible(saved["config"], snapshot)
        metrics.truncate_after(state.step)
    run_end = min(total_steps, int(args.max_steps)) if args.max_steps else total_steps
    if state.step >= run_end:
        LOG.warning("nothing to do: step %d already >= %d (planned %d, --max-steps %s); run_summary.json left untouched",
                    state.step, run_end, total_steps, args.max_steps)
        return EXIT_OK
    state.invocations += 1
    common.write_json({**asdict(state), "config": snapshot, "status": "running", "started": datetime.now().isoformat(timespec="seconds")},
                      run_dir / "run_summary.json")
    max_seconds = float(args.max_hours) * 3600.0 if args.max_hours else None
    log_every = int(args.log_every or pcfg["log_every_steps"])
    eval_every = int(pcfg["eval_every_steps"] if args.eval_every is None else args.eval_every)
    ckpt_every = int(args.checkpoint_every or pcfg["checkpoint_every_steps"])
    keep = int(args.keep_checkpoints)
    LOG.info("run %s: steps %d -> %d of %d planned%s; log every %d, eval every %s, checkpoint every %d (keep %d), VRAM hard cap %.0f GiB",
             run_dir.name, state.step, run_end, total_steps, f", --max-hours {args.max_hours}" if max_seconds else "",
             log_every, eval_every or "never", ckpt_every, keep, hard_cap)

    stop = StopSignal()
    stop.install()
    perm = epoch_permutation(seed, state.epoch, train_store.n)
    LOG.info("data cursor: epoch %d, position %d, next window index %s (permutation seed %d)",
             state.epoch, state.position, int(perm[state.position]) if state.position < train_store.n else "n/a", seed + state.epoch)

    # ------------------------------------------------------------------- the loop
    tokens_per_step = eff_batch * L
    metrics_tail: list[dict[str, Any]] = []
    last_saved_step = state.step if resume_dir is not None else -1
    status, exit_code = "running", EXIT_OK
    nonfinite_streak = 0
    interval_losses: list[float] = []
    interval_steps = 0
    pause_s = 0.0                     # eval + checkpoint time, excluded from tokens/s
    run_tokens = 0
    run_train_s = 0.0

    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    t_loop0 = common.timer()          # starts AFTER checkpoint load: first tokens/s is clean
    t_row = t_loop0
    LOG.info("training loop started (setup took %.1fs)", t_loop0 - t_invocation)

    def write_row(val: tuple[float, float] | None, lr: float, grad_norm: float, loss: float) -> None:
        nonlocal t_row, interval_losses, interval_steps, pause_s, run_tokens, run_train_s
        torch.cuda.synchronize()
        now = common.timer()
        wall = now - t_row
        active = max(1e-6, wall - pause_s)
        tok_s = interval_steps * tokens_per_step / active
        run_tokens += interval_steps * tokens_per_step
        run_train_s += active
        vram = common.vram_stats()
        free, total = torch.cuda.mem_get_info()
        state.peak_reserved_gb = max(state.peak_reserved_gb, vram["max_reserved_gb"])
        state.peak_allocated_gb = max(state.peak_allocated_gb, vram["max_allocated_gb"])
        elapsed = now - t_loop0
        state.train_seconds_total += wall
        row = {
            "step": state.step, "epoch": (state.step - 1) // steps_per_epoch + 1,
            "epoch_frac": round(state.step / steps_per_epoch, 4),
            "loss": round(loss, 5), "loss_avg": round(sum(interval_losses) / max(1, len(interval_losses)), 5),
            "lr": f"{lr:.6e}", "grad_norm": round(grad_norm, 4), "tokens_per_s": round(tok_s),
            "max_alloc_gb": round(vram["max_allocated_gb"], 3), "max_reserved_gb": round(vram["max_reserved_gb"], 3),
            "device_used_gb": round((total - free) / GIB, 3), "elapsed_s": round(elapsed, 1),
            "train_seconds_total": round(state.train_seconds_total, 1),
            "val_mlm_loss": None if val is None else round(val[0], 5), "val_mlm_acc": None if val is None else round(val[1], 5),
        }
        metrics.append(row)
        metrics_tail.append(row)
        del metrics_tail[:-50]
        remaining = run_end - state.step
        eta = remaining * (active / max(1, interval_steps))
        LOG.info("step %d/%d ep %d | loss %.4f (avg %.4f) | lr %.2e | gnorm %.2f | %s tok/s | alloc %.2f GiB res %.2f GiB used %.2f GiB"
                 " | elapsed %s | ETA %s%s",
                 state.step, run_end, row["epoch"], loss, row["loss_avg"], lr, grad_norm, f"{tok_s:,.0f}",
                 vram["max_allocated_gb"], vram["max_reserved_gb"], row["device_used_gb"], fmt_dur(elapsed), fmt_dur(eta),
                 "" if val is None else f" | val_mlm_loss {val[0]:.4f} val_acc {100 * val[1]:.2f}%")
        t_row = now
        interval_losses, interval_steps, pause_s = [], 0, 0.0
        torch.cuda.reset_peak_memory_stats()

    def checkpoint(reason: str) -> None:
        nonlocal last_saved_step, pause_s
        if last_saved_step == state.step:
            return
        t0 = common.timer()
        save_checkpoint(run_dir, model, opt, sched, state, snapshot, mask_gen, metrics_tail, keep, reason)
        last_saved_step = state.step
        pause_s += common.timer() - t0

    while state.step < run_end:
        lr = sched.lr_at(state.step + 1)
        for g in opt.param_groups:
            g["lr"] = lr
        micro_losses = []
        phase = "forward/backward"
        try:
            for _ in range(accum):
                idx = perm[state.position:state.position + micro_batch]
                state.position += micro_batch
                ids = train_store.batch(idx, device)
                inputs, labels, pos = mlm_mask(ids, pool_like, mcfg.vocab_size, p_mask, replace_probs, mask_gen)
                with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                    logits = model(inputs, attention_mask=None, masked_positions=pos)
                    loss = F.cross_entropy(logits, labels)
                (loss / accum).backward()
                micro_losses.append(float(loss.detach()))
            step_loss = sum(micro_losses) / accum
            if not math.isfinite(step_loss):
                nonfinite_streak += 1
                opt.zero_grad(set_to_none=True)
                LOG.error("non-finite loss %r at step %d (streak %d); optimizer step skipped", step_loss, state.step + 1, nonfinite_streak)
                if nonfinite_streak >= 5:
                    status, exit_code = "diverged", EXIT_NAN
                    break
                state.step += 1
            else:
                nonfinite_streak = 0
                phase = "optimizer step"
                grad_norm = float(torch.nn.utils.clip_grad_norm_(model.parameters(), float(pcfg["grad_clip"])))
                opt.step()
                opt.zero_grad(set_to_none=True)
                state.step += 1
                state.last_loss = step_loss
            state.tokens_seen += tokens_per_step
            interval_losses.append(step_loss)
            interval_steps += 1

            epoch_end = state.step % steps_per_epoch == 0
            if epoch_end:
                LOG.info("epoch %d/%d complete at step %d", state.epoch + 1, epochs, state.step)
                state.epoch += 1
                state.position = 0
                perm = epoch_permutation(seed, state.epoch, train_store.n)

            val = None
            finished = state.step >= run_end
            if eval_every and (state.step % eval_every == 0 or finished) and len(val_idx):
                phase = "eval"
                t0 = common.timer()
                val = evaluate_mlm(model, val_store, val_idx, micro_batch, device, seed + VAL_MASK_SEED_OFFSET, pool_like,
                                   mcfg.vocab_size, p_mask, replace_probs)
                state.last_val_loss, state.last_val_acc = val
                pause_s += common.timer() - t0
        except torch.OutOfMemoryError as e:
            # Another process grabbing VRAM (desktop spike, a second job) OOMs this one even
            # though its own reserved peak is under the cap. In the forward/backward phase the
            # weights and optimizer state are intact, so the partial step is discarded; after a
            # completed optimizer step (eval phase) the step counts. Either way the data cursor
            # is rewound to the last completed step, a checkpoint is saved and the resume
            # continues from exactly there. An OOM inside opt.step itself is the one case where
            # weights may be partially updated; it is logged as such.
            ids = inputs = labels = pos = logits = loss = None   # drop the half-built graph before saving
            opt.zero_grad(set_to_none=True)
            torch.cuda.empty_cache()
            model.train()
            state.position = (state.step % steps_per_epoch) * eff_batch
            free, total = torch.cuda.mem_get_info()
            LOG.error("CUDA OOM during the %s of step %d (%s); device has %.2f of %.2f GiB in use across all processes, this process "
                      "reserved %.2f GiB. Saving a checkpoint at step %d and stopping (exit %d).%s", phase, state.step + (phase != "eval"),
                      str(e).splitlines()[0][:160], (total - free) / GIB, total / GIB, torch.cuda.memory_reserved() / GIB, state.step,
                      EXIT_VRAM, " WARNING: the OOM hit inside opt.step; weights may be partially updated." if phase == "optimizer step" else "")
            status, exit_code = "oom", EXIT_VRAM
            break
        if state.step % log_every == 0 or val is not None or finished or stop.name or epoch_end:
            write_row(val, lr, grad_norm if math.isfinite(step_loss) else float("nan"), step_loss)

        reserved_peak = torch.cuda.max_memory_reserved() / GIB
        if reserved_peak > hard_cap:
            state.peak_reserved_gb = max(state.peak_reserved_gb, reserved_peak)
            LOG.error("VRAM guard: reserved peak %.2f GiB exceeds the %.0f GiB hard cap; saving and stopping (exit %d). "
                      "Resume with a smaller --micro-batch or --grad-checkpointing.", reserved_peak, hard_cap, EXIT_VRAM)
            status, exit_code = "vram_guard", EXIT_VRAM
            break
        if finished:
            status = "complete" if state.step >= total_steps else "max_steps"
            break
        if stop.name:
            status = f"signal:{stop.name}"
            break
        if max_seconds is not None and (common.timer() - t_loop0) >= max_seconds:
            LOG.warning("--max-hours %.2f reached at step %d; saving a checkpoint and exiting 0 (resume with --resume)", args.max_hours, state.step)
            status = "max_hours"
            break
        if epoch_end:
            checkpoint(f"epoch {state.epoch} end")
        elif state.step % ckpt_every == 0:
            checkpoint("periodic")

    # -------------------------------------------------------------------- wrap up
    if interval_steps:
        write_row(None, lr, float("nan"), state.last_loss if state.last_loss is not None else float("nan"))
    checkpoint(status)
    run_wall = common.timer() - t_loop0
    mean_tok_s = run_tokens / max(1e-6, run_train_s)
    summary = {
        **asdict(state), "status": status, "exit_code": exit_code, "run_end_step": run_end,
        "this_invocation": {"steps": run_tokens // tokens_per_step, "tokens": run_tokens, "loop_seconds": round(run_wall, 1),
                            "mean_tokens_per_s": round(mean_tok_s), "setup_seconds": round(t_loop0 - t_invocation, 1)},
        "mean_tokens_per_s": round(mean_tok_s), "config": snapshot, "params": counts,
        "finished": datetime.now().isoformat(timespec="seconds"), "log": current_logfile(),
    }
    common.write_json(summary, run_dir / "run_summary.json")
    print(f"\nTotal training time (this invocation): {fmt_dur(run_wall)} ({run_wall:.1f} s) for {run_tokens // tokens_per_step} steps, "
          f"{mean_tok_s:,.0f} tokens/s mean")
    print(f"Total training time (all invocations): {fmt_dur(state.train_seconds_total)} ({state.train_seconds_total:.1f} s), "
          f"{state.step} steps, {state.tokens_seen / 1e6:.1f}M tokens, peak VRAM reserved {state.peak_reserved_gb:.2f} GiB")
    LOG.info("run %s stopped: %s at step %d/%d (exit %d)", run_dir.name, status, state.step, total_steps, exit_code)
    if status == "complete":
        write_final(run_dir, run_dir / f"step-{state.step:07d}", cfg, status)
    if status in ("complete", "max_steps", "max_hours") and state.step >= RUN_LOG_MIN_STEPS:
        common.append_run_log(run_log_block(run_dir, snapshot, state, status, mean_tok_s, current_logfile(), counts))
        LOG.info("RUN_LOG.md updated")
    else:
        LOG.info("RUN_LOG.md not updated (status %s, step %d; needs a normal stop at >= %d steps)", status, state.step, RUN_LOG_MIN_STEPS)
    return exit_code


# ============================================================================ export / plot
def cmd_export(args: argparse.Namespace, cfg: dict) -> int:
    global LOG
    LOG = common.setup_logging("pretrain", cfg["paths"]["logs"])
    run_dir = Path(cfg["paths"]["checkpoints"]) / "pretrain" / args.run_name
    step_dir = latest_step_dir(run_dir)
    if step_dir is None:
        LOG.error("no step-* checkpoint under %s", run_dir)
        return EXIT_CONFIG
    summary_p = run_dir / "run_summary.json"
    status = common.read_json(summary_p).get("status", "exported") if summary_p.exists() else "exported"
    final = write_final(run_dir, step_dir, cfg, status)
    cfgj = common.read_json(final / "config.json")
    print(json.dumps(cfgj["training"], indent=2))
    print(f"final/: {sorted(p.name for p in final.iterdir())} -> {final}")
    return EXIT_OK


def cmd_plot(args: argparse.Namespace, cfg: dict) -> int:
    global LOG
    LOG = common.setup_logging("pretrain", cfg["paths"]["logs"])
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    run_dir = Path(cfg["paths"]["checkpoints"]) / "pretrain" / args.run_name
    rows = MetricsCSV(run_dir / "metrics.csv").read()
    if not rows:
        LOG.error("no metrics.csv under %s", run_dir)
        return EXIT_CONFIG
    steps = np.array([int(r["step"]) for r in rows])
    loss = np.array([float(r["loss"]) for r in rows])
    lr = np.array([float(r["lr"]) for r in rows])
    val_rows = [r for r in rows if r.get("val_mlm_loss")]
    fig, ax = plt.subplots(figsize=(9, 5))
    ax.plot(steps, loss, color="tab:blue", lw=1.2, label="train MLM loss (per logged step)")
    if val_rows:
        ax.plot([int(r["step"]) for r in val_rows], [float(r["val_mlm_loss"]) for r in val_rows],
                "o-", color="tab:orange", lw=1.5, ms=4, label="val MLM loss")
    ax.axhline(math.log(int(cfg["model"]["vocab_size"])), color="grey", ls=":", lw=1, label=f"ln(V) = {math.log(int(cfg['model']['vocab_size'])):.2f}")
    ax.set_xlabel("optimizer step")
    ax.set_ylabel("MLM cross-entropy")
    ax.set_title(f"Pretraining {args.run_name}: {len(rows)} logged steps, last train loss {loss[-1]:.3f}"
                 + (f", last val loss {float(val_rows[-1]['val_mlm_loss']):.3f}" if val_rows else ""))
    ax.grid(alpha=0.3)
    ax2 = ax.twinx()
    ax2.plot(steps, lr, color="tab:green", lw=1, alpha=0.7, label="learning rate")
    ax2.set_ylabel("learning rate")
    ax2.ticklabel_format(axis="y", style="sci", scilimits=(0, 0))
    h1, l1 = ax.get_legend_handles_labels()
    h2, l2 = ax2.get_legend_handles_labels()
    ax.legend(h1 + h2, l1 + l2, loc="upper right", fontsize=9)
    out = Path(cfg["paths"]["reports"]) / "pretrain_curves.png"
    out.parent.mkdir(parents=True, exist_ok=True)
    fig.tight_layout()
    fig.savefig(out, dpi=150)
    LOG.info("wrote %s (%d rows, %d val points)", out, len(rows), len(val_rows))
    print(out)
    return EXIT_OK


# ==================================================================================== CLI
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--config", default="config.yaml")
    sub = ap.add_subparsers(dest="cmd", required=True)
    tr = sub.add_parser("train", help="MLM pretraining (resumable)")
    tr.add_argument("--run-name", default="large-24L")
    tr.add_argument("--resume", nargs="?", const="latest", default=None, metavar="PATH",
                    help="resume from a step dir / run dir, or 'latest' under the run (default when the flag has no value)")
    tr.add_argument("--epochs", type=int, default=None, help="planned epochs (config.pretrain.epochs); fixes the schedule length")
    tr.add_argument("--max-steps", type=int, default=None, help="stop this invocation once the GLOBAL step reaches N (schedule unchanged)")
    tr.add_argument("--max-hours", type=float, default=None, help="stop this invocation after H hours of training: save, exit 0")
    tr.add_argument("--micro-batch", type=int, default=None, help="default config.pretrain.micro_batch_start (16), or the checkpoint's on resume")
    tr.add_argument("--grad-checkpointing", action="store_true", help="activation checkpointing (config.model.grad_checkpointing)")
    tr.add_argument("--limit-docs", type=int, default=None, help="dev: tokenise only the first N documents (separate cache files)")
    tr.add_argument("--checkpoint-every", type=int, default=None)
    tr.add_argument("--log-every", type=int, default=None)
    tr.add_argument("--eval-every", type=int, default=None, help="0 disables periodic eval")
    tr.add_argument("--keep-checkpoints", type=int, default=3)
    tr.add_argument("--fresh", action="store_true", help="start over in a run dir that already holds checkpoints (deletes them)")
    ex = sub.add_parser("export", help="write final/ (model.safetensors + config.json + tokenizer.json) from the latest checkpoint")
    ex.add_argument("--run-name", required=True)
    pl = sub.add_parser("plot", help="reports/pretrain_curves.png from metrics.csv")
    pl.add_argument("--run-name", required=True)
    return ap.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    cfg = common.load_config(args.config)
    if args.cmd == "train":
        return cmd_train(args, cfg)
    if args.cmd == "export":
        return cmd_export(args, cfg)
    return cmd_plot(args, cfg)


if __name__ == "__main__":
    sys.exit(main())