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
```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)