stabilityguard 0.3.0

Advanced PyTorch training stability platform — RLHF monitoring, distributed training support, mixed precision stability, and comprehensive logging.

⚡ StabilityGuard

Advanced PyTorch training stability with predictive spike detection.
Catch gradient explosions before they happen. Auto-tune thresholds. Recover gracefully.

---

The Problem

Training large neural networks fails catastrophically when gradient spikes or NaN explosions occur. A single corrupted gradient can poison your optimizer's state (momentum buffers, Adam's second moments), forcing you to restart from the last checkpoint—often thousands of steps back.

The failure cascade:

1. A gradient spike hits one layer (10-1000× larger than normal)
2. Optimizer state gets corrupted with NaN/Inf values
3. Subsequent steps propagate NaN through the entire model
4. By the time PyTorch throws RuntimeError: Function returned nan, it's too late—your checkpoint is poisoned

What PyTorch tells you:

RuntimeError: Function 'AddmmBackward0' returned nan values in its 0th output

What you actually need to know:

• Which layer caused the spike?
• When did it start (exact step number)?
• What was the gradient magnitude vs. normal baseline?
• Can we catch it before it corrupts the optimizer?

The cost: For a 7B LLM on 8×A100 GPUs, a single NaN explosion at 40% training wastes ~72 GPU-hours. You lose days of compute because one gradient went haywire.

StabilityGuard solves this by monitoring every layer's gradients in real-time and catching spikes before they corrupt your training run.

The Fix

pip install stabilityguard

Before StabilityGuard (standard PyTorch):

from torch.optim import AdamW

optimizer = AdamW(model.parameters(), lr=2e-4)

for batch in dataloader:
    loss = model(batch)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

After StabilityGuard (one-line change):

from torch.optim import AdamW
from stabilityguard import GuardedOptimizer  # ← the only change

base_opt = AdamW(model.parameters(), lr=2e-4)
optimizer = GuardedOptimizer(base_opt, model,
    spike_threshold=10.0,      # gradient norm ratio to trigger alert
    nan_action="skip",         # "skip" | "rollback" | "raise"
    log_every=50               # steps between diagnostic summaries
)

for batch in dataloader:
    loss = model(batch)
    loss.backward()
    optimizer.step()           # GuardedOptimizer intercepts here
    optimizer.zero_grad()

🚀 What's New in v0.3.0

StabilityGuard v0.3.0 introduces 4 major feature categories that transform it into a comprehensive training stability platform:

1. 🤖 RLHF/PPO Stability Guard

Industry-first monitoring for Reinforcement Learning from Human Feedback training.

from stabilityguard.rlhf import RLHFGuard

rlhf_guard = RLHFGuard(
    model=policy_model,
    ref_model=reference_model,
    kl_threshold=0.1,           # KL divergence limit
    reward_collapse_threshold=0.01,  # Entropy threshold
    value_divergence_threshold=10.0  # Critic stability
)

# During PPO training
rlhf_guard.check_step(
    logits=policy_logits,
    ref_logits=ref_logits,
    rewards=rewards,
    values=values,
    old_logprobs=old_logprobs,
    new_logprobs=new_logprobs,
    step=step
)

Features:

• KL Divergence Monitor: Automatic penalty adjustment when policy drifts
• Reward Collapse Detector: Identifies degenerate reward model behavior
• Value Divergence Monitor: Tracks critic network stability
• PPO Ratio Monitor: Analyzes policy update magnitudes

2. 🌐 Distributed Training Support

Multi-GPU gradient spike detection with rank attribution.

from stabilityguard.distributed import DistributedGuardedOptimizer

optimizer = DistributedGuardedOptimizer(
    base_optimizer,
    model,
    spike_threshold=10.0,
    backend="nccl"  # or "gloo"
)

# Automatically coordinates spike detection across all GPUs
# Identifies which rank caused the spike

Features:

• Distributed Spike Detector: All-reduce coordination across ranks
• FSDP Stability Guard: Specialized monitoring for Fully Sharded Data Parallel
• DeepSpeed Stability Guard: Support for ZeRO-1, ZeRO-2, ZeRO-3
• Rank Attribution: Identifies which GPU caused instability

3. 🎯 Mixed Precision Stability

Comprehensive FP16/BF16/FP8 training monitoring.

from stabilityguard.precision import MixedPrecisionGuard

precision_guard = MixedPrecisionGuard(
    model=model,
    precision="fp16",  # or "bf16", "fp8"
    initial_scale=2**16
)

# Scale loss
scaled_loss = precision_guard.scale_loss(loss)
scaled_loss.backward()

# Check for issues and adjust scale
precision_guard.update_scale(model, optimizer, step)

Features:

• Precision Monitor: Overflow/underflow detection
• Adaptive Loss Scaler: Stability-aware scaling
• Dynamic Range Tracking: Real-time precision analysis
• Automatic Recommendations: Actionable precision advice

4. 📊 Advanced Logging System

Comprehensive training diagnostics beyond gradient monitoring.

from stabilityguard.logging import AdvancedLogger

logger = AdvancedLogger(
    log_dir="./logs",
    enable_gradient_flow=True,
    enable_activation_stats=True,
    enable_weight_updates=True,
    enable_checkpoint_scoring=True
)

# Log comprehensive statistics
log_data = logger.log_step(
    step=step,
    loss=loss.item(),
    model=model,
    optimizer=optimizer
)

# Get checkpoint health score (0-100)
score = logger.score_checkpoint("checkpoint.pt", training_history)

Features:

• Gradient Flow Tracker: Per-layer gradient monitoring
• Activation Stats Logger: Dead neuron detection
• Weight Update Tracker: Parameter change analysis
• Checkpoint Health Scorer: Stability-based checkpoint ranking

---

🔧 v0.2.0 Features (Still Available)

1. 🔮 Edge of Stability (Predictive Detection)

Predict spikes 10-50 steps before they happen using Hessian spectral radius estimation.

optimizer = GuardedOptimizer(base_opt, model,
    enable_edge_of_stability=True,
    eos_check_interval=10
)

2. 🔄 SPAM Optimizer (Momentum Reset)

Automatically reset momentum buffers when spikes are detected, preventing corruption propagation.

optimizer = GuardedOptimizer(base_opt, model,
    enable_spam=True,
    spam_lr_reduction=0.5,
    spam_recovery_steps=100
)

3. 🎯 Auto-Calibration (Zero Manual Tuning)

Eliminate manual threshold tuning by learning optimal thresholds from your data.

optimizer = GuardedOptimizer(base_opt, model,
    enable_auto_calibration=True,
    auto_calibration_warmup_steps=100
)

4. ✂️ HELENE Clipping (Adaptive Per-Layer)

Per-layer adaptive gradient clipping based on local Hessian conditioning.

optimizer = GuardedOptimizer(base_opt, model,
    enable_helene=True,
    helene_base_clip=1.0
)

All features are opt-in (default False) for backward compatibility.

What You See

When a spike hits at step 847:

╔══════════════════════════════════════════════════════════════╗
║  ⚠ STABILITYGUARD — SPIKE DETECTED @ step 847                ║
╠══════════════════════════════════════════════════════════════╣
║  Trigger layer  : transformer.h.11.mlp.c_proj                ║
║  Grad norm      : 847.3  (baseline: 1.2, ratio: 706.1x)      ║
║  Action taken   : optimizer.step() SKIPPED                   ║
║  Loss (pre-skip): 14.71                                      ║
╚══════════════════════════════════════════════════════════════╝
stabilityguard.log written → ./sg_logs/spike_step847.json

Your model weights are untouched. Training continues.

How It Works

1. Backward hooks on every nn.Module capture per-layer gradient L2 norms
2. EMA baselines track normal gradient behavior (α=0.01, ~100 step lookback)
3. Spike detection fires when current_norm / ema_baseline > threshold
4. NaN/Inf short-circuit catches corrupted gradients via torch.isfinite
5. Actions execute automatically: skip the step, rollback to checkpoint, or raise

Configuration

Parameter	Default	Description
spike_threshold	10.0	Ratio of current norm to EMA baseline that triggers a spike
nan_action	"skip"	Action on spike: "skip", "rollback", or "raise"
log_every	50	Steps between periodic diagnostic summaries
log_dir	"./sg_logs"	Directory for JSON spike reports
ema_alpha	0.01	EMA smoothing factor (smaller = slower adaptation)
warmup_steps	10	Steps before spike detection activates
verbose	True	Print diagnostic summaries to stdout

v0.2.0 Advanced Features

Parameter	Default	Description
enable_edge_of_stability	False	Enable predictive spike detection
enable_spam	False	Enable momentum reset on spikes
enable_auto_calibration	False	Enable automatic threshold tuning
enable_helene	False	Enable adaptive per-layer clipping
eos_check_interval	10	Steps between Edge of Stability checks
eos_power_iterations	20	Accuracy of λ_max estimation
spam_lr_reduction	0.5	LR reduction factor after spike
spam_recovery_steps	100	Steps to recover LR after spike
auto_calibration_warmup_steps	100	Warmup steps for auto-calibration
helene_base_clip	1.0	Base gradient clip value for HELENE

Actions

Action	Behavior
skip	Skip optimizer.step() — corrupted gradients discarded, weights unchanged
rollback	Restore model + optimizer to the last clean checkpoint
raise	Throw GradientSpikeError for interactive debugging

Integrations

Weights & Biases

from stabilityguard.integrations.wandb import WandBBridge
bridge = WandBBridge()
# Metrics logged under sg/ namespace automatically

MLflow

from stabilityguard.integrations.mlflow import MLflowBridge
bridge = MLflowBridge()

HuggingFace Transformers

from stabilityguard.integrations.huggingface import StabilityGuardCallback

trainer = Trainer(
    model=model,
    args=training_args,
    callbacks=[StabilityGuardCallback(spike_threshold=10.0)],
)

Performance

Metric	Value
Per-step overhead (GPU)	Low (estimated 1-5% on modern GPUs)
Per-step overhead (CPU)	High (~60%, not recommended for production)
Per-step overhead (spike detected)	~3ms (includes JSON write)
Recommended for	GPU-based training (production workloads)
External dependencies	0 (only PyTorch)
License	MIT

v0.2.0 Feature Overhead

Feature	Overhead	When
Edge of Stability	~40 backward passes	Every eos_check_interval steps
SPAM	Negligible	Only on spike detection
Auto-Calibration	Negligible	Only during warmup
HELENE	~20 backward passes	Every step (if enabled)

Understanding Spike Logs

StabilityGuard writes detailed JSON logs for every spike detected to ./sg_logs/:

{
  "snapshot": {
    "step": 847,                   // When spike occurred
    "spike_layer": "fc2",          // Which layer caused it
    "spike_ratio": 12.8,           // Severity (current/baseline)
    "action": "skip",              // Action taken
    "loss": "tensor(14.71, ...)"   // Loss at spike time
  },
  "layer_norms": {
    "fc2": 15.36,                  // Current gradient norm
    "fc2.weight": 14.82
  },
  "ema_baselines": {
    "fc2": 1.20,                   // Expected "normal" norm
    "fc2.weight": 1.15
  },
  "nan_layers": []                 // NaN corruption check
}

Spike Severity Scale

spike_ratio	Severity	Action
1.0-2.0	Normal	No spike
2.0-5.0	Moderate ⚠️	Monitor
5.0-10.0	Severe 🔥	Investigate
>10.0	Critical 💥	Fix immediately

Quick Analysis Commands

# Count total spikes
ls sg_logs/ | wc -l

# Find most problematic layer
grep -h "spike_layer" sg_logs/*.json | sort | uniq -c | sort -rn

# Check for NaN corruption
grep -l "nan_layers.*\[" sg_logs/*.json

Install from source

git clone https://github.com/ashwinmsad1/stabilityguard.git
cd stabilityguard
pip install -e ".[test]"
pytest tests/ -v

Documentation

v0.3.0 Documentation

• v0.3.0 Implementation Complete - Implementation summary
• v0.3.0 Roadmap - Technical specifications and architecture

v0.2.0 Documentation

• v0.2.0 Release Summary - Detailed feature overview
• v0.2.0 Roadmap - Technical specifications
• Hessian-Vector Products Explained - Computational cost analysis

General Documentation

• Long-term Roadmap - Vision through v2.0.0
• Changelog - Version history

License

MIT — use it anywhere, no restrictions.