# Wave 8.10: TFT GPU Memory Profile Report **Date**: 2025-10-15 **Agent**: Wave 8.10 **Task**: Measure TFT GPU memory usage and validate <500MB target **Status**: ❌ **FAILED - Memory Optimization Required** --- ## Executive Summary The TFT (Temporal Fusion Transformer) model **FAILS** the <500MB memory target for GPU inference. With batch_size=32 (F32 precision), the model consumes **2,952MB** during forward pass, which is **~6x over budget**. ### Key Findings | Component | Memory (F32) | Budget | Status | |-----------|--------------|--------|--------| | TFT Base Model | 72MB | <300MB | ✅ PASS | | Forward Activations | 2,880MB | <200MB | ❌ FAIL (14x over) | | Backward Gradients | 0MB | <200MB | ✅ PASS | | Optimizer State (est) | 144MB | <200MB | ✅ PASS | | **Peak Training** | **3,096MB** | **<1000MB** | ❌ **FAIL (3.1x over)** | ### Critical Issue The TFT forward pass allocates **2,880MB** for activations with batch_size=32, making it **impossible to fit in 4GB GPU** alongside other models (DQN 6MB + PPO 145MB + MAMBA-2 164MB). --- ## Test Configuration ### Hardware - **GPU**: NVIDIA RTX 3050 Ti (4GB VRAM) - **CUDA**: Available (DeviceId(1)) - **Baseline Memory**: 103MB used, 3,669MB free ### Model Configuration ```rust TFTConfig { input_dim: 256, hidden_dim: 64, // Reduced from typical 128-256 num_heads: 4, num_layers: 2, prediction_horizon: 5, sequence_length: 60, // 60 timesteps num_quantiles: 9, num_static_features: 5, num_known_features: 10, num_unknown_features: 241, batch_size: 32, // Standard ensemble batch size mixed_precision: false, // F32 precision } ``` ### Test Methodology 1. **Baseline Memory**: Measure GPU memory before model creation 2. **Model Initialization**: Create TFT model on GPU 3. **Forward Pass**: Batch_size=32 inference 4. **Backward Pass**: Compute quantile loss + gradients 5. **Training Simulation**: Estimate Adam optimizer state (2x model params) --- ## Detailed Memory Breakdown ### Step 1: Model Initialization ``` Baseline: 103MB After Init: 175MB Model Memory: +72MB ``` ✅ **Model Parameters**: 72MB fits well under 300MB budget **Analysis**: - TFT architecture: Variable Selection Networks + LSTM + Attention + Quantile Output - Hidden_dim=64 keeps parameter count low - 2 layers minimize depth overhead ### Step 2: Forward Pass (Inference) ``` After Forward: 3,055MB Forward Memory: +2,952MB ❌ CRITICAL ISSUE ``` ❌ **Activation Memory**: 2,952MB is **14.7x over 200MB budget** **Root Cause Analysis**: 1. **Sequence Length**: 60 timesteps × 241 unknown features = 14,460 values per sample 2. **Batch Size**: 32 samples × 60 × 241 = 463,200 activations baseline 3. **Multi-Component Architecture**: - 3 Variable Selection Networks (static, historical, future) - LSTM encoder/decoder (hidden states for 60 timesteps) - Multi-head attention (4 heads × 60 × 60 attention matrices) - Quantile output layer (9 quantiles × 5 horizons) **Memory Amplification**: ``` Input: 32 × 60 × 241 = 463,200 values × 4 bytes = 1.85MB VSN: 32 × 60 × 64 = 122,880 values × 4 bytes = 0.49MB (×3 = 1.47MB) LSTM: 32 × 60 × 64 = 122,880 values × 4 bytes = 0.49MB (×2 states = 0.98MB) Attention: 32 × 4 × 60 × 60 = 460,800 values × 4 bytes = 1.84MB Quantile: 32 × 5 × 9 = 1,440 values × 4 bytes = 0.006MB ────────────────────────────────────────────────────────────────── Theoretical Total: ~4.8MB Actual Measured: 2,952MB ────────────────────────────────────────────────────────────────── Overhead Factor: 615x ❌ ``` **Hypothesis**: Candle framework is holding intermediate tensors in GPU memory during forward pass, possibly for gradient computation. The 615x overhead suggests aggressive memory retention for backpropagation. ### Step 3: Backward Pass (Gradients) ``` After Backward: 3,055MB Backward Memory: +0MB ✅ No additional memory ``` ✅ **Gradient Storage**: No additional memory allocated (gradients likely stored in existing activation buffers) ### Step 4: Training Epoch Simulation ``` Optimizer State (est): +144MB (2x model params for Adam) Training Peak (est): 3,096MB total ``` ❌ **Training Peak**: 3,096MB exceeds 1GB budget by **3.1x** ### Step 5: Ensemble Budget Check ``` Model Training Peak ───────────────────────────────── DQN 6MB PPO 145MB MAMBA-2 164MB TFT 3,096MB ❌ FAILS ───────────────────────────────── Total Ensemble: 3,411MB GPU Capacity: 4,096MB Free Memory: 685MB (16.7% remaining) ``` ❌ **Ensemble Fit**: While TFT *technically* fits with other models, leaving only 685MB free is **operationally unviable**: - No memory for concurrent inference - No headroom for memory spikes - Risk of OOM during training --- ## Optimization Strategies ### Priority 1: Mixed Precision (FP16) - 50% Reduction **Implementation**: ```rust TFTConfig { mixed_precision: true, // Enable FP16 // ... } ``` **Expected Savings**: - Model Parameters: 72MB → 36MB (50%) - Forward Activations: 2,880MB → 1,440MB (50%) - Backward Gradients: 0MB → 0MB (no change) - Optimizer State: 144MB → 72MB (50%) - **Training Peak**: 3,096MB → **1,548MB** ✅ **UNDER 2GB** **Pros**: - Easiest to implement (single config flag) - Compatible with RTX 3050 Ti Tensor Cores - Minimal accuracy loss for inference tasks **Cons**: - May need loss scaling for stable training - Requires FP16-compatible Candle build ### Priority 2: Gradient Checkpointing - Trade Compute for Memory **Implementation**: ```rust // Recompute activations during backward instead of storing them TFTConfig { memory_efficient: true, gradient_checkpointing: true, } ``` **Expected Savings**: - Forward Activations: 2,880MB → ~720MB (75% reduction) - **Training Peak**: 3,096MB → **936MB** ✅ **UNDER 1GB** **Pros**: - Dramatic memory reduction - No precision loss **Cons**: - 30-50% slower training (recomputation overhead) - Requires framework support (check Candle docs) ### Priority 3: Reduce Batch Size - Linear Memory Reduction **Implementation**: ```rust TFTConfig { batch_size: 8, // Reduce from 32 → 8 } ``` **Expected Savings**: - Forward Activations: 2,880MB → 720MB (75%) - **Training Peak**: 3,096MB → **774MB** ✅ **UNDER 1GB** **Pros**: - Guaranteed to work - No code changes **Cons**: - 4x longer training time - Less stable gradients (smaller batch size) ### Priority 4: Architecture Simplification **Options**: - Reduce sequence_length: 60 → 30 timesteps (50% reduction) - Reduce num_quantiles: 9 → 5 quantiles (44% reduction) - Reduce attention heads: 4 → 2 heads (50% reduction) **Expected Savings**: - Sequence length 60 → 30: 2,880MB → 1,440MB - Quantiles 9 → 5: Minimal impact (~10MB) - Attention heads 4 → 2: ~200MB reduction **Pros**: - Direct control over memory usage **Cons**: - Degrades model capability (shorter context, less uncertainty quantification) - Requires model retraining --- ## Recommended Action Plan ### Phase 1: Immediate (1 day) 1. **Enable Mixed Precision (FP16)**: - Set `TFTConfig.mixed_precision = true` - Expected: 3,096MB → 1,548MB ✅ - Test GPU memory usage with updated config 2. **Validate FP16 Accuracy**: - Run TFT E2E training test - Compare loss convergence (F32 vs FP16) - Acceptable loss degradation: <5% ### Phase 2: Short-term (2-3 days) 1. **Implement Gradient Checkpointing**: - Research Candle framework support - Add `gradient_checkpointing` config flag - Expected: 3,096MB → 936MB ✅ 2. **Benchmark Training Speed**: - Measure epochs/second (F32 baseline vs FP16 vs gradient checkpointing) - Acceptable slowdown: <2x ### Phase 3: Medium-term (1 week) 1. **Optimize Batch Size**: - Profile memory usage with batch_size=[8, 16, 24] - Find optimal batch_size for <500MB inference - Update ensemble training coordinator 2. **Architecture Tuning**: - Reduce sequence_length if batch_size reductions insufficient - Test model performance with shorter context windows --- ## Success Criteria (Revised) ### Target Metrics | Component | Current (F32) | Target (FP16) | Status | |-----------|---------------|---------------|--------| | Model Parameters | 72MB | <50MB | ✅ | | Inference Memory | 2,952MB | <500MB | ❌ → ✅ (with FP16+checkpointing) | | Training Peak | 3,096MB | <1000MB | ❌ → ✅ (with FP16+checkpointing) | | Ensemble Total | 3,411MB | <2500MB | ❌ → ✅ (with optimizations) | ### Acceptance Criteria 1. ✅ TFT inference memory <500MB (FP16 + gradient checkpointing) 2. ✅ TFT training peak <1GB (FP16 + gradient checkpointing) 3. ✅ All 4 models (DQN + PPO + MAMBA-2 + TFT) fit concurrently in 4GB GPU 4. ✅ >1GB free memory for inference headroom 5. ✅ <5% accuracy degradation from FP16 conversion 6. ✅ <2x training slowdown from gradient checkpointing --- ## Test Results Summary ### Memory Profiling Test Output ``` 🧪 E2E Test: TFT GPU Memory Profiling Device: Cuda(CudaDevice(DeviceId(1))) 📊 GPU Memory Baseline: Total: 4096MB, Used: 103MB, Free: 3669MB 🏗️ Step 1: Model Initialization ✓ Model created Memory after init: 175MB (model: +72MB) 🔍 Step 2: Forward Pass (Inference) ✓ Forward pass complete Memory after forward: 3055MB (peak: +2952MB) ❌ CRITICAL 🔙 Step 3: Backward Pass (Gradients) ✓ Loss computed: 0.947735 Memory after backward: 3055MB (peak: +2952MB) 🚀 Step 4: Training Epoch Simulation Estimated optimizer memory: +144MB Estimated training peak: 3096MB ❌ OVER BUDGET 📈 Memory Profile Summary: Component Memory (F32) ───────────────────── ──────────── TFT Base Model ~72MB Forward Activations ~2880MB Backward Gradients ~0MB Optimizer State (est) ~144MB Peak Training (est) ~3096MB ✅ Validation: ❌ Forward memory should be <500MB, got 2952MB ❌ Training peak should be <1GB, got 3096MB ❌ Total ensemble (3411MB) leaves insufficient headroom ``` ### Test Verdict **FAILED**: TFT requires urgent memory optimization before production deployment. --- ## Next Steps ### Immediate Actions (Today) 1. **Enable FP16 Mixed Precision**: ```bash # Update TFTConfig in tft_e2e_training.rs mixed_precision: true, ``` 2. **Re-run Memory Profiling**: ```bash cargo test -p ml --test tft_e2e_training test_tft_gpu_memory_profiling -- --test-threads=1 --nocapture ``` 3. **Create Follow-up Task**: - Wave 8.11: FP16 Mixed Precision Implementation - Wave 8.12: Gradient Checkpointing Integration ### Documentation Updates - Update `CLAUDE.md` with TFT memory limitations - Add `ml/TFT_MEMORY_OPTIMIZATION.md` with detailed guide - Update ensemble training coordinator with memory constraints --- ## Appendix: Raw Test Data ### GPU Memory Measurements ``` Measurement Point Used Free Delta ──────────────────────────────────────────────────── Baseline 103MB 3669MB - After Model Init 175MB 3597MB +72MB After Forward Pass 3055MB 717MB +2952MB After Backward Pass 3055MB 717MB +0MB Training Peak (est) 3199MB 573MB +3096MB ``` ### System Configuration ``` GPU: NVIDIA RTX 3050 Ti VRAM: 4096MB CUDA: Available Device: Cuda(CudaDevice(DeviceId(1))) ``` ### Test Command ```bash cargo test -p ml --test tft_e2e_training test_tft_gpu_memory_profiling -- --test-threads=1 --nocapture ``` --- ## Conclusion The TFT model's **2,952MB forward activation memory** makes it **unsuitable for 4GB GPU deployment** without aggressive optimization. The recommended path forward is: 1. **FP16 Mixed Precision** (50% reduction → 1,548MB) ✅ Feasible 2. **Gradient Checkpointing** (75% reduction → 774MB) ✅ Ideal 3. **Batch Size Tuning** (backup strategy) ✅ Guaranteed With these optimizations, TFT can achieve the <500MB inference target and enable concurrent ensemble deployment on RTX 3050 Ti 4GB GPU. **Action Owner**: Wave 8.11 Agent **Due Date**: 2025-10-16 **Priority**: **CRITICAL** (blocks production deployment)