Harden Apple MPS runs against command-buffer recovery

[charlesmartin14]

Root cause

The attached failure is not a primary WeightWatcher/SVD problem. Metal reported asynchronous command-buffer recovery twice, the second recovery was followed by non-finite train/validation metrics, and the later NumPy SVD failure was only the first CPU operation that rejected the contaminated matrices.

Changes

• isolate every Apple-MPS optimizer/seed replicate in a fresh worker process;
• keep ordinary SGD/AdamW/Muon and the opt-in MuonClip extension compatible with that worker boundary;
• make one fresh-process resume attempt by default after a nonzero MPS worker exit;
• resume only from checkpoint_latest.pt and only after the updated loader verifies that the checkpoint is finite;
• copy complete model and optimizer state to CPU before checkpoint replacement;
• reject any checkpoint containing non-finite floating-point or complex tensors;
• keep checkpoint writes atomic so a failed save cannot replace the prior verified checkpoint;
• release MPS/CUDA caches between programmatic replicates;
• provide --mps-retries and a debugging-only --no-mps-isolation override;
• document the failure mode and recovery policy.

Scientific behavior

The patch does not change the model, optimizer equations, learning-rate schedules, batch size, data sampling, evaluation probes, WeightWatcher settings, or result-directory layout. A retry restores model, optimizer, training-generator, Python/NumPy/Torch, and MPS RNG state from the last finite atomic checkpoint. It does not silently switch devices.

Expected behavior

On MPS, a multi-seed command now launches one sequential worker process per optimizer/seed. If Metal recovery corrupts a worker, current finite-metric checks terminate it before WeightWatcher. The supervisor waits briefly and resumes once from the last verified checkpoint in a new Metal process. A repeated failure remains visible and stops the experiment.

Tests

New tests cover:

• refusal to overwrite a checkpoint when optimizer state contains NaN;
• ordinary versus MuonClip worker-module selection;
• stripping --overwrite/--no-resume on a recovery attempt;
• fresh-process retry from an existing latest checkpoint.