diff --git a/baseline/nanogpt_one_head/CONTINUE_ON_ERROR.md b/baseline/nanogpt_one_head/CONTINUE_ON_ERROR.md new file mode 100644 index 0000000..da672a1 --- /dev/null +++ b/baseline/nanogpt_one_head/CONTINUE_ON_ERROR.md @@ -0,0 +1,39 @@ +# Continue after a failed MPS replicate + +The Apple-MPS supervisor runs every optimizer/seed replicate in a fresh worker +process. A failed worker is retried from the last verified finite checkpoint. +If all configured attempts fail, the supervisor records the incomplete run and +continues to the next optimizer/seed instead of terminating the entire batch. + +This behavior is deliberately explicit rather than silent: + +- the incomplete run receives `run_failed.json`; +- the batch receives `_batch_status.json` in the results root; +- the terminal prints a final completed/failed summary; +- the failed run never receives `run_complete.json` and is therefore excluded + from normal completed-run analysis. + +Use `--fail-fast` when debugging to restore immediate termination on the first +replicate that exhausts its retries. + +Example: + +```bash +rg-onehead-train \ + --config configs/reference.yaml \ + --optimizer muon \ + --seeds 1337,2027,4099,5003,6007,7013,8017,9011 \ + --device auto \ + --mps-retries 1 +``` + +If one seed fails twice, the remaining seeds still run. Inspect the latest batch +record with: + +```bash +cat /tmp/rg-nanogpt-one-head/results/_batch_status.json +``` + +A failed seed may be resumed later by running that seed alone without +`--no-resume`; the runner will use its last verified finite +`checkpoint_latest.pt` when available. 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 7e6421c..85a78ae 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,13 +2,16 @@ import argparse from copy import deepcopy +from datetime import datetime, timezone import gc +import json import os from pathlib import Path +import shutil import subprocess import sys import time -from typing import Sequence +from typing import Any, Sequence import torch @@ -223,19 +226,113 @@ def _mps_worker_command( return command +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _atomic_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text( + json.dumps(payload, indent=2, sort_keys=True, default=str), + encoding="utf-8", + ) + temporary.replace(path) + + +def _clear_failure_marker(run_dir: Path) -> None: + (run_dir / "run_failed.json").unlink(missing_ok=True) + + +def _record_failed_replicate( + *, + run_dir: Path, + optimizer_name: str, + seed: int, + attempts: int, + exit_code: int, + latest_checkpoint: Path, +) -> dict[str, Any]: + record: dict[str, Any] = { + "failed": True, + "completed": False, + "optimizer": str(optimizer_name), + "seed": int(seed), + "attempts": int(attempts), + "last_exit_code": int(exit_code), + "last_verified_checkpoint": ( + str(latest_checkpoint) if latest_checkpoint.is_file() else None + ), + "failure_policy": "skip_replicate_and_continue_batch", + "recorded_at_utc": _utc_now(), + } + _atomic_json(run_dir / "run_failed.json", record) + return record + + +def _prepare_retry_without_checkpoint(run_dir: Path) -> None: + """Restart cleanly when a failed worker never wrote a finite checkpoint.""" + + if not run_dir.exists(): + return + archive = run_dir.with_name( + run_dir.name + + ".failed-no-checkpoint." + + datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + ) + if archive.exists(): + shutil.rmtree(archive) + run_dir.replace(archive) + print( + "[one-head-mps] no verified checkpoint exists; archived the partial " + f"run at {archive} and will restart from initialization", + flush=True, + ) + + +def _write_batch_status( + *, + results_root: Path, + started_at_utc: str, + optimizers: Sequence[str], + seeds: Sequence[int], + completed: list[dict[str, Any]], + failed: list[dict[str, Any]], +) -> Path: + status_path = results_root / "_batch_status.json" + payload = { + "schema_version": 1, + "started_at_utc": started_at_utc, + "finished_at_utc": _utc_now(), + "accelerator": "mps", + "requested_optimizers": [str(value) for value in optimizers], + "requested_seeds": [int(value) for value in seeds], + "requested_replicates": int(len(optimizers) * len(seeds)), + "completed_replicates": int(len(completed)), + "failed_replicates": int(len(failed)), + "all_completed": not failed, + "completed": completed, + "failed": failed, + } + _atomic_json(status_path, payload) + return status_path + + def _run_isolated_mps_workers( *, args: argparse.Namespace, cfg: dict, seeds: Sequence[int], -) -> None: +) -> dict[str, Any]: 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") + raise RuntimeError( + "internal MPS worker supervisor selected a non-MPS device" + ) prepare_fineweb_edu(cfg, data_root) optimizers = ( @@ -244,13 +341,19 @@ def _run_isolated_mps_workers( else (str(args.optimizer),) ) max_attempts = 1 + int(args.mps_retries) + fail_fast = bool(getattr(args, "fail_fast", False)) environment = os.environ.copy() environment.setdefault("PYTHONUNBUFFERED", "1") environment.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") + started_at_utc = _utc_now() + completed: list[dict[str, Any]] = [] + failed: list[dict[str, Any]] = [] for optimizer_name in optimizers: for seed in seeds: run_dir = run_directory(results_root, optimizer_name, int(seed)) + succeeded = False + last_return_code = 0 for attempt in range(1, max_attempts + 1): first_attempt = attempt == 1 command = _mps_worker_command( @@ -272,35 +375,108 @@ def _run_isolated_mps_workers( env=environment, check=False, ) + last_return_code = int(result.returncode) if result.returncode == 0: + _clear_failure_marker(run_dir) + completed.append( + { + "optimizer": str(optimizer_name), + "seed": int(seed), + "attempts": int(attempt), + "run_dir": str(run_dir), + } + ) print( "[one-head-mps] worker complete " f"optimizer={optimizer_name} seed={seed}", flush=True, ) + succeeded = 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" + if attempt < max_attempts: + _release_accelerator(resolved_device) + time.sleep(2.0) + if latest.is_file(): + print( + "[one-head-mps] worker failed; Metal will restart in " + "a fresh process and resume from the last finite " + f"atomic checkpoint: {latest}", + flush=True, + ) + else: + _prepare_retry_without_checkpoint(run_dir) + continue + + failure = _record_failed_replicate( + run_dir=run_dir, + optimizer_name=optimizer_name, + seed=int(seed), + attempts=max_attempts, + exit_code=last_return_code, + latest_checkpoint=latest, + ) + failed.append(failure) + checkpoint_note = ( + f"last verified checkpoint: {latest}" + if latest.is_file() + else "no verified checkpoint was written" + ) + message = ( + "isolated MPS worker exhausted its retries with exit code " + f"{last_return_code} for optimizer={optimizer_name} " + f"seed={seed}; {checkpoint_note}" + ) + if fail_fast: + status_path = _write_batch_status( + results_root=results_root, + started_at_utc=started_at_utc, + optimizers=optimizers, + seeds=seeds, + completed=completed, + failed=failed, ) raise RuntimeError( - "isolated MPS worker failed with exit code " - f"{result.returncode} for optimizer={optimizer_name} " - f"seed={seed}; {checkpoint_note}" + message + f"; batch status: {status_path}" ) print( - "[one-head-mps] worker failed; allowing Metal to reset, " - "then resuming from the last finite atomic checkpoint: " - f"{latest}", + "[one-head-mps] SKIPPING failed replicate and continuing " + f"to the next optimizer/seed: {message}", flush=True, ) + break + + if not succeeded: _release_accelerator(resolved_device) - time.sleep(2.0) + + status_path = _write_batch_status( + results_root=results_root, + started_at_utc=started_at_utc, + optimizers=optimizers, + seeds=seeds, + completed=completed, + failed=failed, + ) + print( + "[one-head-mps] batch finished: " + f"completed={len(completed)} failed={len(failed)} " + f"status={status_path}", + flush=True, + ) + if failed: + print( + "[one-head-mps] PARTIAL SUCCESS: failed replicates remain " + "incomplete and are excluded from completed-run analysis; inspect " + f"{status_path}", + flush=True, + ) + return { + "completed": completed, + "failed": failed, + "status_path": status_path, + } def main() -> None: @@ -335,6 +511,14 @@ def main() -> None: "default: 1" ), ) + parser.add_argument( + "--fail-fast", + action="store_true", + help=( + "stop the MPS batch after a replicate exhausts its retries; by " + "default the failed replicate is recorded and the next seed runs" + ), + ) parser.add_argument( "--no-mps-isolation", action="store_true", diff --git a/baseline/nanogpt_one_head/tests/test_mps_recovery.py b/baseline/nanogpt_one_head/tests/test_mps_recovery.py index c9d367e..19319c3 100644 --- a/baseline/nanogpt_one_head/tests/test_mps_recovery.py +++ b/baseline/nanogpt_one_head/tests/test_mps_recovery.py @@ -1,6 +1,7 @@ from __future__ import annotations import argparse +import json from pathlib import Path from types import SimpleNamespace @@ -23,11 +24,41 @@ def _worker_args(tmp_path: Path, **overrides) -> argparse.Namespace: "overwrite": False, "no_resume": False, "mps_retries": 1, + "fail_fast": False, } values.update(overrides) return argparse.Namespace(**values) +def _mock_mps_environment( + tmp_path: Path, + monkeypatch, +) -> tuple[Path, Path]: + data_root = tmp_path / "data" + results_root = tmp_path / "results" + 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) + return data_root, results_root + + 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) @@ -115,31 +146,9 @@ def test_mps_supervisor_retries_from_latest_checkpoint( monkeypatch, ) -> None: args = _worker_args(tmp_path, mps_retries=1) - data_root = tmp_path / "data" - results_root = tmp_path / "results" + _, results_root = _mock_mps_environment(tmp_path, monkeypatch) 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)) @@ -152,7 +161,7 @@ def fake_run(command, env, check): monkeypatch.setattr(training.subprocess, "run", fake_run) - training._run_isolated_mps_workers( + result = training._run_isolated_mps_workers( args=args, cfg={"training": {"seeds": [1337]}}, seeds=(1337,), @@ -161,3 +170,84 @@ def fake_run(command, env, check): assert len(calls) == 2 assert "--mps-worker" in calls[0] assert "--mps-worker" in calls[1] + assert len(result["completed"]) == 1 + assert result["failed"] == [] + assert result["status_path"].is_file() + + +def test_mps_supervisor_skips_failed_seed_and_continues( + tmp_path, + monkeypatch, +) -> None: + args = _worker_args( + tmp_path, + mps_retries=0, + fail_fast=False, + ) + _, results_root = _mock_mps_environment(tmp_path, monkeypatch) + calls: list[list[str]] = [] + + def fake_run(command, env, check): + del env, check + calls.append(list(command)) + seed = command[command.index("--seeds") + 1] + return SimpleNamespace(returncode=1 if seed == "1337" else 0) + + monkeypatch.setattr(training.subprocess, "run", fake_run) + + result = training._run_isolated_mps_workers( + args=args, + cfg={"training": {"seeds": [1337, 2027]}}, + seeds=(1337, 2027), + ) + + assert len(calls) == 2 + assert len(result["completed"]) == 1 + assert result["completed"][0]["seed"] == 2027 + assert len(result["failed"]) == 1 + assert result["failed"][0]["seed"] == 1337 + + failure_path = ( + results_root / "muon" / "seed_1337" / "run_failed.json" + ) + failure = json.loads(failure_path.read_text(encoding="utf-8")) + assert failure["failure_policy"] == "skip_replicate_and_continue_batch" + assert failure["last_exit_code"] == 1 + + batch = json.loads( + (results_root / "_batch_status.json").read_text(encoding="utf-8") + ) + assert batch["requested_replicates"] == 2 + assert batch["completed_replicates"] == 1 + assert batch["failed_replicates"] == 1 + assert batch["all_completed"] is False + + +def test_mps_supervisor_fail_fast_remains_available( + tmp_path, + monkeypatch, +) -> None: + args = _worker_args( + tmp_path, + mps_retries=0, + fail_fast=True, + ) + _, results_root = _mock_mps_environment(tmp_path, monkeypatch) + calls: list[list[str]] = [] + + def fake_run(command, env, check): + del env, check + calls.append(list(command)) + return SimpleNamespace(returncode=9) + + monkeypatch.setattr(training.subprocess, "run", fake_run) + + with pytest.raises(RuntimeError, match="exhausted its retries"): + training._run_isolated_mps_workers( + args=args, + cfg={"training": {"seeds": [1337, 2027]}}, + seeds=(1337, 2027), + ) + + assert len(calls) == 1 + assert (results_root / "_batch_status.json").is_file()