#!/usr/bin/env python3
"""
train_hypernix3_mini.py — pretrain hypernix.3-mini from scratch: a 48.7M-parameter
hyperNix0x-v2 model, 1.21B tokens, on a single GTX 1080 (8 GB, Pascal sm_61),
inside a one-week wall-clock budget.

    Blaze Industries / ray0rf1re
    model:     hypernix.3-mini
    arch:      hyperNix0x-v2  (hypernix.training.brewer)
    runtime:   hypernix.models.neo_oven  (NeoOven + brewer_adapter)
    optimizer: AdamW by default; Agedcookerv5 (Pascal-tuned PressureCookerV5)
               and PressureCookerV6 available via --optimizer

WHAT IT DOES
------------
  1. tokenizer  train a *fully custom* 32k byte-level BPE vocab from a sample of
                the streamed corpus (no borrowed tokenizer, no vocab surgery)
  2. pack       stream the corpus, tokenize it, append to a flat uint16 memmap
                (~2.42 GB on disk for 1.21B tokens — the raw text is never kept)
  3. train      pretrain from scratch against that memmap, resumable, with a
                VRAM escalation ladder and a wall-clock deadline governor
  4. sample     load the result back through NeoOven and generate from it

    python train_hypernix3_mini.py all             # 1 → 2 → 3
    python train_hypernix3_mini.py train --resume  # after a reboot
    python train_hypernix3_mini.py benchmark       # measure tok/s, project ETA
    python train_hypernix3_mini.py sample --sample-prompt "The reason"

Every stage takes --local-text ./my_corpus to use text you already have instead
of streaming from the Hub.

Everything from every stage is appended to ./runs/<name>/train.log, in a format
`hypernix.models.neo_oven.parse_training_log` and the tvtop/tvtop++ monitors
already parse (`step N/M loss=X ppl=Y lr=Z`).

THE BUDGET (defaults, all overridable)
--------------------------------------
  params        48,706,048  (8L x d512, 8 heads / 2 KV heads, d_ff 2203, tied)
  tokens        1,210,000,000  (~24.8 tokens/param — Chinchilla-shaped)
  context       1024
  batch         65,536 tokens/step  (micro 8 x ctx 1024 x accum 8)
  steps         18,464
  disk          ~2.5 GB tokens + ~1.5 GB temporary tokenizer sample
                (streamed, never fully downloaded: ~4 GB pulled over the wire,
                nowhere near the 100 GB ceiling)
  host RAM      < 3 GB — the corpus is a memmap, the stream is bounded
  VRAM          ~4.5 GB at the default micro-batch, with a documented ladder
                down to ~1.8 GB if anything OOMs
  wall clock    3.8e17 FLOPs of work. On a 8.87 TFLOPS card that is 40 h at 30%
                MFU, 60 h at 20%, 80 h at 15% — so the expected finish is 2-3.5
                days and the one-week deadline only fails below ~7% MFU. The
                script prints this table at startup, then replaces it at step
                100 with the throughput it actually measured; --enforce-deadline
                makes it trim the token budget if that measurement says the run
                would overrun.

WHY FP32 ON PASCAL
------------------
GP104 runs FP16 arithmetic at 1/64 of its FP32 rate, and has no bf16 and no
FlashAttention. AMP is therefore a *slowdown* here, not a speedup; this script
trains in FP32 and spends the memory it saves elsewhere (bigger micro-batch).
`--amp` is wired up (autocast + GradScaler) for anyone running this on a
newer card, but it is off by default and warns when enabled here.

TWO UPSTREAM ISSUES THIS SCRIPT WORKS AROUND (see patch_brewer_attention)
-------------------------------------------------------------------------
1. BROKEN CAUSALITY WITH SLIDING-WINDOW ATTENTION. brewer.BrewerAttention
   combines the causal mask and the sliding-window mask with `torch.maximum`.
   For additive masks (0 = keep, -inf = drop) that keeps the *less* masked of
   the two, so on odd-numbered layers every query can see up to
   `sliding_window_size - 1` tokens into its own future. The model trains to a
   beautiful loss and generates garbage. Verify it yourself:

       cfg = BrewerConfig(..., use_sliding_window=True, sliding_window_size=8)
       # perturb token 10 of the input; outputs at positions 3..9 move.

   `maximum` should be `minimum`. This script defaults to
   use_sliding_window=False *and* installs a corrected attention forward, so
   `--sliding-window` is safe to turn on here.
2. SDPA IS FED AN EXPLICIT MASK EVEN WHEN THE MASK IS PLAIN CAUSAL, which
   forces the math backend and materializes B*H*T*T scores (268 MB/layer at
   micro-batch 8). The patched forward passes `is_causal=True` instead and lets
   PyTorch pick its memory-efficient kernel (supported on sm_61).

Both patches are opt-out (`--no-patch`), and neither touches your installed
hypernix — they are runtime monkeypatches scoped to this process.

REQUIREMENTS
------------
    hypernix (>= the release carrying models/brewer_adapter.py)
    torch built with sm_61 support   # checked at startup, see require_pascal()
    tokenizers, datasets, numpy
    transformers            # only for the NeoOven generation path

Written for Python 3.12, no venv activation needed:
    python train_hypernix3_mini.py all
"""

from __future__ import annotations

import argparse
import json
import logging
import math
import os
import random
import shutil
import sys
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Iterator

# Keep BLAS from fighting the dataloader for the i7-7700HQ's 4 cores.
os.environ.setdefault("OMP_NUM_THREADS", "4")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "true")
# Fragmentation is what actually kills long runs on an 8 GB card.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import numpy as np

# ---------------------------------------------------------------------------
# Defaults — the whole recipe lives here
# ---------------------------------------------------------------------------

RUN_NAME = "hypernix.3-mini"

# -- model ------------------------------------------------------------------
VOCAB_SIZE = 32_000
N_LAYERS = 8
N_HEADS = 8
N_KV_HEADS = 2          # GQA 4:1
D_MODEL = 512
D_FF = 2203             # chosen so the exact parameter count lands on 48.7M
CTX = 1024
ROPE_THETA = 100_000.0
SLIDING_WINDOW = 512    # only used when --sliding-window is passed

# -- data -------------------------------------------------------------------
TOKEN_BUDGET = 1_210_000_000
VAL_TOKENS = 2_000_000
TOK_SAMPLE_BYTES = 1_500_000_000    # text sampled to train the BPE vocab
DISK_CEILING_GB = 100.0             # hard stop; the plan needs ~4

#: Streamed sources, mixed by weight. Nothing is downloaded in full — with a
#: 1.21B-token budget this pulls roughly 4 GB over the wire out of a 28.5 GB
#: shard. `text_key: None` auto-detects among text/content/code/raw_content.
SOURCES: list[dict[str, Any]] = [
    {"repo": "HuggingFaceFW/fineweb-edu", "config": "sample-10BT",
     "split": "train", "weight": 1.0, "text_key": "text"},
    # Add code/math by uncommenting — weights are normalised for you:
    # {"repo": "open-web-math/open-web-math", "config": "default",
    #  "split": "train", "weight": 0.10, "text_key": "text"},
    # {"repo": "bigcode/the-stack-smol", "config": "default",
    #  "split": "train", "weight": 0.10, "text_key": "content"},
]

# -- optimisation -----------------------------------------------------------
TOKENS_PER_STEP = 65_536        # 8 x 1024 x 8
MICRO_BATCH = 8
PEAK_LR = 6e-4
MIN_LR_FRAC = 0.1
WARMUP_FRAC = 0.02
WEIGHT_DECAY = 0.1
BETAS = (0.9, 0.95)
GRAD_CLIP = 1.0

# -- schedule / bookkeeping -------------------------------------------------
DEADLINE_HOURS = 168.0           # exactly one week
CALIBRATE_AT_STEP = 100          # re-plan the budget once, from measured tok/s
LOG_EVERY = 10
EVAL_EVERY = 500
EVAL_BATCHES = 20
SAVE_EVERY = 1000
SAVE_EVERY_MINUTES = 30.0
SAMPLE_EVERY = 2000
KEEP_CHECKPOINTS = 3

SPECIALS = ["<|endoftext|>", "<|pad|>", "<fim_prefix>", "<fim_middle>", "<fim_suffix>"]
EOS = "<|endoftext|>"

LOG = logging.getLogger("brew")


# ---------------------------------------------------------------------------
# Paths + logging
# ---------------------------------------------------------------------------

@dataclass
class Paths:
    root: Path

    @property
    def tokenizer(self) -> Path: return self.root / "tokenizer"
    @property
    def tok_json(self) -> Path: return self.tokenizer / "tokenizer.json"
    @property
    def tok_sample(self) -> Path: return self.root / "tokenizer_sample.txt"
    @property
    def bin(self) -> Path: return self.root / "tokens.u16.bin"
    @property
    def pack_state(self) -> Path: return self.root / "pack_state.json"
    @property
    def ckpt_dir(self) -> Path: return self.root / "checkpoints"
    @property
    def latest(self) -> Path: return self.ckpt_dir / "latest.pt"
    @property
    def final(self) -> Path: return self.root / "final"
    @property
    def log(self) -> Path: return self.root / "train.log"
    @property
    def report(self) -> Path: return self.root / "run_report.json"

    def mkdirs(self) -> None:
        for d in (self.root, self.tokenizer, self.ckpt_dir):
            d.mkdir(parents=True, exist_ok=True)


def setup_logging(paths: Paths, stage: str) -> None:
    paths.mkdirs()
    fmt = logging.Formatter(
        fmt="%(asctime)s [" + stage + "] %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
    )
    LOG.setLevel(logging.INFO)
    LOG.handlers.clear()
    fh = logging.FileHandler(paths.log, mode="a", encoding="utf-8")
    fh.setFormatter(fmt)
    sh = logging.StreamHandler(sys.stdout)
    sh.setFormatter(fmt)
    LOG.addHandler(fh)
    LOG.addHandler(sh)
    LOG.propagate = False


def human(n: float, unit: str = "") -> str:
    for suffix, scale in (("T", 1e12), ("B", 1e9), ("M", 1e6), ("K", 1e3)):
        if abs(n) >= scale:
            return f"{n / scale:.3g}{suffix}{unit}"
    return f"{n:.0f}{unit}"


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


def free_disk_gb(path: Path) -> float:
    p = path
    while not p.exists():
        p = p.parent
    return shutil.disk_usage(p).free / 1e9


# ---------------------------------------------------------------------------
# hypernix import
# ---------------------------------------------------------------------------

def import_hypernix(extra_src: str | None) -> None:
    if extra_src:
        sys.path.insert(0, str(Path(extra_src).expanduser().resolve()))
    try:
        import hypernix  # noqa: F401
    except ModuleNotFoundError as exc:  # pragma: no cover
        raise SystemExit(
            "hypernix is not importable. Install it (`pip install hypernix`) or "
            "point at a checkout with --hypernix-src /path/to/HyperNix-pip/src"
        ) from exc


def require_pascal(strict: bool) -> "object":
    """Fail loudly *now* rather than 40 minutes into the first epoch.

    Recent PyTorch wheels have been dropping older compute capabilities; a
    torch that cannot emit sm_61 will happily import, see the GTX 1080, and
    then die at the first kernel launch with 'no kernel image is available'.
    """
    import torch

    if not torch.cuda.is_available():
        msg = "no CUDA device visible — this script is written for a GTX 1080"
        if strict:
            raise SystemExit(msg)
        LOG.warning("%s; continuing on CPU (development mode)", msg)
        return torch.device("cpu")

    cap = torch.cuda.get_device_capability(0)
    name = torch.cuda.get_device_name(0)
    arches = torch.cuda.get_arch_list()
    LOG.info("gpu=%s capability=sm_%d%d torch=%s arch_list=%s",
             name, cap[0], cap[1], torch.__version__, ",".join(arches))
    tag = f"sm_{cap[0]}{cap[1]}"
    major_tag = f"sm_{cap[0]}0"  # same-major cubins (sm_60) are forward compatible with sm_61+ under CUDA binary-compat rules
    if tag not in arches and major_tag not in arches and "compute_" not in " ".join(arches):
        raise SystemExit(
            f"this torch build has no {tag} kernels (arch list: {', '.join(arches)}).\n"
            f"Install a wheel that still ships Pascal, e.g.\n"
            f"  pip install torch --index-url https://download.pytorch.org/whl/cu126"
        )
    if cap >= (7, 0):
        LOG.warning("capability %s is newer than Pascal — the FP32-only choices "
                    "below are leaving performance on the table", tag)
    return torch.device("cuda")


# ---------------------------------------------------------------------------
# Attention patch — see the module docstring
# ---------------------------------------------------------------------------

_MASK_CACHE: dict[tuple, Any] = {}


def _causal_window_mask(T: int, window: int | None, device, dtype):
    key = (T, window, device, dtype)
    cached = _MASK_CACHE.get(key)
    if cached is not None:
        return cached
    import torch

    idx = torch.arange(T, device=device)
    dist = idx.unsqueeze(1) - idx.unsqueeze(0)          # query i, key j -> i - j
    allowed = dist >= 0                                  # causal: j <= i
    if window is not None:
        allowed &= dist < window                         # and within the window
    mask = torch.zeros(T, T, device=device, dtype=dtype)
    mask.masked_fill_(~allowed, torch.finfo(dtype).min)
    _MASK_CACHE[key] = mask
    return mask


def patch_brewer_attention() -> None:
    """Correct mask composition + use the causal fast path.

    Replaces BrewerAttention.forward only; projections, RoPE, GQA expansion and
    the config are all the upstream ones, so a checkpoint written here loads
    into an unpatched hypernix unchanged.
    """
    import torch
    import torch.nn.functional as F
    from hypernix.training import brewer

    def forward(self, x, attn_mask=None):
        B, T, _ = x.shape
        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)

        cos, sin = self._get_rope(T, x.device, q.dtype)
        q, k = brewer._apply_rope(q, k, cos, sin)

        if self.n_rep > 1:
            k = k.repeat_interleave(self.n_rep, dim=1)
            v = v.repeat_interleave(self.n_rep, dim=1)

        drop = self.dropout_p if self.training else 0.0
        if attn_mask is None and not self.use_sliding_window:
            out = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=drop)
        else:
            window = self.sliding_window_size if self.use_sliding_window else None
            mask = _causal_window_mask(T, window, x.device, q.dtype)
            if attn_mask is not None:
                mask = mask + attn_mask
            out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=drop)

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.o_proj(out)

    brewer.BrewerAttention.forward = forward
    torch.backends.cudnn.benchmark = True
    LOG.info("patched BrewerAttention.forward "
             "(causal/window mask composition + is_causal fast path)")


# ---------------------------------------------------------------------------
# Stage 1 — a fully custom vocabulary
# ---------------------------------------------------------------------------

def stream_local(pattern: str, skip_docs: int = 0) -> Iterator[str]:
    """Yield text files matching a glob, for corpora you already have on disk."""
    root = Path(pattern).expanduser()
    files = sorted(root.rglob("*")) if root.is_dir() else sorted(Path().glob(pattern))
    files = [f for f in files if f.is_file()]
    if not files:
        raise SystemExit(f"--local-text {pattern!r} matched no files")
    LOG.info("local corpus: %d files", len(files))
    seen = 0
    while True:                       # loop the corpus if it is smaller than the budget
        for f in files:
            seen += 1
            if seen <= skip_docs:
                continue
            try:
                yield f.read_text(encoding="utf-8", errors="replace")
            except OSError as exc:
                LOG.warning("skipping %s: %s", f, exc)


def stream_texts(sources: list[dict[str, Any]], seed: int, skip_docs: int = 0,
                 local: str | None = None) -> Iterator[str]:
    """Weighted interleave of streamed HF datasets, yielding raw document text."""
    if local:
        yield from stream_local(local, skip_docs)
        return
    from datasets import interleave_datasets, load_dataset

    streams, weights, keys = [], [], []
    for src in sources:
        ds = load_dataset(src["repo"], src.get("config"), split=src.get("split", "train"),
                          streaming=True)
        streams.append(ds)
        weights.append(float(src.get("weight", 1.0)))
        keys.append(src.get("text_key"))

    if len(streams) == 1:
        merged, key = streams[0], keys[0]
    else:
        total = sum(weights)
        merged = interleave_datasets(streams, probabilities=[w / total for w in weights],
                                     seed=seed, stopping_strategy="all_exhausted")
        key = None

    if skip_docs:
        LOG.info("skipping %s already-consumed documents to resume the stream "
                 "(re-reads, does not re-tokenize)", human(skip_docs))
        merged = merged.skip(skip_docs)

    candidates = ("text", "content", "code", "raw_content", "document")
    for record in merged:
        if key and key in record:
            text = record[key]
        else:
            text = next((record[c] for c in candidates
                         if isinstance(record.get(c), str)), None)
        if isinstance(text, str) and text.strip():
            yield text


def stage_tokenizer(paths: Paths, args: argparse.Namespace) -> None:
    """Train a 32k byte-level BPE from scratch on a sample of the corpus."""
    from tokenizers import Tokenizer, decoders, models, pre_tokenizers, processors, trainers

    if paths.tok_json.exists() and not args.force:
        LOG.info("tokenizer already at %s (use --force to retrain)", paths.tok_json)
        return

    target = args.tok_sample_bytes
    if not paths.tok_sample.exists() or paths.tok_sample.stat().st_size < target * 0.95:
        LOG.info("sampling %s of text to train the vocabulary on", human(target, "B"))
        written, docs, t0 = 0, 0, time.time()
        with paths.tok_sample.open("w", encoding="utf-8") as fh:
            for text in stream_texts(args.sources, args.seed, local=args.local_text):
                fh.write(text)
                fh.write("\n")
                written += len(text) + 1
                docs += 1
                if docs % 20_000 == 0:
                    rate = written / max(1e-9, time.time() - t0)
                    LOG.info("  sample %s / %s  (%s docs, %s/s)",
                             human(written, "B"), human(target, "B"),
                             human(docs), human(rate, "B"))
                if written >= target:
                    break
        LOG.info("sample written: %s in %s docs", human(written, "B"), human(docs))

    LOG.info("training byte-level BPE — vocab_size=%d, specials=%s",
             args.vocab_size, ", ".join(SPECIALS))
    tok = Tokenizer(models.BPE(unk_token=None))
    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tok.decoder = decoders.ByteLevel()
    tok.post_processor = processors.ByteLevel(trim_offsets=False)
    trainer = trainers.BpeTrainer(
        vocab_size=args.vocab_size,
        special_tokens=SPECIALS,
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
        min_frequency=2,
        show_progress=True,
    )
    t0 = time.time()
    tok.train([str(paths.tok_sample)], trainer)
    tok.save(str(paths.tok_json))
    LOG.info("vocabulary trained in %s -> %s (%d tokens)",
             hms(time.time() - t0), paths.tok_json, tok.get_vocab_size())

    # transformers-side metadata, so NeoOven's _load_tokenizer finds a fast
    # tokenizer here instead of silently falling back to bytes.
    (paths.tokenizer / "tokenizer_config.json").write_text(json.dumps({
        "tokenizer_class": "PreTrainedTokenizerFast",
        "model_max_length": args.ctx,
        "bos_token": EOS, "eos_token": EOS, "unk_token": None, "pad_token": "<|pad|>",
        "clean_up_tokenization_spaces": False,
        "added_tokens_decoder": {
            str(i): {"content": s, "special": True, "lstrip": False, "rstrip": False,
                     "normalized": False, "single_word": False}
            for i, s in enumerate(SPECIALS)
        },
    }, indent=2))
    (paths.tokenizer / "special_tokens_map.json").write_text(json.dumps({
        "bos_token": EOS, "eos_token": EOS, "pad_token": "<|pad|>",
    }, indent=2))

    sanity = "def fibonacci(n):\n    return n if n < 2 else fibonacci(n-1)+fibonacci(n-2)\n"
    ids = tok.encode(sanity).ids
    assert tok.decode(ids) == sanity, "byte-level BPE failed to round-trip"
    LOG.info("round-trip ok — %d chars -> %d tokens (%.2f chars/token)",
             len(sanity), len(ids), len(sanity) / len(ids))


# ---------------------------------------------------------------------------
# Stage 2 — pack the corpus into a flat uint16 memmap
# ---------------------------------------------------------------------------

def stage_pack(paths: Paths, args: argparse.Namespace) -> None:
    """Stream -> tokenize -> append. Bounded RAM, resumable, uint16 on disk."""
    from tokenizers import Tokenizer

    if not paths.tok_json.exists():
        raise SystemExit("no tokenizer yet — run the 'tokenizer' stage first")

    budget = args.tokens
    need_gb = budget * 2 / 1e9
    free_gb = free_disk_gb(paths.root)
    LOG.info("packing %s tokens -> %.2f GB (uint16); %.1f GB free on this volume",
             human(budget), need_gb, free_gb)
    if need_gb > args.disk_ceiling_gb:
        raise SystemExit(f"{need_gb:.1f} GB exceeds the --disk-ceiling-gb of "
                         f"{args.disk_ceiling_gb:.0f} GB")
    if free_gb < need_gb * 1.2:
        raise SystemExit(f"need ~{need_gb * 1.2:.1f} GB free, have {free_gb:.1f} GB")

    state = {"tokens": 0, "docs": 0}
    if paths.pack_state.exists() and not args.force:
        state = json.loads(paths.pack_state.read_text())
    elif paths.bin.exists():
        paths.bin.unlink()

    if state["tokens"] >= budget:
        LOG.info("already packed %s tokens — nothing to do", human(state["tokens"]))
        return

    # Truncate to exactly what the state claims, so a kill -9 mid-write can
    # never leave a half token at the tail.
    if paths.bin.exists():
        with paths.bin.open("r+b") as fh:
            fh.truncate(state["tokens"] * 2)
        LOG.info("resuming: %s tokens already on disk", human(state["tokens"]))

    tok = Tokenizer.from_file(str(paths.tok_json))
    eos_id = tok.token_to_id(EOS)
    assert eos_id is not None

    batch: list[str] = []
    tokens, docs = state["tokens"], state["docs"]
    t0, last_log, last_tokens = time.time(), time.time(), tokens

    def flush(fh, batch: list[str]) -> int:
        # encode_batch releases the GIL and threads inside Rust, so this is the
        # parallel path without a single Python worker process.
        encs = tok.encode_batch(batch)
        ids: list[int] = []
        for enc in encs:
            ids.extend(enc.ids)
            ids.append(eos_id)
        np.asarray(ids, dtype=np.uint16).tofile(fh)
        return len(ids)

    with paths.bin.open("ab") as fh:
        for text in stream_texts(args.sources, args.seed, skip_docs=docs,
                                 local=args.local_text):
            batch.append(text)
            docs += 1
            if len(batch) < args.pack_batch:
                continue
            tokens += flush(fh, batch)
            batch.clear()

            if time.time() - last_log > 30:
                fh.flush()
                paths.pack_state.write_text(json.dumps({"tokens": tokens, "docs": docs}))
                rate = (tokens - last_tokens) / max(1e-9, time.time() - last_log)
                eta = (budget - tokens) / max(1.0, rate)
                LOG.info("  packed %s / %s tokens (%.1f%%)  %s tok/s  eta %s",
                         human(tokens), human(budget), 100 * tokens / budget,
                         human(rate), hms(eta))
                last_log, last_tokens = time.time(), tokens

            if tokens >= budget:
                break
        if batch and tokens < budget:
            tokens += flush(fh, batch)

    paths.pack_state.write_text(json.dumps({"tokens": tokens, "docs": docs}))
    size_gb = paths.bin.stat().st_size / 1e9
    LOG.info("packed %s tokens from %s documents in %s",
             human(tokens), human(docs), hms(time.time() - t0))
    LOG.info("dataset footprint: %.2f GB on disk, ceiling %.0f GB — clear",
             size_gb, args.disk_ceiling_gb)
    if not args.keep_sample and paths.tok_sample.exists():
        paths.tok_sample.unlink()
        LOG.info("removed the temporary tokenizer sample")


# ---------------------------------------------------------------------------
# Data: memmap sampler (host RAM stays flat regardless of corpus size)
# ---------------------------------------------------------------------------

class TokenBin:
    def __init__(self, path: Path, ctx: int, val_tokens: int, seed: int) -> None:
        self.path = path
        self.ctx = ctx
        total = path.stat().st_size // 2
        if total < ctx * 64:
            raise SystemExit(f"{path} holds only {total} tokens — run 'pack' first")
        self.data = np.memmap(path, dtype=np.uint16, mode="r", shape=(total,))
        self.total = total
        self.val_start = max(ctx + 1, total - val_tokens)
        self.rng = np.random.default_rng(seed)
        LOG.info("token bin: %s tokens (%.2f GB), val tail = %s tokens",
                 human(total), path.stat().st_size / 1e9, human(total - self.val_start))

    def _batch(self, lo: int, hi: int, micro: int, rng: np.random.Generator):
        import torch

        starts = rng.integers(lo, hi - self.ctx - 1, size=micro)
        # int64 copy of (micro, ctx+1) — 8*1025*8 B, trivial and page-cache backed.
        block = np.stack([self.data[s:s + self.ctx + 1] for s in starts]).astype(np.int64)
        t = torch.from_numpy(block)
        return t[:, :-1].contiguous(), t[:, 1:].contiguous()

    def train_batch(self, micro: int):
        return self._batch(0, self.val_start, micro, self.rng)

    def val_batches(self, micro: int, n: int):
        rng = np.random.default_rng(1234)   # same val set every time
        for _ in range(n):
            yield self._batch(self.val_start, self.total, micro, rng)


# ---------------------------------------------------------------------------
# Model / oven construction
# ---------------------------------------------------------------------------

def build_config(args: argparse.Namespace):
    from hypernix.training.brewer import BrewerConfig

    cfg = BrewerConfig(
        name=args.name,
        vocab_size=args.vocab_size,
        n_layers=args.n_layers,
        n_heads=args.n_heads,
        n_kv_heads=args.n_kv_heads,
        d_model=args.d_model,
        d_ff=args.d_ff,
        max_seq_len=args.ctx,
        rope_theta=ROPE_THETA,
        norm_eps=1e-5,
        dropout=0.0,
        tie_embeddings=True,
        use_sliding_window=args.sliding_window,
        sliding_window_size=SLIDING_WINDOW,
        attention_type="gqa",
    )
    # NeoOven._max_context reads max_position_embeddings and otherwise assumes
    # 2048; BrewerConfig has no such field, so generation would extrapolate RoPE
    # past what we trained. Dataclass, not slots — this sticks.
    cfg.max_position_embeddings = args.ctx
    return cfg


def build_oven(cfg, device, paths: Paths):
    """A fresh hyperNix0x-v2 in a NeoOven, with our own tokenizer attached.

    This is what neo_oven.new_brewed does, minus its preset-name-only argument:
    wrap the BrewerModel in brewer_adapter (which is what gives it NeoOven's
    `model(ids, labels=...) -> {"loss": ...}` convention) and hand it to NeoOven.
    """
    import torch
    from hypernix.models import brewer_adapter
    from hypernix.models.generate import _load_tokenizer
    from hypernix.models.neo_oven import NeoOven
    from hypernix.training.brewer import BrewerModel

    model = brewer_adapter.wrap(BrewerModel(cfg)).to(device, dtype=torch.float32)
    tok, kind = _load_tokenizer(paths.tokenizer)
    if kind != "hf":
        LOG.warning("tokenizer at %s did not load as a fast tokenizer (kind=%r) — "
                    "generation samples will be wrong", paths.tokenizer, kind)
    oven = NeoOven(model=model, tokenizer=tok, tokenizer_kind=kind, device=device,
                   dtype=torch.float32, model_dir=paths.final, repo_id=cfg.name)
    return oven


def report_size(cfg, model) -> int:
    params = model.num_params()
    LOG.info("built %s — %s parameters (%.4f M) | %dL d_model=%d heads=%d/%d "
             "d_ff=%d ctx=%d vocab=%d tied=%s sliding_window=%s",
             cfg.name, f"{params:,}", params / 1e6, cfg.n_layers, cfg.d_model,
             cfg.n_heads, cfg.n_kv_heads, cfg.d_ff, cfg.max_seq_len,
             cfg.vocab_size, cfg.tie_embeddings, cfg.use_sliding_window)
    embed = cfg.vocab_size * cfg.d_model
    LOG.info("  %.1f%% of parameters are the (tied) embedding matrix; "
             "%s are 'body' parameters", 100 * embed / params, f"{params - embed:,}")
    return params


# ---------------------------------------------------------------------------
# Forward + loss, with a VRAM ladder
# ---------------------------------------------------------------------------

@dataclass
class MemoryLadder:
    """What to give up, and in what order, when the 8 GB runs out.

    Each rung costs some speed and buys VRAM. We only ever climb, and every
    climb is logged so the run report says what the card actually needed.
    """
    micro_batch: int
    accum: int
    ce_chunk: int = 0          # 0 = one shot; >0 = checkpointed chunks of rows
    grad_ckpt: bool = False
    rung: int = 0

    def climb(self) -> bool:
        if self.rung == 0:
            self.ce_chunk, self.rung = 2048, 1
            LOG.warning("VRAM ladder -> rung 1: checkpointed cross-entropy "
                        "(chunk=%d rows); ~1 GB saved, ~10%% slower", self.ce_chunk)
            return True
        if self.rung == 1:
            self.grad_ckpt, self.rung = True, 2
            LOG.warning("VRAM ladder -> rung 2: gradient checkpointing on all "
                        "blocks; ~1.5 GB saved, ~30%% slower")
            return True
        if self.micro_batch > 1:
            self.micro_batch //= 2
            self.accum *= 2
            self.rung += 1
            LOG.warning("VRAM ladder -> rung %d: micro_batch=%d accum=%d "
                        "(tokens/step unchanged)", self.rung, self.micro_batch, self.accum)
            return True
        return False


def forward_loss(inner, ids, targets, ladder: MemoryLadder):
    """Loss for one micro-batch, bypassing BrewerModel.forward's monolithic head.

    Same computation, same shift as brewer_adapter's `loss` — the difference is
    that the (B, T, 32000) logits tensor is optional. At rung >= 1 the LM head
    and cross-entropy run in checkpointed row-chunks, so only one chunk's
    logits exist at a time and the rest is recomputed in the backward pass.
    """
    import torch
    import torch.nn.functional as F
    from torch.utils.checkpoint import checkpoint

    h = inner.embed(ids)
    for block in inner.blocks:
        h = checkpoint(block, h, use_reentrant=False) if ladder.grad_ckpt else block(h)
    h = inner.norm_out(h)

    flat_h = h.reshape(-1, h.shape[-1])
    flat_t = targets.reshape(-1)

    if ladder.ce_chunk <= 0:
        logits = inner.lm_head(flat_h)
        return F.cross_entropy(logits.float(), flat_t)

    def chunk_loss(hh, tt):
        return F.cross_entropy(inner.lm_head(hh).float(), tt, reduction="sum")

    total = flat_t.new_zeros((), dtype=torch.float32)
    n = flat_t.numel()
    for i in range(0, n, ladder.ce_chunk):
        hh, tt = flat_h[i:i + ladder.ce_chunk], flat_t[i:i + ladder.ce_chunk]
        total = total + checkpoint(chunk_loss, hh, tt, use_reentrant=False)
    return total / n


def is_oom(exc: BaseException) -> bool:
    return isinstance(exc, RuntimeError) and "out of memory" in str(exc).lower()


def project_walltime(params: int, cfg, tokens: int, tflops: float, mfu: float) -> float:
    """Analytic pre-flight estimate, so the deadline is a number before launch.

    6*N covers forward+backward on every weight that participates in a matmul
    (the tied LM head does, the embedding lookup does not); the second term is
    the attention score/value work, which scales with context rather than
    parameters and is ~8%% of the total at ctx=1024.
    """
    body = params - cfg.vocab_size * cfg.d_model      # everything but the embedding
    head = cfg.vocab_size * cfg.d_model               # tied, but still a real GEMM
    span = cfg.sliding_window_size if cfg.use_sliding_window else cfg.max_seq_len
    attn = 3 * 2 * 2 * cfg.n_layers * (span / 2) * cfg.d_model
    per_token = 6 * (body + head) + attn
    return tokens * per_token / (tflops * 1e12 * mfu)


# ---------------------------------------------------------------------------
# LR schedule
# ---------------------------------------------------------------------------

def lr_at(step: int, total: int, peak: float, warmup: int, min_frac: float) -> float:
    if step < warmup:
        return peak * (step + 1) / max(1, warmup)
    progress = (step - warmup) / max(1, total - warmup)
    progress = min(1.0, progress)
    cos = 0.5 * (1.0 + math.cos(math.pi * progress))
    return peak * (min_frac + (1.0 - min_frac) * cos)


def build_optimizer(name: str, params, args, total_steps: int, warmup: int):
    import torch

    name = name.lower()
    if name == "adamw":
        # foreach=True is the multi-tensor path; fused=True is sm_70+ only.
        return torch.optim.AdamW(params, lr=args.lr, betas=BETAS,
                                 weight_decay=WEIGHT_DECAY, eps=1e-8, foreach=True), True
    if name in ("agedcookerv5", "pcv5"):
        from hypernix.optimizers.optimizer_framework import ScheduleConfig
        from hypernix.optimizers.pressure_cooker_v5 import Agedcookerv5
        LOG.info("using Agedcookerv5 (PressureCookerV5, Pascal-enforced): it owns "
                 "its own warmup/plateau/cooldown schedule, so the cosine below "
                 "is disabled for this run")
        sched = ScheduleConfig(lr=args.lr, warmup_steps=warmup,
                               plateau_steps=max(1, int(total_steps * 0.6)),
                               cooldown_steps=max(1, total_steps - warmup
                                                  - int(total_steps * 0.6)),
                               min_lr=args.lr * MIN_LR_FRAC)
        return Agedcookerv5(params, schedule=sched, lr=args.lr,
                            weight_decay=WEIGHT_DECAY, grad_clip=GRAD_CLIP), False
    if name in ("pcv6", "pressurecookerv6"):
        from hypernix.optimizers.pressure_cooker_v6 import PressureCookerV6
        return PressureCookerV6(params, lr=args.lr, weight_decay=WEIGHT_DECAY), True
    raise SystemExit(f"unknown optimizer {name!r}")


# ---------------------------------------------------------------------------
# Checkpoints
# ---------------------------------------------------------------------------

def save_checkpoint(path: Path, oven, cfg, opt, step: int, tokens: int,
                    ladder: MemoryLadder, best: float, elapsed: float) -> None:
    import torch

    inner = oven.model.inner
    tmp = path.with_suffix(".tmp")
    torch.save({
        "arch": "hypernix0x-v2",
        "version": 2,
        "config": cfg.to_dict(),
        "model_state_dict": inner.state_dict(),   # brewer's own layout
        "optimizer_state_dict": opt.state_dict(),
        "step": step,
        "tokens": tokens,
        "best_val": best,
        "elapsed": elapsed,
        "ladder": asdict(ladder),
        "rng": {
            "torch": torch.get_rng_state(),
            "cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
            "numpy": np.random.get_state(),
            "python": random.getstate(),
        },
    }, tmp)
    tmp.replace(path)


def prune_checkpoints(ckpt_dir: Path, keep: int) -> None:
    snaps = sorted(ckpt_dir.glob("step_*.pt"))
    for old in snaps[:-keep] if len(snaps) > keep else []:
        old.unlink(missing_ok=True)


def export_final(paths: Paths, oven, cfg) -> None:
    """Directory layout that brewer_adapter.load / neo_oven.preheat_brewed read."""
    import torch

    paths.final.mkdir(parents=True, exist_ok=True)
    cfg.save(paths.final / "config.json")
    torch.save(oven.model.inner.state_dict(), paths.final / "model.pt")
    for fname in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"):
        src = paths.tokenizer / fname
        if src.exists():
            shutil.copy2(src, paths.final / fname)
    LOG.info("final model exported -> %s "
             "(load it with neo_oven.preheat_brewed(path, tokenizer_source=path))",
             paths.final)


# ---------------------------------------------------------------------------
# Stage 3 — train
# ---------------------------------------------------------------------------

def stage_train(paths: Paths, args: argparse.Namespace, benchmark_only: bool = False) -> None:
    import torch

    device = require_pascal(strict=not args.allow_cpu)
    if not args.no_patch:
        patch_brewer_attention()

    torch.manual_seed(args.seed)
    random.seed(args.seed)
    np.random.seed(args.seed)

    cfg = build_config(args)
    oven = build_oven(cfg, device, paths)
    params = report_size(cfg, oven.model)

    if abs(params - args.target_params) > args.target_params * 0.01:
        LOG.warning("parameter count %s is more than 1%% from the %s target — "
                    "adjust --d-ff", f"{params:,}", human(args.target_params))

    from hypernix.models import neo_oven as neo

    stats = neo.parameter_stats(oven.model)
    LOG.info("parameter memory: %s", stats)

    data = TokenBin(paths.bin, args.ctx, args.val_tokens, args.seed)

    tokens_per_step = args.micro_batch * args.ctx * args.accum
    total_steps = math.ceil(args.tokens / tokens_per_step)
    warmup = max(10, int(total_steps * WARMUP_FRAC))
    ladder = MemoryLadder(micro_batch=args.micro_batch, accum=args.accum)

    LOG.info("plan: %s tokens = %d steps x %s tokens/step "
             "(micro=%d ctx=%d accum=%d) | peak_lr=%.2e warmup=%d | deadline=%.0fh",
             human(args.tokens), total_steps, human(tokens_per_step),
             ladder.micro_batch, args.ctx, ladder.accum, args.lr, warmup,
             args.deadline_hours)
    LOG.info("pre-flight (GTX 1080 = 8.87 TFLOPS fp32): %s at 30%% MFU, %s at 15%%, "
             "%s at 8%% — the budget needs better than %.1f%% to land inside %.0fh. "
             "The real number replaces this at step %d.",
             hms(project_walltime(params, cfg, args.tokens, 8.87, 0.30)),
             hms(project_walltime(params, cfg, args.tokens, 8.87, 0.15)),
             hms(project_walltime(params, cfg, args.tokens, 8.87, 0.08)),
             100 * project_walltime(params, cfg, args.tokens, 8.87, 1.0)
             / (args.deadline_hours * 3600),
             args.deadline_hours, args.calibrate_at)

    opt, use_our_schedule = build_optimizer(args.optimizer, oven.model.parameters(),
                                            args, total_steps, warmup)

    step, tokens_done, best_val, prior_elapsed = 0, 0, float("inf"), 0.0
    if args.resume and paths.latest.exists():
        ck = torch.load(paths.latest, map_location=device, weights_only=False)
        oven.model.inner.load_state_dict(ck["model_state_dict"])
        opt.load_state_dict(ck["optimizer_state_dict"])
        step, tokens_done = ck["step"], ck["tokens"]
        best_val, prior_elapsed = ck.get("best_val", float("inf")), ck.get("elapsed", 0.0)
        ladder = MemoryLadder(**ck["ladder"])
        rng = ck.get("rng") or {}
        if rng.get("torch") is not None:
            torch.set_rng_state(rng["torch"].cpu() if hasattr(rng["torch"], "cpu") else rng["torch"])
        if rng.get("numpy") is not None:
            np.random.set_state(rng["numpy"])
        if rng.get("python") is not None:
            random.setstate(rng["python"])
        LOG.info("resumed at step %d/%d (%s tokens, %s already spent)",
                 step, total_steps, human(tokens_done), hms(prior_elapsed))

    # A disabled GradScaler is a transparent passthrough, so the loop below has
    # one code path whether or not --amp is set.
    scaler = torch.amp.GradScaler(device.type, enabled=args.amp)
    # PressureCookerV5 clips inside its own step(); clipping twice would quietly
    # change its update rule.
    external_clip = args.optimizer.lower() not in ("agedcookerv5", "pcv5")

    oven.model.train()
    t_start = time.time()
    last_save = time.time()
    window_tokens, window_t0 = 0, time.time()
    loss_accum, loss_n = 0.0, 0
    deadline_s = args.deadline_hours * 3600
    replanned = False

    def vram() -> str:
        d = neo.vram_stats()          # {} on CPU, MB on CUDA
        if not d:
            return "cpu"
        return (f"{d['allocated_mb'] / 1024:.2f}/{d['reserved_mb'] / 1024:.2f}GB"
                f" peak {d['max_allocated_mb'] / 1024:.2f}GB")

    while step < total_steps:
        if use_our_schedule:
            lr_now = lr_at(step, total_steps, args.lr, warmup, MIN_LR_FRAC)
            for group in opt.param_groups:
                group["lr"] = lr_now
        else:
            lr_now = opt.param_groups[0].get("lr", args.lr)

        opt.zero_grad(set_to_none=True)
        micro_loss = 0.0
        micro = 0
        while micro < ladder.accum:
            ids, targets = data.train_batch(ladder.micro_batch)
            ids = ids.to(device, non_blocking=True)
            targets = targets.to(device, non_blocking=True)
            loss = None
            try:
                with torch.autocast(device_type=device.type, dtype=torch.float16,
                                    enabled=args.amp):
                    loss = forward_loss(oven.model.inner, ids, targets, ladder) / ladder.accum
                scaler.scale(loss).backward()
            except RuntimeError as exc:
                if not is_oom(exc):
                    raise
                LOG.warning("CUDA OOM at step %d (micro %d/%d)", step, micro, ladder.accum)
                loss = ids = targets = None
                opt.zero_grad(set_to_none=True)
                if device.type == "cuda":
                    torch.cuda.empty_cache()
                if not ladder.climb():
                    raise SystemExit("out of VRAM at micro_batch=1 with every "
                                     "memory option enabled — lower --ctx") from exc
                micro, micro_loss = 0, 0.0
                continue
            micro_loss += float(loss.detach()) * ladder.accum
            micro += 1

        scaler.unscale_(opt)                       # no-op when AMP is off
        if GRAD_CLIP and external_clip:
            torch.nn.utils.clip_grad_norm_(oven.model.parameters(), GRAD_CLIP)
        scaler.step(opt)
        scaler.update()

        step += 1
        batch_tokens = ladder.micro_batch * args.ctx * ladder.accum
        tokens_done += batch_tokens
        window_tokens += batch_tokens
        loss_accum += micro_loss / ladder.accum
        loss_n += 1

        if step % args.log_every == 0:
            elapsed = prior_elapsed + (time.time() - t_start)
            tok_s = window_tokens / max(1e-9, time.time() - window_t0)
            window_tokens, window_t0 = 0, time.time()
            avg = loss_accum / max(1, loss_n)
            loss_accum, loss_n = 0.0, 0
            eta = (total_steps - step) * batch_tokens / max(1.0, tok_s)
            # tvtop / parse_training_log compatible:
            LOG.info("step %d/%d loss=%.4f ppl=%.2f lr=%.3e | %s tok/s | "
                     "%s done | eta %s | vram %s | rung %d",
                     step, total_steps, avg, math.exp(min(avg, 20)), lr_now,
                     human(tok_s), human(tokens_done), hms(eta), vram(), ladder.rung)

            if elapsed + eta > deadline_s:
                LOG.warning("projected finish %s exceeds the %.0fh deadline",
                            hms(elapsed + eta), args.deadline_hours)

        # One re-plan, from measured throughput rather than a guess.
        if not replanned and step >= args.calibrate_at and use_our_schedule:
            replanned = True
            tok_s = tokens_done / max(1e-9, prior_elapsed + time.time() - t_start)
            projected = args.tokens / max(1.0, tok_s)
            LOG.info("calibration: %s tok/s sustained -> %s for the full %s tokens "
                     "(deadline %s, %s headroom)",
                     human(tok_s), hms(projected), human(args.tokens),
                     hms(deadline_s), hms(deadline_s - projected))
            if projected > deadline_s:
                if args.enforce_deadline:
                    affordable = int(tok_s * deadline_s * 0.97)
                    total_steps = max(step + 1, affordable // batch_tokens)
                    LOG.warning("--enforce-deadline: trimming the run to %d steps "
                                "(%s tokens) so the cosine still lands at min_lr "
                                "inside one week", total_steps,
                                human(total_steps * batch_tokens))
                else:
                    LOG.warning("this run will overrun the deadline. Options: "
                                "--enforce-deadline (trim tokens), --ctx 512 "
                                "(~1.7x faster), or fewer tokens via --tokens")

        if step % args.eval_every == 0 or step == total_steps:
            oven.model.eval()
            with torch.no_grad():
                vals = []
                for ids, targets in data.val_batches(ladder.micro_batch, args.eval_batches):
                    ids, targets = ids.to(device), targets.to(device)
                    vals.append(float(forward_loss(oven.model.inner, ids, targets, ladder)))
            oven.model.train()
            v = sum(vals) / len(vals)
            best_val = min(best_val, v)
            LOG.info("eval step %d/%d loss=%.4f ppl=%.2f (best %.4f)",
                     step, total_steps, v, math.exp(min(v, 20)), best_val)

        if step % args.sample_every == 0 and oven.tokenizer_kind == "hf":
            try:
                oven.model.eval()
                text = oven.complete(args.sample_prompt, max_new_tokens=48,
                                     temperature=0.8, top_k=50)
                oven.model.train()
                LOG.info("sample @ %d: %r", step, text[:300])
            except Exception as exc:  # never let a sample kill a week-long run
                LOG.warning("sample failed at step %d: %s", step, exc)

        due = (step % args.save_every == 0
               or (time.time() - last_save) > args.save_every_minutes * 60
               or step == total_steps)
        if due:
            elapsed = prior_elapsed + (time.time() - t_start)
            save_checkpoint(paths.latest, oven, cfg, opt, step, tokens_done,
                            ladder, best_val, elapsed)
            if step % args.save_every == 0:
                save_checkpoint(paths.ckpt_dir / f"step_{step:07d}.pt", oven, cfg, opt,
                                step, tokens_done, ladder, best_val, elapsed)
                prune_checkpoints(paths.ckpt_dir, args.keep_checkpoints)
            last_save = time.time()
            LOG.info("checkpoint @ step %d (%s tokens, %s elapsed)",
                     step, human(tokens_done), hms(elapsed))

        if benchmark_only and step >= args.calibrate_at:
            LOG.info("benchmark complete — stopping before the real run")
            return

        elapsed = prior_elapsed + (time.time() - t_start)
        if elapsed > deadline_s:
            LOG.warning("wall-clock deadline of %.0fh reached at step %d — stopping "
                        "cleanly; --resume will continue if you want more",
                        args.deadline_hours, step)
            break

    elapsed = prior_elapsed + (time.time() - t_start)
    export_final(paths, oven, cfg)
    paths.report.write_text(json.dumps({
        "name": cfg.name,
        "params": params,
        "config": cfg.to_dict(),
        "tokens_trained": tokens_done,
        "steps": step,
        "total_steps": total_steps,
        "best_val_loss": best_val,
        "best_val_ppl": math.exp(min(best_val, 20)),
        "elapsed_seconds": elapsed,
        "elapsed_human": hms(elapsed),
        "deadline_hours": args.deadline_hours,
        "within_deadline": elapsed <= deadline_s,
        "final_ladder": asdict(ladder),
        "dataset_bytes": paths.bin.stat().st_size,
        "optimizer": args.optimizer,
        "sources": args.sources,
    }, indent=2, default=str))
    LOG.info("DONE — %s tokens in %s (deadline %.0fh, %s). best val ppl %.2f",
             human(tokens_done), hms(elapsed), args.deadline_hours,
             "within budget" if elapsed <= deadline_s else "OVER BUDGET",
             math.exp(min(best_val, 20)))


# ---------------------------------------------------------------------------
# Stage 4 — generate from the result, through NeoOven
# ---------------------------------------------------------------------------

def stage_sample(paths: Paths, args: argparse.Namespace) -> None:
    from hypernix.models.neo_oven import preheat_brewed

    if not args.no_patch:
        patch_brewer_attention()
    source = paths.final if paths.final.exists() else paths.latest
    if not source.exists():
        raise SystemExit("nothing trained yet")
    oven = preheat_brewed(source, tokenizer_source=paths.tokenizer,
                          device=None, dtype="float32")
    out = oven.complete(args.sample_prompt, max_new_tokens=args.max_new_tokens,
                        temperature=args.temperature, top_k=50)
    LOG.info("prompt=%r", args.sample_prompt)
    LOG.info("output=%r", out)
    print(out)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__.split("\n\n")[0],
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument("stage", choices=["all", "tokenizer", "pack", "train",
                                     "benchmark", "sample"])
    p.add_argument("--name", default=RUN_NAME)
    p.add_argument("--out", default="./runs", help="run directory root")
    p.add_argument("--hypernix-src", default=None,
                   help="path to a HyperNix-pip/src checkout, if not pip-installed")
    p.add_argument("--seed", type=int, default=1337)
    p.add_argument("--force", action="store_true", help="redo a completed stage")

    g = p.add_argument_group("model")
    g.add_argument("--vocab-size", type=int, default=VOCAB_SIZE)
    g.add_argument("--n-layers", type=int, default=N_LAYERS)
    g.add_argument("--n-heads", type=int, default=N_HEADS)
    g.add_argument("--n-kv-heads", type=int, default=N_KV_HEADS)
    g.add_argument("--d-model", type=int, default=D_MODEL)
    g.add_argument("--d-ff", type=int, default=D_FF)
    g.add_argument("--ctx", type=int, default=CTX)
    g.add_argument("--target-params", type=int, default=48_700_000)
    g.add_argument("--sliding-window", action="store_true",
                   help="enable SWA on odd layers (safe only with the patch)")

    g = p.add_argument_group("data")
    g.add_argument("--tokens", type=int, default=TOKEN_BUDGET)
    g.add_argument("--val-tokens", type=int, default=VAL_TOKENS)
    g.add_argument("--tok-sample-bytes", type=int, default=TOK_SAMPLE_BYTES)
    g.add_argument("--pack-batch", type=int, default=1000)
    g.add_argument("--disk-ceiling-gb", type=float, default=DISK_CEILING_GB)
    g.add_argument("--keep-sample", action="store_true")
    g.add_argument("--local-text", default=None,
                   help="train on a local directory/glob of .txt instead of streaming HF")

    g = p.add_argument_group("training")
    g.add_argument("--micro-batch", type=int, default=MICRO_BATCH)
    g.add_argument("--accum", type=int, default=0, help="0 = derive from tokens/step")
    g.add_argument("--tokens-per-step", type=int, default=TOKENS_PER_STEP)
    g.add_argument("--lr", type=float, default=PEAK_LR)
    g.add_argument("--optimizer", default="adamw",
                   choices=["adamw", "agedcookerv5", "pcv5", "pcv6"])
    g.add_argument("--deadline-hours", type=float, default=DEADLINE_HOURS)
    g.add_argument("--enforce-deadline", action="store_true",
                   help="trim the token budget if calibration says it won't fit")
    g.add_argument("--calibrate-at", type=int, default=CALIBRATE_AT_STEP)
    g.add_argument("--resume", action="store_true")
    g.add_argument("--amp", action="store_true",
                   help="FP16 autocast — a pessimisation on Pascal, see the docstring")
    g.add_argument("--no-patch", action="store_true",
                   help="do not fix BrewerAttention's mask composition")
    g.add_argument("--allow-cpu", action="store_true")

    g = p.add_argument_group("bookkeeping")
    g.add_argument("--log-every", type=int, default=LOG_EVERY)
    g.add_argument("--eval-every", type=int, default=EVAL_EVERY)
    g.add_argument("--eval-batches", type=int, default=EVAL_BATCHES)
    g.add_argument("--save-every", type=int, default=SAVE_EVERY)
    g.add_argument("--save-every-minutes", type=float, default=SAVE_EVERY_MINUTES)
    g.add_argument("--sample-every", type=int, default=SAMPLE_EVERY)
    g.add_argument("--keep-checkpoints", type=int, default=KEEP_CHECKPOINTS)
    g.add_argument("--sample-prompt", default="The first thing to understand about")
    g.add_argument("--max-new-tokens", type=int, default=128)
    g.add_argument("--temperature", type=float, default=0.8)

    args = p.parse_args(argv)
    args.sources = SOURCES
    if args.accum <= 0:
        per_micro = args.micro_batch * args.ctx
        args.accum = max(1, round(args.tokens_per_step / per_micro))
    if args.amp:
        LOG.warning("--amp on Pascal: GP104 runs FP16 GEMMs at 1/64 of FP32. "
                    "This will be slower, not faster.")
    return args


def main(argv: list[str] | None = None) -> None:
    args = parse_args(argv)
    paths = Paths(Path(args.out).expanduser() / args.name)
    setup_logging(paths, args.stage)
    LOG.info("=" * 78)
    LOG.info("hypernix0xv2 pretrain — stage=%s run=%s", args.stage, paths.root)
    import_hypernix(args.hypernix_src)

    if args.stage in ("all", "tokenizer"):
        stage_tokenizer(paths, args)
    if args.stage in ("all", "pack"):
        stage_pack(paths, args)
    if args.stage in ("all", "train"):
        stage_train(paths, args)
    if args.stage == "benchmark":
        stage_train(paths, args, benchmark_only=True)
    if args.stage == "sample":
        stage_sample(paths, args)


if __name__ == "__main__":
    main()