[Mamba] Support TP>1 with quantization for mamba2 mixer in case n_groups % tp_size == 0

[tomeras91]

Purpose

PR #13660 built on #10909 and fixed TP loading issues of mamba2 mixer weights. Concretely, it dealt with duplication of weights needed in cases where n_groups % tp_size != 0. This logic is a bit complicated and required creating custom weight loaders, which means using quantized layers is not possible since their weight loading logic can't be overridden. Following that, PR #14617 introduced a special assertion to make sure mamba2 mixer isn't run with quantization and TP>1.

Yet, due to the complexity/impact tradeoff, PR #13660 didn't support all values of n_groups % tp_size, but rather only 2 cases - n_groups % tp_size == 0 and n_groups == 1. The custom weight loading logic is needed only for the latter case. In the former, no weight duplication is needed and the MergedColumnParallelLinear class can be used.

This PR splits the weight creation of mamba2 mixer to two code paths

1. The current path in main is now used only if n_groups % tp_size != 0 (i.e. n_groups==1)
2. Use of MergedColumnParallelLinear if n_groups % tp_size == 0
   This enables quantization of the mamba2 mixer block in case n_groups % tp_size == 0, which is the more common case.

Test Plan

1. Make sure TP>1 results of unquantized models stay the same. This is done by running the following on main and on this branch:

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

2. Make sure a quantized mamba2 model can be loaded with TP>1, and get equivalent results as in TP=1.

Test Result

1. Got the following GSM8K scores:
   main:

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

PR:

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

Results are identical.

2. Validated with an internal quantized mamba2 model

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

`

[tomeras91]

CC @fabianlim

[gemini-code-assist]

Code Review

This pull request introduces a significant enhancement by enabling quantization for the Mamba2 mixer with tensor parallelism when n_groups is divisible by tp_size. This is achieved by adding a new code path using MergedColumnParallelLinear. The changes are well-structured. However, I've identified a critical logic error in the condition that selects between the new and old code paths, which would lead to incorrect behavior for the n_groups=1 case with TP>1. A fix is suggested to ensure the correct path is chosen based on the original model configuration.

[fabianlim]

oic was this a bug?

[tomeras91]

no.. this is just some renaming since I now create self.groups_ssm_state_size = self.n_groups * self.ssm_state_size during init

[fabianlim]

LGTM

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @tomeras91.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[tdoublep]

LGTM - thanks