# Gradient Checkpointing Validation Guide

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

---

## Quick Start

### Test 1: Baseline (No Checkpointing)
```bash
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
```bash
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)
```bash
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**:
```bash
# 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**:
```bash
# 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**:
```bash
# 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**:
```bash
# 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
```bash
# 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:

```bash
# 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):
   ```bash
   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):
   ```bash
   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)