Correction to TP logic for Mamba Mixer 2 when Num Groups not divisible by TP Size

[fabianlim]

The current logic is incorrect for the case when n_groups cannot can be divided by TP size, this requires more change to the kernels. This is because it currently uses this simple ratio logic to map the head to the correct group. However, one can imagine in the more general case it is more complicated.

For example, n_groups=3, n_heads=15, and TP_size=5, then in this case we end up with this kind of splitting of the heads (in the below, the numbers 0 - 2 represent which each of the 15 heads map to )

000 | 001 | 111 | 122 | 222

In this case, we will end up with a very heterogenous situation, were we need to

• duplicate to at least 2 groups per shard,
• cannot rely on the ratio logic if we want to stay with 2 groups per shard.

Of course, if we duplicated groups to the extent that they equal heads, then its possible, but not sure if there is an more efficient method.

Current Strategy in this PR

For now, maybe its easier to patch it to specicially only support the following two special cases

1. If TP size divides n_groups,
2. If TP_size does not divide n_groups, but n_groups == 1.

These two scenarios support existing models such as Codestral, Bamba, Zamba, etc.., where n_groups are either 1 or some power of 2

cc: @tlrmchlsmth @yury-tokpanov

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[DarkLight1337]

Thanks for fixing, can we add some tests to avoid future regressions?

[fabianlim]

yes currently we have a TP test for selected models, but AFAIK the tests for mamba2 are not yet automated, and there is some problem running them in a suite because they are not setup properly using torch multiprocessing, but let me discuss this with @tlrmchlsmth

[tlrmchlsmth]

The current logic is incorrect for the case when n_groups cannot be divided by TP size

I'm a little confused by this because we're seeing poor gsm8k results on TP size 2 on Mamba Codestral, which has n_groups == 8, so n_groups is divisible in this case.

For reference, these are GSM8k results on Codestral-7B from @yury-tokpanov
TP=1

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.4731|±  |0.0138|
|     |       |strict-match    |     5|exact_match|↑  |0.4610|±  |0.0137|

vs TP=2

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.0167|±  |0.0035|
|     |       |strict-match    |     5|exact_match|↑  |0.0076|±  |0.0024|

[tlrmchlsmth]

For now, maybe its easier to patch it to specicially only support the following two special cases

1. If TP size divides n_groups,
2. If TP_size does not divide n_groups, but n_groups == 1.

These two scenarios support existing models such as Codestral, Bamba, Zamba, etc.., where n_groups are either 1 or some power of 2

This does make sense to me, and I'm on board with this approach

[tlrmchlsmth]

Should this be:

assert (n_groups % self.tp_size) == 0 or n_groups == 1, ...)

[tlrmchlsmth]

I hit this assert when running

lm_eval --model vllm --model_args pretrained=mistralai/Mamba-Codestral-7B-v0.1,gpu_memory_utilization=0.8,max_model_len=4096,tensor_parallel_size=2 --batch_size auto --trust_remote_code --cache_requests true --tasks gsm8k

but with my suggestion, the TP==2 gsm8k results look good:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.4701|±  |0.0137|
|     |       |strict-match    |     5|exact_match|↑  |0.4549|±  |0.0137|

[fabianlim]

oh sorry you are right

[fabianlim]

I'm a little confused by this because we're seeing poor gsm8k results on TP size 2 on Mamba Codestral, which has

@tlrmchlsmth sorry that was a typo, i fixed it in the main description

[yury-tokpanov]

I think originally groups were introduced to support easy TP for mamba2, right? So, in the original logic n_groups should always be divisible by TP size.

I'm wondering, whether the general case for n_groups being not divisible by TP size would be quite exotic. Probably supporting the special case of n_groups == 1 would be enough for now.

[yury-tokpanov]

Thanks for the fix and all the work on TP for mamba2!

[tlrmchlsmth]

Thanks a lot for the fix!