Repository: huggingface/transformers (Python)
Date: 2024-03-09 22:35:02+00:00
Task type: Bug Fix

Issue:
Cannot propagate gradients in Mamba
### System Info

- `transformers` version: 4.39.0.dev0
- Platform: macOS-14.2.1-arm64-arm-64bit
- Python version: 3.11.7
- Huggingface_hub version: 0.21.4
- Safetensors version: 0.4.2
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.2.0 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no

### Who can help?

@ArthurZucker @younesbelkada 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python
from transformers import MambaConfig, MambaForCausalLM, AutoTokenizer
import torch

tokenizer = AutoTokenizer.from_pretrained("state-spaces/mamba-130m-hf")
model = MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf")
inputs = tokenizer("Hey how are you doing?", return_tensors="pt")

def loss_fn(logits, target):
    return torch.nn.functional.cross_entropy(logits[:, -1, :], target)

outputs = model(**inputs)
loss = loss_fn(outputs.logits, torch.randn([1, 50280]))
loss.backward()
```

This produces the following error:

```bash
{
	"name": "RuntimeError",
	"message": "one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [1, 1536, 16]], which is output 0 of torch::autograd::CopyBackwards, is at version 1; expected version 0 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).",
	"stack": "---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[1], line 13
     11 outputs = model(**inputs)
     12 loss = loss_fn(outputs.logits, torch.randn([1, 50280]))
---> 13 loss.backward()

File ~/Documents/projects/inseq/.venv/lib/python3.11/site-packages/torch/_tensor.py:522, in Tensor.backward(self, gradient, retain_graph, create_graph, inputs)
    512 if has_torch_function_unary(self):
    513     return handle_torch_function(
    514         Tensor.backward,
    515         (self,),
   (...)
    520         inputs=inputs,
    521     )
--> 522 torch.autograd.backward(
    523     self, gradient, retain_graph, create_graph, inputs=inputs
    524 )

File ~/Documents/projects/inseq/.venv/lib/python3.11/site-packages/torch/autograd/__init__.py:266, in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
    261     retain_graph = create_graph
    263 # The reason we repeat the same comment below is that
    264 # some Python versions print out the first line of a multi-line function
    265 # calls in the traceback and some print out the last line
--> 266 Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
    267     tensors,
    268     grad_tensors_,
    269     retain_graph,
    270     create_graph,
    271     inputs,
    272     allow_unreachable=True,
    273     accumulate_grad=True,
    274 )

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [1, 1536, 16]], which is output 0 of torch::autograd::CopyBackwards, is at version 1; expected version 0 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True)."
}
```

### Expected behavior

I would expect gradient backpropagation to work normally in the Mamba model, i.e. not to obtain an error upon calling `loss.backward()`


Issue discussion:
Hi @gsarti, thanks for reporting! 

Looking at the error message, it's likely due to an in place operation in the model implementation. Would you like to open a PR to fix this? 
Pretty sure ~setting `use_cache=False` fixes it, let me check~ let's fix it (It's only for the slow version, which I tried but not on CPU)! I do not recommend training with the slow forward
Hi @ArthurZucker,

Thanks for your reply, I'm not sure if I understand correctly: does the `use_cache=False` fix the behavior? My use case is actually gradient-based feature attribution with [Inseq](https://github.com/inseq-team/inseq), which should be much less intensive than a full training run.

Solution patch:
diff --git a/src/transformers/models/mamba/modeling_mamba.py b/src/transformers/models/mamba/modeling_mamba.py
--- a/src/transformers/models/mamba/modeling_mamba.py
+++ b/src/transformers/models/mamba/modeling_mamba.py
@@ -211,7 +211,7 @@ def slow_forward(self, input_states, cache_params=None):
 
         # 2. Convolution sequence transformation
         if cache_params is not None:
-            ssm_state = cache_params.ssm_states[self.layer_idx]
+            ssm_state = cache_params.ssm_states[self.layer_idx].clone()
             if cache_params.seqlen_offset > 0:
                 conv_state = cache_params.conv_states[self.layer_idx]                   # [batch, intermediate_size, conv_kernel_size]
                 conv_state = torch.roll(conv_state, shifts=-1, dims=-1)


Test patch:
diff --git a/tests/models/mamba/test_modeling_mamba.py b/tests/models/mamba/test_modeling_mamba.py
--- a/tests/models/mamba/test_modeling_mamba.py
+++ b/tests/models/mamba/test_modeling_mamba.py
@@ -170,7 +170,7 @@ def create_and_check_mamba_model(self, config, input_ids, *args):
         self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))
         self.parent.assertEqual(len(result.hidden_states), config.num_hidden_layers + 1)
 
-    def create_and_check_causl_lm(self, config, input_ids, *args):
+    def create_and_check_causal_lm(self, config, input_ids, *args):
         model = MambaForCausalLM(config)
         model.to(torch_device)
         model.eval()
@@ -197,7 +197,30 @@ def create_and_check_state_equivalency(self, config, input_ids, *args):
         self.parent.assertTrue(torch.allclose(torch.cat([output_one, output_two], dim=1), output_whole, atol=1e-5))
         # TODO the orignal mamba does not support decoding more than 1 token neither do we
 
-    def create_and_check_forward_and_backwards(self, config, input_ids, *args, gradient_checkpointing=False):
+    def create_and_check_mamba_cached_slow_forward_and_backwards(
+        self, config, input_ids, *args, gradient_checkpointing=False
+    ):
+        model = MambaModel(config)
+        model.to(torch_device)
+        if gradient_checkpointing:
+            model.gradient_checkpointing_enable()
+
+        # create cache
+        cache = model(input_ids, use_cache=True).cache_params
+        cache.seqlen_offset = 0
+
+        # use cache
+        token_emb = model.embeddings(input_ids)
+        outputs = model.layers[0].mixer.slow_forward(token_emb, cache)
+
+        loss = torch.log(1 + torch.abs(outputs.sum()))
+        self.parent.assertEqual(loss.shape, ())
+        self.parent.assertEqual(outputs.shape, (self.batch_size, self.seq_length, self.hidden_size))
+        loss.backward()
+
+    def create_and_check_mamba_lm_head_forward_and_backwards(
+        self, config, input_ids, *args, gradient_checkpointing=False
+    ):
         model = MambaForCausalLM(config)
         model.to(torch_device)
         if gradient_checkpointing:
@@ -304,12 +327,20 @@ def test_mamba_model(self):
 
     def test_mamba_lm_head_model(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_causl_lm(*config_and_inputs)
+        self.model_tester.create_and_check_causal_lm(*config_and_inputs)
 
     def test_state_equivalency(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_state_equivalency(*config_and_inputs)
 
+    def test_mamba_cached_slow_forward_and_backwards(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_mamba_cached_slow_forward_and_backwards(*config_and_inputs)
+
+    def test_mamba_lm_head_forward_and_backwards(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_mamba_lm_head_forward_and_backwards(*config_and_inputs)
+
     def test_initialization(self):
         config, _ = self.model_tester.prepare_config_and_inputs_for_common()