Harden Apple MPS runs against command-buffer recovery

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.

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),