align / code /scripts /olmoe_sft_smoke.py
Add files using upload-large-folder tool
bcc8d35 verified 15 days ago
57.5 kB
#!/usr/bin/env python
"""OLMoE-1B-7B-0125 post-training smoke on 4x A10G: FSDP2, bf16, adamw8bit, MoE telemetry.
    .venv/bin/python scripts/olmoe_sft_smoke.py prepare            # CPU: cached SFT rows -> OLMoE token ids
    .venv/bin/python scripts/olmoe_sft_smoke.py drive              # reference run + SIGTERM/resume run + report
    torchrun --standalone --nproc_per_node=4 scripts/olmoe_sft_smoke.py train --config configs/olmoe_sft_smoke.toml --run-dir RUN
    torchrun --standalone --nproc_per_node=4 scripts/olmoe_sft_smoke.py probe --run-dir RUN   # FSDP cost vs tokens
What this settles (docs/MEASUREMENTS.md, "Track M: OLMoE-1B-7B-0125"):
  * does the real 6.92B model fit and step under FSDP2 + adamw8bit on 23GB cards, and with what headroom;
  * tokens/s and step time, and how much of a micro-step is FSDP traffic (the "is EP needed?" gate);
  * does the LBL-off post-training default leave the router balanced (MaxVio / utilization per layer);
  * the checkpoint-save host-RSS peak (the 45GB-floor question) for a 13.8GB bf16 weight file;
  * does a SIGTERM'd, supervisor-relaunched run continue where it stopped.
DATA. The repo's cached SFT corpus is the PRM800K math-solution shards built for Mistral
(`nano-mistral-o1/data/shards/sft/train`, 11,739 rows, Mistral-v0.3 token ids). OLMoE has a different tokenizer
(50,279 BPE + specials), so `prepare` DECODES a seeded slice with the vendored Mistral tokenizer back to text,
splits it at the stored `prompt_len`, and re-tokenizes prompt and response with the OLMoE HF tokenizer in the
legacy venv (a subprocess; `transformers`/`tokenizers` are not align deps). The label mask is the prompt/response
boundary, an EOS closes each row. This is real SFT data, re-tokenized -- NOT synthetic. Rows are kept only if
their OLMoE length is in [MIN_LEN, seq_len], and every micro-batch is padded to exactly seq_len so each
micro-step costs the same 1024 tokens of compute (padded tokens run through the MoE; router telemetry masks them).
IMPORT ORDER is load-bearing (apply_resource_guards sets env torch reads at import): this module imports only the
stdlib at top level; torch and every align module that pulls it in are imported inside functions, after the guard.
LOADING FOUR 6.9B RANKS WITHOUT A HOST OOM. Each rank materialises the bf16 model on CPU (13.8GB) plus
`hf_to_align`'s stacked copy (another 13.8GB at the peak) before moving it to its GPU. Four ranks at once is
~110GB of a 186GB no-swap box -- the 2026-08-12 shape. Ranks therefore load ONE AT A TIME (barrier-ordered), and the
Trainer's resume reads go through `SerialWeightFormat` for the same reason.
"""

from __future__ import annotations

import argparse
import gc
import glob
import json
import math
import os
import re
import signal
import subprocess
import sys
import threading
import time

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
LEGACY_PY = "/apps/development/rlhf/nano-mistral-o1/.venv/bin/python"
HF_HOME = "/apps/development/rlhf/hf_cache"
CKPT_GLOB = os.path.join(HF_HOME, "hub", "models--allenai--OLMoE-1B-7B-0125", "snapshots", "*")
SFT_SHARDS = "/apps/development/rlhf/nano-mistral-o1/data/shards/sft/train"
OUT_ROOT = os.path.join(REPO, ".tmp", "olmoe")
DATA_DIR = os.path.join(OUT_ROOT, "sft_data")
SMOKE_RUNS = os.path.join(OUT_ROOT, "smoke")
DEFAULT_CONFIG = os.path.join(REPO, "configs", "olmoe_sft_smoke.toml")

SEQ_LEN = 1024
MIN_LEN = 640
N_ROWS = 320  # 20 optimizer steps x accum 4 x world 4 x micro 1
WORLD = 4
GRAD_ACCUM = 4
N_OPT_STEPS = 20
SAVE_AT_STEP = 10
SAVE_EVERY_MICRO = SAVE_AT_STEP * GRAD_ACCUM
PAD_ID = 1  # OLMoE pad_token_id
EOS_ID = 50279

# ---------------------------------------------------------------------------
# prepare: cached Mistral-tokenized SFT rows -> OLMoE-tokenized rows
# ---------------------------------------------------------------------------

_LEGACY_TOKENIZE = r'''
import json, sys
import numpy as np
from safetensors.numpy import save_file
from transformers import AutoTokenizer
ckpt, rows_path, out_path, seq_len, min_len, n_rows, eos = sys.argv[1:8]
seq_len, min_len, n_rows, eos = int(seq_len), int(min_len), int(n_rows), int(eos)
tok = AutoTokenizer.from_pretrained(ckpt)
ids_all, offs, plens, keep = [], [0], [], []
rejected = {"short": 0, "long": 0}
for r in json.load(open(rows_path)):
    p = tok(r["prompt"], add_special_tokens=False)["input_ids"]
    q = tok(r["response"], add_special_tokens=False)["input_ids"] + [eos]
    n = len(p) + len(q)
    if n > seq_len:
        rejected["long"] += 1; continue
    if n < min_len:
        rejected["short"] += 1; continue
    ids_all.extend(p + q); offs.append(len(ids_all)); plens.append(len(p)); keep.append(r["src_index"])
    if len(plens) == n_rows:
        break
if len(plens) < n_rows:
    raise SystemExit(f"only {len(plens)} rows survived the length filter (wanted {n_rows}); rejected={rejected}")
save_file({"ids": np.asarray(ids_all, dtype=np.int32), "offsets": np.asarray(offs, dtype=np.int64),
           "prompt_len": np.asarray(plens, dtype=np.int32), "src_index": np.asarray(keep, dtype=np.int64)}, out_path)
print(json.dumps({"rows": len(plens), "tokens": len(ids_all), "rejected": rejected, "vocab": len(tok)}))
'''


def default_ckpt() -> str:
    hits = sorted(glob.glob(CKPT_GLOB))
    if not hits:
        raise FileNotFoundError(f"no OLMoE-1B-7B-0125 snapshot under {CKPT_GLOB}")
    return hits[-1]


def data_paths(data_dir: str = DATA_DIR) -> tuple[str, str]:
    return os.path.join(data_dir, "sft_ids.safetensors"), os.path.join(data_dir, "sft_meta.json")


def split_response(full_text: str, prompt_text: str) -> str:
    """The response text of a decoded row: `full_text` minus its `prompt_text` prefix.
    Decoding the response tokens ALONE loses a leading space (SentencePiece decodes the first
    piece without its prefix marker), which would silently change the first response token's
    OLMoE id. Decoding the whole row and cutting at the prompt's decoded length keeps it. If the
    prompt text is not a prefix of the full text the decode did not round-trip and the row is
    unusable: raise, never guess a boundary.
    """
    if not full_text.startswith(prompt_text):
        raise ValueError("decoded prompt is not a prefix of the decoded row; cannot place the prompt/response boundary")
    return full_text[len(prompt_text):]


def prepare(ckpt: str, shard_dir: str, data_dir: str, n_rows: int, seed: int) -> dict:
    import numpy as np

    sys.path.insert(0, REPO)
    from align.data.store import ShardDataset
    from align.tokenizer import Tokenizer

    os.makedirs(data_dir, exist_ok=True)
    ds = ShardDataset(shard_dir)
    lens = ds.lengths()
    # Mistral length ~ OLMoE length within ~15%; over-select, the legacy step applies the exact filter.
    cand = np.nonzero((lens >= int(MIN_LEN * 0.85)) & (lens <= SEQ_LEN))[0]
    rng = np.random.default_rng(seed)
    rng.shuffle(cand)
    cand = cand[: int(n_rows * 2.5)]
    mistral = Tokenizer.from_pretrained(os.path.join(REPO, "assets", "tokenizer", "mistral-7b-v0.3"))
    rows = []
    for i in cand.tolist():
        r = ds[i]
        pl = int(r["prompt_len"])
        ids = r.input_ids.tolist()
        full, prompt = mistral.decode(ids), mistral.decode(ids[:pl])
        # Mistral's BOS/EOS decode as the literal text "<s>" / "</s>"; they are not OLMoE text.
        full, prompt = full.removeprefix("<s>").removesuffix("</s>"), prompt.removeprefix("<s>")
        try:
            resp = split_response(full, prompt)
        except ValueError:
            continue
        rows.append({"src_index": i, "prompt": prompt, "response": resp})
    rows_path = os.path.join(data_dir, "rows_text.json")
    with open(rows_path, "w") as f:
        json.dump(rows, f)
    out_path, meta_path = data_paths(data_dir)
    env = {"PATH": os.environ.get("PATH", ""), "HOME": os.environ.get("HOME", ""), "HF_HOME": HF_HOME,
           "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", "TOKENIZERS_PARALLELISM": "false",
           "TMPDIR": os.path.join(REPO, ".tmp"), "OMP_NUM_THREADS": "4"}
    p = subprocess.run([LEGACY_PY, "-c", _LEGACY_TOKENIZE, ckpt, rows_path, out_path, str(SEQ_LEN), str(MIN_LEN), str(n_rows), str(EOS_ID)],
                       capture_output=True, text=True, timeout=900, env=env, cwd=os.path.join(REPO, ".tmp"))
    if p.returncode != 0:
        raise RuntimeError(f"tokenize failed:\n{p.stderr[-2000:]}")
    meta = json.loads(p.stdout.strip().splitlines()[-1])
    meta.update({"source": shard_dir, "source_receipt": "nano-mistral-o1/data/sft.receipt.json (tasksource/PRM800K @ 547b1950)",
                 "seed": seed, "seq_len": SEQ_LEN, "min_len": MIN_LEN, "pad_id": PAD_ID, "eos_id": EOS_ID,
                 "synthetic": False, "retokenized_from": "mistral-7b-v0.3 -> OLMoE-1B-7B-0125 tokenizer"})
    with open(meta_path, "w") as f:
        json.dump(meta, f, indent=2)
    os.remove(rows_path)
    return meta


# ---------------------------------------------------------------------------
# pure analysis helpers (numpy/stdlib only; unit-tested)
# ---------------------------------------------------------------------------


def read_jsonl(path: str) -> list[dict]:
    """Rows of a jsonl file; a torn FINAL line (a kill mid-append) is dropped, any other bad line raises."""
    if not os.path.exists(path):
        return []
    lines = [l for l in open(path).read().split("\n") if l.strip()]
    out = []
    for i, l in enumerate(lines):
        try:
            out.append(json.loads(l))
        except json.JSONDecodeError:
            if i != len(lines) - 1:
                raise
    return out


def dedupe_by_micro(rows: list[dict]) -> dict[int, dict]:
    """micro_step -> row, the LAST occurrence winning (a relaunch may re-log the boundary it resumed at)."""
    return {int(r["micro_step"]): r for r in rows}


def mean_loss_per_micro(per_rank: list[dict[int, dict]]) -> dict[int, float]:
    """Mean over ranks of each rank's micro-batch loss, for micro-steps EVERY rank logged."""
    common = set.intersection(*[set(d) for d in per_rank]) if per_rank else set()
    return {m: sum(d[m]["loss"] for d in per_rank) / len(per_rank) for m in sorted(common)}


def compare_loss_curves(ref: dict[int, float], sup: dict[int, float], stop_micro: int, floor_mult: float = 2.0,
                        floor_pad: float = 5e-3) -> dict:
    """Does the resumed segment of `sup` (micro > stop_micro) stay as close to the uninterrupted `ref`
    as the pre-stop segment (micro <= stop_micro), which is the SAME computation and so measures pure
    run-to-run nondeterminism (NCCL order, SDPA backward, stochastic rounding).
    Verdict: resumed max |dloss| <= floor_mult * pre-stop max |dloss| + floor_pad. The pad keeps a
    near-zero floor (a lucky run) from turning ordinary noise into a failure. No overlap between the segments
    raises (a vacuous pass is the failure mode).
    WHAT THIS DOES NOT PROVE (reviewer finding, 2026-10-12). The loss-curve tolerance (1e-2) has never been shown
    to move when the resume is WRONG: no mutant resume (zeroed 8-bit moments, base weights instead of the
    checkpoint) was run on the GPUs. A lost optimizer state at lr 1e-5 over 36 micro-steps could plausibly stay
    inside it. Provenance is therefore checked separately and exactly, by `check_resume_provenance` (the relaunch
    loaded THE checkpoint at the stop micro-step, no history row was re-logged from zero) and `compare_batches`
    (every rank saw the same row at every micro-step as the uninterrupted run). This function is the
    continuity check on top of those, not a substitute for them.
    """
    common = sorted(set(ref) & set(sup))
    pre = [abs(ref[m] - sup[m]) for m in common if m <= stop_micro]
    post = [abs(ref[m] - sup[m]) for m in common if m > stop_micro]
    if not pre or not post:
        raise ValueError(f"need steps on both sides of the stop (stop={stop_micro}); pre={len(pre)} post={len(post)}")
    floor, dev = max(pre), max(post)
    tol = floor_mult * floor + floor_pad
    return {
        "stop_micro": stop_micro, "n_pre": len(pre), "n_post": len(post),
        "pre_stop_max_abs": floor, "pre_stop_mean_abs": sum(pre) / len(pre),
        "resumed_max_abs": dev, "resumed_mean_abs": sum(post) / len(post),
        "tolerance": tol, "passed": dev <= tol,
    }


def step_time_stats(rows: list[dict], accum: int, save_every_micro: int, drop_first: int = 2) -> dict:
    """Micro/optimizer-step wall time from consecutive `t_wall` stamps of ONE rank's rows.
    Intervals dropped: the first `drop_first` (allocator warmup, first-use kernels) and any interval that
    contains a checkpoint save (the micro-step after micro % save_every == 0 pays for the write). The
    median is reported with the spread; a mean would be dominated by whatever a neighbour job did.
    """
    rows = sorted(rows, key=lambda r: r["micro_step"])
    dts = []
    for a, b in zip(rows, rows[1:]):
        if b["micro_step"] != a["micro_step"] + 1:
            continue
        if a["micro_step"] % save_every_micro == 0 or a["micro_step"] < drop_first:
            continue
        dts.append(b["t_wall"] - a["t_wall"])
    if not dts:
        raise ValueError("no usable intervals")
    dts.sort()
    med = dts[len(dts) // 2]
    return {"n": len(dts), "micro_median_s": med, "micro_min_s": dts[0], "micro_max_s": dts[-1],
            "optimizer_step_s": med * accum}


# ---------------------------------------------------------------------------
# train / probe (torchrun targets)
# ---------------------------------------------------------------------------


def _vm_kb(field: str) -> float:
    with open("/proc/self/status") as f:
        for line in f:
            if line.startswith(field + ":"):
                return float(line.split()[1])
    return float("nan")


def _mem_available_gb() -> float:
    with open("/proc/meminfo") as f:
        for line in f:
            if line.startswith("MemAvailable:"):
                return float(line.split()[1]) / 1024**2
    return float("nan")


def trim_host_heap() -> bool:
    """Give freed glibc heap pages back to the OS (`malloc_trim(0)`); True if anything was released.
    MEASURED NEED (docs/MEASUREMENTS.md, Track M): after the rank-serial load, ranks 1-3 each kept a 13.4GiB
    anonymous `[heap]` -- the freed bf16 state dict (tensors under glibc's dynamic mmap threshold, 4-5MB per
    expert, land in the brk heap, and a live small chunk above them stops the heap from shrinking). Three
    ranks x 13.4GiB = 40GiB of dead host memory for the whole run, on a box with no swap, and rank 0 kept ~11GiB
    after its first checkpoint save for the same reason. `malloc_trim` walks the arenas and `madvise`s the free
    pages away. A no-op (False) off glibc.
    """
    try:
        import ctypes

        return bool(ctypes.CDLL("libc.so.6").malloc_trim(0))
    except (OSError, AttributeError):
        return False


def _reset_hwm() -> bool:
    """Reset this process's peak RSS (VmHWM) so a window's high-water mark is its own."""
    try:
        with open("/proc/self/clear_refs", "w") as f:
            f.write("5")
        return True
    except OSError:
        return False


class RssWindow:
    """Peak host memory over a window: own RSS (VmHWM after a reset, and a 0.1s sampler) and system
    MemAvailable (min). MemAvailable counts clean page cache as available, so its dip is anonymous
    memory the window really added -- the quantity the 45GB floor is about."""

    def __init__(self, period: float = 0.1):
        self.period = period

    def __enter__(self):
        self.hwm_reset = _reset_hwm()
        self.rss0 = _vm_kb("VmRSS") / 1024**2
        self.avail0 = _mem_available_gb()
        self.rss_peak = self.rss0
        self.avail_min = self.avail0
        self._stop = threading.Event()
        self.t0 = time.perf_counter()

        def loop():
            while not self._stop.is_set():
                self.rss_peak = max(self.rss_peak, _vm_kb("VmRSS") / 1024**2)
                self.avail_min = min(self.avail_min, _mem_available_gb())
                self._stop.wait(self.period)

        self._th = threading.Thread(target=loop, daemon=True)
        self._th.start()
        return self

    def __exit__(self, *exc):
        self._stop.set()
        self._th.join()
        self.dur_s = time.perf_counter() - self.t0
        self.hwm = _vm_kb("VmHWM") / 1024**2
        return False

    def record(self) -> dict:
        return {"dur_s": round(self.dur_s, 2), "rss_start_gb": round(self.rss0, 2), "rss_peak_sampled_gb": round(self.rss_peak, 2),
                "vmhwm_gb": round(self.hwm, 2), "hwm_reset": self.hwm_reset,
                "mem_available_start_gb": round(self.avail0, 1), "mem_available_min_gb": round(self.avail_min, 1),
                "mem_available_drop_gb": round(self.avail0 - self.avail_min, 1)}


def _load_dataset(data_dir: str):
    from safetensors.numpy import load_file

    arrs = load_file(data_paths(data_dir)[0])
    return arrs["ids"], arrs["offsets"], arrs["prompt_len"]


def _make_sampler_cls():
    """Defined lazily: it needs torch, which is imported only after the resource guards."""
    import torch

    class RowBatch(list):
        """A list of row indices (what the Trainer logs as `batch`) that also carries the collated
        tensors the adapter and `MoETrainModule` read (`input_ids`, `labels`, `attention_mask`)."""

    class RankSampler:
        """Micro-batch 1. Epoch order is a pure function of (seed, epoch) -- the contract the Trainer's
        resumable stream relies on -- and rank r takes elements r, r+W, r+2W of it, so the W ranks of one
        micro-step see W DISTINCT rows (dense `scripts/sft.py` does not shard; see the report)."""

        def __init__(self, data_dir, rank, world, seed, device, seq_len=SEQ_LEN, n_rows=None):
            self.ids, self.offs, self.plen = _load_dataset(data_dir)
            self.n = len(self.plen) if n_rows is None else n_rows
            self.rank, self.world, self.seed, self.epoch = rank, world, seed, 0
            self.device, self.seq_len = device, seq_len

        def set_epoch(self, e):
            self.epoch = int(e)

        def row(self, i):
            return self.ids[self.offs[i]:self.offs[i + 1]], int(self.plen[i])

        def __iter__(self):
            g = torch.Generator().manual_seed(self.seed + self.epoch)
            order = torch.randperm(self.n, generator=g).tolist()
            usable = (self.n // self.world) * self.world
            for j in range(self.rank, usable, self.world):
                i = order[j]
                x, p = self.row(i)
                L = len(x)
                ids = torch.full((1, self.seq_len), PAD_ID, dtype=torch.long)
                ids[0, :L] = torch.from_numpy(x.astype("int64"))
                labels = torch.full((1, self.seq_len), -100, dtype=torch.long)
                labels[0, p:L] = ids[0, p:L]
                mask = torch.zeros((1, self.seq_len), dtype=torch.long)
                mask[0, :L] = 1
                b = RowBatch([i])
                b.input_ids, b.labels, b.attention_mask = ids.to(self.device), labels.to(self.device), mask.to(self.device)
                b.n_valid = L
                yield b

        def __len__(self):
            return (self.n // self.world)

    return RankSampler, RowBatch


def _serial_format_cls():
    from align.train.ckpt_format import OlmoeFormat

    class SerialWeightFormat(OlmoeFormat):
        """`OlmoeFormat` whose `load_weights` runs on ONE rank at a time.
        `Trainer.resume_from` has every rank read the full 13.8GB state and `hf_to_align` stacks a second
        copy at its peak; four ranks at once is ~110GB. Barrier-ordering the reads bounds the transient to
        one rank's, plus the full dicts the already-loaded ranks still hold until their local shard is cut.
        """

        def load_weights(self, path):
            from align.utils.dist import barrier, rank, world_size

            out = None
            for r in range(world_size()):
                if rank() == r:
                    out = super().load_weights(path)
                barrier()
            return out

    return SerialWeightFormat


def build_sharded_model(ckpt: str, ocfg, backend: str, attn: str, device):
    """Rank-serial load -> GPU -> activation checkpointing -> FSDP2, as `stage_smoke.build_stage_model` does
    for dense, with the host-RAM serialisation above."""
    import torch

    from align.model import olmoe as om
    from align.model import olmoe_checkpoint as oc
    from align.model.checkpoint import load_hf_safetensors
    from align.train import fsdp
    from align.utils.dist import barrier, rank, world_size

    model = None
    peaks = {}
    for r in range(world_size()):
        if rank() == r:
            with RssWindow() as w:
                sd = oc.hf_to_align(load_hf_safetensors(ckpt, dtype=torch.bfloat16), ocfg)
                model = om.OlmoeForCausalLM.from_state_dict(ocfg, sd, attn_implementation=attn, moe_backend=backend)
                del sd
                model = model.to(device)
            peaks = w.record()
            gc.collect()
            peaks["trimmed"] = trim_host_heap()
            peaks["rss_after_trim_gb"] = round(_vm_kb("VmRSS") / 1024**2, 2)
        barrier()
    fsdp.apply_activation_checkpointing(model)
    fsdp.apply_fsdp(model)
    return model, peaks


def _write_json(path: str, obj) -> None:
    tmp = path + ".partial"
    with open(tmp, "w") as f:
        json.dump(obj, f, indent=2, default=str)
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, path)


def train_main(a) -> int:
    from align.config import load_config
    from align.utils.guard import apply_resource_guards

    cfg = load_config(a.config, list(a.set or []))
    apply_resource_guards(cpu_threads=8, heavy=True, label="olmoe-sft-smoke")
    import torch

    from align.model.olmoe_checkpoint import read_olmoe_config
    from align.train.moe_module import build_moe_train_module
    from align.train.offload_adamw import make_optimizer
    from align.train.sft_loss import make_sft_loss_adapter
    from align.train.stage_gates import assert_distinct_devices
    from align.train.trainer import MANDATORY_METRICS, Trainer
    from align.utils.dist import barrier, init_distributed, local_rank, rank as get_rank, seed_everything, world_size as get_world
    from align.utils.tracking import Tracker

    run_dir = os.path.abspath(a.run_dir)
    os.makedirs(run_dir, exist_ok=True)
    init_distributed(run_dir=run_dir)
    assert_distinct_devices()
    R, W = get_rank(), get_world()
    dev = torch.device("cuda", local_rank())
    torch.cuda.set_device(dev)
    with open(os.path.join(run_dir, f"rank{R}.pid"), "w") as f:
        f.write(str(os.getpid()))
    seed_everything(cfg.train.seed)

    ckpt = a.ckpt or default_ckpt()
    ocfg = read_olmoe_config(ckpt)
    model, load_peaks = build_sharded_model(ckpt, ocfg, a.moe_backend, "sdpa", dev)
    torch.cuda.synchronize()
    load_gpu_peak = torch.cuda.max_memory_allocated() / 2**30
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()

    RankSampler, _ = _make_sampler_cls()
    sampler = RankSampler(a.data_dir, R, W, cfg.train.seed, dev, n_rows=a.n_rows)
    sft = make_sft_loss_adapter(require_fused=True)

    def adapter(m, batch):
        return sft(m, (batch.input_ids, batch.labels))

    opt = make_optimizer(cfg, [p for p in model.parameters() if p.requires_grad])
    module = build_moe_train_module(cfg, model, adapter, opt, stage="sft", olmoe_cfg=ocfg)
    fmt = _serial_format_cls()(ocfg, tokenizer_dir=ckpt)
    tracker = Tracker(run_dir if R == 0 else os.path.join(run_dir, f"rank{R}"), "sft",
                      mandatory_keys=MANDATORY_METRICS, run_name=cfg.run_name)
    trainer = Trainer(cfg=cfg, model=model, sampler=sampler, loss_adapter=adapter, run_dir=run_dir, stage="sft",
                      train_module=module, weight_format=fmt, tracker=tracker)

    # -- telemetry: per-step memory/time, and the host-RSS window around every checkpoint save -------------
    orig_log = tracker.log

    def log(step, metrics=None, **kw):
        torch.cuda.synchronize()
        m = {**(dict(metrics) if metrics else {}), **kw}
        m["t_wall"] = time.time()
        m["mem/alloc_peak_gb"] = torch.cuda.max_memory_allocated() / 2**30
        m["mem/reserved_peak_gb"] = torch.cuda.max_memory_reserved() / 2**30
        m["rss_gb"] = _vm_kb("VmRSS") / 1024**2
        orig_log(step, m)

    tracker.log = log
    orig_save = trainer.save_checkpoint
    ckpt_log = os.path.join(run_dir, f"ckpt_rss_rank{R}.jsonl")

    def save():
        micro = trainer.micro_step
        with RssWindow() as w:
            out = orig_save()
        gc.collect()
        trimmed = trim_host_heap()  # after the window: the window reports the save's own peak, untrimmed
        rec = {"trimmed": trimmed, "rss_after_trim_gb": round(_vm_kb("VmRSS") / 1024**2, 2), "micro_step": micro, "optimizer_step": trainer.optimizer_step, "rank": R, **w.record()}
        with open(ckpt_log, "a") as f:
            f.write(json.dumps(rec) + "\n")
        return out

    trainer.save_checkpoint = save

    resumed = trainer.resume_from_latest()
    if resumed:
        # MEASURED (smoke/sup/ckpt_rss_rank1.jsonl): after the resume read, ranks 1-3 sat at 13.7GB RSS until the
        # first save trimmed them -- every rank materialises the full bf16 state dict, and the freed copy is
        # stranded in glibc's brk heap exactly as after the initial load. Same fix, same reason.
        gc.collect()
        trim_host_heap()
    _write_json(os.path.join(run_dir, f"launch_rank{R}_{int(time.time())}.json"), {
        "rank": R, "resumed_from": resumed, "resumed_micro_step": trainer.micro_step, "load": load_peaks,
        "load_gpu_peak_gb": load_gpu_peak, "gpu_alloc_after_wrap_gb": torch.cuda.memory_allocated() / 2**30,
        "optimizer": getattr(trainer.optimizer, "banner_tag", type(trainer.optimizer).__name__),
        "n_params": sum(p.numel() for p in model.parameters()), "moe_backend": a.moe_backend,
        "lb_weight": module.lb_weight, "z_weight": module.z_weight, "world": W, "cfg_run": cfg.run_name,
    })
    summary = trainer.train()
    barrier()
    if R == 0:
        _write_json(os.path.join(run_dir, f"summary_{int(time.time())}.json"), {k: v for k, v in summary.items() if isinstance(v, (int, float, str, bool, type(None)))})
    if trainer.stop_exit_code:
        print(f"{trainer.stop_reason} stop: checkpointed, exiting {trainer.stop_exit_code}", file=sys.stderr)
        return trainer.stop_exit_code
    return 0


def _time_median(fn, n: int, warm: int = 2) -> float:
    import torch

    for _ in range(warm):
        fn()
    torch.cuda.synchronize()
    ts = []
    for _ in range(n):
        t0 = time.perf_counter()
        fn()
        torch.cuda.synchronize()
        ts.append(time.perf_counter() - t0)
    ts.sort()
    return ts[len(ts) // 2]


def probe_collectives(ocfg, dev, W: int) -> dict:
    """Isolated FSDP2 traffic at real sizes, un-overlapped. NCCL algbw = full tensor bytes / time."""
    import torch
    import torch.distributed as dist

    from align.utils.dist import barrier

    d, i, e = ocfg.hidden_size, ocfg.intermediate_size, ocfg.num_experts
    block_params = 4 * d * d + 2 * d + 2 * d + e * d + e * 3 * d * i  # attn + norms + router + experts
    sweep = []
    for frac in (1 / 16, 1 / 4, 1.0):
        n = int(block_params * frac) // W * W
        ag_out = torch.empty(n, dtype=torch.bfloat16, device=dev)
        ag_in = torch.randn(n // W, dtype=torch.bfloat16, device=dev)
        rs_in = torch.randn(n, dtype=torch.float32, device=dev)
        rs_out = torch.empty(n // W, dtype=torch.float32, device=dev)
        barrier()
        t_ag = _time_median(lambda: dist.all_gather_into_tensor(ag_out, ag_in), 10)
        barrier()
        t_rs = _time_median(lambda: dist.reduce_scatter_tensor(rs_out, rs_in), 10)
        sweep.append({"params": n, "all_gather_bf16_s": t_ag, "reduce_scatter_fp32_s": t_rs,
                      "ag_algbw_GBps": n * 2 / t_ag / 1e9, "rs_algbw_GBps": n * 4 / t_rs / 1e9})
        del ag_out, ag_in, rs_in, rs_out
    full = sweep[-1]
    # Expert-parallel dispatch/combine: every token's k expert-slots travel as d_model bf16 activations;
    # with EP=W each rank keeps 1/W of them locally. One all-to-all moves T*k*d*2 bytes per rank, and a
    # layer issues 2 per forward, 2 in the activation-checkpoint recompute, 2 in the backward.
    a2a = []
    for tokens in (1024, 4096):
        n = tokens * ocfg.top_k * d
        n -= n % W
        x_in = torch.randn(n, dtype=torch.bfloat16, device=dev)
        x_out = torch.empty_like(x_in)
        barrier()
        t = _time_median(lambda: dist.all_to_all_single(x_out, x_in), 10)
        a2a.append({"tokens_per_rank": tokens, "bytes_per_rank": n * 2, "all_to_all_s": t, "algbw_GBps": n * 2 / t / 1e9,
                    "ep_comm_per_micro_s": ocfg.n_layers * 6 * t})
        del x_in, x_out
    return {"block_params": block_params, "sweep": sweep, "all_to_all": a2a, "n_blocks": ocfg.n_layers,
            "nccl_p2p_level": os.environ.get("NCCL_P2P_LEVEL"),
            # AG fwd + AG in the activation-checkpoint recompute + RS, per block, fully exposed (no overlap).
            "exposed_per_micro_s": ocfg.n_layers * (2 * full["all_gather_bf16_s"] + full["reduce_scatter_fp32_s"])}


def probe_compute(ocfg, dev, seq_lens: list[int], backend: str) -> dict:
    """Compute-only fwd+bwd of a TRUNCATED random-init model per GPU, no FSDP, no collectives: t(n_layers, L)
    = fixed + n_layers * per_layer, so the full 16-layer compute is extrapolated, not guessed."""
    import torch

    from align.model import olmoe as om
    from align.train import fsdp
    from align.train.sft_loss import make_sft_loss_adapter

    sft = make_sft_loss_adapter(require_fused=True)
    out = []
    for n_layers in (2, 4):
        m = om.OlmoeForCausalLM(ocfg.replace(n_layers=n_layers), attn_implementation="sdpa", moe_backend=backend).to(dev, torch.bfloat16)
        fsdp.apply_activation_checkpointing(m)
        for L in seq_lens:
            x = torch.randint(2, ocfg.vocab_size - 32, (1, L), device=dev)

            def step():
                o = sft(m, (x, x))
                o.loss.backward()
                m.zero_grad(set_to_none=True)

            try:
                torch.cuda.reset_peak_memory_stats()
                out.append({"n_layers": n_layers, "seq_len": L, "median_s": _time_median(step, 3, warm=1),
                            "alloc_peak_gb": torch.cuda.max_memory_allocated() / 2**30})
            except torch.OutOfMemoryError:
                m.zero_grad(set_to_none=True)
                torch.cuda.empty_cache()
                out.append({"n_layers": n_layers, "seq_len": L, "oom": True})
        del m
        torch.cuda.empty_cache()
    res = {"points": out, "extrapolated_16_layer_s": {}}
    for L in seq_lens:
        a2 = next((p for p in out if p["n_layers"] == 2 and p["seq_len"] == L and "median_s" in p), None)
        a4 = next((p for p in out if p["n_layers"] == 4 and p["seq_len"] == L and "median_s" in p), None)
        if a2 and a4:
            per_layer = (a4["median_s"] - a2["median_s"]) / 2
            fixed = a2["median_s"] - 2 * per_layer
            res["extrapolated_16_layer_s"][str(L)] = {"per_layer_s": per_layer, "fixed_s": fixed, "total_16_s": fixed + 16 * per_layer}
    return res


def probe_main(a) -> int:
    """FSDP cost vs tokens. Three independent parts (`--only`): `collectives` (no model: NCCL AG/RS bandwidth at
    the real block size), `compute` (truncated model per GPU, no FSDP: per-layer compute vs tokens), `scan`
    (the real FSDP2 model: micro-step time at several L; FSDP's per-micro traffic is independent of L, so
    t(L) = comm + b*L, and a flat t(L) means communication-bound)."""
    from align.config import load_config
    from align.utils.guard import apply_resource_guards

    cfg = load_config(a.config, list(a.set or []))
    apply_resource_guards(cpu_threads=8, heavy=True, label="olmoe-sft-probe")
    import torch

    from align.model.olmoe_checkpoint import read_olmoe_config
    from align.train.sft_loss import make_sft_loss_adapter
    from align.train.stage_gates import assert_distinct_devices
    from align.utils.dist import barrier, init_distributed, local_rank, rank as get_rank, seed_everything, world_size as get_world

    run_dir = os.path.abspath(a.run_dir)
    os.makedirs(run_dir, exist_ok=True)
    init_distributed(run_dir=run_dir)
    assert_distinct_devices()
    R, W = get_rank(), get_world()
    dev = torch.device("cuda", local_rank())
    torch.cuda.set_device(dev)
    seed_everything(0)
    ckpt = a.ckpt or default_ckpt()
    ocfg = read_olmoe_config(ckpt)
    seq_lens = [int(x) for x in a.seq_lens.split(",")]
    out = {"world": W, "moe_backend": a.moe_backend, "only": a.only}
    if a.only in ("all", "collectives"):
        out["collectives"] = probe_collectives(ocfg, dev, W)
        barrier()
    if a.only in ("all", "compute"):
        out["compute"] = probe_compute(ocfg, dev, seq_lens, a.moe_backend)
        barrier()
    if a.only in ("all", "scan"):
        model, _ = build_sharded_model(ckpt, ocfg, a.moe_backend, "sdpa", dev)
        torch.cuda.empty_cache()
        sft = make_sft_loss_adapter(require_fused=True)
        ids_all, _, _ = _load_dataset(a.data_dir)
        pts = []
        for L in seq_lens:
            torch.cuda.reset_peak_memory_stats()
            reps = []
            try:
                for it in range(a.warmup + a.reps):
                    start = (it * 7919 + R * 104729) % (len(ids_all) - L)
                    x = torch.from_numpy(ids_all[start:start + L].astype("int64"))[None].to(dev)
                    torch.cuda.synchronize()
                    barrier()
                    t0 = time.perf_counter()
                    o = sft(model, (x, x))
                    o.loss.backward()
                    torch.cuda.synchronize()
                    dt = time.perf_counter() - t0
                    model.zero_grad(set_to_none=True)
                    if it >= a.warmup:
                        reps.append(dt)
                reps.sort()
                pt = {"seq_len": L, "micro_median_s": reps[len(reps) // 2], "micro_min_s": reps[0], "micro_max_s": reps[-1],
                      "alloc_peak_gb": torch.cuda.max_memory_allocated() / 2**30, "reserved_peak_gb": torch.cuda.max_memory_reserved() / 2**30}
            except torch.OutOfMemoryError:
                model.zero_grad(set_to_none=True)
                torch.cuda.empty_cache()
                pt = {"seq_len": L, "oom": True}
            pts.append(pt)
            barrier()
        out["scan"] = pts
    if a.only == "gradvar":
        out["gradvar"] = probe_gradvar(a, cfg, ckpt, ocfg, dev, R, W)
    if R == 0:
        _write_json(os.path.join(run_dir, "probe.json"), out)
    barrier()
    return 0


def probe_gradvar(a, cfg, ckpt, ocfg, dev, R: int, W: int) -> dict:
    """Is the pre-clip global grad norm of ONE accumulation window reproducible?
    The training runs logged a step-1 pre-clip norm of 5.44, 8.09 and 11.65 in three runs with identical weights,
    rows and seed (losses equal to 1e-3), converging to within ~5% from step 2. This repeats the exact step-1
    window (the first `grad_accum` micro-batches of this rank's sampler stream, attention mask included) several
    times in ONE process on frozen weights and reports each global norm: if they differ, the nondeterminism is
    in forward/backward itself (bf16 atomics, router near-ties), not in the optimizer or the first-step path."""
    import torch
    import torch.distributed as dist

    from align.train.sft_loss import make_sft_loss_adapter
    from align.utils.dist import barrier

    model, _ = build_sharded_model(ckpt, ocfg, a.moe_backend, "sdpa", dev)
    trim_host_heap()
    torch.cuda.empty_cache()
    sft = make_sft_loss_adapter(require_fused=True)
    RankSampler, _ = _make_sampler_cls()
    it = iter(RankSampler(a.data_dir, R, W, cfg.train.seed, dev))
    for _ in range(a.skip_micro):
        next(it)
    window = [next(it) for _ in range(cfg.optimizer.grad_accum)]

    class WithMask:
        def __init__(self, mask):
            self.mask = mask

        def __call__(self, **kw):
            return model(attention_mask=self.mask, **kw)

    norms, losses = [], []
    for rep in range(a.grad_reps):
        tot = 0.0
        for b in window:
            o = sft(WithMask(b.attention_mask), (b.input_ids, b.labels))
            (o.loss / len(window)).backward()
            tot += float(o.loss.detach()) / len(window)
        sq = torch.zeros((), dtype=torch.float64, device=dev)
        for p in model.parameters():
            if p.grad is not None:
                g = p.grad.to_local() if hasattr(p.grad, "to_local") else p.grad
                sq += g.double().pow(2).sum()
        dist.all_reduce(sq)
        norms.append(float(sq.sqrt()))
        losses.append(tot)
        model.zero_grad(set_to_none=True)
        barrier()
    return {"rows": [b[0] for b in window], "grad_norms": norms, "rank_mean_losses": losses}


# ---------------------------------------------------------------------------
# drive: reference run, SIGTERM + supervisor resume run, report
# ---------------------------------------------------------------------------


def _gpu_used_mib() -> list[int]:
    p = subprocess.run(["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"], capture_output=True, text=True, timeout=30)
    return [int(x) for x in p.stdout.split()]


class HostMonitor:
    """1 Hz sampler of GPU memory (nvidia-smi view: includes the CUDA context + NCCL), MemAvailable and load,
    to <run_dir>/monitor.jsonl. The torch-allocator peaks in history.jsonl are the other half."""

    def __init__(self, path: str, period: float = 1.0, pid_dir: str | None = None):
        self.path, self.period, self.pid_dir = path, period, pid_dir

    def _rank_memory(self) -> list[dict]:
        """Per-rank RSS split (anonymous / file-backed / shmem, GiB) from smaps_rollup: which part of a rank's
        host footprint is heap that could be returned, and which is mapped files."""
        out = []
        for r in range(WORLD):
            try:
                pid = int(open(os.path.join(self.pid_dir, f"rank{r}.pid")).read())
                d = {}
                for line in open(f"/proc/{pid}/smaps_rollup"):
                    k, _, v = line.partition(":")
                    if k in ("Rss", "Anonymous", "Shmem", "Pss_File", "Pss_Anon", "Pss_Shmem"):
                        d[k] = round(float(v.split()[0]) / 1024**2, 2)
                out.append({"rank": r, "pid": pid, **d})
            except Exception:  # noqa: BLE001 -- the process may be between attempts
                pass
        return out

    def __enter__(self):
        self._stop = threading.Event()
        self.peak = [0, 0, 0, 0]
        self.min_avail = float("inf")
        self.max_load = 0.0

        def loop():
            with open(self.path, "a") as f:
                while not self._stop.is_set():
                    try:
                        g = _gpu_used_mib()
                    except Exception:
                        g = []
                    av, ld = _mem_available_gb(), os.getloadavg()[0]
                    n_tick = getattr(self, "_n", 0)
                    self._n = n_tick + 1
                    rm = self._rank_memory() if (self.pid_dir and n_tick % 10 == 0) else None
                    self.peak = [max(p, x) for p, x in zip(self.peak, g)] if g else self.peak
                    self.min_avail, self.max_load = min(self.min_avail, av), max(self.max_load, ld)
                    rec = {"t": time.time(), "gpu_mib": g, "mem_available_gb": round(av, 1), "load1": ld}
                    if rm:
                        rec["ranks"] = rm
                    f.write(json.dumps(rec) + "\n")
                    f.flush()
                    self._stop.wait(self.period)

        self._th = threading.Thread(target=loop, daemon=True)
        self._th.start()
        return self

    def __exit__(self, *exc):
        self._stop.set()
        self._th.join()
        return False


def host_preflight_ok() -> tuple[bool, dict]:
    sys.path.insert(0, REPO)
    from align.utils.guard import host_state

    st = host_state()
    used = _gpu_used_mib()
    st["gpu_used_mib"] = used
    ok = (st["ram_available_gb"] >= 45 and st["load1"] < 2 * st["n_cores"] and st["disk_in_flight"] < 64
          and st["d_state_procs"] <= 5 and len(used) >= WORLD and all(u < 2048 for u in used[:WORLD]))
    return ok, st


def torchrun_cmd(sub: str, run_dir: str, extra: list[str] | None = None) -> list[str]:
    return [os.path.join(REPO, ".venv", "bin", "torchrun"), "--standalone", f"--nproc_per_node={WORLD}",
            os.path.abspath(__file__), sub, "--run-dir", run_dir] + (extra or [])


def _child_env() -> dict:
    env = dict(os.environ)
    for k in ("TMPDIR", "HF_HOME", "HF_HUB_CACHE", "XDG_CACHE_HOME", "TRITON_CACHE_DIR"):
        env.pop(k, None)  # let apply_resource_guards set the /apps defaults
    env.update({"TMPDIR": os.path.join(REPO, ".tmp"), "PYTHONPATH": REPO, "OMP_NUM_THREADS": "8",
                "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", "HF_HOME": HF_HOME,
                "HF_HUB_DISABLE_XET": "1", "ALIGN_PG_TIMEOUT_S": "900", "NCCL_DEBUG": "WARN"})
    return env


def run_to_completion(cmd: list[str], run_dir: str, timeout: float, log_name: str) -> int:
    """Run `cmd` in its own session; on timeout SIGTERM then SIGKILL ITS OWN process group (never by pattern)."""
    os.makedirs(run_dir, exist_ok=True)
    log = open(os.path.join(run_dir, log_name), "ab")
    p = subprocess.Popen(cmd, env=_child_env(), stdout=log, stderr=subprocess.STDOUT, start_new_session=True, cwd=REPO)
    t0 = time.time()
    while p.poll() is None:
        if time.time() - t0 > timeout:
            os.killpg(p.pid, signal.SIGTERM)
            try:
                p.wait(60)
            except subprocess.TimeoutExpired:
                os.killpg(p.pid, signal.SIGKILL)
            return 124
        time.sleep(1.0)
    return p.returncode


def checkpoint_complete(run_dir: str, micro: int) -> bool:
    sys.path.insert(0, REPO)
    d = os.path.join(run_dir, "checkpoints", f"micro_{micro:08d}")
    if not os.path.isdir(d):
        return False
    from align.train.trainer import checkpoint_problem

    return checkpoint_problem(d) is None


def run_supervised_with_sigterm(run_dir: str, sigterm_after_micro: int, timeout: float) -> dict:
    """supervise.py -> torchrun. Once the step-10 checkpoint is complete on disk, SIGTERM ONE rank's pid (rank 1):
    the control plane must stop every rank at the same optimizer-step boundary after a checkpoint, and the
    supervisor must relaunch because progress was made and the exit is not a HALT."""
    os.makedirs(run_dir, exist_ok=True)
    cmd = [os.path.join(REPO, ".venv", "bin", "python"), os.path.join(REPO, "scripts", "supervise.py"), "--run-dir", run_dir,
           "--max-retries", "2", "--backoff", "10", "--"] + torchrun_cmd("train", run_dir, ["--config", DEFAULT_CONFIG])
    log = open(os.path.join(run_dir, "supervisor.log"), "ab")
    p = subprocess.Popen(cmd, env=_child_env(), stdout=log, stderr=subprocess.STDOUT, start_new_session=True, cwd=REPO)
    t0, sent = time.time(), None
    while p.poll() is None:
        if time.time() - t0 > timeout:
            os.killpg(p.pid, signal.SIGTERM)
            try:
                p.wait(90)
            except subprocess.TimeoutExpired:
                os.killpg(p.pid, signal.SIGKILL)
            return {"returncode": 124, "timed_out": True, "sigterm": sent}
        if sent is None and checkpoint_complete(run_dir, sigterm_after_micro):
            try:
                pid = int(open(os.path.join(run_dir, "rank1.pid")).read())
                with open(f"/proc/{pid}/cmdline", "rb") as f:
                    assert b"olmoe_sft_smoke" in f.read()  # the pid file is not stale
                os.kill(pid, signal.SIGTERM)
                sent = {"pid": pid, "t_since_start_s": round(time.time() - t0, 1), "after_checkpoint_micro": sigterm_after_micro,
                        "history_micro_at_send": (read_jsonl(os.path.join(run_dir, "history.jsonl")) or [{}])[-1].get("micro_step")}
            except Exception as e:  # noqa: BLE001
                sent = {"error": repr(e)}
        time.sleep(0.2)
    return {"returncode": p.returncode, "timed_out": False, "sigterm": sent, "elapsed_s": round(time.time() - t0, 1)}


def load_run(run_dir: str) -> dict:
    per_rank = [dedupe_by_micro(read_jsonl(os.path.join(run_dir if r == 0 else os.path.join(run_dir, f"rank{r}"), "history.jsonl")))
                for r in range(WORLD)]
    return {"per_rank": per_rank, "loss": mean_loss_per_micro(per_rank)}


def moe_curves(per_rank0: dict[int, dict], accum: int) -> dict:
    """Per-layer MaxVio / utilization at the optimizer-step boundaries (the exact-window values)."""
    steps = sorted(m for m in per_rank0 if m % accum == 0)
    layers = sorted({int(k.split("/")[1][1:]) for k in per_rank0[steps[0]] if re.match(r"moe/l\d+/max_vio", k)}) if steps else []
    return {
        "steps": steps,
        "max_vio": {l: [per_rank0[m][f"moe/l{l}/max_vio"] for m in steps] for l in layers},
        "utilization": {l: [per_rank0[m][f"moe/l{l}/utilization"] for m in steps] for l in layers},
        "global": {k: [per_rank0[m][k] for m in steps] for k in ("moe/max_vio", "moe/max_vio_worst", "moe/utilization", "moe/utilization_worst",
                                                                  "moe/frac_dead_worst", "moe/router_entropy", "moe/load_entropy")},
    }


def drive(a) -> int:
    os.makedirs(SMOKE_RUNS, exist_ok=True)
    ref_dir, sup_dir = os.path.join(SMOKE_RUNS, "ref"), os.path.join(SMOKE_RUNS, "sup")
    if not os.path.isfile(data_paths(a.data_dir)[0]):
        print("preparing data ...", flush=True)
        print(prepare(default_ckpt(), SFT_SHARDS, a.data_dir, N_ROWS, 0), flush=True)
    report = {"started": time.strftime("%Y-%m-%d %H:%M:%S"), "config": DEFAULT_CONFIG}

    ok, st = host_preflight_ok()
    report["preflight_ref"] = st
    if not ok:
        print("preflight FAILED, not starting:", st, flush=True)
        return 83
    if not a.skip_ref:
        with HostMonitor(os.path.join(SMOKE_RUNS, "monitor_ref.jsonl"), pid_dir=ref_dir) as mon:
            rc = run_to_completion(torchrun_cmd("train", ref_dir, ["--config", DEFAULT_CONFIG]), ref_dir, a.timeout, "launcher.log")
        report["ref"] = {"returncode": rc, "gpu_nvidia_smi_peak_mib": mon.peak, "mem_available_min_gb": mon.min_avail, "load1_max": mon.max_load}
        print("reference run exit", rc, flush=True)
        if rc != 0:
            _write_json(os.path.join(SMOKE_RUNS, "report.json"), report)
            return rc
    ok, st = host_preflight_ok()
    report["preflight_sup"] = st
    if not ok:
        print("preflight FAILED before supervised run:", st, flush=True)
        _write_json(os.path.join(SMOKE_RUNS, "report.json"), report)
        return 83
    with HostMonitor(os.path.join(SMOKE_RUNS, "monitor_sup.jsonl"), pid_dir=sup_dir) as mon:
        sup = run_supervised_with_sigterm(sup_dir, SAVE_EVERY_MICRO, a.timeout)
    sup.update({"gpu_nvidia_smi_peak_mib": mon.peak, "mem_available_min_gb": mon.min_avail, "load1_max": mon.max_load})
    report["sup"] = sup
    print("supervised run:", sup, flush=True)
    _write_json(os.path.join(SMOKE_RUNS, "report.json"), report)
    return 0 if sup.get("returncode") == 0 else 1


def detect_stop_micro(attempts: list[dict]) -> int:
    """The micro-step the interrupted attempt left its newest complete checkpoint at, from supervisor.jsonl.
    The first `exit` event with a non-zero child exit that `progressed`. A log with no failed attempt means the
    run was never interrupted and a resume comparison would be vacuous: raise rather than pick a number."""
    for r in attempts:
        if r.get("event") == "exit" and r.get("child_exit") not in (0, None):
            if not r.get("progressed"):
                raise ValueError(f"attempt {r.get('attempt')} failed without progress; the supervisor would not have relaunched")
            return int(r["ckpt_micro_step"])
    raise ValueError("no failed attempt in the supervisor log: the run was never interrupted, so there is no resume to verify")


def _rank_dir(run_dir: str, r: int) -> str:
    return run_dir if r == 0 else os.path.join(run_dir, f"rank{r}")


def load_launches(run_dir: str, world: int = WORLD) -> dict[int, list[dict]]:
    """rank -> that rank's `launch_rank{r}_{unix_ts}.json` records, oldest first (one per attempt)."""
    out: dict[int, list[dict]] = {}
    for r in range(world):
        files = sorted(glob.glob(os.path.join(run_dir, f"launch_rank{r}_*.json")),
                       key=lambda f: int(re.search(r"_(\d+)\.json$", f).group(1)))
        out[r] = [json.load(open(f)) for f in files]
    return out


def load_raw_rows(run_dir: str, world: int = WORLD) -> dict[int, list[dict]]:
    """rank -> history rows in FILE order (no de-duplication: a restart-from-zero must stay visible)."""
    return {r: read_jsonl(os.path.join(_rank_dir(run_dir, r), "history.jsonl")) for r in range(world)}


def check_resume_provenance(launches: dict[int, list[dict]], raw_rows: dict[int, list[dict]], stop_micro: int) -> dict:
    """Was the second attempt a RESUME of the stop checkpoint, on every rank? Exact, not statistical.
    THE HOLE THIS CLOSES. `analyze` de-duplicates history by micro-step (last row wins) and then compares loss
    curves. A relaunch that silently IGNORED the checkpoint and trained from micro 0 re-logs micro 1..80; the
    de-duplicated curve is then a complete second run of the reference computation, the 'pre-stop' and
    'resumed' segments both equal the reference to run-to-run noise, and the verdict PASSES. A resume that
    did not happen is the failure mode this smoke exists to catch, so it is checked directly:
      * every rank has >= 2 launch records; the first started fresh (`resumed_micro_step == 0`, no path), the
        last loaded `micro_{stop:08d}` and continued from exactly `stop_micro`;
      * each rank's raw history starts at micro 1 and advances by exactly 1 per row (a repeat is allowed only
        at `stop_micro` itself: the boundary a relaunch may re-log). A row going back to a lower micro-step
        is a restart.
    """
    problems: list[str] = []
    want = f"micro_{stop_micro:08d}"
    for r, ls in sorted(launches.items()):
        if len(ls) < 2:
            problems.append(f"rank {r}: {len(ls)} launch record(s); an interrupted-and-resumed run has >= 2")
            continue
        first, last = ls[0], ls[-1]
        if first.get("resumed_micro_step") != 0 or first.get("resumed_from"):
            problems.append(f"rank {r}: first attempt did not start fresh ({first.get('resumed_from')!r}, micro {first.get('resumed_micro_step')})")
        if last.get("resumed_micro_step") != stop_micro:
            problems.append(f"rank {r}: last attempt resumed at micro {last.get('resumed_micro_step')}, expected {stop_micro}")
        if not str(last.get("resumed_from") or "").rstrip("/").endswith(want):
            problems.append(f"rank {r}: last attempt loaded {last.get('resumed_from')!r}, expected a path ending in {want}")
    for r, rows in sorted(raw_rows.items()):
        seq = [int(x["micro_step"]) for x in rows]
        if not seq:
            problems.append(f"rank {r}: empty history")
            continue
        if seq[0] != 1:
            problems.append(f"rank {r}: history starts at micro {seq[0]}, not 1")
        for a_, b_ in zip(seq, seq[1:]):
            if b_ == a_ + 1 or (b_ == a_ and a_ == stop_micro):
                continue
            problems.append(f"rank {r}: history goes micro {a_} -> {b_} (a restart or a gap, not a continuation)")
            break
    return {"ok": not problems, "problems": problems, "stop_micro": stop_micro}


def compare_batches(ref_rows: dict[int, list[dict]], sup_rows: dict[int, list[dict]]) -> dict:
    """Did every rank see the SAME row at every micro-step in the interrupted run as in the uninterrupted one?
    The sampler is a pure function of (seed, epoch, rank) and the Trainer's stream is resumable, so a correct
    resume reproduces the reference batch sequence EXACTLY; a data position that drifted (a skipped or replayed
    window) shows up here as an integer mismatch, with none of the loss curve's noise tolerance. Micro-steps
    present on one side only are mismatches too: a run that stopped early is not a match."""
    mismatches = []
    n = 0
    for r in sorted(ref_rows):
        a = {int(x["micro_step"]): x.get("batch") for x in ref_rows[r]}
        b = {int(x["micro_step"]): x.get("batch") for x in sup_rows.get(r, [])}
        for m in sorted(set(a) | set(b)):
            n += 1
            if a.get(m, "MISSING") != b.get(m, "MISSING"):
                mismatches.append({"rank": r, "micro_step": m, "ref": a.get(m, "MISSING"), "sup": b.get(m, "MISSING")})
    return {"ok": not mismatches and n > 0, "compared": n, "mismatches": mismatches[:10], "n_mismatches": len(mismatches)}


def analyze(a) -> int:
    root = getattr(a, "runs_dir", None) or SMOKE_RUNS
    ref_dir, sup_dir = os.path.join(root, "ref"), os.path.join(root, "sup")
    ref, sup = load_run(ref_dir), load_run(sup_dir)
    sup_attempts = read_jsonl(os.path.join(sup_dir, "supervisor.jsonl"))
    stop_micro = int(a.stop_micro) if a.stop_micro is not None else detect_stop_micro(sup_attempts)
    cmp = compare_loss_curves(ref["loss"], sup["loss"], stop_micro)
    prov = check_resume_provenance(load_launches(sup_dir), load_raw_rows(sup_dir), stop_micro)
    batches = compare_batches(load_raw_rows(ref_dir), load_raw_rows(sup_dir))
    steps = {"ref": step_time_stats(list(ref["per_rank"][0].values()), GRAD_ACCUM, SAVE_EVERY_MICRO),
             "sup": step_time_stats(list(sup["per_rank"][0].values()), GRAD_ACCUM, SAVE_EVERY_MICRO)}
    out = {"loss_compare": cmp, "resume_provenance": prov, "batch_compare": batches, "step_time": steps,
           "supervisor_log": sup_attempts, "moe": moe_curves(ref["per_rank"][0], GRAD_ACCUM)["global"]}
    _write_json(os.path.join(root, "analysis.json"), out)
    print(json.dumps({k: v for k, v in out.items() if k not in ("supervisor_log", "moe")}, indent=2))
    return 0 if (cmp["passed"] and prov["ok"] and batches["ok"]) else 1


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    sub = ap.add_subparsers(dest="cmd", required=True)
    p = sub.add_parser("prepare")
    p.add_argument("--ckpt")
    p.add_argument("--shards", default=SFT_SHARDS)
    p.add_argument("--data-dir", default=DATA_DIR)
    p.add_argument("--n-rows", type=int, default=N_ROWS)
    p.add_argument("--seed", type=int, default=0)
    for name in ("train", "probe"):
        t = sub.add_parser(name)
        t.add_argument("--config", default=DEFAULT_CONFIG)
        t.add_argument("--run-dir", required=True)
        t.add_argument("--ckpt")
        t.add_argument("--data-dir", default=DATA_DIR)
        t.add_argument("--n-rows", type=int)
        t.add_argument("--moe-backend", default="loop", choices=("loop", "grouped_mm"))
        t.add_argument("--set", action="append", help="config override key=value (dotted)")
        if name == "probe":
            t.add_argument("--seq-lens", default="512,1024,2048,4096")
            t.add_argument("--only", default="all", choices=("all", "collectives", "compute", "scan", "gradvar"),
                           help="collectives = no model load (NCCL bandwidth only); compute = truncated model, no FSDP; "
                                "scan = FSDP t(L); gradvar = repeat the step-1 accumulation window and compare global grad norms")
            t.add_argument("--grad-reps", type=int, default=4)
            t.add_argument("--skip-micro", type=int, default=0,
                           help="gradvar: start the window this many micro-steps into the epoch (with --ckpt <run>/checkpoints/micro_N "
                                "this replays the window that followed checkpoint N)")
            t.add_argument("--warmup", type=int, default=2)
            t.add_argument("--reps", type=int, default=5)
    d = sub.add_parser("drive")
    d.add_argument("--data-dir", default=DATA_DIR)
    d.add_argument("--timeout", type=float, default=2700.0)
    d.add_argument("--skip-ref", action="store_true")
    an = sub.add_parser("analyze")
    an.add_argument("--stop-micro", type=int, help="default: read from <sup>/supervisor.jsonl")
    an.add_argument("--runs-dir", help=f"directory holding ref/ and sup/ (default {SMOKE_RUNS})")
    a = ap.parse_args(argv)
    sys.path.insert(0, REPO)
    if a.cmd == "prepare":
        print(json.dumps(prepare(a.ckpt or default_ckpt(), a.shards, a.data_dir, a.n_rows, a.seed), indent=2))
        return 0
    if a.cmd == "train":
        from align.exit_codes import run_with_exit_codes

        return run_with_exit_codes(lambda: train_main(a))
    if a.cmd == "probe":
        return probe_main(a)
    if a.cmd == "drive":
        return drive(a)
    return analyze(a)


if __name__ == "__main__":
    raise SystemExit(main())