aac0597cd2
- Fixed DQN early stopping checkpoint naming bug (Option B)
- Added is_final: bool parameter to checkpoint callback signature
- Trainer now distinguishes final checkpoints from regular epoch checkpoints
- Final checkpoints use 'dqn_final_epoch{N}' naming convention
- Regular checkpoints use 'dqn_epoch_{N}' naming convention
- Completed comprehensive TFT OOM investigation
- Spawned 3 parallel agents for memory analysis
- Identified 16.4GB memory leak (29.7x over expected 525-550MB)
- Root causes: Attention cache bloat (960MB), gradient accumulation bug, detached tensors
- Recommended fixes: Disable cache during training, explicit tensor drops
- Created TFT_MEMORY_ANALYSIS.md, TFT_MEMORY_LEAK_ANALYSIS.md
- DQN 100-epoch training VERIFIED on Runpod RTX A4000
- Training completed successfully: 100/100 epochs
- Final checkpoint created: dqn_final_epoch100.safetensors
- Training speed: 4.8 sec/epoch (3.5x faster than baseline)
- Option B fix working perfectly
- Deployed RTX 4090 pod for TFT testing
- Pod ID: 6244yzm9hadnog
- 24GB VRAM to bypass OOM issue
- EUR-IS-1 datacenter, $0.59/hr
Files modified:
- ml/examples/train_dqn.rs (checkpoint callback signature)
- ml/src/trainers/dqn.rs (callback signature + is_final parameter)
- CLAUDE.md (compacted to ~11k chars)
Generated reports:
- TFT_MEMORY_ANALYSIS.md (15-section memory breakdown)
- TFT_MEMORY_QUICK_SUMMARY.md (executive summary)
- TFT_MEMORY_LEAK_ANALYSIS.md (5 critical leaks identified)
Co-Authored-By: Claude <noreply@anthropic.com>
2.8 KiB
2.8 KiB
GRAD-B3: TFT Encoder Gradient Checkpointing - Quick Reference
Status: ✅ ALREADY IMPLEMENTED Date: 2025-10-25
Quick Facts
| Metric | Value |
|---|---|
| Status | ✅ Production Ready (Implemented) |
| Memory Reduction | 63-71% (exceeds 30-40% target) |
| Training Overhead | ~20% (acceptable) |
| Compilation | 0 errors, 0 warnings |
| Backward Compatible | Yes (default: disabled) |
| QAT Compatible | ❌ No (workaround exists) |
Usage
Enable Checkpointing
cargo run -p ml --example train_tft_parquet --release --features cuda -- \
--parquet-file test_data/ES_FUT_180d.parquet \
--use-gradient-checkpointing \
--epochs 50
Disable Checkpointing (Default)
cargo run -p ml --example train_tft_parquet --release --features cuda -- \
--parquet-file test_data/ES_FUT_180d.parquet \
--epochs 50
Checkpointed Layers
- ✅ Static Encoder (GRN Stack)
- ✅ Historical Encoder (GRN Stack)
- ✅ Future Encoder (GRN Stack)
- ✅ LSTM Encoder (Temporal)
- ✅ LSTM Decoder (Temporal)
- ✅ Temporal Attention (Self-Attention)
Memory Impact
| Configuration | VRAM Usage | Reduction |
|---|---|---|
| No Checkpointing | 420-530 MB | - |
| With Checkpointing | 105-155 MB | 63-71% |
| Checkpointing + INT8 | 50-75 MB | 75-80% |
When to Use
✅ Use gradient checkpointing when:
- Training on 4GB GPU (RTX 3050 Ti)
- Experiencing OOM errors
- Batch size > 32
- Memory > compute priority
❌ Don't use when:
- GPU has >8GB VRAM
- Speed is critical
- Using QAT mode (incompatible)
Code Locations
| Component | File | Line |
|---|---|---|
| Forward Method | ml/src/tft/mod.rs |
529 |
| Config Field | ml/src/trainers/tft.rs |
434 |
| CLI Flag | ml/examples/train_tft_parquet.rs |
- |
| Static Encoder | ml/src/tft/mod.rs |
569 |
| Historical Encoder | ml/src/tft/mod.rs |
575 |
| Future Encoder | ml/src/tft/mod.rs |
581 |
| LSTM Encoder | ml/src/tft/mod.rs |
593 |
| LSTM Decoder | ml/src/tft/mod.rs |
599 |
| Temporal Attention | ml/src/tft/mod.rs |
616 |
Documentation
- Implementation Guide:
GRADIENT_CHECKPOINTING_IMPLEMENTATION.md - Quick Reference:
GRADIENT_CHECKPOINTING_QUICK_REFERENCE.md - QAT Workaround:
QAT_GRADIENT_CHECKPOINTING_WORKAROUND.md - This Report:
AGENT_GRAD_B3_ENCODER_CHECKPOINTING_REPORT.md
QAT Limitation
⚠️ Gradient checkpointing does NOT work with QAT mode
# This will print a warning and disable checkpointing
--use-qat --use-gradient-checkpointing # ← Checkpointing ignored
Workaround: Use 2-phase training (see QAT_GRADIENT_CHECKPOINTING_WORKAROUND.md)
Agent Status
GRAD-B3: ✅ COMPLETE - NO ACTION REQUIRED
Implementation already exists from previous wave. Skip to next agent.
Updated: 2025-10-25