foxhunt/AGENT_205_SMOKE_TEST_RESULTS.md

Agent 205: MAMBA-2 3-Epoch Smoke Test Results

Date: 2025-10-15 Status: ❌ FAILED - Shape mismatch in training loop Duration: ~9 seconds (failed at first batch)

---

Executive Summary

The smoke test FAILED with a shape mismatch error during the first training batch. The error occurs in prepare_scan_input_with_gradients function where batch matrix multiplication is not properly handling 3D tensors.

Root Cause: Missing batch dimension broadcast in gradient-enabled forward pass.

---

Test Execution

Command

cargo run -p ml --example train_mamba2_dbn --release --features cuda -- --epochs 3

Success Checkpoints

✅ Compilation successful (0.37s) ✅ CUDA initialization successful (RTX 3050 Ti detected) ✅ Data loading successful (7,223 messages from 4 DBN files) ✅ Feature extraction successful (72 sequences created) ✅ Train/val split successful (57 train, 15 val) ✅ Model initialization successful (211,200 parameters) ❌ FAILED at first training batch

---

Error Analysis

Error Message

Error: Training failed

Caused by:
    Model error: Candle error: shape mismatch in matmul, lhs: [32, 60, 512], rhs: [512, 16]

Stack Trace

ml::mamba::Mamba2SSM::forward_with_gradients
ml::mamba::Mamba2SSM::train_batch

Tensor Shapes

• Input to prepare_scan_input_with_gradients: [32, 60, 512] (batch, seq, d_inner)
• B matrix: [16, 512] (d_state, d_inner)
• B.t(): [512, 16] (d_inner, d_state)
• Attempted matmul: [32, 60, 512] × [512, 16] → FAILS

Root Cause

The prepare_scan_input_with_gradients function (line 1186) uses:

let Bu = input.matmul(&B.t()?)?;

This works for 2D tensors but fails for 3D batch tensors because Candle's matmul doesn't automatically broadcast the batch dimension.

The inference version (prepare_scan_input, line 707) correctly broadcasts B:

let batch_size = input.dim(0)?;
let B_t = B.t()?.contiguous()?;
let d_inner = B_t.dim(0)?;
let d_state = B_t.dim(1)?;
let B_broadcasted = B_t.unsqueeze(0)?.broadcast_as((batch_size, d_inner, d_state))?;
let Bu = input.matmul(&B_broadcasted)?;

The training version is missing this broadcast logic.

---

Configuration

Model Parameters

• d_model: 256
• d_state: 16
• d_inner: 512 (d_model × expand = 256 × 2)
• seq_len: 60
• batch_size: 32
• num_layers: 6
• total_parameters: 211,200

Data

• Symbols: 6E.FUT (Euro FX futures)
• Files: 4 DBN files (2024-01-02 to 2024-01-05)
• Total messages: 7,223 OHLCV bars
• Sequences: 72 total (57 train, 15 val)
• Memory: ~4MB

Hardware

• GPU: RTX 3050 Ti (CUDA enabled)
• Device ID: 2

---

Required Fix

Location

ml/src/mamba/mod.rs:1179-1188

Current Code (BROKEN)

fn prepare_scan_input_with_gradients(
    &self,
    input: &Tensor,
    _A: &Tensor,
    B: &Tensor,
) -> Result<Tensor, MLError> {
    // Multiply input by B matrix for state transition
    let Bu = input.matmul(&B.t()?)?;  // ❌ FAILS: [32,60,512] × [512,16]
    Ok(Bu)
}

Fixed Code (REQUIRED)

fn prepare_scan_input_with_gradients(
    &self,
    input: &Tensor,
    _A: &Tensor,
    B: &Tensor,
) -> Result<Tensor, MLError> {
    // FIXED: Broadcast B to match batch dimension
    // input: [batch, seq, d_inner], B: [d_state, d_inner]
    // B.t(): [d_inner, d_state] → broadcast to [batch, d_inner, d_state]
    let batch_size = input.dim(0)?;
    let B_t = B.t()?.contiguous()?;
    let d_inner = B_t.dim(0)?;
    let d_state = B_t.dim(1)?;
    let B_broadcasted = B_t.unsqueeze(0)?.broadcast_as((batch_size, d_inner, d_state))?;
    let Bu = input.matmul(&B_broadcasted)?;  // ✅ WORKS: [32,60,512] × [32,512,16] = [32,60,16]
    Ok(Bu)
}

---

Impact Assessment

Blocking Issues

1. Critical: Training loop completely blocked by shape mismatch
2. Scope: Affects all MAMBA-2 training (DQN/PPO/TFT unaffected)
3. Workaround: None - must fix before training

Non-Blocking Success

• Data loading: 100% functional (0.70ms for 1,674 bars)
• Model initialization: 100% functional (211K parameters)
• CUDA integration: 100% functional (RTX 3050 Ti)
• Feature extraction: 100% functional (72 sequences)

---

Next Steps

Immediate (Agent 206)

1. Apply the fix to prepare_scan_input_with_gradients (5 lines)
2. Recompile with cargo build -p ml --release --features cuda
3. Rerun smoke test: cargo run -p ml --example train_mamba2_dbn --release --features cuda -- --epochs 3
4. Verify first epoch completes without errors
5. Monitor for finite loss (not NaN)

Expected After Fix

• ✅ First batch processes successfully
• ✅ Loss is finite (expect ~0.01-1.0 range initially)
• ✅ Gradients computed correctly
• ✅ Model weights update
• ✅ Validation loss computed
• ✅ Checkpoint saved after each epoch

Timeline

• Fix duration: 2 minutes (5 lines of code)
• Recompile: ~30 seconds
• Smoke test: ~5-10 minutes (3 epochs)
• Total: ~12 minutes to validate fix

---

Lessons Learned

Code Quality Issues

1. Inconsistency: Inference path has correct broadcast logic, training path doesn't
2. Missing tests: No unit tests caught this before integration
3. Code duplication: prepare_scan_input and prepare_scan_input_with_gradients should share logic

Testing Gaps

• No shape validation tests for batch matmul operations
• No smoke tests run before Agent 205
• Missing unit tests for SSM operations with batch dimensions

Recommendations

1. Run smoke tests before declaring "compilation success"
2. Add unit tests for all SSM tensor operations
3. Refactor to eliminate code duplication between inference/training paths
4. Add shape assertions at function boundaries

---

Conclusion

The smoke test successfully identified a critical shape mismatch bug in the training loop that would have blocked all MAMBA-2 training. The fix is simple (5 lines) and mirrors existing working code from the inference path.

Agent 206 should apply this fix immediately and rerun the smoke test.

---

Test Log: /tmp/mamba2_smoke_test.log Agent: 205 Predecessor: Agent 204 (compilation) Successor: Agent 206 (apply fix, retest)