foxhunt/WAVE_8_10_TFT_GPU_MEMORY_PROFILE.md

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

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:

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:

// 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:

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:

   # Update TFTConfig in tft_e2e_training.rs
   mixed_precision: true,
   

2. Re-run Memory Profiling:

   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

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)