foxhunt/TFT_LSTM_VARMAP_BUG_FIX.md

TFT LSTM Encoder VarMap Bug Fix

Issue

Critical Bug: ml/src/tft/lstm_encoder.rs:337 created a local VarMap inside the LSTMEncoder::new() constructor, causing LSTM parameters to NOT be saved in TFT checkpoints.

This is identical to the MAMBA-2 VarMap bug that was previously fixed.

Root Cause

// BUGGY CODE (Line 337)
pub fn new(..., device: &Device) -> Result<Self, MLError> {
    let varmap = VarMap::new();  // ❌ Local VarMap - params won't be saved!
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, device);

    for i in 0..num_layers {
        let layer = LSTMLayer::new(..., vs.pp(format!("layer_{}", i)))?;
    }
}

Problem: The local VarMap goes out of scope when the constructor returns, so LSTM layer parameters are never registered in the parent model's VarMap. This means:

• LSTM weights are NOT saved in checkpoints
• Model loading will fail or use random weights
• Training progress is lost across sessions

Solution Applied

1. Fixed Constructor Signature

Changed from accepting Device to accepting parent's VarBuilder:

// FIXED CODE
pub fn new(
    num_layers: usize,
    input_size: usize,
    hidden_size: usize,
    vb: VarBuilder<'_>,  // Use parent's VarBuilder
) -> Result<Self, MLError> {
    let mut layers = Vec::new();

    for i in 0..num_layers {
        let layer_input_size = if i == 0 { input_size } else { hidden_size };
        let layer = LSTMLayer::new(layer_input_size, hidden_size, vb.pp(&format!("layer_{}", i)))?;
        layers.push(layer);
    }

    Ok(Self { layers, num_layers, hidden_size })
}

Key Changes:

• Removed local VarMap creation
• Changed parameter from device: &Device to vb: VarBuilder<'_>
• LSTM layers now register in parent's VarMap via vb.pp(&format!("layer_{}", i))

2. Updated Call Sites

Updated test files to create VarMap and pass VarBuilder:

// BEFORE
let device = Device::Cpu;
let lstm = LSTMEncoder::new(2, 64, 128, &device)?;

// AFTER
use candle_nn::{VarBuilder, VarMap};
use candle_core::DType;

let device = Device::Cpu;
let varmap = VarMap::new();
let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
let lstm = LSTMEncoder::new(2, 64, 128, vb.pp("lstm_encoder"))?;

3. Files Modified

• /home/jgrusewski/Work/foxhunt/ml/src/tft/lstm_encoder.rs - Fixed constructor
• /home/jgrusewski/Work/foxhunt/ml/src/tft/quantized_lstm.rs - Updated 2 test cases

Validation Plan

1. ✅ Build: cargo build -p ml --release - PASSED (no warnings)

2. ⏳ Tests: cargo test -p ml lstm - Running

3. 📋 Expected Results:

   • Tests will FAIL initially (parameter signature changed)
   • Tests in these files need updating:
     • ml/tests/tft_lstm_int8_quantization_test.rs (10 calls)
     • ml/tests/tft_lstm_encoder_unit_test.rs (8 calls)
     • ml/tests/tft_varmap_regression_test.rs (1 call)

4. 📊 Checkpoint Size Verification:

   • After fixing tests, train a TFT model with LSTM encoder
   • Checkpoint size should increase by 50-200KB (LSTM params)
   • Load checkpoint and verify LSTM weights are non-random

Impact

✅ Benefits

• LSTM parameters now properly saved in checkpoints
• Model state fully recoverable across sessions
• Training progress preserved
• Same fix pattern as MAMBA-2 (proven approach)

⚠️ Breaking Changes

• API signature change: LSTMEncoder::new() now requires VarBuilder instead of Device
• All call sites must be updated to create VarMap and VarBuilder
• Test files require systematic updates

📝 Notes

• Main TFT model in mod.rs uses simplified Linear layers, not LSTMEncoder
• LSTMEncoder is used in:
  • Quantized LSTM (quantized_lstm.rs)
  • Test suites (3 files, 19 total calls)
• No production code paths broken (only tests need fixing)

Next Steps

1. Wait for test results
2. Fix failing test files systematically
3. Verify checkpoint integrity with integration test
4. Document in CLAUDE.md

Related Fixes

• MAMBA-2 VarMap bug (fixed in previous wave)
• Same root cause, same solution pattern
• Validates fix approach is correct