Views
No views yet
Fully autonomous debugging and error recovery for Hugging Face TRL trainers. Add one callback, wrap withSelfHealingTrainer, and cut debugging costs to near zero.
┌─────────────────────────────────────────┐
│ LAYER 4: ORCHESTRATION │
│ SelfHealingTrainer retry loop │
│ while not converged: try → recover │
├─────────────────────────────────────────┤
│ LAYER 3: RECOVERY │
│ HealingActions: rollback, halve LR, │
│ halve batch, reclip, clear cache │
├─────────────────────────────────────────┤
│ LAYER 2: DIAGNOSIS │
│ Root-cause classifier: NaN/divergence/ │
│ OOM/data/API — with literature refs │
├─────────────────────────────────────────┤
│ LAYER 1: DETECTION │
│ SelfHealingCallback: loss, gradients, │
│ memory, ZClip adaptive clipping │
└─────────────────────────────────────────┘pip install git+https://huggingface.co/ScottzillaSystems/self-healing-training1from self_healing import SelfHealingTrainer, HealingConfig
2from trl import SFTTrainer, SFTConfig
3
4# Your normal training setup
5trainer = SFTTrainer(
6 model=model,
7 args=SFTConfig(
8 output_dir="./output",
9 learning_rate=2e-5,
10 per_device_train_batch_size=4,
11 ),
12 train_dataset=dataset,
13 tokenizer=tokenizer,
14)
15
16# Wrap with self-healing — that's it!
17sh = SelfHealingTrainer(
18 trainer,
19 HealingConfig(
20 max_recovery_attempts=5,
21 zclip_enabled=True,
22 ),
23)
24
25# Optional: dry-run to catch config errors before full training
26sh.dry_run(num_steps=2)
27
28# Train with full autonomy
29result = sh.train()| Failure | Detection | Recovery | Paper |
|---|---|---|---|
| NaN loss | math.isnan(loss) after each step | Rollback → halve LR → enable grad clip | ZClip arxiv:2504.02507 |
| CUDA OOM | on_exception catches OutOfMemoryError | Halve batch (preserve effective via GA) → gradient checkpointing → clear cache | Unicron arxiv:2401.00134 |
| Loss spike | Loss > 5× running mean over window | ZClip adaptive gradient clipping → emergency checkpoint | ZClip arxiv:2504.02507 |
| Divergence | Loss increasing for N consecutive steps | Rollback → halve LR | Pioneer Agent arxiv:2604.09791 |
| Gradient explosion | grad_norm > 100 | ZClip → enable max_grad_norm=1.0 | AdaGC arxiv:2502.11034 |
| DPO plateau | loss ≈ 0.693 (random chance) | Increase LR 2-5× → check data quality | Rafailov et al. (2023) |
| Overfitting | eval_loss - train_loss > 2.0 | Alert with actionable recommendation | Standard practice |
| API errors | Exception with "api/network/timeout" | Exponential backoff (30s → 60s → 120s → ...) | Standard pattern |
| Data errors | Exception with "shape/dimension/index" | Skip batch → log bad sample | Deep Researcher arxiv:2604.05854 |
| Crash postmortem | Always | postmortem.json with exit reason, last step, metrics, recovery history | PTT pattern |
postmortem.json:1{
2 "exit_reason": "exception",
3 "exception_type": "OutOfMemoryError",
4 "last_step": 847,
5 "timestamp": "2026-04-30T15:26:04Z",
6 "final_metrics": {"loss": 2.15, "grad_norm": 42.3},
7 "recovery_actions": [
8 {
9 "failure": "oom",
10 "diagnosis": "CUDA Out of Memory. Batch size exceeds GPU capacity.",
11 "actions": ["halve_batch_size", "enable_gradient_checkpointing", "clear_cache"]
12 }
13 ],
14 "running_time_seconds": 1847.3
15}report_to="trackio" in your training args. SHTS emits:healing/recovery_attempts, healing/nan_count, healing/loss_spike_ratio, healing/eval_gapzclip/raw_grad_norm, zclip/clipped_grad_norm, zclip/z_score, zclip/total_clipshttps://huggingface.co/spaces/<username>/<trackio-space>1# Aggressive — for unstable training, low tolerance
2config = HealingConfig.aggressive()
3# nan_patience=1, zclip_z_threshold=2.0, max_recovery_attempts=10
4
5# Conservative — only intervene on clear failures
6config = HealingConfig.conservative()
7# nan_patience=10, loss_spike_factor=10.0, zclip_z_threshold=4.0, max_recovery_attempts=2
8
9# Custom
10config = HealingConfig(
11 nan_patience=5,
12 loss_spike_factor=8.0,
13 divergence_patience=100,
14 max_recovery_attempts=3,
15 zclip_enabled=True,
16 zclip_z_threshold=3.0,
17)| Trainer | Status | Notes |
|---|---|---|
SFTTrainer (TRL) | ✅ Full | All metrics captured |
DPOTrainer (TRL) | ✅ Full | DPO plateau detection (loss≈0.693) |
GRPOTrainer (TRL) | ✅ Full | Group reward monitoring |
PPOTrainer (TRL) | ✅ Full | KL divergence tracking |
ORPOTrainer (TRL) | ✅ Full | Odds ratio monitoring |
KTOTrainer (TRL) | ✅ Full | Desirable/undesirable logps |
CPOTrainer (TRL) | ✅ Full | Contrastive preference |
Trainer (Transformers) | ✅ Full | Standard ML training |
SelfHealingTrainer.train()
│
├── dry_run() ← Validate setup first
│
└── while not converged:
│
├── trainer.train() ← Run training
│ │
│ ├── on_step_end ← Detect NaN, spikes, divergence
│ ├── on_log ← Monitor gradients (ZClip)
│ ├── on_evaluate ← Check overfitting
│ └── on_exception ← Catch OOM, API, data errors
│
├── [recovery needed?]
│ ├── diagnose ← Classify failure type
│ ├── heal ← Apply recovery actions
│ └── retry ← resume_from_checkpoint=True
│
└── [converged] ← Done!| Paper | ID | Contribution |
|---|---|---|
| Unicron | arxiv:2401.00134 | Cost-aware self-healing at cluster scale, error taxonomy (4 types), elastic scaling |
| ZClip | arxiv:2504.02507 | Z-score adaptive gradient clipping, eliminates catastrophic loss spikes |
| AdaGC | arxiv:2502.11034 | Per-tensor adaptive gradient clipping, optimizer-agnostic |
| Pioneer Agent | arxiv:2604.09791 | Structured decision tree by score buckets for autonomous iteration |
| Deep Researcher | arxiv:2604.05854 | Dry-run validation, zero-cost monitoring, constant-size memory |
| CheckFree | arxiv:2506.15461 | Pipeline-parallel recovery via neighbor averaging |
| DPO | Rafailov et al. (2023) | DPO plateau at 0.693 = random chance (Section 4.2) |
| PTT | post-training-toolkit | DiagnosticsCallback + postmortem pattern |