Repository navigation
[Bugfix][Model][CPU] Fix Mamba2 quantized in_proj weight and scale loading for TP >1 - #58083
Merged
bigPYJ1151 merged 2 commits intoSep 30, 2026
Merged
Conversation
… for TP > 1 In Mamba2 (MambaMixer2), in_proj concatenates gate, up, B, C, and dt projections into an asymmetric dimension. When tensor_parallel_size > 1, groups B and C are replicated across ranks, handled by the custom mamba_v2_sharded_weight_loader. However, mamba_loader was only assigned to in_proj.weight, leaving in_proj.weight_scale (created by compressed-tensors and other quantizers) with the default ColumnParallelLinear loader. This caused Rank 1 to slice past the checkpoint dimension size, crashing with: RuntimeError: start (3352) + length (3352) exceeds dimension size (6448). This fix: 1. Applies mamba_loader to weight, weight_scale, and input_scale. 2. Uses view_as in mamba_v2_sharded_weight_loader to ensure 1D and 2D scale slices copy seamlessly into parameter data. Co-authored-by: Antigravity <antigravity@google.com> Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Akashcodes732
changed the title
Akashcodes732
changed the title
Contributor
Author
|
Hi @bigPYJ1151 |
Contributor
Author
|
Hi @bigPYJ1151 |
Contributor
Author
|
Hi @bigPYJ1151 Can you please take a look at the changes made in this PR ? |
bigPYJ1151
approved these changes
Sep 30, 2026
Member
|
/ci run |
|
✅ Triggered Buildkite CI #92082 for commit |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Fix
RuntimeError: start (...) + length (...) exceeds dimension size (...)when serving quantized Mamba-2 models (e.g.granite-4.0-h-tiny-quantized.w8a8) with TP > 1.MambaMixer2,self.in_projgroups B and C are replicated whenngroups < tp_size. Whileweightusedmamba_loader, quantized parameters (weight_scale,input_scale) used default linear slicing, causing out-of-bounds indexing on Rank 1+.mamba_loaderto all parameters ofself.in_projand align sliced scale dimensions withtarget_slice.Test Plan