Subsystem 03

OOM Auto-Recovery & Self-Healing Engine

Zero-loss training recovery that intercepts CUDA out-of-memory errors and loss explosions in-flight without aborting the run.

Why Unattended Local Training Fails

Machine learning fine-tuning jobs on consumer GPUs are inherently fragile. Even when pre-flight estimates indicate safe headroom, two runtime anomalies routinely terminate processes:

Burst Sequence Outliers (CUDA OOM): A mini-batch containing conversations with outlier token lengths (e.g. 3,800 tokens instead of the mean 1,200 tokens) causes quadratic attention memory consumption ($O(N^2)$), instantly crashing the PyTorch allocator.
Gradient Explosion / NaN Divergence: Mixed-precision (FP16/BF16) training on small datasets occasionally encounters extreme gradient norms, causing weights to turn into NaNs or training loss to spike by >300%.

The 5-Level Autonomous Escalation Ladder

MoroAI replaces the standard unhandled exception loop with a progressive Autonomous Healer Protocol. When an anomaly occurs, the engine intercepts the signal and executes the lowest-impact resolution:

Level 1: Allocator Cache Evacuation Latency: ~400ms

Invokes torch.cuda.empty_cache() and forces Python cyclic garbage collection. This defragments the CUDA memory allocator, reclaiming up to 1.5GB of stale buffer blocks without altering training state.

Level 2: Dynamic Micro-Batch Halving Mathematical Gradient Equivalence

Divides the per-device micro-batch size by 2 while doubling gradient accumulation steps (e.g., micro-batch 4 & accum 4 → micro-batch 2 & accum 8). Effective batch size (B_eff = b_micro * N_accum) is mathematically preserved.

Level 3: Full Gradient Checkpointing Enforcement Memory Reduction: ~60%

Dynamically enables PyTorch activation recomputation across all attention heads and MLP projection layers. Rather than storing intermediate forward activations, they are recomputed during backward pass, reducing activation VRAM by 60% with a modest 15-20% compute overhead.

Level 4: Sequence Length Batch Clamping Targeted Batch Recovery

If an isolated mini-batch triggers an OOM due to extreme sequence length, the healer isolates the spike batch, clamps token lengths to the 95th percentile, backpropagates safely, and logs the truncation event to the lineage ledger.

Level 5: Checkpoint Rollback on Loss Divergence Loss Spike Protection

A moving exponential loss tracker monitors gradient stability. If step loss exceeds 3x moving average loss (EMA) or evaluates to NaN, training automatically rolls back to the last certified checkpoint, lowers the learning rate by 20%, and resumes.

CLI Invocation

Launch training with auto-heal enabled:

# Launch unattended training with the 5-step healer active
moro train --config ./moro.yaml --auto-heal

Live Recovery Log Stream

Step 420/1200 | Loss: 1.12 | VRAM: 7.9GB
[WARN] CUDA OOM at step 421 (token length 3,840)
[HEAL] Level 1: Flushed CUDA memory pool (Freed 820MB)
[HEAL] Level 2: Micro-batch 2 -> 1, Accum 8 -> 16
[RESUME] Checkpoint 400 reloaded
Step 421/1200 | Loss: 1.11 | VRAM: 6.1GB (STABLE)
Step 422/1200 | Loss: 1.09 | VRAM: 6.2GB
[SUCCESS] Full training finished with 0 data loss

← VRAM Recipe Predictor Multi-Layer Eval Harness