diff --git "a/code/align/train/trainer.py" "b/code/align/train/trainer.py" new file mode 100644--- /dev/null +++ "b/code/align/train/trainer.py" @@ -0,0 +1,2405 @@ +"""The one Trainer for SFT and DPO, plus the resume machinery it depends on. + +The stage scripts (sections 11/12) differ only in loss adapter, dataset, and +config; the loop, the checkpoint cadence, the schedule, the probes, and the +tripwires are shared so a fix to any of them cannot land in one stage and miss +the other. + +M-A (course-correction, 2026-08-22): the OLMo-core Trainer/TrainModule split. +Everything that touches the model -- forward, backward, the gated optimizer +step, FSDP2 state movement -- now lives in `align.train.train_module.TrainModule`; +this class owns the loop, the schedule, the checkpoint ORCHESTRATION (cadence, +idempotence, barriers, prune), the instruments and history.jsonl. `model`, +`optimizer` and `loss_adapter` remain readable and assignable on the Trainer +(scripts and tests use both), delegating to the module. + +Three properties here are load-bearing, and each of them is a documented +predecessor failure rather than a design preference: + + * THE SCHEDULE IS COUNTED IN OPTIMIZER STEPS. Under grad-accum 4, a horizon + counted in dataloader iterations decays the learning rate four times too + fast, and the run ends at a learning rate its config never asked for. + * THE CHECKPOINT CADENCE IS COUNTED IN MICRO-STEPS. That keeps "a kill loses + at most N micro-steps of work" true when someone changes accumulation. The + host on this box seizes without warning; a 3-hour run that checkpoints on + an accumulation-dependent cadence loses an unpredictable amount. + * RESUME IS LOSSLESS AND SEEDED. The batch stream's offset is the number of + batches CONSUMED, advanced before the yield -- see ResumableBatchStream. + +Checkpoints are bf16-weights-only, which is safe precisely because the fp32 +optimizer master is restored alongside them and is the authoritative copy. + +G0 control plane (2026-10-12). Everything below is a DECISION that more than +one rank must take identically, and the failure it prevents is the same each +time: one rank acts alone while its peers sit in a collective, so the run does +not crash, it hangs until the NCCL watchdog fires (the 5-hour stall of +2026-08-12 had exactly this shape). Hence one rule: a rank-local fact (a +signal, a wall clock, a tripwire verdict, a non-finite loss) becomes an action +only after a MAX all-reduce on the gloo group, at a program point every rank +reaches. + + * stop (SIGTERM/SIGINT, a supervisor's STOP_REQUEST file, + `train.max_wall_seconds`) -- agreed at each OPTIMIZER-STEP boundary, saved + there, never mid-accumulation; + * HALT tripwires and the arm-time self-test -- rank-reduced before raising; + * the skip cap -- computed from reduced values, so every rank counts the same; + * checkpoint integrity -- fsync before the atomic rename, and a completeness + check so a truncated directory is skipped by `latest_checkpoint`. +""" + +from __future__ import annotations + +import json +import logging +import math +import os +import re +import shutil +import signal +import struct +import threading +import time +import zipfile +from typing import Optional + +import torch + +from ..exit_codes import ( + EXIT_NONFINITE_HALT, + EXIT_TRIPWIRE_HALT, + EXIT_WALL_BUDGET, + ExitCodeError, + write_halt_marker, +) +from ..utils.supervisor import ( + DRAIN_READY_NAME, + STOP_MARKER_NAME, + STOP_REQUEST_NAME, + clear_run_marker, + read_run_marker, + write_run_marker, +) +from ..utils.dist import ( + all_gather_objects, + all_reduce_max, + any_rank, + capture_rng_state, + restore_rng_state, + world_size as _dist_world_size, +) +from ..utils.guard import assert_paths_off_root +from ..utils.tracking import Tracker +from . import sharded_ckpt +from .ckpt_format import DenseFormat +from .offload_adamw import make_optimizer +from .probes import ( + GradNormTripwire, + LengthDriftTripwire, + LogpsChosenTripwire, + TrunkMovementProbe, + aggregate_micro_step_times, + assert_tripwires_armed, + throughput_report, +) + +# Loss surface lives with the TrainModule now; re-exported here because the +# loss adapters and tests import them from this module. +from .train_module import ( # noqa: F401 + LOSS_NORMALIZATIONS, + LossAdapter, + LossOutput, + TrainModule, + reduce_window_tokens, +) + +logger = logging.getLogger(__name__) + +NAN = float("nan") + +#: Every log record must carry these. A missing key raises in the Tracker -- +#: undefined values are NaN, never absent and never 0. +MANDATORY_METRICS = ("loss", "lr", "micro_step", "optimizer_step", "grad_norm") + +#: THE CANONICAL MARGIN KEY, repo-wide. This is the name section-05's DPO +#: kernel emits, and the Trainer reads it DIRECTLY rather than through an +#: alias. One name, one authority: the previous arrangement had the kernel +#: emitting `rewards/margins` while the Trainer seeded and read `margin`, which +#: disarmed the chosen-likelihood-collapse tripwire completely. +MARGIN_KEY = "rewards/margins" +LOGPS_CHOSEN_KEY = "logps/chosen" +LENGTHS_CHOSEN_KEY = "lengths/chosen" +GRAD_NORM_KEY = "grad_norm" + +#: Tripwire names the Trainer knows how to arm. +TRIPWIRE_BUILDERS = ("logps_chosen", "length_drift", "grad_norm") + +#: Measured on this box (tmp/offload_bench.json): 2.56s CPU AdamW per 1.81B +#: rank shard + ~0.82s PCIe transfer per rank per optimizer step. At accum 4 +#: that is 20-30% step-time overhead; at accum 1 it is 50-100%. +MIN_GRAD_ACCUM_FULL_FT = 4 + + +#: Skip-cap defaults (read with getattr off `cfg.train`, because config.py is +#: another lane's file; a config that carries the field overrides these, and +#: `None` disables that cap). +DEFAULT_MAX_CONSECUTIVE_SKIPPED_STEPS = 3 +DEFAULT_MAX_SKIPPED_FRACTION = 0.05 +#: The fraction cap is not evaluated before this many optimizer-step +#: boundaries: 1 skip in the first 5 steps is 20% and says nothing. +SKIP_FRACTION_MIN_STEPS = 100 + + +class HaltError(ExitCodeError): + """A run-ending decision that every rank raises together. + + Subclasses RuntimeError (via ExitCodeError) so every pre-existing + `pytest.raises(RuntimeError, match=...)` on a tripwire HALT keeps matching. + `exit_code` is the process exit status the stage scripts map it to + (`align.exit_codes.run_with_exit_codes`) and that the supervisor refuses to + retry; `Trainer.train` also writes it to `/HALT.json`, because + torchrun reports every worker failure as exit 1. + """ + + exit_code = EXIT_TRIPWIRE_HALT + + +class NonFiniteHaltError(HaltError): + """The skip cap (consecutive / fractional) fired: the run is not training.""" + + exit_code = EXIT_NONFINITE_HALT + + +class _StopSignals: + """SIGTERM/SIGINT -> a flag, for the duration of one `train()` call. + + THE SIGTERM-MID-SAVE DEFECT (REFERENCE_GAP_ANALYSIS gap 1; scripts/sft.py + documents a torchrun SIGTERM landing mid-save). Python's default SIGTERM + action kills the process wherever it is -- including inside a 13.5GB + checkpoint write or between a rank's gather and its peers'. The handler does + the only thing a handler may safely do (set a flag); the loop acts on it at + an optimizer-step boundary, after all ranks have agreed. + + A SECOND SIGINT raises KeyboardInterrupt: an operator whose first Ctrl-C is + being honoured slowly (a save in flight) still needs a way out, and the + save path already deletes its `.partial` on any BaseException. + Handlers can only be installed from the main thread; elsewhere (a test + harness thread) this degrades to "no handlers", loudly. + """ + + SIGNALS = (signal.SIGTERM, signal.SIGINT) + + def __init__(self) -> None: + self.signum: Optional[int] = None + self._previous: dict = {} + self.installed = False + + def _handle(self, signum, frame) -> None: # noqa: ARG002 + if self.signum is not None and signum == signal.SIGINT: + raise KeyboardInterrupt + if self.signum is None: + self.signum = int(signum) + + def __enter__(self) -> "_StopSignals": + if threading.current_thread() is not threading.main_thread(): + logger.warning( + "Trainer.train() is not on the main thread: SIGTERM/SIGINT handlers " + "cannot be installed, so no coordinated stop on signal" + ) + return self + for sig in self.SIGNALS: + self._previous[sig] = signal.signal(sig, self._handle) + self.installed = True + return self + + def __exit__(self, *exc) -> None: + for sig, prev in self._previous.items(): + signal.signal(sig, signal.SIG_DFL if prev is None else prev) + self._previous.clear() + self.installed = False + + +# --------------------------------------------------------------------------- +# checkpoint durability and completeness +# --------------------------------------------------------------------------- + + +def _fsync_path(path: str) -> None: + """fsync a file OR a directory by path (a directory fsync is what makes a + rename/create inside it durable).""" + fd = os.open(path, os.O_RDONLY) + try: + os.fsync(fd) + finally: + os.close(fd) + + +def fsync_tree(root: str) -> None: + """fsync every file under `root`, then every directory (leaves first). + + THE ATOMIC-RENAME-IS-NOT-DURABLE DEFECT. `os.rename(staging, final)` is + atomic against a process crash, not against a power loss or a host seizure + (the box has seized before): without an fsync the rename can reach the + journal before the file DATA reaches disk, leaving a `micro_N` directory of + zero-length files that passes every `os.path.isdir` check -- a half-written + checkpoint that looks complete. Data first, then the directory entries, then + (caller) the rename, then the parent directory. + """ + dirs = [] + for dirpath, _dirnames, filenames in os.walk(root): + dirs.append(dirpath) + for f in filenames: + _fsync_path(os.path.join(dirpath, f)) + for d in reversed(dirs): + _fsync_path(d) + + +def _torch_save_synced(obj, path: str) -> None: + """`torch.save` to `path` and fsync it before returning.""" + with open(path, "wb") as f: + torch.save(obj, f) + f.flush() + os.fsync(f.fileno()) + + +def _safetensors_problem(path: str) -> Optional[str]: + """None if `path` is a structurally complete safetensors file, else why not. + + Reads ONLY the header: 8-byte little-endian length, then that many bytes of + JSON whose `data_offsets` say how long the tensor buffer must be. A file cut + anywhere short of that fails the size check, so truncation is detected + without reading 13GB. + """ + if not os.path.isfile(path): + return f"missing weights file {os.path.basename(path)}" + size = os.path.getsize(path) + if size < 8: + return f"{os.path.basename(path)} is {size} bytes (no safetensors header)" + with open(path, "rb") as f: + (n,) = struct.unpack(" size - 8 or n > 100 * 2**20: + return f"{os.path.basename(path)} header claims {n} bytes of a {size}-byte file" + try: + header = json.loads(f.read(n)) + except ValueError: + return f"{os.path.basename(path)} header is not valid JSON (truncated?)" + end = 0 + for k, v in header.items(): + if k == "__metadata__": + continue + end = max(end, int(v["data_offsets"][1])) + if size < 8 + n + end: + return ( + f"{os.path.basename(path)} is truncated: {size} bytes on disk, header " + f"requires {8 + n + end}" + ) + return None + + +def _zip_problem(path: str) -> Optional[str]: + """None if `path` is a readable torch.save zip archive (central directory + intact -- it is written LAST, so any truncation destroys it).""" + if not os.path.isfile(path): + return f"missing {os.path.basename(path)}" + if os.path.getsize(path) == 0: + return f"{os.path.basename(path)} is empty" + try: + with zipfile.ZipFile(path) as z: + if not any(n.endswith("data.pkl") for n in z.namelist()): + return f"{os.path.basename(path)} has no data.pkl" + except (zipfile.BadZipFile, OSError) as e: + return f"{os.path.basename(path)} is not a readable torch archive ({e})" + return None + + +def checkpoint_problem(path: str) -> Optional[str]: + """`_checkpoint_problem` made TOTAL: it returns a reason, it never raises. + + THE VERDICT-THAT-CRASHES DEFECT. The completeness check parses files a dying + host may have cut anywhere, so a half-written manifest can be syntactically + valid but structurally wrong (a `per_rank_files` layout with no + `world_size`, a safetensors header that is JSON but not an object). Those + surfaced as a KeyError/AttributeError out of the check, and every caller + that asks "is this the newest GOOD checkpoint?" -- `latest_checkpoint` + (resume), the supervisor's progress rule -- died on the exact directory it + exists to step over, instead of skipping it and falling back. Unknown is + not complete: any exception is the finding. + """ + try: + return _checkpoint_problem(path) + except Exception as e: # noqa: BLE001 -- see docstring + return f"checkpoint metadata is malformed ({type(e).__name__}: {e})" + + +def _checkpoint_problem(path: str) -> Optional[str]: + """None iff the checkpoint directory is COMPLETE; otherwise the first reason. + + "Complete" is what resume will actually need, checked cheaply (headers and + file presence, no 13GB reads): + + * weights: `model.safetensors` -- or every shard named by + `model.safetensors.index.json` -- present with a header whose declared + tensor extent fits inside the file; + * `resume.pt`: a readable archive whose top-level object is a dict with + its loop state (`micro_step`, `stream`, `scheduler`); + * when its manifest says `per_rank_files`: one readable + `optimizer_rank{r}.pt` for EVERY rank the manifest names (the SAVED + world size, not the current one -- a rank-count mismatch is refused + later with its own message and must not be disguised as "incomplete"). + + The rename-after-fsync protocol means a complete directory is the norm; this + exists for the hosts that seize between the rename and the data reaching + disk, for a checkpoint copied by hand, and for a disk that filled mid-copy. + """ + if not os.path.isdir(path): + return "not a directory" + if sharded_ckpt.is_sharded(path): + # Sharded layout (train.checkpoint_format="sharded"): the manifest names + # one weights file per rank. The directory's own manifest decides, not + # the current config, exactly as for per_rank_files below. + problem = sharded_ckpt.sharded_problem(path, _safetensors_problem) + if problem: + return problem + shards = [] + else: + index = os.path.join(path, "model.safetensors.index.json") + if os.path.isfile(index): + try: + with open(index) as f: + shards = sorted(set(json.load(f)["weight_map"].values())) + except (ValueError, KeyError, OSError) as e: + return f"unreadable model.safetensors.index.json ({e})" + else: + shards = ["model.safetensors"] + for shard in shards: + problem = _safetensors_problem(os.path.join(path, shard)) + if problem: + return problem + + resume = os.path.join(path, "resume.pt") + problem = _zip_problem(resume) + if problem: + return problem + try: + # mmap: the gather-mode resume.pt carries every rank's optimizer state + # (tens of GB); we want the manifest, not the bytes. + try: + state = torch.load(resume, weights_only=True, map_location="cpu", mmap=True) + except RuntimeError: + state = torch.load(resume, weights_only=True, map_location="cpu") + except Exception as e: # noqa: BLE001 -- any failure to load IS the finding + return f"resume.pt does not load ({type(e).__name__}: {e})" + if not isinstance(state, dict) or not {"micro_step", "stream", "scheduler"} <= state.keys(): + return "resume.pt is missing its loop state (micro_step/stream/scheduler)" + layout = state.get("optimizer_state") + if isinstance(layout, dict) and layout.get("layout") == "per_rank_files": + for r in range(int(layout["world_size"])): + problem = _zip_problem(os.path.join(path, f"optimizer_rank{r}.pt")) + if problem: + return f"per-rank optimizer state incomplete: {problem}" + return None + + +# --------------------------------------------------------------------------- +# the batch stream +# --------------------------------------------------------------------------- + + +class ResumableBatchStream: + """A seeded, resumable iterator over (epoch, batch) with an exact offset. + + THE OFF-BY-ONE. `offset` is the number of batches CONSUMED, and it is + advanced BEFORE the yield. Advancing it after the yield looks equivalent + and is not: a generator does not execute past its yield until the consumer + asks for the next item, so a checkpoint taken while holding batch i would + record offset i and the resumed run would replay batch i. One duplicated + batch is one step of stale gradient -- nothing crashes, no metric moves, + and the run stops being the run its seed describes. + + A sibling lane shipped exactly this bug: their trainer-level resume test + passed while the stream offset was off by one, and only a dedicated + offset test found it. + """ + + def __init__(self, sampler, epochs: int = 1): + self.sampler = sampler + self.epochs = int(epochs) + self.epoch = 0 + self.offset = 0 + self.consumed_total = 0 + + def __iter__(self): + while self.epoch < self.epochs: + self.sampler.set_epoch(self.epoch) + batches = list(self.sampler) + for i in range(self.offset, len(batches)): + self.offset = i + 1 + self.consumed_total += 1 + yield batches[i] + self.epoch += 1 + self.offset = 0 + + def unconsume(self) -> None: + """Hand back the batch most recently yielded, because it was never trained. + + The loop must pull a batch before it can decide the run is over, so a + horizon stop leaves the stream one batch ahead of the work actually + done. Without this the end-of-training checkpoint records an offset + (and `consumed_total`) of micro_step + 1: the resumed stream would skip + a batch no step ever saw, which is the off-by-one the class docstring + describes, arriving from the other side. Exact because a generator is + suspended AT its yield: `offset == i + 1` and the epoch rollover has not + run yet. + """ + if self.offset < 1 or self.consumed_total < 1: + raise RuntimeError("unconsume() with nothing consumed") + self.offset -= 1 + self.consumed_total -= 1 + + def state_dict(self) -> dict: + return {"epoch": self.epoch, "offset": self.offset, "consumed_total": self.consumed_total} + + def load_state_dict(self, sd: dict) -> None: + self.epoch = int(sd["epoch"]) + self.offset = int(sd["offset"]) + self.consumed_total = int(sd["consumed_total"]) + + +# --------------------------------------------------------------------------- +# schedule +# --------------------------------------------------------------------------- + + +def scheduler_horizon_optimizer_steps(rows: int, global_batch: int, epochs: int) -> int: + """`epochs * ceil(rows / global_batch)`, in OPTIMIZER steps. + + The unit is the whole point. Counting this in dataloader iterations under + accum 4 decays the LR four times too fast, and the symptom -- a run that + plateaus early -- looks like a data or capacity problem, not a scheduler + bug. + """ + if rows <= 0 or global_batch <= 0 or epochs <= 0: + raise ValueError( + f"horizon needs positive rows/global_batch/epochs, got {rows}/{global_batch}/{epochs}" + ) + return epochs * math.ceil(rows / global_batch) + + +class WarmupCosineSchedule: + """Linear warmup then cosine decay, stepped once per OPTIMIZER step.""" + + def __init__( + self, base_lr: float, total_steps: int, warmup_frac: float = 0.03, min_lr_frac: float = 0.0 + ): + if total_steps <= 0: + raise ValueError("total_steps must be positive") + self.base_lr = float(base_lr) + self.total_steps = int(total_steps) + self.warmup_steps = max(1, int(round(warmup_frac * total_steps))) + self.min_lr_frac = float(min_lr_frac) + self.step_count = 0 + + def lr(self) -> float: + t = self.step_count + if t < self.warmup_steps: + return self.base_lr * (t + 1) / self.warmup_steps + # Clamped at the horizon: a scheduler stepped past its total must not + # produce a negative or oscillating LR. + progress = min(1.0, (t - self.warmup_steps) / max(1, self.total_steps - self.warmup_steps)) + cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) + return self.base_lr * (self.min_lr_frac + (1.0 - self.min_lr_frac) * cosine) + + def step(self) -> None: + self.step_count += 1 + + def state_dict(self) -> dict: + return { + "step_count": self.step_count, + "total_steps": self.total_steps, + "warmup_steps": self.warmup_steps, + "base_lr": self.base_lr, + } + + def load_state_dict(self, sd: dict) -> None: + self.step_count = int(sd["step_count"]) + self.total_steps = int(sd["total_steps"]) + self.warmup_steps = int(sd["warmup_steps"]) + self.base_lr = float(sd["base_lr"]) + + +# --------------------------------------------------------------------------- +# the Trainer +# --------------------------------------------------------------------------- + + +class Trainer: + """Micro-step loop, accumulation, gated optimizer step, cadence, resume.""" + + def __init__( + self, + cfg, + model, + sampler, + loss_adapter: LossAdapter, + run_dir: str, + optimizer=None, + scheduler=None, + tracker=None, + stage: str = "sft", + train_module: Optional[TrainModule] = None, + weight_format=None, + allow_resume_mismatch: bool = False, + ): + self.cfg = cfg + self.sampler = sampler + self.stage = stage + + accum = cfg.optimizer.grad_accum + # RULED (team-lead, pre-launch): keyed to `offload_adamw`, because a + # guard must be keyed to the FAILURE IT PREVENTS, not to the + # circumstances that usually accompany it. This rule's failure mode is + # offload amortization -- the measured 2.56s CPU optimizer step and + # 0.82s PCIe transfer that only that optimizer pays. Under adamw8bit + # the state lives on-GPU and that failure mode is ABSENT, so the guard + # is inapplicable rather than weakened. Firing it there would have + # forced two pre-registration edits (global batch size and a halved + # horizon) to satisfy a constraint with nothing behind it. + if ( + cfg.train.full_ft + and cfg.optimizer.name == "offload_adamw" + and accum < MIN_GRAD_ACCUM_FULL_FT + ): + raise ValueError( + f"grad_accum={accum} is below {MIN_GRAD_ACCUM_FULL_FT}, which " + "full-FT with offloaded AdamW requires. Measured on this box: " + "the CPU optimizer step is 2.56s per 1.81B-param rank shard " + "and PCIe transfer is ~0.82s per rank per optimizer step. " + "Amortised over 4 accumulation micro-steps that is 20-30% " + "step-time overhead; at accum 1 it is 50-100% and the run " + "spends more wall-clock moving state than computing. Raise " + "grad_accum, or set train.full_ft=false for a small model. " + "If host RAM forced the adamw8bit fallback, set " + "optimizer.name=adamw8bit -- its state is on-GPU and this " + "rule does not apply to it." + ) + self.accum = accum + #: U01-5: resume refuses a checkpoint whose seed / grad_accum / data fingerprint differ + #: from this run's unless the operator says so (kwarg, or `train.allow_resume_mismatch`). + self.allow_resume_mismatch = bool( + allow_resume_mismatch or getattr(cfg.train, "allow_resume_mismatch", False) + ) + #: U01-1: micro-batches accumulated since the last optimizer boundary, and what became of + #: a partial tail window ({"micros", "action"}), None when the run ended on a boundary. + self._pending_micros = 0 + self.tail_window: Optional[dict] = None + + # THE MID-ACCUMULATION SAVE DEFECT (align port review, 2026-10-12). A + # checkpoint holds weights, optimizer state, schedule and stream offset -- + # NOT the gradients accumulated so far in the current window. A cadence + # that does not land on an optimizer-step boundary therefore saves a + # half-filled window: the resumed run restarts the counter at the saved + # micro_step, applies its first step after only (accum - k) micro-batches + # (each already scaled by 1/accum), and silently trains a different run. + # Measured on the tiny harness (accum 2, save_every 3): max weight + # divergence 3e-3 at lr 1e-3 against the uninterrupted run, no error, no + # metric out of line. The RL stage config already refuses this (its saves + # are weight syncs); the shared Trainer is where SFT/DPO/pretrain need it. + save_every = cfg.train.save_every_micro_steps + if save_every % accum != 0: + raise ValueError( + f"train.save_every_micro_steps={save_every} is not a multiple of " + f"optimizer.grad_accum={accum}: a checkpoint taken mid-accumulation " + "drops the window's accumulated gradients, so a resume from it is " + "NOT lossless. Pick a cadence that lands on an optimizer-step " + f"boundary (e.g. {max(1, save_every // accum) * accum})." + ) + + assert_paths_off_root([run_dir]) + self.run_dir = run_dir + self.ckpt_root = os.path.join(run_dir, "checkpoints") + os.makedirs(self.ckpt_root, exist_ok=True) + self._discard_partials() + + self.stream = ResumableBatchStream(sampler, epochs=cfg.train.epochs) + # The module is the ONLY holder of model/loss/optimizer; the Trainer's + # `model`/`optimizer`/`loss_adapter` attributes below are delegating + # properties, kept assignable because scripts free the optimizer + # before eval (`trainer.optimizer = None`) and tests swap adapters. + self.train_module = train_module or TrainModule( + model=model, + loss_adapter=loss_adapter, + optimizer=optimizer + or make_optimizer(cfg, [p for p in model.parameters() if p.requires_grad]), + grad_accum=accum, + ) + + # U03-2: "micro_mean" (default) is the original loop, untouched. "global_token_mean" buffers + # each accumulation window, all-reduces its target-token count and scales every micro-batch + # by n_micro / N_window. Refuse, at construction, every combination that cannot honour it. + self.loss_normalization = getattr(cfg.train, "loss_normalization", "micro_mean") + if self.loss_normalization not in LOSS_NORMALIZATIONS: + raise ValueError( + f"train.loss_normalization={self.loss_normalization!r}; expected one of {LOSS_NORMALIZATIONS}" + ) + if self.loss_normalization == "global_token_mean": + if type(self.train_module).train_batch is not TrainModule.train_batch: + raise ValueError( + f"train.loss_normalization='global_token_mean' needs the stock TrainModule.train_batch; " + f"{type(self.train_module).__name__} overrides it (RL/PPO/MoE modules own their loss scaling)" + ) + if not callable(getattr(self.train_module.loss_adapter, "count_tokens", None)): + raise ValueError( + "train.loss_normalization='global_token_mean' needs a loss adapter exposing " + "`count_tokens(batch)` (make_sft_loss_adapter provides it)" + ) + self._window_tokens: Optional[int] = None + + # Which converter writes/reads the weight file (align/train/ckpt_format.py). + # Explicit argument > the module's own (an MoE module knows its family) + # > dense from cfg.model, the behaviour every pre-MoE caller has. + self.weight_format = ( + weight_format + or getattr(self.train_module, "weight_format", None) + or DenseFormat(cfg.model) + ) + + self._validate_checkpoint_format() + + total_micro = len(sampler) * cfg.train.epochs + horizon = max(1, math.ceil(total_micro / self.accum)) + self.scheduler = scheduler or WarmupCosineSchedule( + base_lr=cfg.optimizer.lr, total_steps=horizon, warmup_frac=cfg.train.warmup_frac + ) + + self.tracker = tracker or Tracker( + run_dir, stage, mandatory_keys=MANDATORY_METRICS, run_name=cfg.run_name + ) + # Remember what the probe was built on: resume rebuilds it, and it must be + # rebuilt on the SAME view (a PPO run passes an adapter-only view because + # the frozen base can never move), not on the raw model. + self._probe_target = model + self.trunk_probe = TrunkMovementProbe(model) + self.trunk_report: dict = {} + tw = cfg.tripwires + # Only ARMED tripwires are constructed. An unarmed stage does not get + # a guard it cannot feed -- SFT has no preference margin, and arming + # the collapse tripwire there would either halt every run or (the old + # behaviour) quietly feed it NaN forever. + unknown = set(tw.armed) - set(TRIPWIRE_BUILDERS) + if unknown: + raise ValueError( + f"tripwires.armed names {sorted(unknown)}, which the Trainer " + f"cannot build; known: {list(TRIPWIRE_BUILDERS)}" + ) + self.tripwires: dict = {} + if "logps_chosen" in tw.armed: + self.tripwires["logps_chosen"] = LogpsChosenTripwire(window=tw.logps_window) + if "grad_norm" in tw.armed: + self.tripwires["grad_norm"] = GradNormTripwire() + if "length_drift" in tw.armed: + self.tripwires["length_drift"] = LengthDriftTripwire( + baseline=tw.chosen_longer_baseline, + band=tw.length_band, + min_pairs=tw.length_min_pairs, + ) + self.tripwire_alarms: list[dict] = [] + self._tripwire_self_test_done = False + + self.micro_step = 0 + self.optimizer_step = 0 + self._last_applied_lr: Optional[float] = None + #: Applied optimizer steps SINCE the trunk snapshot was taken. Not the + #: same as optimizer_step: after a resume the total is already nonzero + #: while the fresh snapshot has seen no movement yet, and probing on + #: the total would fire on every resumed run. + self._applied_since_snapshot = 0 + self.checkpoints_written: list[int] = [] + #: R5 (2026-08-23 audit): checkpoints pinned against pruning. Only the + #: final two checkpoints of each destroyed phase-1 run survived + #: (keep-2), so nothing from before the instability existed to rewind + #: to. The FIRST fire of any tripwire pins the NEXT save. + self.pinned_checkpoints: set[int] = set() + self._pin_next_save = False + self.step_times: list[float] = [] + self._probe_done = False + + # -- G0 control plane ------------------------------------------------ + tr = cfg.train + self.max_wall_seconds = getattr(tr, "max_wall_seconds", None) + if self.max_wall_seconds is not None and not float(self.max_wall_seconds) > 0: + raise ValueError( + f"train.max_wall_seconds={self.max_wall_seconds!r} must be > 0 (or None to disable)" + ) + self.max_consecutive_skipped = getattr( + tr, "max_consecutive_skipped_steps", DEFAULT_MAX_CONSECUTIVE_SKIPPED_STEPS + ) + self.max_skipped_fraction = getattr(tr, "max_skipped_fraction", DEFAULT_MAX_SKIPPED_FRACTION) + if self.max_consecutive_skipped is not None and int(self.max_consecutive_skipped) < 0: + raise ValueError( + f"train.max_consecutive_skipped_steps={self.max_consecutive_skipped!r} must be >= 0" + ) + if self.max_skipped_fraction is not None and not 0.0 <= float(self.max_skipped_fraction) <= 1.0: + raise ValueError( + f"train.max_skipped_fraction={self.max_skipped_fraction!r} must be in [0, 1]" + ) + self._stop_signals = _StopSignals() + self.stop_reason: Optional[str] = None + self._train_t0: Optional[float] = None + #: On-device OR of "this micro-batch's loss was non-finite" since the + #: last optimizer-step boundary; read once per boundary (one sync). + self._nonfinite_loss_acc: Optional[torch.Tensor] = None + #: Replicated on every rank (updated from reduced values only) and + #: persisted in resume.pt, so a resume does not forgive a bad run. + self.skip_stats = {"consecutive": 0, "bad_total": 0, "boundaries": 0} + + # -- delegation to the TrainModule --------------------------------------- + + @property + def model(self): + return self.train_module.model + + @model.setter + def model(self, value) -> None: + self.train_module.model = value + if value is None: + # U01-11: clearing the model is the post-train handoff; the probe's own references + # (`_probe_target`, `trunk_probe._model`) would otherwise keep the replica alive. + self._drop_probe_refs() + + def _drop_probe_refs(self) -> None: + self._probe_target = None + probe = getattr(self, "trunk_probe", None) + if probe is not None: + probe.release() + + def release(self) -> None: + """Free everything the Trainer holds on the training replica (model, optimizer, probe + snapshots). Call before building the post-train eval model; the Trainer cannot resume or + probe afterwards.""" + self.model = None + self.optimizer = None + self._drop_probe_refs() + + @property + def optimizer(self): + return self.train_module.optimizer + + @optimizer.setter + def optimizer(self, value) -> None: + self.train_module.optimizer = value + + @property + def loss_adapter(self) -> LossAdapter: + return self.train_module.loss_adapter + + @loss_adapter.setter + def loss_adapter(self, value: LossAdapter) -> None: + self.train_module.loss_adapter = value + + # -- the loop ----------------------------------------------------------- + + def train(self) -> dict: + self.stop_reason = None + with self._stop_signals: + self._train_t0 = time.monotonic() + self._train_wall_t0 = time.time() + self._announce_drain(True) + try: + return self._train_loop() + except HaltError as e: + # The verdict must survive the launcher: torchrun exits 1 for + # any worker failure, so the supervisor reads this file instead + # (align.exit_codes). Identical on every rank; rank 0 writes. + self._write_halt_marker(e.exit_code, str(e)) + raise + finally: + self._announce_drain(False) + + def _announce_drain(self, on: bool) -> None: + """Rank 0 tells a supervisor, through `DRAIN_READY.json`, that this + process polls `STOP_REQUEST` at every boundary -- so a stop should be + REQUESTED and waited for rather than signalled and killed (G3 audit + U01-1: a signal stop needs a boundary plus a save, longer than the + launcher's SIGTERM->SIGKILL window). Present only while `train()` runs: + before it there is no boundary to wait for, after it the post-train + eval is not interruptible this way.""" + from ..utils.dist import is_rank0 + + if not is_rank0(): + return + if on: + write_run_marker(self.run_dir, DRAIN_READY_NAME, pid=os.getpid()) + else: + clear_run_marker(self.run_dir, DRAIN_READY_NAME) + + def _stop_requested_by_file(self) -> bool: + """True iff a supervisor's `STOP_REQUEST` written during THIS train() + call exists. Read by rank 0 only (a shared run dir); the boundary + all-reduce spreads it, so a rank with a stale view cannot stop alone. + A request older than this call is a leftover and is ignored.""" + from ..utils.dist import is_rank0 + + if not is_rank0(): + return False + since = getattr(self, "_train_wall_t0", 0.0) + return read_run_marker(self.run_dir, STOP_REQUEST_NAME, since) is not None + + def _write_stop_marker(self) -> None: + """Leave `STOP.json` (code 128+signum) so the supervisor can tell an + operator's stop from a crash: torchrun reports both as exit 1 (G3 + audit U01-2). Written when the stop is AGREED, before the save, so a + launcher that SIGKILLs the save still leaves the verdict.""" + from ..utils.dist import is_rank0 + + if is_rank0(): + code = 128 + int(self._stop_signals.signum or signal.SIGTERM) + write_run_marker( + self.run_dir, STOP_MARKER_NAME, code=code, reason="coordinated stop (signal)" + ) + + def _write_halt_marker(self, code: int, reason: str) -> None: + from ..utils.dist import is_rank0 + + if is_rank0(): + write_halt_marker(self.run_dir, code, reason[:2000]) + + @property + def stop_exit_code(self) -> int: + """The process exit status the stage script should return after `train()`. + + 0 for a run that ended on its own terms (data exhausted, horizon, step + cap). A COORDINATED STOP is not a finished stage: continuing into the + post-train eval/gate of a half-trained model (what the scripts did) + would pass or fail it on weights nobody meant to evaluate. A wall-clock + stop is `EXIT_WALL_BUDGET` (the budget is spent; the supervisor must not + relaunch it); a signal stop is the shell convention 128+signum. + """ + if self.stop_reason == "wall_clock": + return EXIT_WALL_BUDGET + if self.stop_reason == "signal": + return 128 + int(self._stop_signals.signum or signal.SIGTERM) + return 0 + + def _train_loop(self) -> dict: + cfg = self.cfg + limit = cfg.train.max_micro_steps + self.train_module.set_train() + + stop_reason = "data_exhausted" + global_mean = self.loss_normalization == "global_token_mean" + it = iter(self.stream) + lookahead: list = [] # global_token_mean only: the rest of the current window, already pulled + while True: + if lookahead: + batch = lookahead.pop(0) + else: + try: + batch = next(it) + except StopIteration: + break + # THE ZOMBIE-RUN GUARD. Past the scheduler's horizon the cosine + # clamps, and with min_lr_frac=0 it clamps to EXACTLY 0 -- the + # loop then runs forward, backward, step, log and checkpoint on + # frozen weights, with every metric looking normal. DPO Arm A did + # that for 287 optimizer steps (~24 GPU-hours) because nothing + # stopped it. There is no learning left past the horizon by + # construction, so the run ends and says why. + if self.optimizer_step >= self.scheduler.total_steps: + stop_reason = "horizon" + logger.info( + "stopping at the scheduler horizon: %d optimizer steps " + "completed of %d planned. Past this point the cosine " + "schedule clamps the learning rate to %.3g and further " + "steps would not change the weights.", + self.optimizer_step, + self.scheduler.total_steps, + self.scheduler.lr(), + ) + # The batch just pulled is not going to be trained on: give it + # back so the end-of-training checkpoint's stream offset is the + # number of batches TRAINED, not trained + 1. + self.stream.unconsume() + break + + t0 = time.perf_counter() + if global_mean and self.micro_step % self.accum == 0: + # Window start: pull the rest of the window, count its target tokens, reduce over + # ranks. Saves/stops only happen on window boundaries, so the stream being up to + # accum-1 batches ahead mid-window is never observable in a checkpoint. + window = [batch] + while len(window) < self.accum: + try: + window.append(next(it)) + except StopIteration: + break + lookahead = window[1:] + count = self.train_module.loss_adapter.count_tokens + self._window_tokens = reduce_window_tokens(sum(int(count(b)) for b in window)) + self.micro_step += 1 + + if global_mean: + out = self.train_module.train_batch(batch, window_tokens=self._window_tokens) + else: + out = self.train_module.train_batch(batch) + self._note_loss(out.loss) + self._pending_micros += 1 + + stepped = False + stop: Optional[str] = None + if self.micro_step % self.accum == 0: + stepped, stop = self._optimizer_boundary() + # U01-8: the armed grad_norm HALT wire must inspect EVERY optimizer step. `_log` + # feeds it only on logged micros, so with log_every=10 / grad_accum=2 four steps in + # five were never seen. Feed the boundaries `_log` will not cover (a logged boundary + # is fed there, from the same value: one observation per step, never two). + if ( + stepped + and "grad_norm" in self.tripwires + and self.micro_step % cfg.train.log_every_micro_steps != 0 + ): + self._check_tripwires( + {GRAD_NORM_KEY: self.train_module.last_grad_norm}, names=("grad_norm",) + ) + self.step_times.append(time.perf_counter() - t0) + + if self.micro_step % cfg.train.log_every_micro_steps == 0: + self._log(batch, out, stepped) + if self.micro_step % cfg.train.save_every_micro_steps == 0: + self.save_checkpoint() + if stop is not None: + # Before the `limit` test: a stop request must win, because the + # capped-run branch below skips the final save and a SIGTERM + # that discarded the window's work is the failure this prevents. + stop_reason = stop + if stop == "signal": + self._write_stop_marker() + logger.warning( + "coordinated stop (%s) at optimizer step %d (micro %d)", + stop, + self.optimizer_step, + self.micro_step, + ) + break + if limit is not None and self.micro_step >= limit: + stop_reason = "max_micro_steps" + break + + stop_reason = self._finish_partial_window(stop_reason) + + # A COORDINATED STOP always saves: the boundary is the one point where + # the checkpoint is lossless (no half-filled accumulation window), and + # the cadence save may be up to save_every micro-steps away. + if stop_reason in ("signal", "wall_clock"): + self.save_checkpoint() + + # END-OF-TRAINING SAVE. The cadence saves at multiples of save_every; + # an epoch whose total is not divisible leaves the tail only in + # process memory. Stage-11's first full epoch ended at micro 2250 + # with the last save at 2000, then the post-epoch eval crashed -- the + # final 125 optimizer steps existed nowhere on disk. Only for + # uncapped runs: a capped probe must never spend minutes on a 14.5GB + # gather or leave a probe artifact that keep-N pruning could count + # against real checkpoints. + elif limit is None and self.micro_step % cfg.train.save_every_micro_steps != 0: + self.save_checkpoint() + + self.stop_reason = stop_reason + if stop_reason == "wall_clock": + self._write_halt_marker(EXIT_WALL_BUDGET, "train.max_wall_seconds reached; checkpointed") + summary = self.summary() + summary["stop_reason"] = stop_reason + logger.info( + "training finished (%s): %d micro-steps, %d optimizer steps", + stop_reason, + self.micro_step, + self.optimizer_step, + ) + return summary + + # -- G0 control plane ----------------------------------------------------- + + def _finish_partial_window(self, stop_reason: str) -> str: + """U01-1: decide what happens to a PARTIAL accumulation window when the loop ends. + + THE DROPPED-TAIL DEFECT (Codex audit U01-1). `train.epochs * len(sampler)` micro-batches + that are not a multiple of `grad_accum` left the last `< accum` micro-batches accumulated + but never stepped: their .grad stayed on the params (a later forward/eval, or an optimizer + reused after the run, would see it), the scheduler horizon (ceil) planned a step that was + never taken, and nothing recorded the loss. + + End of DATA (`data_exhausted`): the tail is a real, if short, window -- step it, with the + gradient rescaled to a mean over the `k` micro-batches it holds (each was scaled 1/accum), + so the horizon (ceil) is consistent with the steps taken. Under `global_token_mean` the + window is already normalised by its own token count (see `train_batch`), so no rescale. + A module that overrides `train_batch` owns its loss scaling, so the correct factor is + unknown: its tail is DROPPED, loudly, rather than stepped with a wrong scale. Any other + ending (`max_micro_steps` cap) drops the tail the same way. Dropping always zeroes the + grads and records `tail_window` in the summary. + """ + k = int(self._pending_micros) + if k == 0: + return stop_reason + tm = self.train_module + if stop_reason == "data_exhausted": + stock = type(tm).train_batch is TrainModule.train_batch + if stock: + if self.loss_normalization == "micro_mean" and k < self.accum: + factor = float(self.accum) / float(k) + with torch.no_grad(): + for prm in self.model.parameters(): + if prm.grad is not None: + prm.grad.mul_(factor) + logger.warning( + "data exhausted mid-window: stepping the %d-micro-batch tail (grad_accum=%d)", + k, + self.accum, + ) + _stepped, stop = self._optimizer_boundary() + self.tail_window = {"micros": k, "action": "stepped"} + if stop == "signal": + self._write_stop_marker() + if stop is not None: + return stop + return stop_reason + tm.zero_grads() + self._pending_micros = 0 + self._nonfinite_loss_acc = None + self.tail_window = {"micros": k, "action": "dropped"} + logger.warning( + "run ended (%s) with a partial accumulation window of %d micro-batch(es); its gradients " + "were DROPPED and zeroed (recorded in the summary as tail_window)", + stop_reason, + k, + ) + return stop_reason + + def _optimizer_boundary(self) -> tuple[bool, Optional[str]]: + """The optimizer-step boundary: apply the LR, step, advance the schedule, agree on stop/skip. + + Returns (stepped, stop_reason). Shared by the in-loop boundary and the end-of-data tail + window (U01-1) so both take exactly the same path. + """ + self._pending_micros = 0 + # THE UNAPPLIED-SCHEDULE DEFECT (found in the align port, present + # in the legacy trainer too). The scheduler was stepped and its LR + # was LOGGED, but nothing ever wrote it into the optimizer: the + # optimizer was built once with lr=cfg.optimizer.lr, so every + # SFT/DPO/GRPO run trained at a CONSTANT lr (no warmup, no decay) + # while history.jsonl reported a decaying one. Measured on the tiny + # harness: scheduler.lr() 1e-3 -> 7.5e-4 -> 2.5e-4 while + # optimizer.param_groups[0]["lr"] stayed 1e-3. The logged value and + # the applied value must be the same number, so apply it here, + # immediately before the step it governs. Idempotent with modules + # that set their own lr (PPO). + self._apply_scheduled_lr() + skipped_before = self.train_module.skipped_steps + stepped = self.train_module.optim_step() + # A step that was not applied is only a SKIP if the optimizer's + # grad gate said so (its `skipped_steps` counter moved). A + # module may decline to step by design -- PPO's critic-only + # warmup iterations return False every time -- and counting + # those would halt a healthy run on its fourth warmup window. + gate_skipped = self.train_module.skipped_steps > skipped_before + if stepped: + # The schedule counts steps that ACTUALLY happened: a run + # whose grads went non-finite for a stretch must not decay + # its LR on schedule while training nothing. + self.scheduler.step() + self.optimizer_step += 1 + self._applied_since_snapshot += 1 + # ONE reduced decision per optimizer step: stop request, skip + # health. Raises (identically on every rank) past the skip cap. + stop = self._boundary_control(gate_skipped) + return stepped, stop + + def _note_loss(self, loss) -> None: + """Remember, on-device, whether any micro-batch loss in this window was + non-finite. + + THE LOSS-ONLY-NAN BLIND SPOT. The skip gate looks at GRADIENTS. A + non-finite loss whose gradient is finite (a decoupled-backward proxy, a + NaN in a logged-only term) steps happily, and a loss that is NaN every + micro-step looks like a healthy run to the gate. The flag is OR-ed on + the loss's own device and read once per optimizer step, so it costs no + extra host sync per micro-step. + """ + bad = ~torch.isfinite(loss.detach()).all() + acc = self._nonfinite_loss_acc + self._nonfinite_loss_acc = bad if acc is None else (acc | bad) + + def _wall_limit_hit(self) -> bool: + if self.max_wall_seconds is None or self._train_t0 is None: + return False + return (time.monotonic() - self._train_t0) >= float(self.max_wall_seconds) + + def _boundary_control(self, gate_skipped: bool) -> Optional[str]: + """The optimizer-step-boundary agreement. Returns the stop reason or None. + + ONE MAX all-reduce carries four rank-local facts: [signalled, + wall-clock expired, a loss was non-finite, the grad gate skipped the step]. + Every rank therefore sees the same four numbers and takes the same + branch -- the only way "any rank signalled => all stop" can be true + without a rank stopping alone while a peer waits in the next + micro-step's FSDP all-gather. Wall clocks are per-rank, so they are + reduced too: ranks disagree about elapsed time by the skew of their + launch. + + Called only on a boundary, so a stop never lands mid-accumulation (a + checkpoint holds no half-filled gradient window; see __init__). + """ + nonfinite = False + if self._nonfinite_loss_acc is not None: + nonfinite = bool(self._nonfinite_loss_acc) + self._nonfinite_loss_acc = None + sig, wall, nf, gate_skipped = all_reduce_max( + [ + int(self._stop_signals.signum is not None or self._stop_requested_by_file()), + int(self._wall_limit_hit()), + int(nonfinite), + int(gate_skipped), + ] + ) + self._record_step_health(bool(nf), bool(gate_skipped)) + if sig: + return "signal" + if wall: + return "wall_clock" + return None + + def _record_step_health(self, nonfinite_loss: bool, gate_skipped: bool) -> None: + """Count bad optimizer steps; HALT past the consecutive or fractional cap. + + THE SILENT-SKIP RUN. The grad gate skips a non-finite step and trains + on; if the cause is persistent (an overflowing loss scale, a poisoned + shard, an LR that has already blown the weights) it skips FOREVER: + every step "succeeds", the schedule stalls (it counts applied steps), + and the run burns GPU-hours learning nothing. `skipped_steps` was + logged but nothing ever read it. A step is BAD if the grad gate skipped it OR + any of its micro-batch losses was non-finite. The inputs are the + already-reduced flags, so every rank holds identical counters and + raises the identical error with no further collective. + """ + st = self.skip_stats + st["boundaries"] += 1 + if nonfinite_loss or gate_skipped: + st["consecutive"] += 1 + st["bad_total"] += 1 + else: + st["consecutive"] = 0 + cause = ( + "non-finite loss" + (" and skipped step" if gate_skipped else "") + if nonfinite_loss + else "skipped step (non-finite gradient)" + ) + cap = self.max_consecutive_skipped + if cap is not None and st["consecutive"] > int(cap): + self._halt_skip_cap( + f"{st['consecutive']} consecutive bad optimizer steps exceed " + f"train.max_consecutive_skipped_steps={cap} (last cause: {cause})", + cause, + ) + frac = self.max_skipped_fraction + if frac is not None and st["boundaries"] >= SKIP_FRACTION_MIN_STEPS: + observed = st["bad_total"] / st["boundaries"] + if observed > float(frac): + self._halt_skip_cap( + f"{st['bad_total']} of {st['boundaries']} optimizer steps were bad " + f"({observed:.1%}), above train.max_skipped_fraction={frac} " + f"(last cause: {cause})", + cause, + ) + + def _halt_skip_cap(self, detail: str, cause: str) -> None: + self.tripwire_alarms.append( + {"step": self.micro_step, "alarm": "skip_cap", "response": f"HALT: {cause}"} + ) + raise NonFiniteHaltError( + f"HALT skip cap at micro_step {self.micro_step} (optimizer step " + f"{self.optimizer_step}): {detail}. The run is not training: fix the " + "cause (loss scale / LR / data) or raise the cap deliberately; the last " + "good checkpoint is on disk." + ) + + def _agree_or_raise( + self, local_exc: Optional[BaseException], what: str, halt_on_failure: bool = False + ) -> None: + """Make a rank-local failure a GLOBAL one before anyone raises. + + THE LONE-RAISER HANG. A tripwire verdict, the arm-time self-test and the + trunk probe all read rank-local numbers (this rank's data, this rank's + shard). A rank that raises alone unwinds out of the loop while its peers + enter the next micro-step's collective and wait for it -- the process + group does not fail, it blocks until the watchdog fires 30 minutes + later. So: reduce the verdict (MAX) first, every rank at the same + program point, then raise everywhere. + + A `HaltError` is raised IDENTICALLY on every rank, built from the + rank-ordered gather of every failing rank's message. Any other + exception (a KeyError from a missing metric) is re-raised as itself on + the rank that hit it, and peers get a `HaltError` naming it. + + `halt_on_failure=True` (the trunk probe and the arm-time self-test) converts ANY failure + into a `HaltError` on the originating rank too (U01-21): those checks are deterministic, + so a plain exception there was an ordinary crash with no HALT.json -- the supervisor + relaunched a run that would fail identically. + """ + if ( + halt_on_failure + and local_exc is not None + and not isinstance(local_exc, HaltError) + ): + converted = HaltError(f"{what} failed: {type(local_exc).__name__}: {local_exc}") + converted.__cause__ = local_exc + local_exc = converted + if _dist_world_size() == 1: + if local_exc is not None: + raise local_exc + return + from ..utils.dist import any_rank + + if not any_rank(local_exc is not None): + return + # (message, exit code) per rank, so the identical HaltError every rank + # raises carries the SAME exit status: a non-finite verdict on one rank + # must not be reported as a tripwire HALT by its peers. + msgs = all_gather_objects( + None + if local_exc is None + else ( + f"{type(local_exc).__name__}: {local_exc}", + int(getattr(local_exc, "exit_code", EXIT_TRIPWIRE_HALT)), + ) + ) + failed = [(r, m) for r, m in enumerate(msgs) if m is not None] + if local_exc is not None and not isinstance(local_exc, HaltError): + raise local_exc + cls = NonFiniteHaltError if any(c == EXIT_NONFINITE_HALT for _, (_, c) in failed) else HaltError + raise cls( + "; ".join(f"[rank {r}] {m}" for r, (m, _) in failed) + + f" (rank-reduced during {what}: all ranks stop together)" + ) + + def _apply_scheduled_lr(self) -> None: + """Write the scheduler's current LR into every optimizer param group. + + A no-op for an optimizer that exposes no `param_groups` (a test stub); + a real optimizer that does expose them gets the scheduled value. + """ + # Remember the value so the log reports what the step USED. The scheduler + # is stepped right after the optimizer step, so `scheduler.lr()` at log + # time is lr_{k+1} on the micro that applied step k: the log ran one + # schedule step ahead of the optimizer (review finding, low). + lr = self.scheduler.lr() + self._last_applied_lr = lr + opt = self.optimizer + groups = getattr(opt, "param_groups", None) + if not groups: + return + for g in groups: + g["lr"] = lr + + def _log(self, batch, out: LossOutput, stepped: bool) -> None: + metrics = { + "loss": float(out.loss.detach()), + # The lr the optimizer last APPLIED; before the first step (or just + # after a resume) the lr the next step will apply. + "lr": ( + self._last_applied_lr + if getattr(self, "_last_applied_lr", None) is not None + else self.scheduler.lr() + ), + "micro_step": self.micro_step, + "optimizer_step": self.optimizer_step, + # Undefined on a micro-step that did not step, and on a skipped + # step. NaN, never 0.0 -- a 0.0 grad norm reads as 'converged'. + "grad_norm": (self.train_module.last_grad_norm if stepped else NAN), + "batch": list(batch), + "n_tokens": out.n_tokens, + "skipped_steps": self.train_module.skipped_steps, + } + # NOTE: tripwire INPUTS are deliberately NOT seeded here. Seeding + # `margin`/`logps/chosen` with NaN is exactly what disarmed the + # collapse tripwire: it made the key always present and always + # unusable, so the guard ran blind and reported itself armed. The loss + # adapter supplies them when they exist; when it does not, the + # arm-time self-test below halts rather than pretending. + metrics.update(out.metrics) + self.tracker.log(self.micro_step, metrics) + + # Fail fast: the trunk probe runs at the first log step AFTER the + # first APPLIED optimizer step. Not earlier -- before any step lands, + # "the trunk has not moved" is trivially true and the probe would fire + # on every run. Not later: the whole point is to fail before + # GPU-hours accumulate behind a frozen trunk. + if not self._probe_done and self._applied_since_snapshot >= 1: + err = None + try: + self.trunk_report = self.trunk_probe.check(step=self.micro_step) + except Exception as e: # noqa: BLE001 -- re-raised by _agree_or_raise + err = e + self._agree_or_raise(err, "trunk-movement probe", halt_on_failure=True) + self._probe_done = True + + # Arm-time self-test, once, at the first log step: every armed + # tripwire's inputs present and finite, or halt before training + # proceeds. This turns metric-name drift into a launch failure. + if not self._tripwire_self_test_done: + err = None + try: + assert_tripwires_armed(metrics, self.tripwires) + except Exception as e: # noqa: BLE001 -- re-raised by _agree_or_raise + err = e + self._agree_or_raise(err, "tripwire arm-time self-test", halt_on_failure=True) + self._tripwire_self_test_done = True + + self._check_tripwires(metrics) + + def _check_tripwires(self, metrics: dict, names: Optional[tuple] = None) -> None: + """Feed each armed tripwire from the CANONICAL metric keys. + + Reads `metrics[...]` directly and lets a KeyError surface: an armed + tripwire whose input has gone missing mid-run must fail loudly, not + fall back to a default. The arm-time self-test catches this at step + one; this is the backstop if a later adapter stops emitting the key. + """ + halt: Optional[BaseException] = None + for name, wire in self.tripwires.items(): + if names is not None and name not in names: + continue + try: + if name == "logps_chosen": + r = wire.update( + metrics[LOGPS_CHOSEN_KEY], metrics[LENGTHS_CHOSEN_KEY], metrics[MARGIN_KEY] + ) + elif name == "grad_norm": + r = wire.update(metrics[GRAD_NORM_KEY]) + elif name == "length_drift": + # `pairs` is the sample the fraction was computed over: the + # tripwire needs it to know whether it has seen enough to + # speak, and it is a key the DPO adapter already emits. + r = wire.update(metrics["chosen_longer_fraction"], n_pairs=metrics["pairs"]) + else: # pragma: no cover + raise RuntimeError(f"no feed defined for tripwire {name!r}") + except Exception as e: # noqa: BLE001 -- a missing input is a rank-local failure + halt = e + break + if r.tripped: + if not any(a["alarm"] == name for a in self.tripwire_alarms): + # R5: first fire of this wire pins the next checkpoint, so + # an early-stop/rewind decision has somewhere to rewind TO. + self._pin_next_save = True + self.tripwire_alarms.append( + {"step": self.micro_step, "alarm": name, "response": r.response} + ) + if getattr(wire, "HALT", False): + # R3 (2026-08-23 audit): a HALT tripwire stops the run at + # the alarm. The recorded rerun trained 221 optimizer steps + # past a 355,830 grad-norm blow-up because every guard only + # logged; a gradient explosion is a defect, not a metric. + halt = HaltError( + f"HALT tripwire {name!r} fired at micro_step " + f"{self.micro_step}: {r.response} (value={r.value})" + ) + break + # The verdict is rank-local (this rank's data, this rank's shard); + # reducing it is what keeps a rank that raised from leaving its peers + # in the next micro-step's collective. Skipped only when NO tripwire is + # armed -- config-derived, hence the same on every rank. + if self.tripwires: + self._agree_or_raise(halt, "tripwire check") + elif halt is not None: # pragma: no cover -- nothing to feed + raise halt + + def summary(self) -> dict: + return { + "micro_step": self.micro_step, + "optimizer_step": self.optimizer_step, + "skipped_steps": self.train_module.skipped_steps, + "checkpoints_written": list(self.checkpoints_written), + "pinned_checkpoints": sorted(self.pinned_checkpoints), + "trunk_report": self.trunk_report, + "tripwire_alarms": list(self.tripwire_alarms), + "skip_stats": dict(self.skip_stats), + "tail_window": None if self.tail_window is None else dict(self.tail_window), + } + + def throughput_probe(self, tokens_per_step: int) -> dict: + """The probe's report, for the launch banner. + + `self.step_times` holds per-MICRO durations; `tokens_per_step` and + the scheduler horizon are per-OPTIMIZER-step. The aggregation is the + unit conversion (task #25) -- without it the banner over-reported + tokens/s and under-projected wall-clock by ~grad_accum x. + """ + return throughput_report( + step_times_s=aggregate_micro_step_times( + self.step_times[: self.cfg.train.probe_steps], self.accum + ), + tokens_per_step=tokens_per_step, + optimizer_steps_planned=self.scheduler.total_steps, + ) + + # -- checkpointing ------------------------------------------------------ + + def full_state_dict(self): + """The module's unsharded CPU gather; kept on the Trainer because + callers (eval handoff, tests) address the Trainer. Incident history + lives on `TrainModule.full_model_state_dict`.""" + return self.train_module.full_model_state_dict() + + def save_checkpoint(self) -> str: + """Write bf16 weights + resume state, atomically. + + Staged under `.partial` and renamed into place, so a host seizure + mid-write leaves nothing that looks resumable. A half-written + checkpoint is worse than none: it is found, loaded, and trusted. + """ + import torch.distributed as dist + + from ..utils.dist import is_rank0 + + # Idempotent BEFORE the collective: a re-save of an existing + # micro-step (resume that re-reaches an already-saved boundary) must + # neither crash on the rename nor delete a good checkpoint. Checked + # on every rank against the shared FS so all ranks skip together and + # nobody enters the state-gather collective alone. + name = f"micro_{self.micro_step:08d}" + # THE RANK-LOCAL PIN DEFECT (G3 audit U01-3). The flag is raised by a + # tripwire verdict, and logps_chosen / length_drift are fed from THIS + # rank's batch, so it can be set on rank 1 alone -- while pruning and + # resume.pt's pin list are rank 0's. Reduce it (MAX) so every rank pins + # or none does. save_checkpoint is entered by every rank at the same + # micro-step (cadence, coordinated stop, end-of-training), and the + # idempotence check just below is already a collective on this path. + if bool(all_reduce_max([int(self._pin_next_save)])[0]): + self._pin_next_save = True + if self._pin_next_save: + # R5: pinned on EVERY rank and on every save path (fresh, re-save, + # idempotent skip) -- prune consults this set on rank 0 only, but + # the flag must clear everywhere or a resume re-pins arbitrarily. + self.pinned_checkpoints.add(self.micro_step) + self._pin_next_save = False + logger.info("checkpoint %s PINNED (first tripwire fire precedes it)", name) + if self._resave_is_redundant(name): + logger.info("checkpoint %s already exists; skipping re-save", name) + if self.micro_step not in self.checkpoints_written: + self.checkpoints_written.append(self.micro_step) + return os.path.join(self.ckpt_root, name) if is_rank0() else "" + + if self._sharded_checkpoints(): + return self._save_checkpoint_sharded(name) + + # Gathered on EVERY rank (it is a collective); written by rank 0 only. + # The same rename race that killed the gate ledger lives here: four + # ranks renaming one `.partial` directory, first one wins, three die. + state = self.full_state_dict() + + # Optimizer state is SHARD-LOCAL per rank; a resume.pt holding only + # rank 0's moments hands ranks 1..N-1 the wrong slices on resume -- + # plausible training on corrupted state, the exact failure the + # world-size guard names, through a different door. Gather every + # rank's CPU-cloned state to rank 0 over gloo (control plane; NCCL + # would bind this to GPU state). + from ..utils.dist import gloo_group + + opt_sd = self.train_module.optim_state_dict() + from ..utils.dist import barrier + + # Every rank's RNG streams (collective, symmetric, a few KB each); only + # rank 0 writes them. Gathered BEFORE the layout dispatch so both layouts + # share one collective sequence on every rank. + rng_all = all_gather_objects(capture_rng_state()) + tw_all = all_gather_objects(self._tripwire_state()) + + if self._per_rank_optimizer_files(): + return self._save_checkpoint_per_rank(name, state, opt_sd, rng_all, tw_all) + + group = gloo_group() + if group is not None: + per_rank: list = [None] * dist.get_world_size() + dist.gather_object(opt_sd, per_rank if is_rank0() else None, dst=0, group=group) + else: + per_rank = [opt_sd] + + if not is_rank0(): + self.checkpoints_written.append(self.micro_step) + # WAIT for rank 0's write. Without this, peers ran ahead into the + # NEXT training step's FSDP all-gathers while rank 0 spent minutes + # writing 14.5GB; rank 0 could not join and the NCCL watchdog + # SIGABRT'd the run at the micro-300 save (Arm A rerun, + # 2026-08-20). The old no-barrier rationale (rank-0 crash mid-write + # strands peers) still ends the same way -- watchdog kills all -- + # so the barrier costs nothing on the failure path and fixes the + # healthy one. + barrier() + return "" + + final = os.path.join(self.ckpt_root, name) + staging = final + ".partial" + if os.path.isdir(staging): + shutil.rmtree(staging) + os.makedirs(staging, exist_ok=True) + + try: + self.weight_format.save_weights(state, staging, torch.bfloat16) + # HF-layout completeness: the eval stage loads AutoTokenizer from + # the checkpoint dir, and a weights-only dir sent it down the + # slow-tokenizer conversion path (no sentencepiece installed) at + # hour 9 of stage-11. Vocab never changes during full-FT, so the + # pinned snapshot's tokenizer files ARE this checkpoint's. + self._copy_tokenizer(staging, name) + payload = self._resume_payload(rng_all, tw_all) + payload["optimizer_per_rank"] = per_rank + _torch_save_synced(payload, os.path.join(staging, "resume.pt")) + self._write_integrity(staging) + self._publish(staging, final) + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + + # NO BARRIER AFTER THE WRITE, and my own 4-rank test is why. I had one + # here; a rank-0 failure during the write (in that run, a bad optimizer + # stub) meant rank 0 raised BEFORE reaching it while three peers sat in + # it -- a 400-second hang instead of a crash. That is the NCCL-timeout + # shape, produced by a barrier added for safety. + # + self.checkpoints_written.append(self.micro_step) + self._prune_checkpoints() + logger.info("checkpoint written: %s", final) + # Release the peers (see the barrier at the non-rank0 return above: + # both sides of the old barrier argument have now happened in + # production, and the watchdog makes the crash-path cost identical + # either way -- the barrier fixes the healthy path). + from ..utils.dist import barrier + + barrier() + return final + + def _resave_is_redundant(self, name: str) -> bool: + """COLLECTIVE. True iff a COMPLETE `name` is already on disk. + + THE TRUSTED-CORRUPTION DEFECT. "A directory called micro_N exists" was + taken as "micro_N is saved". But `latest_checkpoint` skips a micro_N + that fails `checkpoint_problem` (the host seized after the rename, the + disk filled mid-copy) and resumes from the previous good one -- and the + resumed run then walks back up to N, finds the directory, logs "already + exists" and WRITES NOTHING. The broken directory stays the newest + entry, the weights at N exist nowhere on disk, and (worse) it sits in + `checkpoints_written`, so keep-N pruning counts it as one of the two + kept and deletes the last good one when the next save lands. + + So: an existing directory is reused only if it is complete. A broken one + is moved aside (renamed, never deleted: an operator may want to look at + it; the name no longer matches `micro_*`, so neither resume nor the + supervisor sees it) and the save proceeds. + + The verdict is rank 0's alone, reduced (MAX) to every rank. Deciding + from each rank's own view of the filesystem means a rank that saw the + directory (cached attributes, a node-local disk) skips while a peer + enters the state-gather collective alone, which hangs until the + watchdog -- the lone-rank shape every other decision here avoids. + """ + from ..utils.dist import is_rank0 + + redundant = 0 + final = os.path.join(self.ckpt_root, name) + if is_rank0() and os.path.isdir(final): + problem = checkpoint_problem(final) + if problem is None: + redundant = 1 + else: + aside = os.path.join(self.ckpt_root, f"corrupt_{name}.{int(time.time())}") + logger.warning( + "checkpoint %s exists but is INCOMPLETE (%s): moved to %s and re-saving", + name, + problem, + os.path.basename(aside), + ) + os.rename(final, aside) + _fsync_path(self.ckpt_root) + return bool(all_reduce_max([redundant])[0]) + + def _copy_tokenizer(self, staging: str, name: str) -> None: + """Best-effort by design: weights are the irreplaceable part, and eval + has a snapshot fallback. Loud, never silent. WHICH tokenizer is the + weight format's decision (a dense checkpoint gets the pinned Mistral + files; an OLMoE one only an explicitly configured directory).""" + try: + copied = self.weight_format.copy_tokenizer(self.cfg.paths.model_snapshot, staging) + except (FileNotFoundError, OSError): + copied = [] + if not copied: + logger.warning( + "no tokenizer files copied for checkpoint %s (%s); it is weights-only " + "and eval must fall back to the snapshot", + name, + type(self.weight_format).__name__, + ) + + # -- per-rank optimizer-state files --------------------------------------- + + def _per_rank_optimizer_files(self) -> bool: + """Config flag `train.per_rank_optimizer_state` (default OFF). + + Default mode gathers every rank's optimizer state to rank 0 over gloo: + rank 0 holds the whole thing in host RAM at save time (12 B/param for + fp32 AdamW, 2 B/param under adamw8bit) on a shared box with no swap. + The per-rank mode has each rank write its OWN shard + (`optimizer_rank{r}.pt`) and gathers nothing. Read with getattr so a + config without the field (config.py is another lane's file) means OFF. + """ + return bool(getattr(self.cfg.train, "per_rank_optimizer_state", False)) + + @staticmethod + def _write_integrity(staging: str, opt_shas: Optional[dict] = None, include_weights: bool = True) -> None: + """Record sha256 of the payloads the weights manifest does not cover (U01-3): resume.pt, + the gathered layout's safetensors weights, and (per-rank / sharded layouts) every rank's + optimizer file, whose hashes the ranks gathered. Written LAST, before the publish.""" + sha = {} + for fn in sorted(os.listdir(staging)): + if fn == "resume.pt" or (include_weights and fn.endswith(".safetensors")): + sha[fn] = sharded_ckpt.file_sha256(os.path.join(staging, fn)) + sha.update(opt_shas or {}) + sharded_ckpt.write_integrity(staging, sha) + + def _publish(self, staging: str, final: str) -> None: + """Make `staging` durable, rename it into place, make the rename durable. + + fsync every file and directory under `staging` BEFORE the rename (see + `fsync_tree`), then fsync the checkpoint root so the new directory entry + itself survives a power loss. Order is the whole point. + """ + fsync_tree(staging) + os.rename(staging, final) + _fsync_path(os.path.dirname(final)) + + def _data_fingerprint(self) -> dict: + """What identifies the data this run trains on, as primitives (U01-5). `len(sampler)` and + `epochs` fix the horizon; a sampler may add its own `fingerprint()` (content/order hash).""" + fp: dict = {} + sampler = getattr(self, "sampler", None) + try: + fp["len"] = int(len(sampler)) + except (TypeError, AttributeError): + pass + epochs = getattr(getattr(getattr(self, "cfg", None), "train", None), "epochs", None) + if epochs is not None: + fp["epochs"] = int(epochs) + fn = getattr(sampler, "fingerprint", None) + if callable(fn): + fp["sampler"] = str(fn()) + return fp + + def _tripwire_state(self) -> dict: + """{wire name: its state_dict()} for every armed wire that can persist (U01-6).""" + out = {} + for name, wire in (getattr(self, "tripwires", None) or {}).items(): + fn = getattr(wire, "state_dict", None) + if callable(fn): + out[name] = fn() + return out + + def _resume_payload( + self, rng_per_rank: Optional[list] = None, tripwire_per_rank: Optional[list] = None + ) -> dict: + payload = { + "stream": self.stream.state_dict(), + "scheduler": self.scheduler.state_dict(), + "micro_step": self.micro_step, + "optimizer_step": self.optimizer_step, + "seed": self.cfg.train.seed, + # U01-5: what resume must still agree with. Recorded, never silently assumed. + "grad_accum": getattr(self, "accum", None), + "data_fingerprint": self._data_fingerprint(), + # U01-6: tripwire baselines / latches (rank-ordered, like the RNG streams: logps_chosen + # and length_drift are fed from each rank's own batch). + "tripwire_per_rank": ( + tripwire_per_rank if tripwire_per_rank is not None else [self._tripwire_state()] + ), + # So the logged lr is continuous across a resume (the lr the last + # step used, not the next one). + "last_applied_lr": self._last_applied_lr, + # R5 (the review that found this): the pin set lived only in process + # memory, so a resumed run forgot which checkpoints were pinned and + # the first prune after the resume deleted the one pinned + # checkpoint it had been told to keep. + # `or` tolerates the bare Trainer stubs other lanes' tests call + # `Trainer._resume_payload(stub)` with. + "pinned_checkpoints": sorted(getattr(self, "pinned_checkpoints", None) or ()), + "skip_stats": dict(getattr(self, "skip_stats", None) or {}), + # One entry per rank (rank-ordered): RNG streams are per-rank, and + # rank 0 is the only writer of resume.pt. + "rng_per_rank": ( + rng_per_rank if rng_per_rank is not None else [capture_rng_state()] + ), + } + # Generic extension slot: modules with loop-external resumable state + # (tripwire windows, a latest-weights pointer) expose extra_state(). + fn = getattr(self.train_module, "extra_state", None) + if fn is not None: + payload["extra_state"] = fn() + # Sampler-owned state. The stream's (epoch, offset) is enough for a + # sampler whose order is a pure function of (seed, epoch); a sampler + # whose position depends on what happened at run time (GRPO's prompt + # cursor advances on every dead-group retry) must carry its own. + sfn = getattr(self.sampler, "state_dict", None) + if callable(sfn): + payload["sampler_state"] = sfn() + return payload + + def _save_checkpoint_per_rank(self, name: str, state, opt_sd, rng_all=None, tw_all=None) -> str: + """Barriered, atomic per-rank save. Every rank runs this in lockstep. + + Protocol (barrier count identical on all ranks, so none can strand): + rank 0 makes a clean `.partial` dir -> B1 -> every rank writes its own + `optimizer_rank{r}.pt` via tmp + os.replace (a rank removes ONLY its + own tmp/final file on failure) -> B2 -> rank 0 writes weights and + resume.pt (which records the world size: the manifest), renames the + dir, prunes -> B3. A rank that dies before its write leaves B2 + unreached on its peers, and the rename never happens: the checkpoint + does not exist rather than existing incomplete. + """ + import torch.distributed as dist # noqa: F401 (symmetry with the default path) + + from ..utils.dist import barrier, is_rank0, rank as _rank, world_size as _world + + final = os.path.join(self.ckpt_root, name) + staging = final + ".partial" + if is_rank0(): + if os.path.isdir(staging): + shutil.rmtree(staging) + os.makedirs(staging, exist_ok=True) + barrier() + + mine = os.path.join(staging, f"optimizer_rank{_rank()}.pt") + tmp = mine + ".tmp" + try: + _torch_save_synced(opt_sd, tmp) + os.replace(tmp, mine) + my_sha = {os.path.basename(mine): sharded_ckpt.file_sha256(mine)} + except BaseException: + for f in (tmp, mine): + try: + os.remove(f) + except OSError: + pass + raise + barrier() + opt_shas: dict = {} + for d in all_gather_objects(my_sha): + opt_shas.update(d) + + if not is_rank0(): + self.checkpoints_written.append(self.micro_step) + barrier() + return "" + + try: + self.weight_format.save_weights(state, staging, torch.bfloat16) + self._copy_tokenizer(staging, name) + payload = self._resume_payload(rng_all, tw_all) + payload["optimizer_state"] = {"layout": "per_rank_files", "world_size": _world()} + _torch_save_synced(payload, os.path.join(staging, "resume.pt")) + self._write_integrity(staging, opt_shas) + self._publish(staging, final) + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + self.checkpoints_written.append(self.micro_step) + self._prune_checkpoints() + logger.info("checkpoint written (per-rank optimizer files): %s", final) + barrier() + return final + + # -- sharded checkpoints (train.checkpoint_format = "sharded") ------------- + + def _sharded_checkpoints(self) -> bool: + """Config flag `train.checkpoint_format == "sharded"` (default "gathered"). + Read with getattr so a config without the field means gathered.""" + return getattr(self.cfg.train, "checkpoint_format", "gathered") == "sharded" + + def _module_has_sharded_hooks(self) -> bool: + tm = self.train_module + return callable(getattr(tm, "sharded_model_state", None)) and callable( + getattr(tm, "load_sharded_model_state", None) + ) + + def _validate_checkpoint_format(self) -> None: + """Fail at CONSTRUCTION, not at the first save an hour in. + + THE MERGED-NAME DEFECT this prevents: a module that overrides + `full_model_state_dict` (LoRA folds the adapter into merged weights; PPO + has its own gather) defines what "the weights" mean. The base-class + sharded path reads `model.state_dict()` directly, so it would write the + UNMERGED base under the checkpoint's usual name and a resume/eval would + load a model without the adapter -- no error, a worse model. Such a + module must provide `sharded_model_state()` / `load_sharded_model_state()` + or it cannot use the sharded format. + """ + fmt = getattr(self.cfg.train, "checkpoint_format", "gathered") + if fmt not in sharded_ckpt.CHECKPOINT_FORMATS: + raise ValueError( + f"train.checkpoint_format={fmt!r}; expected one of {sharded_ckpt.CHECKPOINT_FORMATS}" + ) + if fmt != "sharded": + return + tm = self.train_module + plain = ( + type(tm).full_model_state_dict is TrainModule.full_model_state_dict + and type(tm).load_model_state_dict is TrainModule.load_model_state_dict + ) + if not plain: + # MoETrainModule overrides both methods, but without expert + # parallelism they are `super()` calls verbatim -- and that is the + # OLMoE case this format was built for (Track M). WITH ep_groups the + # experts live in EP groups, not in the FSDP shards, and the + # state-dict methods do real work: refused. + try: + from .moe_module import MoETrainModule + + plain = ( + isinstance(tm, MoETrainModule) + and getattr(tm, "ep_groups", None) is None + and type(tm).full_model_state_dict is MoETrainModule.full_model_state_dict + and type(tm).load_model_state_dict is MoETrainModule.load_model_state_dict + ) + except ImportError: + plain = False + if not (plain or self._module_has_sharded_hooks()): + raise ValueError( + f"train.checkpoint_format='sharded' with {type(tm).__name__}, which overrides the " + "weight state-dict methods (merged/adapter weights) and provides no " + "sharded_model_state()/load_sharded_model_state(): the sharded path would checkpoint " + "the unmerged base. Use checkpoint_format='gathered' or add the two hooks." + ) + + def _sharded_state(self): + fn = getattr(self.train_module, "sharded_model_state", None) + return fn() if callable(fn) and self._module_has_sharded_hooks() else self.train_module.model.state_dict() + + def _save_checkpoint_sharded(self, name: str) -> str: + """COLLECTIVE. Every rank writes its own shards; nothing is gathered. + + Protocol (every rank runs it; the collective sequence is identical): + rng all-gather -> rank 0 makes a clean `.partial` -> B1 -> each rank + writes `model_rank{r}.safetensors` + `optimizer_rank{r}.pt` (tmp + + fsync + rename) -> any-rank-failed reduction (a full disk on ONE rank + stops every rank now instead of stranding peers in the next gather) -> + all-gather of the per-rank infos -> rank 0 writes manifest, tokenizer, + resume.pt, fsync_tree + rename, prune -> any-rank-failed reduction + (doubles as the closing sync; rank-0 failure raises everywhere). + The peers' wait is the slowest rank's own 1/N write, not rank 0 writing + 14.5 GB while three ranks idle. + """ + from ..utils.dist import barrier, is_rank0, rank as _rank, world_size as _world + + final = os.path.join(self.ckpt_root, name) + staging = final + ".partial" + rng_all = all_gather_objects(capture_rng_state()) + tw_all = all_gather_objects(self._tripwire_state()) + opt_sd = self.train_module.optim_state_dict() + if is_rank0(): + if os.path.isdir(staging): + shutil.rmtree(staging) + os.makedirs(staging, exist_ok=True) + barrier() + + info, err, my_sha = None, None, {} + mine = os.path.join(staging, f"optimizer_rank{_rank()}.pt") + try: + info = sharded_ckpt.save_rank_shard( + self._sharded_state(), staging, rank=_rank(), world_size=_world() + ) + _torch_save_synced(opt_sd, mine + ".tmp") + os.replace(mine + ".tmp", mine) + my_sha = {os.path.basename(mine): sharded_ckpt.file_sha256(mine)} + except BaseException as e: # noqa: BLE001 -- stay in lockstep with the peers, then re-raise + err = e + for f in (mine + ".tmp", mine): + try: + os.remove(f) + except OSError: + pass + if any_rank(err is not None): + if is_rank0(): + shutil.rmtree(staging, ignore_errors=True) + if err is not None: + raise err + raise sharded_ckpt.ShardedCheckpointError( + f"a peer rank failed writing its shard of {name}; checkpoint abandoned" + ) + infos = all_gather_objects(info) + opt_shas: dict = {} + for d in all_gather_objects(my_sha): + opt_shas.update(d) + + err = None + if is_rank0(): + try: + fmt = sharded_ckpt.describe_weight_format(self.weight_format) + manifest = sharded_ckpt.build_manifest( + infos, + # The OPTIMIZER's slot order, not the model's: the optimizer is built over + # `p.requires_grad` params only (Trainer default, qwen3/olmoe/grpo launchers), + # so with a frozen parameter every slot index after it points one name + # earlier in the full list. `reshard` then refused (or, for same-shape + # neighbours, re-cut the wrong tensor's moments under the wrong name). + param_order=[ + n for n, p in self.train_module.model.named_parameters() if p.requires_grad + ], + export=fmt, + ) + sharded_ckpt.write_manifest(staging, manifest) + self._copy_tokenizer(staging, name) + payload = self._resume_payload(rng_all, tw_all) + payload["optimizer_state"] = {"layout": "per_rank_files", "world_size": _world()} + _torch_save_synced(payload, os.path.join(staging, "resume.pt")) + self._write_integrity(staging, opt_shas, include_weights=False) + self._publish(staging, final) + except BaseException as e: # noqa: BLE001 + err = e + shutil.rmtree(staging, ignore_errors=True) + if any_rank(err is not None): + if err is not None: + raise err + raise sharded_ckpt.ShardedCheckpointError(f"rank 0 failed publishing {name}; checkpoint abandoned") + self.checkpoints_written.append(self.micro_step) + if is_rank0(): + self._prune_checkpoints() + logger.info("checkpoint written (sharded, %d rank files): %s", _world(), final) + return final + return "" + + def _load_sharded_weights(self, path: str) -> None: + """COLLECTIVE. Each rank reads ONLY its own shard file. + + Rank-local verdicts (manifest/world size are identical everywhere, but a + sha256 mismatch or a truncated file is one rank's alone) are reduced + before any weight is copied, so a bad shard on rank 2 stops all four + ranks with the same error instead of leaving three of them in the next + collective. + """ + from ..utils.dist import rank as _rank, world_size as _world + + err = None + try: + m = sharded_ckpt.read_manifest(path) + sharded_ckpt.check_world_size(path, m, _world()) + bad = sharded_ckpt.verify_rank_file(path, _rank()) + if bad: + raise sharded_ckpt.ShardedCheckpointError(bad) + except BaseException as e: # noqa: BLE001 + err = e + if any_rank(err is not None): + if err is not None: + raise err + raise sharded_ckpt.ShardedCheckpointError( + f"a peer rank cannot load its shard of {path}; resume abandoned" + ) + if m.get("resharded_from"): + logger.warning( + "checkpoint %s was RESHARDED from world_size=%s: weights and optimizer are exact, but " + "data-order position and RNG streams follow the old partitioning", + path, + m["resharded_from"].get("world_size"), + ) + # THE LONE-READER DEFECT. The pre-checks above are deterministic across ranks; the reads + # below are not (EIO / a stale NFS handle on ONE rank's file, an OOM-killed page-in). That + # rank would raise while its peers moved on to the optimizer load and then into the first + # training collective, where they wait for the watchdog (30 minutes) instead of crashing + # with the cause. Same reduction as the save side: every rank stops together. + err = None + try: + loader = getattr(self.train_module, "load_sharded_model_state", None) + if callable(loader) and self._module_has_sharded_hooks(): + loader(path, rank=_rank(), world_size=_world()) + else: + n = sharded_ckpt.load_rank_shard( + self.train_module.model.state_dict(), path, rank=_rank(), world_size=_world() + ) + logger.info("loaded %d local shards from %s", n, m["files"][str(_rank())]["name"]) + except BaseException as e: # noqa: BLE001 -- stay in lockstep with the peers, then re-raise + err = e + if any_rank(err is not None): + if err is not None: + raise err + raise sharded_ckpt.ShardedCheckpointError( + f"a peer rank failed reading its shard of {path}; resume abandoned" + ) + + def _prune_checkpoints(self) -> None: + # Prune ONLY checkpoints THIS trainer wrote. Name-sorting the whole + # dir let a previous run's stale higher-numbered dirs shadow the live + # run's fresh saves: keep-2 kept the zombies and DELETED the new + # micro_00000100/200 (Arm A rerun, 2026-08-20). Stale dirs are a + # launch-hygiene problem, not this function's to silently absorb. + keep = self.cfg.train.keep_checkpoints + # R5: pinned checkpoints are outside the keep-N budget entirely. + prunable = set(self.checkpoints_written) - self.pinned_checkpoints + mine = sorted(f"micro_{m:08d}" for m in prunable) + for stale in mine[:-keep] if keep > 0 else mine: + shutil.rmtree(os.path.join(self.ckpt_root, stale), ignore_errors=True) + + def _older_checkpoint_steps(self, below: int) -> list[int]: + """Micro-steps of the `micro_NNNNNNNN` directories on disk strictly below `below` (no partials).""" + if not os.path.isdir(self.ckpt_root): + return [] + out = [] + for d in os.listdir(self.ckpt_root): + m = re.fullmatch(r"micro_(\d{8})", d) + if m and int(m.group(1)) < below: + out.append(int(m.group(1))) + return sorted(out) + + def _discard_partials(self) -> None: + """Delete leftover partials at startup. They masquerade as resumable.""" + if not os.path.isdir(self.ckpt_root): + return + for entry in os.listdir(self.ckpt_root): + if entry.endswith(".partial"): + shutil.rmtree(os.path.join(self.ckpt_root, entry), ignore_errors=True) + logger.warning("discarded partial checkpoint %s", entry) + + def latest_checkpoint(self) -> Optional[str]: + """The newest COMPLETE checkpoint directory, or None. + + THE FOUND-LOADED-AND-TRUSTED DEFECT. Name order alone picked the highest + `micro_N` -- including one whose rename survived a host seizure but + whose data did not (zero-length weights, a cut resume.pt, a missing + optimizer shard). Resume then died in the middle of a 13GB load, or + worse, resumed from a partial state. Each candidate passes + `checkpoint_problem` first; an incomplete one is SKIPPED WITH A WARNING + and left on disk (an operator may want to inspect or repair it; deleting + evidence is not this function's job) and the next-newest is tried. + """ + if not os.path.isdir(self.ckpt_root): + return None + dirs = sorted( + ( + d + for d in os.listdir(self.ckpt_root) + if d.startswith("micro_") and not d.endswith(".partial") + ), + reverse=True, + ) + for d in dirs: + path = os.path.join(self.ckpt_root, d) + problem = checkpoint_problem(path) + if problem is None: + return path + logger.warning( + "skipping INCOMPLETE checkpoint %s: %s (left on disk, not deleted)", d, problem + ) + return None + + def resume_from_latest(self) -> Optional[str]: + path = self.latest_checkpoint() + if path is None: + logger.info("no checkpoint to resume from; starting fresh") + return None + self.resume_from(path) + return path + + def _verify_payload_integrity(self, path: str) -> None: + """COLLECTIVE. sha256-verify the payloads this rank is about to read (U01-3). + + Rank 0 checks the shared files (resume.pt and, for the gathered layout, the weights); every + rank checks its own `optimizer_rank{r}.pt`. The verdicts are reduced before anyone raises, + so a bad file on one rank stops all ranks together instead of stranding the peers in the + next collective. A checkpoint with NO `integrity.json` predates the sidecar: it loads with + a warning. A sidecar that exists and disagrees never loads. + """ + from ..utils.dist import is_rank0, rank as _rank + + recorded = None + problem: Optional[str] = None + try: + recorded = sharded_ckpt.read_integrity(path) + if recorded is not None: + names = [f"optimizer_rank{_rank()}.pt"] + if is_rank0(): + names += [ + n + for n in recorded + if n == "resume.pt" or (n.endswith(".safetensors") and not n.startswith("model_rank")) + ] + problem = sharded_ckpt.verify_integrity(path, names, recorded) + except sharded_ckpt.CheckpointIntegrityError as e: + problem = str(e) + if recorded is None and problem is None: + logger.warning( + "checkpoint %s has no integrity.json (it predates payload hashing): resume.pt and " + "the optimizer files are NOT sha256-verified", + path, + ) + if any_rank(problem is not None): + if problem is not None: + raise sharded_ckpt.CheckpointIntegrityError( + f"checkpoint {path} failed its sha256 integrity check: {problem}" + ) + raise sharded_ckpt.CheckpointIntegrityError( + f"checkpoint {path}: a peer rank failed its sha256 integrity check" + ) + + def _check_resume_compat(self, state: dict, path: str) -> None: + """Refuse a resume whose seed / grad_accum / data fingerprint differ from the saved run (U01-5). + + A stateless sampler's order is a pure function of (seed, epoch, data), so changing any of + them resumes at a stream offset that now means a different batch: weights, optimizer and + schedule all load cleanly and nothing looks wrong. A field the checkpoint does not record + (it predates this check) is warned about, never refused. `allow_resume_mismatch` is the + explicit operator override; it logs what it overrode. + """ + diffs: list[str] = [] + missing: list[str] = [] + if "seed" in state: + if int(state["seed"]) != int(self.cfg.train.seed): + diffs.append(f"seed {state['seed']} -> {self.cfg.train.seed}") + else: + missing.append("seed") + if state.get("grad_accum") is not None: + if int(state["grad_accum"]) != int(self.accum): + diffs.append(f"grad_accum {state['grad_accum']} -> {self.accum}") + else: + missing.append("grad_accum") + saved_fp = state.get("data_fingerprint") + if isinstance(saved_fp, dict): + now_fp = self._data_fingerprint() + for k in sorted(set(saved_fp) & set(now_fp)): + if saved_fp[k] != now_fp[k]: + diffs.append(f"data fingerprint {k} {saved_fp[k]!r} -> {now_fp[k]!r}") + else: + missing.append("data fingerprint") + if missing: + logger.warning( + "checkpoint %s predates resume-compatibility recording (%s): cannot check them", + path, + ", ".join(missing), + ) + if not diffs: + return + msg = ( + f"checkpoint {path} was saved under a different configuration than this run: " + + "; ".join(diffs) + + ". The stream offset, schedule and optimizer state would be applied to a different " + "run with no error. Resume with the original seed / grad_accum / data, or pass " + "allow_resume_mismatch=True (train.allow_resume_mismatch) to override deliberately." + ) + if not self.allow_resume_mismatch: + raise RuntimeError(msg) + logger.warning("OVERRIDDEN (allow_resume_mismatch): %s", msg) + + def _restore_tripwire_state(self, state: dict, path: str) -> None: + """Restore tripwire baselines / latches (U01-6). Absent in older checkpoints: warn.""" + from ..utils.dist import rank as _rank + + if not self.tripwires: + return + per_rank = state.get("tripwire_per_rank") + if per_rank is None: + logger.warning( + "checkpoint %s predates tripwire-state persistence: tripwire baselines restart", path + ) + return + if len(per_rank) != _dist_world_size(): + logger.warning( + "checkpoint %s holds tripwire state for %d rank(s) but this run has %d: " + "tripwire baselines restart", + path, + len(per_rank), + _dist_world_size(), + ) + return + for name, sd in per_rank[_rank()].items(): + wire = self.tripwires.get(name) + loader = getattr(wire, "load_state_dict", None) + if not callable(loader): + continue + try: + loader(sd) + except Exception as e: # noqa: BLE001 -- a wire whose layout changed restarts, loudly + logger.warning("tripwire %r state in %s could not be restored (%s: %s)", name, path, type(e).__name__, e) + + def resume_from(self, path: str) -> None: + """Restore weights, then optimizer, then loop state. + + Order matters: the optimizer's fp32 master is the AUTHORITATIVE weight + copy, so it is loaded after the bf16 weights and overwrites them. That + is what makes a bf16-weights-only checkpoint a lossless resume. + """ + + self._verify_payload_integrity(path) + if sharded_ckpt.is_sharded(path): + # Self-describing, like per_rank_files: the directory's manifest, not + # the current flag, decides how weights are read. + self._load_sharded_weights(path) + else: + nmo_state = self.weight_format.load_weights(path) + # The CPU-slice per-tensor load (and its OOM incident history) lives + # on the module; the Trainer only sequences it against the rest of + # the resume payload below. + self.train_module.load_model_state_dict(nmo_state) + + # weights_only=True: a checkpoint is a file on disk, and the resume + # payload is only tensors, dicts, lists and primitives. Unpickling + # arbitrary objects to reload our own state would be a code-execution + # path we do not need. + # map_location="cpu": legacy resume.pt files carry cuda:0 tensors, + # and letting torch.load restore device puts every rank's copy on + # GPU 0 (the tail-resume OOM). Loading lands on CPU; the optimizer's + # own load_state_dict moves state to each param's device. + state = torch.load(os.path.join(path, "resume.pt"), weights_only=True, map_location="cpu") + from ..utils.dist import rank as _rank + + self._check_resume_compat(state, path) + per_rank = state.get("optimizer_per_rank") + layout = state.get("optimizer_state") + if layout is not None and layout.get("layout") == "per_rank_files": + # The checkpoint is self-describing: its manifest, not the current + # flag, decides how to read it. Shards are rank-local slices, so a + # different rank count would hand rank r a slice cut for another + # partitioning -- plausible training on wrong state. Refuse. + from ..utils.dist import world_size as _world + + saved_w = int(layout["world_size"]) + if saved_w != _world(): + raise RuntimeError( + f"checkpoint {path} holds per-rank optimizer state for " + f"world_size={saved_w}, but this run has world_size={_world()}. " + "Per-rank shards cannot be re-sharded; resume with the same " + "number of ranks." + ) + f = os.path.join(path, f"optimizer_rank{_rank()}.pt") + if not os.path.isfile(f): + raise FileNotFoundError(f"per-rank optimizer file missing: {f}") + self.train_module.load_optim_state_dict( + torch.load(f, weights_only=True, map_location="cpu") + ) + elif per_rank is not None: + # THE SILENT-SHARD DEFECT (gather layout). The per-rank-FILES + # layout records its world size and refuses a mismatch; this + # layout recorded nothing, so a 2-rank resume of a 4-rank save + # quietly handed ranks 0/1 the first two shards of a four-way + # partitioning (and a 4-rank resume of a 2-rank save died with an + # IndexError on ranks 2/3 only, leaving ranks 0/1 in the next + # collective). The list's length IS the saved world size. + if len(per_rank) != _dist_world_size(): + raise RuntimeError( + f"checkpoint {path} holds optimizer state for {len(per_rank)} rank(s), " + f"but this run has world_size={_dist_world_size()}. Per-rank optimizer " + "shards cannot be re-sharded; resume with the same number of ranks." + ) + self.train_module.load_optim_state_dict(per_rank[_rank()]) + else: + # Pre-fix checkpoint: only rank 0's shard-local moments were ever + # persisted; the other ranks' are gone. Loading rank 0's into + # every rank would be silently wrong state, so ALL ranks resume + # with FRESH moments -- consistent, honest, and recorded. Weights, + # schedule and stream position remain exact. + logger.warning( + "LEGACY resume.pt (%s) has no per-rank optimizer state; " + "resuming with FRESH optimizer moments on all ranks. This is " + "a recorded deviation, not a lossless resume.", + path, + ) + # U01-2: fresh MOMENTS are the recorded deviation; fresh MASTERS would be a silent weight + # revert (the optimizer's fp32 master is authoritative and is still the launch weights). + if self.train_module.reseed_optimizer_from_params(): + logger.warning( + "re-seeded the optimizer's fp32 master weights from the restored parameters" + ) + if "extra_state" in state: + loader = getattr(self.train_module, "load_extra_state", None) + if loader is None: + logger.warning("checkpoint carries extra_state but the module cannot load it") + else: + loader(state["extra_state"]) + self._restore_tripwire_state(state, path) + sampler_loader = getattr(self.sampler, "load_state_dict", None) + if callable(sampler_loader): + if "sampler_state" not in state: + # Strict on purpose: a stateful sampler resumed without its + # state restarts its data order from the beginning. Weights, + # optimizer and schedule would all be exact, so nothing would + # look wrong -- the run would just train on the wrong data. + raise RuntimeError( + f"checkpoint {path} carries no sampler state, but this run's " + f"sampler ({type(self.sampler).__name__}) keeps run-time " + "position. Resuming would replay its data from the start. " + "The checkpoint predates sampler-state persistence; start " + "the stage fresh or resume from a newer checkpoint." + ) + sampler_loader(state["sampler_state"]) + elif "sampler_state" in state: + raise RuntimeError( + f"checkpoint {path} carries sampler state but this run's sampler " + f"({type(self.sampler).__name__}) cannot load it: the data " + "source changed between save and resume." + ) + self._last_applied_lr = state.get("last_applied_lr") + # Absent in checkpoints that predate G0: leave the in-memory defaults. + if "pinned_checkpoints" in state: + self.pinned_checkpoints = {int(m) for m in state["pinned_checkpoints"]} + if "skip_stats" in state: + self.skip_stats = {k: int(v) for k, v in state["skip_stats"].items()} + self._nonfinite_loss_acc = None + self.scheduler.load_state_dict(state["scheduler"]) + self.stream.load_state_dict(state["stream"]) + self.micro_step = int(state["micro_step"]) + self.optimizer_step = int(state["optimizer_step"]) + # `checkpoints_written` is a PER-PROCESS list: a fresh process starts + # with it empty. Two consumers key off it, and both read an empty list + # as "nothing has ever been saved": + # * the launch scripts' `_latest_weights` (rollout engine weights): + # after a resume the first rollout was generated by the SFT + # SUBSTRATE -- the policy before every update -- while the trainer + # computed ratios against the resumed policy; + # * keep-N pruning: the resumed checkpoint was invisible to it, so it + # was never pruned and `keep_checkpoints` was exceeded by one for + # the rest of the run (13.5GB each). + # The checkpoint we just resumed from IS the newest one on disk by + # construction (resume_from_latest picks it; an explicit older path is + # the operator's call), so register it. Idempotent: a re-resume or an + # idempotent re-save must not list one micro-step twice. + if self.micro_step not in self.checkpoints_written: + self.checkpoints_written.append(self.micro_step) + # PREDECESSORS TOO (Track M, 2026-10-12: the SIGTERM-resumed OLMoE run left micro_40, _44 and _80 on + # disk = 83GB with keep_checkpoints=2, on a disk that is 95% full). The interrupted attempt wrote + # micro_40 and micro_44; the new process knew only 44, so after saving 80 the list was [44, 80] and 40 + # could never be pruned -- one leaked checkpoint per resume. Directories NUMBERED BELOW the one we + # resumed from are this lineage's older saves; adopting them is what keep-N means. Directories numbered + # ABOVE it are deliberately NOT adopted: those are the stale-run zombies `_prune_checkpoints` was + # written to ignore (a higher-numbered zombie shadowing a fresh save). + # OLDEST FIRST: consumers read `checkpoints_written[-1]` as "the newest save" (scripts/ppo.py + # `_latest_weights` hands it to the rollout engine), so the adopted predecessors go in FRONT of the + # resumed checkpoint, never after it. + adopted = [m for m in self._older_checkpoint_steps(self.micro_step) if m not in self.checkpoints_written] + self.checkpoints_written[:0] = adopted + # The trunk snapshot must be retaken: comparing post-resume parameters + # against a pre-resume snapshot would report movement that the resumed + # run did not produce. + self.trunk_probe = TrunkMovementProbe(self._probe_target if self._probe_target is not None else self.model) + self._probe_done = False + self._applied_since_snapshot = 0 + # Last: nothing above may draw from an RNG after this point, or the + # restored stream would already be advanced. + rng = state.get("rng_per_rank") + if rng is None: + logger.warning( + "checkpoint %s predates RNG persistence: RNG streams restart from the seed", path + ) + elif len(rng) != _dist_world_size(): + logger.warning( + "checkpoint %s holds RNG state for %d rank(s) but this run has %d: " + "RNG streams restart from the seed", + path, + len(rng), + _dist_world_size(), + ) + else: + restore_rng_state(rng[_rank()]) + logger.info( + "resumed from %s at micro_step %d (optimizer_step %d)", + path, + self.micro_step, + self.optimizer_step, + )