# 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**: ```rust fn forward_with_gradients(&mut self, input: &Tensor) -> Result { // Enable gradient tracking let input = input.detach(); // ❌ REMOVE THIS LINE ``` **AFTER**: ```rust fn forward_with_gradients(&mut self, input: &Tensor) -> Result { // Gradients already tracked on input if needed ``` --- ## Fix #2: Enable SSM Gradient Tracking (Lines 259-286) **File**: `ml/src/mamba/mod.rs` **BEFORE**: ```rust 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**: ```rust 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)**: ```rust 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)**: ```rust 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)**: ```rust pub fn new(config: Mamba2Config, device: &Device) -> Result { 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)**: ```rust pub fn new(config: Mamba2Config, device: &Device) -> Result { 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**: ```rust 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**: ```rust 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**: ```rust let loss_value = loss.to_scalar::()? as f64; ``` **AFTER**: ```rust let loss_value = loss.to_scalar::()?; ``` --- ## Fix #6: Update Optimizer to Use Layer Keys (Line 1224) **File**: `ml/src/mamba/mod.rs` **BEFORE**: ```rust 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**: ```rust 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: ```bash # 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.