foxhunt/AGENT_219_QUICK_FIX_GUIDE.md

MAMBA-2 Quick Fix Guide (Agent 219)

CRITICAL: 5 bugs prevent ANY training. Apply fixes in order.

---

Fix #1: Remove Gradient Detach (Line 1101)

File: ml/src/mamba/mod.rs

BEFORE:

fn forward_with_gradients(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
    // Enable gradient tracking
    let input = input.detach();  // ❌ REMOVE THIS LINE

AFTER:

fn forward_with_gradients(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
    // Gradients already tracked on input if needed

---

Fix #2: Enable SSM Gradient Tracking (Lines 259-286)

File: ml/src/mamba/mod.rs

BEFORE:

let A = Tensor::randn(0.0, 1.0, (config.d_state, config.d_state), device)?;
let B = Tensor::randn(0.0, 1.0, (config.d_state, d_inner), device)?;
let C = Tensor::randn(0.0, 1.0, (d_inner, config.d_state), device)?;
let delta = Tensor::ones((config.d_model,), DType::F64, device)?;

AFTER:

let A = Tensor::randn(0.0, 1.0, (config.d_state, config.d_state), device)?
    .requires_grad(true)?;
let B = Tensor::randn(0.0, 1.0, (config.d_state, d_inner), device)?
    .requires_grad(true)?;
let C = Tensor::randn(0.0, 1.0, (d_inner, config.d_state), device)?
    .requires_grad(true)?;
let delta = Tensor::ones((config.d_model,), DType::F64, device)?
    .requires_grad(true)?;

---

Fix #3: Store VarMap (Lines 358 & 377)

File: ml/src/mamba/mod.rs

BEFORE (struct definition, ~line 358):

pub struct Mamba2SSM {
    pub config: Mamba2Config,
    pub metadata: Mamba2Metadata,
    pub state: Mamba2State,
    // ... other fields ...
    pub device: Device,

    // Model parameters
    pub input_projection: Linear,

AFTER (add field):

pub struct Mamba2SSM {
    pub config: Mamba2Config,
    pub metadata: Mamba2Metadata,
    pub state: Mamba2State,
    // ... other fields ...
    pub device: Device,

    // Model parameters
    pub var_map: candle_nn::VarMap,  // ✅ ADD THIS
    pub input_projection: Linear,

BEFORE (constructor, ~line 377):

pub fn new(config: Mamba2Config, device: &Device) -> Result<Self, MLError> {
    let vs = candle_nn::VarMap::new();
    let vb = VarBuilder::from_varmap(&vs, DType::F64, device);
    // ... create layers ...

    Ok(Self {
        config,
        metadata,
        state,
        // ... other fields ...
        input_projection,

AFTER (store VarMap):

pub fn new(config: Mamba2Config, device: &Device) -> Result<Self, MLError> {
    let vs = candle_nn::VarMap::new();
    let vb = VarBuilder::from_varmap(&vs, DType::F64, device);
    // ... create layers ...

    Ok(Self {
        config,
        metadata,
        state,
        // ... other fields ...
        var_map: vs,  // ✅ ADD THIS
        input_projection,

---

Fix #4: Extract Gradients After Backward (Line 1185)

File: ml/src/mamba/mod.rs

BEFORE:

fn backward_pass(&mut self, loss: &Tensor, _input: &Tensor, _target: &Tensor) -> Result<(), MLError> {
    let _grad = loss.backward()?;

    self.clip_gradients(self.config.grad_clip)?;

AFTER:

fn backward_pass(&mut self, loss: &Tensor, _input: &Tensor, _target: &Tensor) -> Result<(), MLError> {
    loss.backward()?;

    // Extract gradients from SSM parameters
    self.gradients.clear();
    for (layer_idx, ssm_state) in self.state.ssm_states.iter().enumerate() {
        if let Some(A_grad) = ssm_state.A.grad() {
            self.gradients.insert(format!("A_{}", layer_idx), A_grad);
        }
        if let Some(B_grad) = ssm_state.B.grad() {
            self.gradients.insert(format!("B_{}", layer_idx), B_grad);
        }
        if let Some(C_grad) = ssm_state.C.grad() {
            self.gradients.insert(format!("C_{}", layer_idx), C_grad);
        }
        if let Some(delta_grad) = ssm_state.delta.grad() {
            self.gradients.insert(format!("delta_{}", layer_idx), delta_grad);
        }
    }

    // Extract gradients from Linear layers
    for (name, var) in self.var_map.data().lock().unwrap().iter() {
        if let Some(grad) = var.grad() {
            self.gradients.insert(name.clone(), grad);
        }
    }

    self.clip_gradients(self.config.grad_clip)?;

---

Fix #5: Direct F64 Loss Extraction (Line 1168)

File: ml/src/mamba/mod.rs

BEFORE:

let loss_value = loss.to_scalar::<f32>()? as f64;

AFTER:

let loss_value = loss.to_scalar::<f64>()?;

---

Fix #6: Update Optimizer to Use Layer Keys (Line 1224)

File: ml/src/mamba/mod.rs

BEFORE:

fn optimizer_step(&mut self) -> Result<(), MLError> {
    // ... setup code ...

    // Collect gradients first to avoid borrow checker issues
    let a_grad = self.gradients.get("A").cloned();
    let b_grad = self.gradients.get("B").cloned();
    let c_grad = self.gradients.get("C").cloned();
    let delta_grad = self.gradients.get("delta").cloned();

    // Apply Adam updates to all SSM parameters
    let num_layers = self.state.ssm_states.len();
    for layer_idx in 0..num_layers {
        // Update A matrix
        if let Some(ref A_grad) = a_grad {

AFTER:

fn optimizer_step(&mut self) -> Result<(), MLError> {
    // ... setup code ...

    // Apply Adam updates to all SSM parameters
    let num_layers = self.state.ssm_states.len();
    for layer_idx in 0..num_layers {
        // Collect layer-specific gradients
        let a_grad = self.gradients.get(&format!("A_{}", layer_idx)).cloned();
        let b_grad = self.gradients.get(&format!("B_{}", layer_idx)).cloned();
        let c_grad = self.gradients.get(&format!("C_{}", layer_idx)).cloned();
        let delta_grad = self.gradients.get(&format!("delta_{}", layer_idx)).cloned();

        // Update A matrix
        if let Some(ref A_grad) = a_grad {

---

Testing

After all fixes:

# Run MAMBA-2 training test
cargo test -p ml --test e2e_mamba2_training -- --nocapture

# Should see:
# - Gradients computed ✅
# - Parameters updating ✅
# - Loss decreasing ✅

---

Estimated Time

• Fix #1-3: 30 minutes
• Fix #4-6: 1 hour
• Testing: 30 minutes
• Total: 2 hours

---

Priority

1. 🔴 CRITICAL: Fixes #1-3 (enable gradient tracking)
2. 🔴 CRITICAL: Fix #4 (extract gradients)
3. 🟡 MEDIUM: Fix #5 (precision)
4. 🟡 MEDIUM: Fix #6 (optimizer keys)

Apply in order - each fix depends on previous ones.