foxhunt/docs/archive/wave_d/reports/FINAL_VERIFICATION_REPORT.md

TFT Memory Leak Fix - Final Verification Report

Date: 2025-10-26 Agent: Agent 3 (Final Verification) Test: Integration test with ES_FUT_small.parquet (5 epochs, batch_size=1)

---

Executive Summary

Build Status: ✅ SUCCESS Test Status: ✅ 90/90 PASSING (2 ignored) Integration Test: ❌ FAILED (OOM during validation) Root Cause: Incomplete implementation of Agent 1's fix + missing validation cache clearing

---

Detailed Analysis

1. Build Verification ✅

cargo build -p ml --release

Result: Compiled successfully in 0.41s (cached build) Status: ✅ PASS

---

2. Unit Test Verification ✅

cargo test -p ml --lib tft --release

Results:

• Running 92 tests
• ✅ 90 passed
• ⏭️ 2 ignored (expected)
• ❌ 0 failed
• ⏱️ Completed in 2.48s

Status: ✅ PASS (100% pass rate)

---

3. Integration Test (Attempt 1) ❌

Command:

RUST_LOG=info cargo run -p ml --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_small.parquet \
  --batch-size 1 \
  --epochs 5 \
  --use-gpu

Configuration:

• Training batch size: 1
• Validation batch size: 32 ⚠️ (WRONG - should be 1)
• Epochs: 5
• GPU: NVIDIA GeForce RTX 3050 Ti (4GB)

Memory Profile:

Epoch 0 START:           1291.0MB / 4096.0MB (31.5% utilization)
Epoch 0 AFTER_TRAINING:  1611.0MB / 4096.0MB (39.3% utilization)
Epoch 0 BEFORE_VALIDATION: 1611.0MB / 4096.0MB (optimizer dropped)
Validation START:        1611.0MB / 4096.0MB
[OOM ERROR]

Error: Training OOM after 0 retries (final batch_size=1)

Root Cause: Validation batch size hardcoded to 32 in CLI argument parser, causing 32x memory spike during validation.

---

4. Integration Test (Attempt 2) ❌

Command:

RUST_LOG=info cargo run -p ml --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_small.parquet \
  --batch-size 1 \
  --validation-batch-size 1 \  # ← Fixed manually
  --epochs 5 \
  --use-gpu

Configuration:

• Training batch size: 1
• Validation batch size: 1 ✅ (Corrected via CLI arg)
• Epochs: 5
• GPU: NVIDIA GeForce RTX 3050 Ti (4GB)

Memory Profile:

Epoch 0 START:           1650.0MB / 4096.0MB (40.3% utilization)
Epoch 0 AFTER_TRAINING:  1611.0MB / 4096.0MB (39.3% utilization)
Epoch 0 BEFORE_VALIDATION: 1611.0MB / 4096.0MB (optimizer dropped)
Validation START:        1611.0MB / 4096.0MB
[OOM ERROR]

Error: Training OOM after 0 retries (final batch_size=1)

Root Cause: Validation loop lacks per-batch cache clearing. With 176 validation samples at batch_size=1, attention cache accumulates ~2500MB+ without clearing.

---

Issues Found

Issue 1: Incomplete Agent 1 Fix ⚠️

Agent 1's Objective: Change default validation_batch_size to match batch_size dynamically.

What Was Fixed:

• ✅ ml/src/trainers/tft.rs: Updated default config logic

What Was Missed:

• ❌ ml/examples/train_tft_parquet.rs line 85: Still has hardcoded default_value = "32"

Impact: Users must manually specify --validation-batch-size 1 or face OOM errors.

Fix Required:

// ml/examples/train_tft_parquet.rs, line 84-86
/// Validation batch size (defaults to match training batch_size)
#[arg(long)]  // Remove default_value = "32"
validation_batch_size: Option<usize>,  // Make optional

// Later in code (around line 261):
validation_batch_size: opts.validation_batch_size.unwrap_or(opts.batch_size),

---

Issue 2: Missing Validation Cache Clearing ⚠️

Agent 2's Objective: Add CUDA cache clearing to prevent fragmentation.

What Was Fixed:

• ✅ Added sync_cuda_device() + model.clear_cache() AFTER validation (line 1206-1214)

What Was Missed:

• ❌ No cache clearing INSIDE the validation loop (lines 1464-1496)

Impact: With 176 validation samples at batch_size=1, the attention cache accumulates memory without clearing, causing OOM.

Current Validation Loop:

// ml/src/trainers/tft.rs, line 1464-1496
for batch in val_loader.iter() {
    let predictions = self.model.forward(...)?;  // ← Accumulates attention cache
    // ... compute loss ...
    batch_count += 1;
    // ❌ NO CACHE CLEARING HERE
}

Fix Required:

for (i, batch) in val_loader.iter().enumerate() {
    let predictions = self.model.forward(...)?;
    // ... compute loss ...
    batch_count += 1;

    // Clear cache every N batches during validation
    if i % 10 == 0 && self.device.is_cuda() {
        if let Err(e) = Self::sync_cuda_device(&self.device) {
            warn!("Failed to sync CUDA during validation: {}", e);
        }
    }
}

// Final cache clear after validation loop
if self.device.is_cuda() {
    Self::sync_cuda_device(&self.device).ok();
}

---

Issue 3: QAT Learning Rate Schedule Bug (Known)

Status: Documented in commit message, not fixed yet.

Agent 4's Finding: QAT learning rate schedule updates self.state.learning_rate but not the actual optimizer's LR.

Impact: QAT warmup/cooldown phases use incorrect learning rate throughout training.

Priority: P1 (fix after validation OOM resolved)

---

Fixes Applied (from other agents)

✅ Agent 1: Validation Batch Size (Partial)

• File: ml/src/trainers/tft.rs
• Change: Updated default validation_batch_size in config
• Status: Working in code, but CLI parser not updated

✅ Agent 2: CUDA Cache Clearing (Partial)

• File: ml/src/trainers/tft.rs (lines 1206-1214)
• Change: Added sync_cuda_device() + model.clear_cache() after validation
• Status: Working, but missing in-loop clearing

✅ Agent 4: Optimizer Drop/Restore

• File: ml/src/trainers/tft.rs (lines 1117-1126)
• Change: Drop optimizer before validation, restore after
• Status: ✅ Working correctly (confirmed by logs showing "Dropped optimizer before validation")

✅ Agent 5: Memory Profiling

• File: ml/src/trainers/tft.rs
• Change: Added 9 memory checkpoints throughout training loop
• Status: ✅ Working correctly (detailed memory logs visible)

---

Memory Analysis

Training Phase (704 samples, batch_size=1)

Epoch 0 START:          1650MB (baseline + model + optimizer)
Epoch 0 AFTER_TRAINING: 1611MB (-39MB, within expected variance)

Memory Delta: -39MB (optimizer state stabilized) Status: ✅ No memory leak detected in training loop

Validation Phase (176 samples, batch_size=1)

BEFORE_VALIDATION: 1611MB (optimizer dropped)
Validation START:  1611MB (stable)
[Expected END]:    ~1650MB (after processing 176 batches)
[Actual]:          OOM before completion

Expected Memory: ~50MB increase for 176 forward passes Actual Memory: >2485MB increase (OOM) Root Cause: Attention cache accumulation without clearing

---

Test Results Summary

Component	Status	Details
Compilation	✅ PASS	0.41s, no errors
Unit Tests	✅ PASS	90/90 (100%)
Training Loop	✅ PASS	No OOM, stable memory
Optimizer Drop	✅ PASS	Confirmed via logs
Memory Profiling	✅ PASS	9 checkpoints working
Validation Loop	❌ FAIL	OOM at batch_size=1
CLI Defaults	❌ FAIL	Hardcoded validation_batch_size=32
Integration (5 epochs)	❌ FAIL	0/5 epochs completed

---

Final Verdict

ALL FIXES WORKING: ❌ NO

Remaining Issues:

1. P0 (Critical): Validation loop lacks per-batch cache clearing → OOM with 176 samples
2. P1 (High): CLI parser has hardcoded validation_batch_size=32 → User confusion + OOM
3. P2 (Medium): QAT learning rate schedule bug (known, deferred)

Completion Status: 60%

Working (3/5):

• ✅ Build system
• ✅ Unit tests
• ✅ Training loop memory management

Not Working (2/5):

• ❌ Validation loop OOM
• ❌ CLI defaults

---

Recommendations

Immediate Actions (Required for Pass)

1. Fix Validation Cache Clearing (15 min):

   • Add sync_cuda_device() call every 10 batches in validation loop
   • Add final cache clear after validation completes
   • Test with 176 validation samples at batch_size=1

2. Fix CLI Defaults (5 min):

   • Change validation_batch_size from default_value = "32" to Option<usize>
   • Use unwrap_or(opts.batch_size) to match training batch size
   • Update help text to clarify default behavior

3. Re-run Integration Test (5 min):

   • Should complete 5/5 epochs without OOM
   • Verify memory stays below 2000MB throughout training+validation
   • Confirm final model checkpoint saved

Optional (Future Work)

4. Fix QAT LR Schedule (30 min):

   • Implement proper optimizer LR update in QAT warmup/cooldown
   • Add test to verify LR changes during QAT phases

5. Add Validation Cache Clearing Test (10 min):

   • Unit test: validate_epoch with 200+ samples should not OOM
   • Verify cache is cleared every N batches

---

Logs

Full integration test logs saved to:

• /home/jgrusewski/Work/foxhunt/tft_final_integration_test.log (Attempt 1, validation_batch_size=32)
• /home/jgrusewski/Work/foxhunt/tft_final_integration_test_v2.log (Attempt 2, validation_batch_size=1)

---

Conclusion

The other agents (1, 2, 4, 5) made significant progress on fixing the TFT memory leak:

• Agent 4 & 5: ✅ Completed successfully (optimizer drop/restore, memory profiling)
• Agent 1 & 2: ⚠️ Partially completed (fixes applied but incomplete)

The integration test failed because:

1. Agent 1's fix was only applied to the trainer code, not the example CLI parser
2. Agent 2's fix was only applied after validation, not during the validation loop

To achieve a passing test, the two issues above must be fixed, then the integration test re-run.

Estimated Time to Fix: 20 minutes (15 min validation cache + 5 min CLI defaults)

---

Generated: 2025-10-26T19:23:00Z Agent: Agent 3 (Final Verification) Status: ⚠️ INCOMPLETE - 2 critical issues blocking test pass