foxhunt/GRADIENT_CHECKPOINTING_VALIDATION_GUIDE.md

Gradient Checkpointing Validation Guide

Date: 2025-10-23 Status: Ready for Testing Estimated Time: 30-60 minutes

---

Quick Start

Test 1: Baseline (No Checkpointing)

cargo run --release --example test_gradient_checkpointing --features cuda -- \
  --mode baseline \
  --epochs 2 \
  --batch-size 32

Expected Result: May OOM on 4GB GPU (RTX 3050 Ti)

Test 2: With Gradient Checkpointing

cargo run --release --example test_gradient_checkpointing --features cuda -- \
  --mode checkpointing \
  --epochs 2 \
  --batch-size 32

Expected Result: Should succeed on 4GB GPU

Test 3: Comparison Test (Both Modes)

cargo run --release --example test_gradient_checkpointing --features cuda -- \
  --mode compare \
  --epochs 2 \
  --batch-size 32

Expected Result: Baseline may OOM, checkpointing succeeds, report shows ~20% time overhead

---

Validation Checklist

Pre-Test Setup

• RTX 3050 Ti GPU available (4GB VRAM)
• CUDA enabled (nvidia-smi working)
• Codebase built with --features cuda
• No other GPU processes running (nvidia-smi shows <500 MB usage)

Test Execution

• Test 1 (baseline) runs without OOM OR OOMs as expected
• Test 2 (checkpointing) completes successfully
• Test 3 (compare) shows ~20% time overhead
• Memory profiling logs appear every 100 batches
• No memory leaks detected (delta <500 MB per epoch)

Memory Validation

Run watch -n 1 nvidia-smi in separate terminal during tests:

• Baseline peak VRAM: ~3.8-4.2 GB (close to limit)
• Checkpointing peak VRAM: ~2.5-3.0 GB (comfortable margin)
• Memory reduction: ~30-40% (1.2-1.5 GB saved)

Performance Validation

• Checkpointing overhead: <25% (+20% expected)
• Training stability: Loss decreases normally
• No gradient computation errors

---

Troubleshooting

Issue: Baseline Does Not OOM

Cause: GPU has >4GB VRAM or batch_size too small

Fix:

# Increase batch size to stress VRAM
cargo run --release --example test_gradient_checkpointing --features cuda -- \
  --mode baseline \
  --batch-size 64  # Double batch size

Issue: Checkpointing Also OOMs

Cause: Hidden dimension too large or base model memory too high

Fix:

# Reduce hidden dimension
cargo run --release --example test_gradient_checkpointing --features cuda -- \
  --mode checkpointing \
  --hidden-dim 128  # Half of default 256

Issue: No Memory Profiling Logs

Cause: cuda feature not enabled or MemoryProfiler API changed

Fix:

# Verify CUDA feature enabled
cargo build --release --example test_gradient_checkpointing --features cuda

# Check MemoryProfiler API
rg "MemoryProfiler::new" ml/src

Issue: >25% Overhead

Cause: CPU bottleneck or small batch size amplifying recomputation cost

Fix:

# Increase batch size to amortize overhead
cargo run --release --example test_gradient_checkpointing --features cuda -- \
  --mode compare \
  --batch-size 64

---

Manual Memory Testing (nvidia-smi)

Setup

# Terminal 1: Watch GPU memory
watch -n 1 'nvidia-smi --query-gpu=memory.used,memory.total,utilization.gpu --format=csv'

# Terminal 2: Run training
cargo run --release --example test_gradient_checkpointing --features cuda -- \
  --mode checkpointing

Expected Output

memory.used [MiB], memory.total [MiB], utilization.gpu [%]
500 MiB, 4096 MiB, 0 %          # Idle
1200 MiB, 4096 MiB, 15 %        # Model loading
2800 MiB, 4096 MiB, 85 %        # Training (checkpointing)
2900 MiB, 4096 MiB, 90 %        # Peak usage
2800 MiB, 4096 MiB, 85 %        # Stable training

Validation Criteria

• Peak VRAM: <3.2 GB (with 800 MB safety margin on 4GB GPU)
• Stable VRAM: <3.0 GB (no memory leaks)
• GPU Utilization: >80% (efficient computation)

---

Production Deployment

Once validation passes:

1. Update Training Scripts

Add --use-gradient-checkpointing to all TFT training commands:

# Wave C TFT Training (201 features)
cargo run --release --example train_tft_parquet --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-gradient-checkpointing  # ← ADD THIS FLAG

# Wave D TFT Training (225 features)
cargo run --release --example train_tft_parquet --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-qat \
  --use-gradient-checkpointing  # ← ADD THIS FLAG

2. Update Documentation

• Add to ml/docs/QAT_GUIDE.md: Gradient checkpointing section
• Update CLAUDE.md: GPU Memory Budget with checkpointing stats
• Create ml/docs/GRADIENT_CHECKPOINTING_GUIDE.md

3. Update CI/CD

• Add gradient checkpointing test to CI pipeline
• Set up nightly GPU memory regression tests
• Monitor production VRAM usage metrics

4. Team Communication

• Announce gradient checkpointing availability
• Share validation results (memory savings, overhead)
• Provide training command examples

---

Expected Results Summary

Metric	Baseline	Checkpointing	Improvement
Peak VRAM	3.8-4.2 GB	2.5-3.0 GB	-30-40%
Training Time	100% (baseline)	~120%	+20% overhead
OOM Risk (4GB GPU)	HIGH	LOW	✅ Resolved
Batch Size (Max)	16-24	32-48	2× increase
Model Accuracy	100% (FP32)	100% (FP32)	No loss

---

Next Steps

1. Run Validation Test (~30 min):

   cargo run --release --example test_gradient_checkpointing --features cuda -- \
     --mode compare
   

2. Document Results (~10 min):

   • Record peak VRAM from nvidia-smi
   • Capture training time overhead
   • Save test output to /home/jgrusewski/Work/foxhunt/GRADIENT_CHECKPOINTING_TEST_RESULTS.txt

3. Update Production Scripts (~10 min):

   • Add --use-gradient-checkpointing to training commands
   • Update CLAUDE.md with results

4. Begin TFT-225 Training (4-6 weeks):

   cargo run --release --example train_tft_parquet --features cuda -- \
     --parquet-file test_data/ES_FUT_180d.parquet \
     --epochs 50 \
     --use-gradient-checkpointing
   

---

Total Validation Time: 30-60 minutes Blocker Status: P0 - Unblocks TFT-225 training on 4GB GPU Risk: LOW (implementation already validated in codebase)