mamba generation throughput lower than original due to DecodingCGCache

[y1xia0w]

System Info

Python 3.10.13, CUDA 12.1
GPU = NVIDIA GeForce RTX 2080 Ti. Max memory = 10.747 GB.

torch==2.2.1
torchaudio==2.1.0
torchvision==0.16.0
tokenizers==0.15.2
transformers ==git+https://github.com/huggingface/transformers@dd1c9052159ae824c8acef7c2552f9fad5ca020a
triton==2.2.0
causal_conv1d==git+https://github.com/Dao-AILab/causal-conv1d.git@96456720c00393a5c32872d8352d7a7ec31fb3db#egg=causal_conv1d
mamba_ssm==git+https://github.com/state-spaces/mamba.git@9127d1f47f367f5c9cc49c73ad73557089d02cb8#egg=mamba_ssm

Who can help?

text models: @ArthurZucker and @younesbelkada
generate: @gante

Information

• The official example scripts
• My own modified scripts

Tasks

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

Reproduction

The key model initialization and generation parts are given as below.

Original code repo

In the original code repo

from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-130m")
model.eval()

model.generate(
        input_ids=input_ids,
        max_length=max_length,
        **cg=True**
    )

Then throughput for generating 1K length is

Number of parameters: 129135360
Prompt length: 100, generation length: 1000
Prompt processing + decoding time: 1011 ms

Using the HF library

from transformers import MambaForCausalLM
model = MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf")
model.eval()

model.generate(
        input_ids=input_ids,
        max_length=max_length
    )

Then throughput for generating 1K length is

Number of parameters: 129135360
Prompt length: 100, generation length: 1000
state-spaces/mamba-130m-hf prompt processing + decoding time: 15970ms

Expected behavior

The "cg=True" is confirmed to be the part has a significant impact on the generation performance for mamba.

I have tried:

1. Passing the "use_cache=True" as follows won't affect the results

model = MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf", use_cache=True)
or
model = MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf", cache_params={use_cache: True})
or
model.config.use_cache=True

2. Modifying the mamba model to force the argument "use_cache=True" in the MambaModel, but still not working.

I assume this is related to the #29605, but modifying the argument directly seems not solving the problem.

[javiermcebrian]

Hi! Thanks for raising this issue. I agree. I've realized the same since some weeks ago while testing Mamba, as previously I was using the mamba-ssm repo and was able to to generation much faster using the DecodingCGCache. I've been tracking that other issue too, as although it seems just variables refactor it includes a change in the prepare_inputs_for_generation function passing the kwargs (https://github.com/huggingface/transformers/pull/29605/files#diff-e1d4758c08973fdac2c23a8a3710872d943ce8509035205da4a681bc4dcaf1c3R694). I didn't created any issue as I wondered that PR will be merged soon due to its simplicity, but seems it is not. Also, I'm not sure if that PR solves all the issue.

[ArthurZucker]

Hey! Thanks both I'll dive a bit on this, contributions are also welcome.
The use_cache should be using the cache, and in my test it was, I'll check.

[gante]

Hi folks 👋

It doesn't look like a caching issue, but a compilation one -- under the hood, the cg flag in the original repo triggers compilation (cg probably stands for cuda graphs).

Our implementation is not compatible with fullgraph compilation (the equivalent), which may explain the difference.

@ArthurZucker in case you want to reproduce:

from transformers import MambaForCausalLM
import torch
from time import time

model = MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf", device_map="auto")
model.eval()
model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)

input_ids = torch.arange(100, device="cuda").unsqueeze(0)
max_length = 1000

start = time()
model = torch.compile(model)
out = model.generate(
    input_ids=input_ids,
    max_length=max_length,
)
print(f"Time: {time() - start:.2f}s")
print(out.shape)

[ArthurZucker]

Might be the einsum that were all replaces by normal operation, saw some issue with this recently as well

[y1xia0w]

@gante Nice catch and thank you for your effort to address it!
@ArthurZucker Could you please elaborate a bit more and also share the other related issues? Thank you very much!

[ArthurZucker]

#29544 is what made me think of that!

[javiermcebrian]

Hi! Thanks for the findings :)
@ArthurZucker is there any update on this issue? I saw #29544 is merged, but I'm not sure if the compilation issue described affecting the CG is fixed or not.
Thanks!

[ArthurZucker]

I did not test the compilation with mamba, I right now can't investigate 😢 would love if you can.
Now to be fair, we need to compare greedy generation only, not sampling, not beam search, making sure apple are compared with apple. As said, the compilation is not currently supported

[ArthurZucker]

I'll set this as a feature request: add support for torch.compile for Mamba

[javiermcebrian]

Hi @ArthurZucker :)
I tested using exactly the same generation config, just greedy search with no sampling. Original Mamba repo was faster than HF during generation.
So then, just to understand, the only possible fix would be supporting compilation for Mamba so that de CG works as in the original repo?
Thanks!

[gante]

Added it to our generate + torch.compile tracker, which we are actively tackling :)