Cannot propagate gradients in Mamba
Author: gsarti
Pull request: no

### 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()`

Comments:
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.

I recommend installing the cuda kernels for now!