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**:
```rust
fn forward_with_gradients(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
    // Enable gradient tracking
    let input = input.detach();  // ❌ REMOVE THIS LINE
```

**AFTER**:
```rust
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**:
```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<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)**:
```rust
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**:
```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::<f32>()? as f64;
```

**AFTER**:
```rust
let loss_value = loss.to_scalar::<f64>()?;
```

---

## 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.