diff --git a/baseline/nanogpt_one_head/MPS_RECOVERY.md b/baseline/nanogpt_one_head/MPS_RECOVERY.md new file mode 100644 index 0000000..ab34262 --- /dev/null +++ b/baseline/nanogpt_one_head/MPS_RECOVERY.md @@ -0,0 +1,96 @@ +# MPS recovery and finite-checkpoint policy + +## Failure mode + +Apple's Metal backend can occasionally report a command-buffer recovery such as: + +```text +Discarded (victim of GPU error/recovery) +kIOGPUCommandBufferCallbackErrorInnocentVictim +``` + +The message is emitted by the asynchronous MPS runtime. A Python training loop +may continue temporarily even though a queued GPU operation was discarded. If a +subsequent update consumes corrupted state, training loss and validation loss +can become `NaN`. WeightWatcher then fails later while attempting an SVD of a +non-finite matrix. The SVD exception is downstream; it is not the original +failure. + +## Repository behavior + +The ordinary `rg-onehead-train` and opt-in `rg-onehead-muonclip` launchers now +isolate every MPS optimizer/seed run in a fresh subprocess. This prevents a +sequence of long replicates from sharing one Metal command queue and one +long-lived MPS allocator state. + +For a command such as: + +```bash +rg-onehead-train \ + --config configs/reference.yaml \ + --optimizer muon \ + --seeds 1337,2027,4099 \ + --device auto +``` + +an Apple-Silicon machine runs three sequential worker processes. The scientific +protocol is unchanged: every worker receives the same config, seed, data root, +results root, batch size, evaluation probes, optimizer, and LR schedule. + +If an isolated MPS worker exits nonzero, the supervisor waits briefly for Metal +to reset and makes one fresh-process resume attempt from +`checkpoint_latest.pt`. The checkpoint includes model state, optimizer state, +training-generator state, Python/NumPy/Torch RNG state, and MPS RNG state. + +The default is: + +```text +initial worker + one fresh-process resume attempt +``` + +Change it with: + +```bash +--mps-retries 0 # no automatic restart +--mps-retries 2 # at most two restarts +``` + +For debugging only, the old same-process behavior can be requested with: + +```bash +--no-mps-isolation +``` + +## Finite-state gate + +Before replacing any training checkpoint, the code now: + +1. synchronizes the accelerator; +2. copies the complete model and optimizer state to CPU; +3. verifies that every floating-point or complex tensor is finite; +4. writes to a temporary file; +5. atomically replaces the target checkpoint only after validation succeeds. + +A contaminated update therefore cannot overwrite the last verified +`checkpoint_latest.pt`. Loading also applies the same finite-state validation, +so a legacy checkpoint containing `NaN` or `Inf` is rejected explicitly. + +The training loop already checks finite train/validation metrics and model +parameters before WeightWatcher. The expected failure is now a direct +`FloatingPointError`, not the misleading downstream NumPy SVD error. + +## What is not automatic + +The supervisor does not silently switch from MPS to CPU or TPU. A device change +would alter numerical execution and should be an explicit experimental choice. +If the same step fails again after a fresh-process resume, stop the MPS run and +restart that seed from scratch on CPU or TPU, or investigate the local macOS and +PyTorch MPS versions. + +## Existing runs + +An already-running Python process is not changed by pulling this commit. To use +MPS worker isolation and the finite-checkpoint gate, stop the old launcher, +pull and reinstall the package, and start or resume with the normal command. +Do not resume from a checkpoint that the updated loader identifies as +contaminated. diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/checkpoints.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/checkpoints.py index 956d996..1349ca8 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/checkpoints.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/checkpoints.py @@ -21,6 +21,56 @@ ) +def _nonfinite_tensor_paths(value: Any, path: str) -> list[str]: + bad: list[str] = [] + if torch.is_tensor(value): + if value.is_floating_point() or value.is_complex(): + if not bool(torch.isfinite(value).all()): + bad.append(path) + return bad + if isinstance(value, dict): + for key, item in value.items(): + bad.extend( + _nonfinite_tensor_paths( + item, + f"{path}.{key}", + ) + ) + return bad + if isinstance(value, (list, tuple)): + for index, item in enumerate(value): + bad.extend( + _nonfinite_tensor_paths( + item, + f"{path}[{index}]", + ) + ) + return bad + + +def _require_finite_checkpoint_state( + *, + model_state: dict[str, Any], + optimizer_states: list[dict[str, Any]] | None, + step: int, +) -> None: + bad = _nonfinite_tensor_paths(model_state, "model") + if optimizer_states is not None: + bad.extend( + _nonfinite_tensor_paths( + optimizer_states, + "optimizers", + ) + ) + if bad: + preview = ", ".join(bad[:12]) + suffix = "" if len(bad) <= 12 else f" (+{len(bad) - 12} more)" + raise FloatingPointError( + "refusing to write or load a contaminated checkpoint at " + f"step={int(step)}; non-finite tensors: {preview}{suffix}" + ) + + def _atomic_torch_save(payload: dict[str, Any], path: Path) -> Path: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(path.suffix + ".tmp") @@ -46,12 +96,25 @@ def save_training_checkpoint( ) -> Path: device = model_device(model) synchronize(device) + + # Materialize model and optimizer state on CPU before touching the target + # path. A Metal command-buffer recovery can otherwise leave finite-looking + # Python control flow around corrupted accelerator tensors. The finite-state + # gate ensures checkpoint_latest.pt always remains the last verified state. + model_state = tree_to_cpu(model.state_dict()) + optimizer_states = tree_to_cpu(optimizer_state_dict(handles)) + _require_finite_checkpoint_state( + model_state=model_state, + optimizer_states=optimizer_states, + step=step, + ) + payload: dict[str, Any] = { - "schema_version": 3, - # Always serialize CPU tensors so checkpoints are portable between - # MPS, CUDA, TPU/XLA, and CPU environments. - "model": tree_to_cpu(model.state_dict()), - "optimizers": tree_to_cpu(optimizer_state_dict(handles)), + "schema_version": 4, + # CPU tensors keep checkpoints portable between MPS, CUDA, TPU/XLA, + # and CPU environments. + "model": model_state, + "optimizers": optimizer_states, "step": int(step), "best_validation_loss": float(best_validation_loss), "best_validation_step": int(best_validation_step), @@ -83,6 +146,11 @@ def load_training_checkpoint( raise RuntimeError( "checkpoint protocol fingerprint does not match the requested run" ) + _require_finite_checkpoint_state( + model_state=payload["model"], + optimizer_states=payload["optimizers"], + step=int(payload.get("step", -1)), + ) model.load_state_dict(payload["model"]) load_optimizer_state_dict(handles, payload["optimizers"]) random.setstate(payload["python_random_state"]) @@ -118,9 +186,15 @@ def save_epoch_model_checkpoint( ) device = model_device(model) synchronize(device) + model_state = tree_to_cpu(model.state_dict()) + _require_finite_checkpoint_state( + model_state=model_state, + optimizer_states=None, + step=step, + ) payload = { - "schema_version": 2, - "model": tree_to_cpu(model.state_dict()), + "schema_version": 3, + "model": model_state, "step": int(step), "nominal_epoch": float(nominal_epoch), "actual_epoch": float(actual_epoch), diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/training.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/training.py index ed7bf45..7e6421c 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/training.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/training.py @@ -2,12 +2,22 @@ import argparse from copy import deepcopy +import gc +import os from pathlib import Path +import subprocess +import sys +import time from typing import Sequence import torch -from .config import SUPPORTED_OPTIMIZERS, canonical_seeds, load_config, roots +from .config import ( + SUPPORTED_OPTIMIZERS, + canonical_seeds, + load_config, + roots, +) from .data import prepare_fineweb_edu from .engine import run_one from .run_utils import run_directory, run_is_complete @@ -38,6 +48,50 @@ def _resolve_roots( ) +def _release_accelerator(device: torch.device) -> None: + """Release cached accelerator state after a replicate exits. + + A fresh subprocess is the primary MPS isolation boundary. This cleanup also + protects programmatic callers that execute several replicates in one Python + process. + """ + + if device.type == "cuda": + try: + torch.cuda.synchronize(device) + finally: + gc.collect() + torch.cuda.empty_cache() + return + + if device.type != "mps" or not hasattr(torch, "mps"): + gc.collect() + return + + try: + torch.mps.synchronize() + except Exception as exc: + # A Metal GPU recovery can make synchronization itself fail. Cleanup is + # still attempted, and the original training exception is allowed to + # propagate from the caller. + print( + "[one-head-mps] warning: synchronize during cleanup failed: " + f"{type(exc).__name__}: {exc}", + flush=True, + ) + finally: + gc.collect() + if hasattr(torch.mps, "empty_cache"): + try: + torch.mps.empty_cache() + except Exception as exc: + print( + "[one-head-mps] warning: empty_cache failed: " + f"{type(exc).__name__}: {exc}", + flush=True, + ) + + def run_optimizer_replicates( *, cfg: dict, @@ -63,21 +117,24 @@ def run_optimizer_replicates( ) if prepare_data: prepare_fineweb_edu(cfg, data_path) - run_dirs = [] + run_dirs: list[Path] = [] for seed in selected_seeds: - run_dirs.append( - run_one( - cfg=deepcopy(cfg), - data_root=data_path, - results_root=results_path, - optimizer_name=optimizer_name, - seed=seed, - device=resolved_device, - resume=resume, - overwrite=overwrite, - progress=progress, + try: + run_dirs.append( + run_one( + cfg=deepcopy(cfg), + data_root=data_path, + results_root=results_path, + optimizer_name=optimizer_name, + seed=seed, + device=resolved_device, + resume=resume, + overwrite=overwrite, + progress=progress, + ) ) - ) + finally: + _release_accelerator(resolved_device) return run_dirs @@ -119,6 +176,133 @@ def run_all_replicates( return outputs +def _mps_worker_module(optimizer_name: str) -> str: + # The opt-in MuonClip launcher installs its extension before delegating to + # training.main. A child process must repeat that installation. + return ( + "rg_nanogpt_one_head.muonclip" + if optimizer_name == "muon_clip" + else "rg_nanogpt_one_head.training" + ) + + +def _mps_worker_command( + *, + args: argparse.Namespace, + optimizer_name: str, + seed: int, + data_root: Path, + results_root: Path, + first_attempt: bool, +) -> list[str]: + command = [ + sys.executable, + "-u", + "-m", + _mps_worker_module(optimizer_name), + "--config", + str(Path(args.config).resolve()), + "--optimizer", + optimizer_name, + "--seeds", + str(int(seed)), + "--data-root", + str(data_root.resolve()), + "--results-root", + str(results_root.resolve()), + "--device", + "mps", + "--mps-worker", + "--mps-retries", + "0", + ] + if first_attempt and bool(args.overwrite): + command.append("--overwrite") + if first_attempt and bool(args.no_resume): + command.append("--no-resume") + return command + + +def _run_isolated_mps_workers( + *, + args: argparse.Namespace, + cfg: dict, + seeds: Sequence[int], +) -> None: + data_root, results_root, resolved_device = _resolve_roots( + data_root=args.data_root, + results_root=args.results_root, + device="mps", + ) + if resolved_device.type != "mps": + raise RuntimeError("internal MPS worker supervisor selected a non-MPS device") + + prepare_fineweb_edu(cfg, data_root) + optimizers = ( + tuple(SUPPORTED_OPTIMIZERS) + if args.optimizer == "all" + else (str(args.optimizer),) + ) + max_attempts = 1 + int(args.mps_retries) + environment = os.environ.copy() + environment.setdefault("PYTHONUNBUFFERED", "1") + environment.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") + + for optimizer_name in optimizers: + for seed in seeds: + run_dir = run_directory(results_root, optimizer_name, int(seed)) + for attempt in range(1, max_attempts + 1): + first_attempt = attempt == 1 + command = _mps_worker_command( + args=args, + optimizer_name=optimizer_name, + seed=int(seed), + data_root=data_root, + results_root=results_root, + first_attempt=first_attempt, + ) + print( + "[one-head-mps] starting isolated worker " + f"optimizer={optimizer_name} seed={seed} " + f"attempt={attempt}/{max_attempts}", + flush=True, + ) + result = subprocess.run( + command, + env=environment, + check=False, + ) + if result.returncode == 0: + print( + "[one-head-mps] worker complete " + f"optimizer={optimizer_name} seed={seed}", + flush=True, + ) + break + + latest = run_dir / "checkpoint_latest.pt" + if attempt >= max_attempts or not latest.is_file(): + checkpoint_note = ( + f"last verified checkpoint: {latest}" + if latest.is_file() + else "no verified checkpoint was written" + ) + raise RuntimeError( + "isolated MPS worker failed with exit code " + f"{result.returncode} for optimizer={optimizer_name} " + f"seed={seed}; {checkpoint_note}" + ) + + print( + "[one-head-mps] worker failed; allowing Metal to reset, " + "then resuming from the last finite atomic checkpoint: " + f"{latest}", + flush=True, + ) + _release_accelerator(resolved_device) + time.sleep(2.0) + + def main() -> None: parser = argparse.ArgumentParser( description="Run the one-head FineWeb-Edu optimizer baselines" @@ -142,7 +326,31 @@ def main() -> None: ) parser.add_argument("--overwrite", action="store_true") parser.add_argument("--no-resume", action="store_true") + parser.add_argument( + "--mps-retries", + type=int, + default=1, + help=( + "fresh-process resume attempts after an MPS worker failure; " + "default: 1" + ), + ) + parser.add_argument( + "--no-mps-isolation", + action="store_true", + help="run MPS replicates in the current process (debugging only)", + ) + parser.add_argument( + "--mps-worker", + action="store_true", + help=argparse.SUPPRESS, + ) args = parser.parse_args() + if args.mps_retries < 0: + parser.error("--mps-retries must be nonnegative") + if args.overwrite and args.no_resume: + parser.error("--overwrite and --no-resume are mutually exclusive") + cfg = load_config(args.config) seeds = ( tuple( @@ -153,13 +361,29 @@ def main() -> None: if args.seeds else canonical_seeds(cfg) ) + if not seeds: + parser.error("at least one seed is required") + + resolved_device = choose_device(args.device) + if ( + resolved_device.type == "mps" + and not args.mps_worker + and not args.no_mps_isolation + ): + _run_isolated_mps_workers( + args=args, + cfg=cfg, + seeds=seeds, + ) + return + common = dict( cfg=cfg, config_path=args.config, seeds=seeds, data_root=args.data_root, results_root=args.results_root, - device=args.device, + device=resolved_device, resume=not args.no_resume, overwrite=args.overwrite, ) diff --git a/baseline/nanogpt_one_head/tests/test_mps_recovery.py b/baseline/nanogpt_one_head/tests/test_mps_recovery.py new file mode 100644 index 0000000..c9d367e --- /dev/null +++ b/baseline/nanogpt_one_head/tests/test_mps_recovery.py @@ -0,0 +1,163 @@ +from __future__ import annotations + +import argparse +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from rg_nanogpt_one_head.checkpoints import save_training_checkpoint +from rg_nanogpt_one_head.optimizers import OptimizerHandle +import rg_nanogpt_one_head.training as training + + +def _worker_args(tmp_path: Path, **overrides) -> argparse.Namespace: + config = tmp_path / "config.yaml" + config.write_text("protocol: test\n", encoding="utf-8") + values = { + "config": str(config), + "optimizer": "muon", + "data_root": str(tmp_path / "data"), + "results_root": str(tmp_path / "results"), + "overwrite": False, + "no_resume": False, + "mps_retries": 1, + } + values.update(overrides) + return argparse.Namespace(**values) + + +def test_checkpoint_refuses_nonfinite_optimizer_state(tmp_path) -> None: + model = torch.nn.Linear(3, 2) + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3) + handle = OptimizerHandle( + role="primary", + optimizer=optimizer, + peak_lr=1e-3, + min_lr=1e-4, + ) + + x = torch.ones(2, 3) + model(x).sum().backward() + optimizer.step() + for state in optimizer.state.values(): + state["exp_avg"].fill_(float("nan")) + break + + path = tmp_path / "checkpoint_latest.pt" + with pytest.raises( + FloatingPointError, + match="refusing to write or load a contaminated checkpoint", + ): + save_training_checkpoint( + path, + model=model, + handles=[handle], + step=25, + best_validation_loss=1.0, + best_validation_step=20, + elapsed_seconds=2.0, + fingerprint="unit", + cfg={"protocol": {"name": "unit"}}, + optimizer_name="adamw", + seed=7, + train_generator=torch.Generator().manual_seed(11), + ) + + assert not path.exists() + assert not path.with_suffix(".pt.tmp").exists() + + +def test_mps_worker_command_uses_the_correct_extension_module(tmp_path) -> None: + args = _worker_args( + tmp_path, + overwrite=True, + no_resume=True, + ) + + ordinary = training._mps_worker_command( + args=args, + optimizer_name="muon", + seed=1337, + data_root=tmp_path / "data", + results_root=tmp_path / "results", + first_attempt=True, + ) + retry = training._mps_worker_command( + args=args, + optimizer_name="muon", + seed=1337, + data_root=tmp_path / "data", + results_root=tmp_path / "results", + first_attempt=False, + ) + muonclip = training._mps_worker_command( + args=args, + optimizer_name="muon_clip", + seed=1337, + data_root=tmp_path / "data", + results_root=tmp_path / "results", + first_attempt=True, + ) + + assert "rg_nanogpt_one_head.training" in ordinary + assert "rg_nanogpt_one_head.muonclip" in muonclip + assert "--mps-worker" in ordinary + assert "--overwrite" in ordinary + assert "--no-resume" in ordinary + assert "--overwrite" not in retry + assert "--no-resume" not in retry + + +def test_mps_supervisor_retries_from_latest_checkpoint( + tmp_path, + monkeypatch, +) -> None: + args = _worker_args(tmp_path, mps_retries=1) + data_root = tmp_path / "data" + results_root = tmp_path / "results" + calls: list[list[str]] = [] + + monkeypatch.setattr( + training, + "_resolve_roots", + lambda **kwargs: ( + data_root, + results_root, + torch.device("mps"), + ), + ) + monkeypatch.setattr( + training, + "prepare_fineweb_edu", + lambda cfg, path: None, + ) + monkeypatch.setattr( + training, + "_release_accelerator", + lambda device: None, + ) + monkeypatch.setattr(training.time, "sleep", lambda seconds: None) + + def fake_run(command, env, check): + del env, check + calls.append(list(command)) + if len(calls) == 1: + run_dir = results_root / "muon" / "seed_1337" + run_dir.mkdir(parents=True) + (run_dir / "checkpoint_latest.pt").write_bytes(b"verified") + return SimpleNamespace(returncode=1) + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(training.subprocess, "run", fake_run) + + training._run_isolated_mps_workers( + args=args, + cfg={"training": {"seeds": [1337]}}, + seeds=(1337,), + ) + + assert len(calls) == 2 + assert "--mps-worker" in calls[0] + assert "--mps-worker" in calls[1]