Add presentation bundle (reports, curves, run log, code) for the Colab notebook
439e7b1 verified 15 days ago
41.6 kB
"""Smoke test: prove the training pipeline and measure VRAM / throughput before any real run.
Sequence (every step logs; results land in reports/smoke_test.json and a markdown summary
is printed at the end):
1. device info, seed, config; ``PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`` set before
   CUDA initialises (variable-size probing fragments the caching allocator otherwise).
2. build ``EncoderForMLM`` from config.yaml and print ``count_parameters``.
3. data: packed real pretraining text if data/processed/{pretrain.parquet,tokenizer.json}
   exist, else random token ids. MLM masking exactly as configured (15%, 80/10/10, loss on
   masked positions only, specials never masked).
4. micro-batch finder: per candidate, 1 warmup + 3 timed optimizer steps, each made of
   two accumulated fwd/bwd micro-steps under bf16 autocast + fused AdamW. Two, not one,
   because production accumulates: from the second micro-step on the fp32 gradients
   (1.26 GiB) are resident while the next activation set is built, and that is the real
   high-water mark (a single fwd/bwd/opt step under-reports it by one gradient copy).
   Peak allocated/reserved recorded, OOMs caught; once without and once with gradient
   checkpointing. The choice follows config.vram: the largest candidate whose reserved
   peak sits in [target_gb_low, target_gb_high], else the largest under target_gb_high,
   never above hard_cap_gb. The throughput run's own reserved peak is then checked
   against the same cap and band as a second opinion.
5. throughput at the chosen micro-batch (no checkpointing): CUDA-event-timed micro-steps
   with the real gradient-accumulation pattern (accumulation = round(256 / micro-batch);
   the effective batch actually realised is reported, and a warning is logged when it is
   not exactly 256), an isolated optimizer-step cost, tokens/s, and projected hours per
   epoch / per 3 epochs. Epoch tokens come from --epoch-tokens, else from the data
   manifest's pretraining character count times the tokens/char measured on the packed
   sample, else the 117M estimate. Plus a check of which SDPA backend actually runs on
   the mask-free (Flash) and masked (mem-efficient) paths.
6. mini convergence check from a fresh init: 30 steps at LR 1e-4 with a 10-step warmup.
   The first loss should sit near ln(V) = 10.40; on real text it should fall visibly.
7. inference latency of ``EncoderForClassification`` for one 512-token email: GPU bf16
   (100 samples), CPU fp32 at 512 and 128 tokens (10 samples), versus the instructor's
   0.074 s/sample benchmark.
8. write the JSON report and print the summary.
CLI: python src/smoke_test.py [--steps 30] [--skip-cpu] [--candidates 16,24,12,8,4]
"""

import os

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

import argparse  # noqa: E402
import gc  # noqa: E402
import logging  # noqa: E402
import math  # noqa: E402
import sys  # noqa: E402
import time  # noqa: E402
import warnings  # noqa: E402
from dataclasses import asdict, dataclass  # noqa: E402
from pathlib import Path  # noqa: E402
from typing import Any  # noqa: E402

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

import torch  # noqa: E402
import torch.nn.functional as F  # noqa: E402
from torch import Tensor  # noqa: E402
from torch.nn.attention import SDPBackend, sdpa_kernel  # noqa: E402

import common  # noqa: E402
from model import EncoderConfig, EncoderForClassification, EncoderForMLM, count_parameters  # noqa: E402

GIB = float(2**30)
INSTRUCTOR_BENCHMARK_S = 0.074
DEFAULT_EPOCH_TOKENS = 117e6          # verify_data.md: ~117M tokens/epoch at 4 chars/token
PROBE_ACCUM = 2                       # micro-steps per finder step: grads resident during the 2nd
VRAM_UNIT_NOTE = "config vram.*_gb values are read as GiB and compared with torch.cuda.max_memory_reserved() / 2**30"
LOG = logging.getLogger("smoke_test")


# ------------------------------------------------------------------------------ data
@dataclass
class TokenPool:
    """Fixed-length token sequences the test samples batches from."""

    seqs: Tensor                 # [n, L] int64 on the training device
    special_ids: list[int]       # never masked, never used as MLM targets
    mask_id: int
    source: str                  # "real" | "random"
    n_docs: int
    detail: str
    tokens_per_char: float | None = None   # measured on the packed sample (real data only)

    @property
    def n(self) -> int:
        return int(self.seqs.shape[0])

    def batch(self, size: int, gen: torch.Generator) -> Tensor:
        idx = torch.randint(self.n, (size,), device=self.seqs.device, generator=gen)
        return self.seqs[idx]


def _special_ids_from_config(cfg: dict) -> tuple[list[int], int]:
    specials = cfg["tokenizer"]["special_tokens"]
    return list(range(len(specials))), specials.index("[MASK]")


def _text_columns(names: set[str], path: Path) -> list[str]:
    cols = ["text"] if "text" in names else [c for c in ("subject", "body") if c in names]
    if not cols:
        raise ValueError(f"{path}: no text column among {sorted(names)}")
    return cols


def _texts(batch) -> list[str]:
    """Model input text per row: the ``text`` column, else subject + blank line + body."""
    names = set(batch.schema.names)
    if "text" in names:
        return [t or "" for t in batch.column("text").to_pylist()]
    subj = batch.column("subject").to_pylist() if "subject" in names else [""] * batch.num_rows
    body = batch.column("body").to_pylist()
    return [f"{s or ''}\n\n{b or ''}" for s, b in zip(subj, body)]


def _assert_no_holdout_overlap(texts: list[str], holdout: Path) -> None:
    """The hard no-leak rule, checked on the exact documents this run packs: none may
    equal a held-out TREC-06 email under the verify_data.md sec. 7 normalisation
    (lowercase, [a-z0-9] only). Cheap (one sha1 per row) and independent of whatever
    bookkeeping the data pipeline recorded."""
    if not holdout.exists():
        LOG.warning("no holdout parquet at %s; in-code leak check skipped", holdout)
        return
    import pyarrow.parquet as pq

    pf = pq.ParquetFile(holdout)
    cols = _text_columns(set(pf.schema_arrow.names), holdout)
    hold: set[str] = set()
    for batch in pf.iter_batches(batch_size=4096, columns=cols):
        hold.update(common.sha1_hex(n) for n in map(common.normalize_text, _texts(batch)) if n)
    hits = sum(common.sha1_hex(n) in hold for n in map(common.normalize_text, texts) if n)
    assert hits == 0, f"{hits} pretraining documents match held-out emails in {holdout}"
    LOG.info("no-leak check: 0 of %d packed docs match the %d holdout emails", len(texts), len(hold))


def build_token_pool(cfg: dict, mcfg: EncoderConfig, device: torch.device,
                     n_docs: int, random_seqs: int) -> TokenPool:
    """Real packed text when the data agent's artefacts exist, else random ids.
    Packing follows the pretraining recipe: documents are concatenated with [SEP]
    between them and cut into max_seq_len windows (RoBERTa FULL-SENTENCES), which is
    what keeps the Flash backend eligible (no padding -> no mask).
    """
    L = mcfg.max_seq_len
    processed = Path(cfg["paths"]["data_processed"])
    parquet = processed / "pretrain.parquet"
    tok_path = Path(cfg["paths"]["tokenizer"])
    if parquet.exists() and tok_path.exists():
        import pyarrow.parquet as pq
        from tokenizers import Tokenizer

        pf = pq.ParquetFile(parquet)
        names = set(pf.schema_arrow.names)
        source_col = next((c for c in ("source", "dataset_name", "dataset") if c in names), None)
        cols = _text_columns(names, parquet) + ([source_col] if source_col else [])
        table = next(pf.iter_batches(batch_size=n_docs, columns=cols))
        if source_col:
            # Hard rule: the TREC-06 holdout must never reach pretraining text.
            srcs = {str(s).lower().replace("_", "-") for s in table.column(source_col).to_pylist()}
            assert not any("trec-06" in s or "trec06" in s for s in srcs), f"TREC-06 rows in {parquet}"
        texts = _texts(table)
        _assert_no_holdout_overlap(texts, processed / "test_holdout.parquet")

        tok = Tokenizer.from_file(str(tok_path))
        special_ids = [i for i in (tok.token_to_id(s) for s in cfg["tokenizer"]["special_tokens"]) if i is not None]
        sep_id, mask_id = tok.token_to_id("[SEP]"), tok.token_to_id("[MASK]")
        assert sep_id is not None and mask_id is not None, "tokenizer lacks [SEP]/[MASK]"
        assert tok.get_vocab_size() <= mcfg.vocab_size, "tokenizer vocab larger than model vocab"
        stream: list[int] = []
        for enc in tok.encode_batch(texts, add_special_tokens=False):
            stream.extend(enc.ids)
            stream.append(sep_id)
        n_seq = len(stream) // L
        if n_seq < 1:
            raise ValueError("not enough text to pack a single sequence")
        seqs = torch.tensor(stream[: n_seq * L], dtype=torch.long).view(n_seq, L).to(device)
        n_chars = sum(len(t) for t in texts)
        tpc = len(stream) / n_chars if n_chars else None
        return TokenPool(seqs, special_ids, mask_id, "real", len(texts),
                         f"{len(texts)} docs ({n_chars:,} chars) -> {len(stream):,} tokens -> {n_seq} packed x {L}",
                         tokens_per_char=tpc)

    special_ids, mask_id = _special_ids_from_config(cfg)
    gen = torch.Generator(device="cpu").manual_seed(cfg["seed"])
    seqs = torch.randint(len(special_ids), mcfg.vocab_size, (random_seqs, L), generator=gen).to(device)
    return TokenPool(seqs, special_ids, mask_id, "random", 0,
                     f"{random_seqs} random x {L} ids in [{len(special_ids)}, {mcfg.vocab_size})")


def mlm_mask(ids: Tensor, pool: TokenPool, vocab_size: int, p: float, replace_probs: list[float],
             gen: torch.Generator) -> tuple[Tensor, Tensor, Tensor]:
    """BERT dynamic masking: select ``p`` of the non-special positions; of those, replace
    80% with [MASK], 10% with a random token, leave 10% unchanged. Returns
    ``(inputs, labels_at_masked, masked_positions)`` so the model can gather before the head.
    """
    special = torch.isin(ids, torch.tensor(pool.special_ids, device=ids.device))
    sel = (torch.rand(ids.shape, device=ids.device, generator=gen) < p) & ~special
    if not bool(sel.any()):
        first = (~special).view(-1).nonzero()[0]
        sel.view(-1)[first] = True
    labels = ids[sel]
    inputs = ids.clone()
    r = torch.rand(ids.shape, device=ids.device, generator=gen)
    p_mask, p_rand, _ = replace_probs
    to_mask = sel & (r < p_mask)
    to_rand = sel & (r >= p_mask) & (r < p_mask + p_rand)
    inputs[to_mask] = pool.mask_id
    n_rand = int(to_rand.sum())
    if n_rand:
        inputs[to_rand] = torch.randint(len(pool.special_ids), vocab_size, (n_rand,),
                                        device=ids.device, generator=gen)
    return inputs, labels, sel


# ------------------------------------------------------------------------ training bits
def build_model(mcfg: EncoderConfig, device: torch.device, seed: int) -> EncoderForMLM:
    common.set_seed(seed)
    with torch.device(device):
        model = EncoderForMLM(mcfg)
    return model.train()


def build_optimizer(model: torch.nn.Module, lr: float, pcfg: dict) -> torch.optim.AdamW:
    """AdamW with weight decay only on matrices (biases and LayerNorm excluded), fused."""
    decay = [p for p in model.parameters() if p.requires_grad and p.ndim >= 2]
    no_decay = [p for p in model.parameters() if p.requires_grad and p.ndim < 2]
    groups = [{"params": decay, "weight_decay": pcfg["weight_decay"]},
              {"params": no_decay, "weight_decay": 0.0}]
    kw: dict[str, Any] = dict(lr=lr, betas=tuple(pcfg["betas"]), eps=pcfg["eps"])
    try:
        return torch.optim.AdamW(groups, fused=True, **kw)
    except (TypeError, RuntimeError) as e:  # fused unavailable on this build/device
        LOG.warning("fused AdamW unavailable (%s); falling back to foreach", e)
        return torch.optim.AdamW(groups, **kw)


class Trainer:
    """Minimal MLM step machinery shared by the finder, the throughput run and the
    convergence check, so all three measure the same code path the real trainer will use."""

    def __init__(self, model: EncoderForMLM, opt: torch.optim.AdamW, pool: TokenPool,
                 cfg: dict, mcfg: EncoderConfig, gen: torch.Generator) -> None:
        self.model, self.opt, self.pool, self.gen = model, opt, pool, gen
        self.vocab = mcfg.vocab_size
        p = cfg["pretrain"]
        self.p_mask = p["mlm_probability"]
        self.replace_probs = list(p["mask_replace_probs"])
        self.grad_clip = p["grad_clip"]

    def fwd_bwd(self, micro_batch: int, loss_scale: float = 1.0) -> float:
        ids = self.pool.batch(micro_batch, self.gen)
        inputs, labels, pos = mlm_mask(ids, self.pool, self.vocab, self.p_mask, self.replace_probs, self.gen)
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            logits = self.model(inputs, attention_mask=None, masked_positions=pos)
            loss = F.cross_entropy(logits, labels)
        (loss * loss_scale).backward()
        return float(loss.detach())

    def opt_step(self) -> None:
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip)
        self.opt.step()
        self.opt.zero_grad(set_to_none=True)

    def step(self, micro_batch: int) -> float:
        loss = self.fwd_bwd(micro_batch)
        self.opt_step()
        return loss

    def accum_step(self, micro_batch: int, accum: int) -> float:
        """One optimizer step from ``accum`` accumulated micro-steps: the production
        pattern. From the second micro-step on, the fp32 ``.grad`` tensors of the whole
        model (4 B/param = 1.26 GiB) are resident while a fresh activation set is built,
        so this is the pattern that exposes the true peak; a lone fwd/bwd/opt step
        measures one gradient copy too little. Returns the mean micro-step loss."""
        total = 0.0
        for _ in range(accum):
            total += self.fwd_bwd(micro_batch, 1.0 / accum)
        self.opt_step()
        return total / accum

    def release(self) -> None:
        self.opt.zero_grad(set_to_none=True)
        gc.collect()
        torch.cuda.empty_cache()


# ------------------------------------------------------------------------- step 4: finder
@dataclass
class ProbeResult:
    micro_batch: int
    grad_checkpointing: bool
    ok: bool
    steps: int                   # timed optimizer steps, each of `micro_steps` accumulated fwd/bwd
    micro_steps: int
    max_allocated_gib: float
    max_reserved_gib: float
    device_used_after_gib: float
    ms_per_micro_step: float | None   # includes 1/micro_steps of an optimizer step: rough, see throughput
    tokens_per_s: float | None
    error: str | None


def probe(trainer: Trainer, micro_batch: int, steps: int, seq_len: int, ckpt: bool) -> ProbeResult:
    trainer.model.config.grad_checkpointing = ckpt
    trainer.release()
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    ok, err, ms = True, None, None
    try:
        # One untimed step absorbs cuBLAS/allocator warmup at this shape and allocates the
        # AdamW state; the peak-memory window covers it too. Every step uses the
        # accumulation pattern so the resident-gradient peak is part of the measurement.
        trainer.accum_step(micro_batch, PROBE_ACCUM)
        torch.cuda.synchronize()
        t0 = common.timer()
        for _ in range(steps):
            trainer.accum_step(micro_batch, PROBE_ACCUM)
        torch.cuda.synchronize()
        ms = (common.timer() - t0) * 1000.0 / (steps * PROBE_ACCUM)
    except torch.OutOfMemoryError as e:
        ok = False
        err = str(e).splitlines()[0][:240]
        LOG.warning("B=%d ckpt=%s: OOM -> %s", micro_batch, ckpt, err)
    peak_alloc = torch.cuda.max_memory_allocated() / GIB
    peak_res = torch.cuda.max_memory_reserved() / GIB
    free, total = torch.cuda.mem_get_info()
    trainer.release()
    tok_s = micro_batch * seq_len / (ms / 1000.0) if ms else None
    r = ProbeResult(micro_batch, ckpt, ok, steps, PROBE_ACCUM, round(peak_alloc, 3), round(peak_res, 3),
                    round((total - free) / GIB, 3), round(ms, 1) if ms else None,
                    round(tok_s) if tok_s else None, err)
    LOG.info("probe B=%d ckpt=%s ok=%s peak_alloc=%.2f GiB peak_reserved=%.2f GiB ms/micro-step=%s tok/s=%s",
             r.micro_batch, ckpt, ok, peak_alloc, peak_res, r.ms_per_micro_step, r.tokens_per_s)
    return r


def choose_micro_batch(results: list[ProbeResult], low: float, high: float, cap: float) -> tuple[int | None, str]:
    fit = [r for r in results if r.ok and r.max_reserved_gib <= cap]
    in_band = [r for r in fit if low <= r.max_reserved_gib <= high]
    if in_band:
        b = max(in_band, key=lambda r: r.micro_batch)
        return b.micro_batch, f"largest candidate with reserved peak in [{low}, {high}] GiB ({b.max_reserved_gib} GiB)"
    under = [r for r in fit if r.max_reserved_gib <= high]
    if under:
        b = max(under, key=lambda r: r.micro_batch)
        return b.micro_batch, f"no candidate reached {low} GiB; largest under {high} GiB ({b.max_reserved_gib} GiB)"
    if fit:
        b = min(fit, key=lambda r: r.micro_batch)
        return b.micro_batch, f"WARNING every fitting candidate exceeds {high} GiB; smallest under the {cap} GiB cap ({b.max_reserved_gib} GiB)"
    return None, "no candidate fits under the hard cap"


# --------------------------------------------------------------------- step 5: throughput
def measure_throughput(trainer: Trainer, micro_batch: int, seq_len: int, eff_batch: int,
                       timed_steps: int, warmup_steps: int, epoch_tokens: float) -> dict[str, Any]:
    trainer.model.config.grad_checkpointing = False
    trainer.release()
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    # Nearest integer, not floor: 256 is divisible by neither 24 nor 12, and every number
    # below must describe the batch that is actually realised (floor would silently train
    # at 240 while reporting 256). The realised batch is reported next to the target.
    accum = max(1, round(eff_batch / micro_batch))
    eff_actual = accum * micro_batch
    if eff_actual != eff_batch:
        LOG.warning("effective batch %d (= %d x %d) differs from the configured %d because %d is not a divisor; "
                    "tokens/s, steps/epoch and hours below use the realised %d",
                    eff_actual, accum, micro_batch, eff_batch, micro_batch, eff_actual)
    scale = 1.0 / accum

    def micro(i: int) -> None:
        trainer.fwd_bwd(micro_batch, scale)
        if (i + 1) % accum == 0:
            trainer.opt_step()

    for i in range(warmup_steps):
        micro(i)
    trainer.release()
    torch.cuda.synchronize()
    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
    start.record()
    for i in range(timed_steps):
        micro(i)
    end.record()
    torch.cuda.synchronize()
    elapsed_ms = start.elapsed_time(end)
    opt_steps_in_window = timed_steps // accum
    trainer.release()

    # Isolated optimizer-step cost (clip + fused AdamW + zero_grad), amortised over `accum`.
    opt_ms = []
    for _ in range(3):
        trainer.fwd_bwd(micro_batch, scale)
        torch.cuda.synchronize()
        s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
        s.record()
        trainer.opt_step()
        e.record()
        torch.cuda.synchronize()
        opt_ms.append(s.elapsed_time(e))
    trainer.release()
    ms_opt = sum(opt_ms) / len(opt_ms)

    ms_fwd_bwd = (elapsed_ms - opt_steps_in_window * ms_opt) / timed_steps
    ms_micro_raw = elapsed_ms / timed_steps
    tokens_per_micro = micro_batch * seq_len
    tok_s_raw = tokens_per_micro / (ms_micro_raw / 1000.0)
    ms_per_opt_step = accum * ms_fwd_bwd + ms_opt
    tok_s_eff = eff_actual * seq_len / (ms_per_opt_step / 1000.0)
    steps_per_epoch = epoch_tokens / (eff_actual * seq_len)
    hours_epoch = epoch_tokens / tok_s_eff / 3600.0
    vram = common.vram_stats()
    out = {
        "micro_batch": micro_batch, "grad_accumulation": accum, "effective_batch": eff_actual,
        "effective_batch_target": eff_batch, "effective_batch_matches_target": eff_actual == eff_batch,
        "peak_reserved_gib": round(vram["max_reserved_gb"], 3), "peak_allocated_gib": round(vram["max_allocated_gb"], 3),
        "timed_micro_steps": timed_steps, "warmup_micro_steps": warmup_steps,
        "ms_per_micro_step_raw": round(ms_micro_raw, 2), "ms_fwd_bwd": round(ms_fwd_bwd, 2),
        "ms_optimizer_step": round(ms_opt, 2), "ms_per_optimizer_step_total": round(ms_per_opt_step, 1),
        "tokens_per_s_raw": round(tok_s_raw), "tokens_per_s_effective": round(tok_s_eff),
        "epoch_tokens": epoch_tokens, "optimizer_steps_per_epoch": round(steps_per_epoch),
        "hours_per_epoch": round(hours_epoch, 2), "hours_3_epochs": round(3 * hours_epoch, 2),
        "vram": vram,
    }
    LOG.info("throughput: %s", out)
    return out


def resolve_epoch_tokens(override: float | None, pool: TokenPool, processed: Path) -> tuple[float, str]:
    """Tokens per pretraining epoch for the ETA: an explicit --epoch-tokens wins; else the
    data manifest's pretraining character count scaled by the tokens/char measured on the
    packed sample (the only place the real tokenizer and the real corpus meet); else the
    research estimate."""
    if override is not None:
        return float(override), "--epoch-tokens"
    manifest = processed / "splits_manifest.json"
    if pool.tokens_per_char and manifest.exists():
        m = common.read_json(manifest)
        chars = int(m["counts"]["totals"]["pretrain"]["chars"])
        note = " [dev subsample]" if m.get("dev_subsample") else ""
        return chars * pool.tokens_per_char, f"splits_manifest.json {chars:,} chars x {pool.tokens_per_char:.4f} tok/char{note}"
    return DEFAULT_EPOCH_TOKENS, "verify_data.md estimate (117M tokens at 4 chars/token)"


def check_sdpa_backends(model: EncoderForMLM, seq_len: int, vocab: int) -> dict[str, Any]:
    """Which fused kernel actually runs. Flash must take the mask-free (packed) path and
    memory-efficient the bool-mask path; a silent math fallback would store the LxL
    softmax and blow the activation budget."""
    model.config.grad_checkpointing = False
    ids = torch.randint(5, vocab, (2, seq_len), device="cuda")
    mask = torch.ones(2, seq_len, dtype=torch.bool, device="cuda")
    mask[1, seq_len // 2:] = False
    cases = {
        "flash_no_mask": ([SDPBackend.FLASH_ATTENTION], None),
        "efficient_bool_mask": ([SDPBackend.EFFICIENT_ATTENTION], mask),
        "flash_bool_mask_expected_fail": ([SDPBackend.FLASH_ATTENTION], mask),
    }
    out: dict[str, Any] = {}
    for name, (backends, m) in cases.items():
        try:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                with sdpa_kernel(backends), torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                    h = model.encoder(ids, m)
                h.float().mean().backward()
            out[name] = "ran"
        except RuntimeError as e:
            out[name] = f"unavailable: {str(e).splitlines()[0][:100]}"
        model.zero_grad(set_to_none=True)
    out["flags"] = {
        "flash_sdp_enabled": torch.backends.cuda.flash_sdp_enabled(),
        "mem_efficient_sdp_enabled": torch.backends.cuda.mem_efficient_sdp_enabled(),
        "math_sdp_enabled": torch.backends.cuda.math_sdp_enabled(),
    }
    LOG.info("sdpa backends: %s", out)
    return out


# -------------------------------------------------------------------- step 6: convergence
def convergence_check(cfg: dict, mcfg: EncoderConfig, pool: TokenPool, device: torch.device,
                      micro_batch: int, steps: int, lr: float, warmup: int) -> dict[str, Any]:
    model = build_model(mcfg, device, cfg["seed"] + 1)
    opt = build_optimizer(model, lr, cfg["pretrain"])
    gen = torch.Generator(device=device).manual_seed(cfg["seed"] + 1)
    tr = Trainer(model, opt, pool, cfg, mcfg, gen)
    losses: list[tuple[int, float]] = []
    for step in range(1, steps + 1):
        for g in opt.param_groups:
            g["lr"] = lr * min(1.0, step / max(1, warmup))
        loss = tr.step(micro_batch)
        if not math.isfinite(loss):
            LOG.error("non-finite loss at step %d", step)
            losses.append((step, loss))
            break
        if step == 1 or step % 5 == 0 or step == steps:
            losses.append((step, round(loss, 4)))
            LOG.info("convergence step %3d/%d  loss %.4f  lr %.2e", step, steps, loss, opt.param_groups[0]["lr"])
    tr.release()
    first, last = losses[0][1], losses[-1][1]
    out = {
        "micro_batch": micro_batch, "steps": steps, "lr": lr, "warmup_steps": warmup,
        "ln_vocab": round(math.log(mcfg.vocab_size), 4), "first_loss": first, "last_loss": last,
        "drop": round(first - last, 4), "losses": losses, "data": pool.source,
        "vram": common.vram_stats(),
    }
    del tr, opt, model
    gc.collect()
    torch.cuda.empty_cache()
    return out


# ------------------------------------------------------------------------ step 7: latency
def _time_gpu(fn, n: int, warmup: int) -> dict[str, float]:
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    times = []
    for _ in range(n):
        s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
        s.record()
        fn()
        e.record()
        torch.cuda.synchronize()
        times.append(s.elapsed_time(e))
    times.sort()
    return {"n": n, "avg_ms": round(sum(times) / n, 2), "p50_ms": round(times[n // 2], 2), "max_ms": round(times[-1], 2)}


def _time_cpu(fn, n: int, warmup: int) -> dict[str, float]:
    for _ in range(warmup):
        fn()
    times = []
    for _ in range(n):
        t0 = time.perf_counter()
        fn()
        times.append((time.perf_counter() - t0) * 1000.0)
    times.sort()
    return {"n": n, "avg_ms": round(sum(times) / n, 1), "p50_ms": round(times[n // 2], 1), "max_ms": round(times[-1], 1)}


def inference_latency(cfg: dict, mcfg: EncoderConfig, device: torch.device, seq_len: int,
                      gpu_samples: int, cpu_samples: int, cpu_threads: int, skip_cpu: bool) -> dict[str, Any]:
    common.set_seed(cfg["seed"])
    with torch.device(device):
        clf = EncoderForClassification(mcfg)
    clf.eval()
    out: dict[str, Any] = {"instructor_benchmark_ms": INSTRUCTOR_BENCHMARK_S * 1000.0, "params": count_parameters(clf)["total"]}
    ids = torch.randint(5, mcfg.vocab_size, (1, seq_len), device=device)
    full_mask = torch.ones(1, seq_len, dtype=torch.bool, device=device)
    with torch.inference_mode(), torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        out[f"gpu_bf16_{seq_len}_no_mask"] = _time_gpu(lambda: clf(ids), gpu_samples, 10)
        out[f"gpu_bf16_{seq_len}_bool_mask"] = _time_gpu(lambda: clf(ids, full_mask), gpu_samples, 10)
    LOG.info("GPU latency: %s", {k: v for k, v in out.items() if k.startswith("gpu")})
    if not skip_cpu:
        clf = clf.cpu().float()
        torch.set_num_threads(cpu_threads)
        out["cpu_threads"] = torch.get_num_threads()
        with torch.inference_mode():
            for L in (seq_len, 128):
                x = ids[:, :L].cpu()
                out[f"cpu_fp32_{L}"] = _time_cpu(lambda: clf(x), cpu_samples, 2)
                LOG.info("CPU fp32 L=%d: %s", L, out[f"cpu_fp32_{L}"])
    del clf
    gc.collect()
    torch.cuda.empty_cache()
    return out


# ------------------------------------------------------------------------------ report
def _probe_table(rows: list[ProbeResult]) -> str:
    head = "| micro-batch | ckpt | result | peak alloc GiB | peak reserved GiB | device used GiB | ms/micro-step | tok/s |\n|---|---|---|---|---|---|---|---|"
    lines = [head]
    for r in rows:
        res = "ok" if r.ok else "OOM"
        lines.append(f"| {r.micro_batch} | {'on' if r.grad_checkpointing else 'off'} | {res} | {r.max_allocated_gib:.2f} | "
                     f"{r.max_reserved_gib:.2f} | {r.device_used_after_gib:.2f} | {r.ms_per_micro_step or '-'} | "
                     f"{f'{r.tokens_per_s:,}' if r.tokens_per_s else '-'} |")
    return "\n".join(lines)


def markdown_summary(rep: dict[str, Any]) -> str:
    d = rep["device"]
    vr = rep["vram_policy"]
    lines = [
        "## Smoke test summary",
        f"- device: {d.get('device', 'cpu')} (sm {d.get('sm', '?')}, {d.get('vram_gib', '?')} GiB, bf16={d.get('bf16', '?')}); "
        f"torch {d.get('torch')} / python {d.get('python')}; device used before model build: {rep['baseline_device_used_gib']} GiB",
        f"- model: {rep['param_counts']['mlm']['total']:,} params (MLM, tied) / {rep['param_counts']['classifier']['total']:,} (classifier)",
        f"- data: {rep['data']['source']} -- {rep['data']['detail']}",
        f"- VRAM policy: target [{vr['target_gb_low']}, {vr['target_gb_high']}] GiB reserved, hard cap {vr['hard_cap_gb']} GiB, alloc_conf={vr['alloc_conf']} ({VRAM_UNIT_NOTE})",
        "", f"### Micro-batch finder (no checkpointing; {rep['finder']['probe_steps']} timed optimizer steps of "
            f"{PROBE_ACCUM} accumulated micro-steps after 1 warmup step per candidate; grads resident during the 2nd micro-step)",
        _probe_table([ProbeResult(**r) for r in rep['finder']['no_ckpt']]),
        f"\n**chosen micro-batch: {rep['finder']['chosen_micro_batch']}** -- {rep['finder']['reason']}",
    ]
    fin = rep["finder"]
    if fin.get("throughput_peak_reserved_gib") is not None:
        verdict = ("WITHIN the hard cap" if fin["throughput_within_hard_cap"] else "**BREACHES the hard cap**")
        band = "in the target band" if fin["throughput_in_target_band"] else "outside the target band"
        lines.append(f"- second opinion from the throughput run (real accumulation at the chosen batch): "
                     f"{fin['throughput_peak_reserved_gib']} GiB reserved peak -- {verdict}, {band}")
    lines += [
        "", "### Micro-batch finder (gradient checkpointing on, ~x1.33 compute)", _probe_table([ProbeResult(**r) for r in rep['finder']['ckpt']]),
        f"\ncheckpointing choice: {rep['finder']['chosen_micro_batch_ckpt']} -- {rep['finder']['reason_ckpt']}",
    ]
    t = rep.get("throughput")
    if t:
        eff_note = "" if t["effective_batch_matches_target"] else f" (**target {t['effective_batch_target']} not divisible by {t['micro_batch']}; all numbers use {t['effective_batch']}**)"
        lines += [
            "", "### Throughput (chosen micro-batch, no checkpointing)",
            f"- micro-batch {t['micro_batch']} x accumulation {t['grad_accumulation']} = effective {t['effective_batch']} sequences{eff_note}",
            f"- fwd+bwd {t['ms_fwd_bwd']} ms/micro-step, optimizer step {t['ms_optimizer_step']} ms -> {t['ms_per_optimizer_step_total']} ms per optimizer step",
            f"- **{t['tokens_per_s_effective']:,} tokens/s** effective ({t['tokens_per_s_raw']:,} raw over {t['timed_micro_steps']} timed micro-steps); peak reserved {t['peak_reserved_gib']} GiB",
            f"- {t['epoch_tokens']/1e6:.1f}M-token epoch ({t.get('epoch_tokens_source', '?')}) = {t['optimizer_steps_per_epoch']} optimizer steps -> "
            f"**{t['hours_per_epoch']} h/epoch, {t['hours_3_epochs']} h for 3 epochs**",
        ]
    s = rep.get("sdpa_backends")
    if s:
        lines += ["", "### SDPA backends", *(f"- {k}: {v}" for k, v in s.items())]
    c = rep.get("convergence")
    if c:
        lines += [
            "", "### Mini convergence check",
            f"- {c['steps']} steps at micro-batch {c['micro_batch']}, LR {c['lr']:.0e} with {c['warmup_steps']}-step warmup, data={c['data']}",
            f"- ln(V) = {c['ln_vocab']}; first loss {c['first_loss']}, last loss {c['last_loss']} (drop {c['drop']})",
            "- trace: " + ", ".join(f"{st}:{lo}" for st, lo in c["losses"]),
        ]
    lat = rep.get("latency")
    if lat:
        lines += ["", "### Inference latency (batch 1), instructor benchmark 74 ms/sample",
                  "| setting | n | avg ms | p50 ms | max ms | vs 74 ms |", "|---|---|---|---|---|---|"]
        for k, v in lat.items():
            if isinstance(v, dict) and "avg_ms" in v:
                ratio = v["avg_ms"] / lat["instructor_benchmark_ms"]
                lines.append(f"| {k} | {v['n']} | {v['avg_ms']} | {v['p50_ms']} | {v['max_ms']} | {ratio:.2f}x |")
    lines += ["", f"total runtime: {rep['runtime_s']} s; report: {rep['report_path']}"]
    return "\n".join(lines)


# -------------------------------------------------------------------------------- main
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")
    ap.add_argument("--steps", type=int, default=30, help="optimizer steps for the convergence check")
    ap.add_argument("--skip-cpu", action="store_true", help="skip the CPU fp32 latency samples")
    ap.add_argument("--candidates", default=None,
                    help="comma list of micro-batches to probe (default: config pretrain.micro_batch_candidates "
                         "if present, else micro_batch_start,24,12,8,4)")
    ap.add_argument("--ckpt-candidates", default=None, help="micro-batches for the checkpointing sweep (default: same)")
    ap.add_argument("--probe-steps", type=int, default=3)
    ap.add_argument("--timed-steps", type=int, default=20)
    ap.add_argument("--n-docs", type=int, default=3000, help="real documents to pack when the parquet exists")
    ap.add_argument("--random-seqs", type=int, default=1024, help="random sequences when no real data exists")
    ap.add_argument("--epoch-tokens", type=float, default=None,
                    help="tokens per pretraining epoch for the ETA (default: data manifest chars x measured tokens/char, else 117e6)")
    ap.add_argument("--conv-lr", type=float, default=1e-4)
    ap.add_argument("--conv-warmup", type=int, default=10)
    ap.add_argument("--cpu-threads", type=int, default=16)
    ap.add_argument("--latency-samples", type=int, default=None, help="GPU latency samples (default config.eval.latency_samples)")
    return ap.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    t_start = common.timer()
    cfg = common.load_config(args.config)
    logger = common.setup_logging("smoke_test", cfg["paths"]["logs"])
    global LOG
    LOG = logger
    device_info = common.print_device_info()
    if not torch.cuda.is_available():
        LOG.error("CUDA is required for the smoke test")
        return 2
    LOG.info("PYTORCH_CUDA_ALLOC_CONF=%s", os.environ.get("PYTORCH_CUDA_ALLOC_CONF"))
    common.set_seed(cfg["seed"])
    device = torch.device("cuda")
    mcfg = EncoderConfig.from_dict(cfg["model"])
    L = mcfg.max_seq_len
    vr = cfg["vram"]
    pcfg = cfg["pretrain"]
    free, total = torch.cuda.mem_get_info()
    baseline_used = round((total - free) / GIB, 2)
    LOG.info("device used before model build (desktop + others): %.2f GiB", baseline_used)

    # 2. model
    model = build_model(mcfg, device, cfg["seed"])
    mlm_counts = count_parameters(model)
    LOG.info("EncoderForMLM params: %s", mlm_counts)
    print(f"EncoderForMLM: total {mlm_counts['total']:,} | embeddings {mlm_counts['embeddings']:,} | "
          f"layers {mlm_counts['encoder_layers']:,} | final_norm {mlm_counts['final_norm']:,} | head {mlm_counts['head']:,}")
    with torch.device("meta"):
        cls_counts = count_parameters(EncoderForClassification(mcfg))
    LOG.info("EncoderForClassification params: %s", cls_counts)
    LOG.info("after model build: %s", common.vram_stats())

    # 3. data
    pool = build_token_pool(cfg, mcfg, device, args.n_docs, args.random_seqs)
    LOG.info("token pool: %s -> %s", pool.source, pool.detail)

    # 4. finder
    opt = build_optimizer(model, pcfg["lr"], pcfg)
    gen = torch.Generator(device=device).manual_seed(cfg["seed"])
    trainer = Trainer(model, opt, pool, cfg, mcfg, gen)
    start = int(pcfg["micro_batch_start"])
    default_cands = ([int(b) for b in pcfg["micro_batch_candidates"]] if pcfg.get("micro_batch_candidates")
                     else [start] + [b for b in (24, 12, 8, 4) if b != start])
    cands = [int(x) for x in args.candidates.split(",")] if args.candidates else default_cands
    ckpt_cands = [int(x) for x in args.ckpt_candidates.split(",")] if args.ckpt_candidates else cands
    cands = [b for b in cands if b <= pool.n] or [pool.n]
    ckpt_cands = [b for b in ckpt_cands if b <= pool.n] or [pool.n]
    LOG.info("finder candidates: %s (no ckpt), %s (ckpt)", cands, ckpt_cands)
    # probe() flips grad_checkpointing on model.config, which IS mcfg (same dataclass
    # instance); remember the config.yaml value so later model builds inherit it, not the
    # last probe's setting.
    ckpt_default = bool(mcfg.grad_checkpointing)
    no_ckpt = [probe(trainer, b, args.probe_steps, L, ckpt=False) for b in cands]
    chosen, reason = choose_micro_batch(no_ckpt, vr["target_gb_low"], vr["target_gb_high"], vr["hard_cap_gb"])
    LOG.info("chosen micro-batch (no ckpt): %s -- %s", chosen, reason)
    with_ckpt = [probe(trainer, b, args.probe_steps, L, ckpt=True) for b in ckpt_cands]
    chosen_ckpt, reason_ckpt = choose_micro_batch(with_ckpt, vr["target_gb_low"], vr["target_gb_high"], vr["hard_cap_gb"])
    LOG.info("chosen micro-batch (ckpt): %s -- %s", chosen_ckpt, reason_ckpt)
    mcfg.grad_checkpointing = ckpt_default

    report: dict[str, Any] = {
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        "device": device_info,
        "alloc_conf": os.environ.get("PYTORCH_CUDA_ALLOC_CONF"),
        "baseline_device_used_gib": baseline_used,
        "model_config": mcfg.to_dict(),
        "param_counts": {"mlm": mlm_counts, "classifier": cls_counts},
        "data": {"source": pool.source, "detail": pool.detail, "n_sequences": pool.n, "seq_len": L,
                 "mlm_probability": pcfg["mlm_probability"], "mask_replace_probs": pcfg["mask_replace_probs"]},
        "vram_policy": vr,
        "vram_policy_unit": VRAM_UNIT_NOTE,
        "finder": {"candidates": cands, "ckpt_candidates": ckpt_cands, "probe_steps": args.probe_steps,
                   "micro_steps_per_probe_step": PROBE_ACCUM,
                   "no_ckpt": [asdict(r) for r in no_ckpt], "ckpt": [asdict(r) for r in with_ckpt],
                   "chosen_micro_batch": chosen, "reason": reason,
                   "chosen_micro_batch_ckpt": chosen_ckpt, "reason_ckpt": reason_ckpt},
    }

    if chosen is None:
        LOG.error("no micro-batch fits; skipping throughput and convergence")
    else:
        # 5. throughput + backend check
        epoch_tokens, epoch_src = resolve_epoch_tokens(args.epoch_tokens, pool, Path(cfg["paths"]["data_processed"]))
        LOG.info("epoch tokens for the ETA: %.3g (%s)", epoch_tokens, epoch_src)
        tp = measure_throughput(trainer, chosen, L, int(pcfg["effective_batch_sequences"]), args.timed_steps, 3, epoch_tokens)
        tp["epoch_tokens_source"] = epoch_src
        report["throughput"] = tp
        # The throughput run is the chosen batch under the real accumulation pattern over
        # many more micro-steps than the probe; its reserved peak is the second opinion on
        # the finder's decision against the same cap and band.
        fin = report["finder"]
        fin["throughput_peak_reserved_gib"] = tp["peak_reserved_gib"]
        fin["throughput_within_hard_cap"] = tp["peak_reserved_gib"] <= vr["hard_cap_gb"]
        fin["throughput_in_target_band"] = vr["target_gb_low"] <= tp["peak_reserved_gib"] <= vr["target_gb_high"]
        if not fin["throughput_within_hard_cap"]:
            LOG.error("micro-batch %d breaches the %s GiB hard cap under real accumulation (%.2f GiB reserved); "
                      "do not train at this batch", chosen, vr["hard_cap_gb"], tp["peak_reserved_gib"])
        elif not fin["throughput_in_target_band"]:
            LOG.warning("micro-batch %d: %.2f GiB reserved under real accumulation is outside the [%s, %s] GiB band",
                        chosen, tp["peak_reserved_gib"], vr["target_gb_low"], vr["target_gb_high"])
        report["sdpa_backends"] = check_sdpa_backends(model, L, mcfg.vocab_size)

    trainer.release()
    del trainer, opt, model
    gc.collect()
    torch.cuda.empty_cache()

    if chosen is not None:
        # 6. convergence from a fresh init
        report["convergence"] = convergence_check(cfg, mcfg, pool, device, chosen, args.steps, args.conv_lr, args.conv_warmup)

    # 7. latency
    del pool
    gc.collect()
    torch.cuda.empty_cache()
    n_lat = args.latency_samples or int(cfg["eval"]["latency_samples"])
    report["latency"] = inference_latency(cfg, mcfg, device, L, n_lat, 10, args.cpu_threads, args.skip_cpu)

    # 8. report
    report["runtime_s"] = round(common.timer() - t_start, 1)
    report_path = Path(cfg["paths"]["reports"]) / "smoke_test.json"
    report["report_path"] = str(report_path)
    common.write_json(report, report_path)
    LOG.info("wrote %s (%.1f s total)", report_path, report["runtime_s"])
    print()
    print(markdown_summary(report))
    return 0


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