Add torch.compile Support For Mamba
Author: zhenglongjiepheonix
Pull request: yes
Labels: name: run-slow

`torch.compile` support for mamba! Closes #31246 

Comments:
It seems that the mamba cache is not compatible with the current cache design used in `generate`, but we have to initialize the cache before we step into `model.forward` in order to make dynamo happy, and I think a specific conditional check for mamba in `get_cache` might not be what we want because it's too specific a patch. We can let user specify and create a mamba cache when using `torch.compile`, or is there any way to solve this so that we set the cache for users in `gerenate`? @ArthurZucker 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_31247). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Could you share benchmark results? 

> Could you share benchmark results?

Sure, with `mamba-1.4b` on a single `A100-SXM4-80GB` in `float16, batch_size=1` inference mode, time is for per token generation:

- slow_foward
<img width="1084" alt="Screen Shot 2024-06-07 at 20 19 37" src="https://github.com/huggingface/transformers/assets/32992656/3568757f-db34-4763-9de9-f8db309958c2">

- cuda_kernel_forward
<img width="1084" alt="Screen Shot 2024-06-07 at 20 20 43" src="https://github.com/huggingface/transformers/assets/32992656/3557790b-3df9-4ee0-bf73-abe739bc040d">

- compile
<img width="1084" alt="Screen Shot 2024-06-07 at 20 22 11" src="https://github.com/huggingface/transformers/assets/32992656/854b71ac-90ad-476a-91ec-0a463745c53f">

And here is the function I used for benchmarking
```python
@torch.no_grad
def perf():
    tokenizer = AutoTokenizer.from_pretrained("state-spaces/mamba-1.4b-hf")
    tokenizer.pad_token = tokenizer.eos_token
    inputs = tokenizer("Hey how are you doing today ? " * 100, return_tensors="pt", padding=True).to('cuda')

    model = MambaForCausalLM.from_pretrained("state-spaces/mamba-1.4b-hf", torch_dtype=torch.float16)
    model.config.use_cache = True
    model.to('cuda')    

    input_ids = inputs.input_ids
    cache = MambaCache(model.config, 1, device=input_ids.device)
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    logits = model(input_ids, cache_params = cache).logits
    next_token = torch.argmax(logits[:, -1], dim=-1)[:, None]

    model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)
    torch.cuda.synchronize()
    for i in range(10):
        start.record()
        logits = model(next_token.clone(), cache_params = cache).logits
        next_token = torch.argmax(logits[:, -1], dim=-1)[:, None]

        end.record()
        torch.cuda.synchronize()
        print(f'Step {i}, Total time: {start.elapsed_time(end)} ms, next_token = {next_token.int()}')

```

As we can see from the results above, it takes a lot of time for the first and second decoding step, and actually in my script I skipped compile for the very first prefilling stage because it takes forever(nearly one hour) to compile, so if we only focus on the decoding phase, then we get a steady `8x` speedup
even if comparing with cuda kernel implementation

@ArthurZucker @zhenglongjiepheonix 

Regarding the need to initialize the cache before the forward pass/outside `generate`: I've hit the same issue in the end-to-end `generate` compilation PR. On top of this, we are getting an increasing number of caches with different kwargs, pushing all complexity into `generate`. I'd really like for us to initialize caches outside `generate` by default, it is causing many needless headaches :D We can (and should), however, create an external function to initialize the cache and add it in a way that automagically handles it in the `pipelines` 

> @ArthurZucker @zhenglongjiepheonix
> 
> Regarding the need to initialize the cache before the forward pass/outside `generate`: I've hit the same issue in the end-to-end `generate` compilation PR. On top of this, we are getting an increasing number of caches with different kwargs, pushing all complexity into `generate`. I'd really like for us to initialize caches outside `generate` by default, it is causing many needless headaches :D We can (and should), however, create an external function to initialize the cache and add it in a way that automagically handles it in the `pipelines`

I agree, I think a model agonistic PR should be opened addressing the cache setting issue if we want compile on `generate`, currently I think all caches are set within `generate`? 

Another concern regarding centralized cache initialization is, if we can avoid doing specific condition checks like
```python
if model is mamba:
    cache = MambaCache(...)
```
because sometimes caches that different models need are inherently different and can't be unified, currently I am doing this for mamba by initializing cache inside its own `prepare_inputs_for_generation`,
I think this decentralized way of initializing caches can divide the complexity of cache initialization in a sense, but it might be against the whole `generate` compile stuff @ArthurZucker @gante 

> currently I think all caches are set within generate?

Not all caches -- for instance, `SinkCache` is initialized as you exemplified above (outside `generate`). Passing a cache instance to `generate` is the API that's guaranteed to work all the time, even with custom classes :)

__________________________

Regarding centralization: if we centralize, we force `generate` to hold a large if/else functions, including checks for specific models like mamba. We could break it up by making special models overwrite the cache initialization function. However, caches are growing so quickly in so many directions that I'm wary of pushing all the logic into `generate`, at least for now 👀 Having the caches being externally initialized also makes the code more modular, and much easier to add new caches (remember: the field is growing fast!)

Mmmm TBH it's quite simple to just have a mapping from `cache_implementation` to `CacheCls`. For example for mamba you set `cache_implementation = "mamba"` / automatically infer or just have it in the MambaConfig. 

Let's think simple here, cache config was already introduce, we just need to init with the general cache config object, and pass in the model's config. 

no? 

Maybe let's forget about cache API refactor in this PR because the main point is about compile, I will try setting mamba cache by doing specific conditional check in `generate`, again related with `cache_position` we might need some shape-specific rather than value-specific checks in order to know whether we are in prefill or decoding phase if we don't want boolen flags @ArthurZucker @gante 

There is no need for a refactor here! You already have this : 
```python 
def _get_cache(self, cache_implementation: str, max_batch_size: int, max_cache_len: int) -> Cache:
        """
        Sets a cache for `generate`, that will persist across calls. A new cache will only be initialized a
        new `generate` call requires a larger cache.

        Returns the resulting cache object.
        """
        cache_cls: Cache = NEED_SETUP_CACHE_CLASSES_MAPPING[cache_implementation]
        if cache_implementation == "sliding_window":
            max_cache_len = min(self.config.sliding_window, max_cache_len)

```
so the mamba cache should simply fit there no? 

> There is no need for a refactor here! You already have this :
> 
> ```python
> def _get_cache(self, cache_implementation: str, max_batch_size: int, max_cache_len: int) -> Cache:
>         """
>         Sets a cache for `generate`, that will persist across calls. A new cache will only be initialized a
>         new `generate` call requires a larger cache.
> 
>         Returns the resulting cache object.
>         """
>         cache_cls: Cache = NEED_SETUP_CACHE_CLASSES_MAPPING[cache_implementation]
>         if cache_implementation == "sliding_window":
>             max_cache_len = min(self.config.sliding_window, max_cache_len)
> ```
> 
> so the mamba cache should simply fit there no?

No, `max_cache_len` does not make sense for mamba, and another issue is  we need `cache_params` instead of `past_key_values`

Looks good! 

Congrats on the merge! 🔥