GitHub Repo
Transformers Community Reports Mamba2 Grouping Gap in Normalization; Non-Fused Path May Differ from Training Core Computation
An October 12 report to the Transformers community says non-fused Mamba2 RMSNorm does not normalize by configured groups, potentially producing different outputs and gradients from the fused training path. The report includes CPU and CUDA comparisons and a local patch check, but no official fix or full-model quality evaluation.
On October 12, the Transformers community raised a Mamba2 computation consistency issue: when n_groups > 1, non-fused gated RMSNorm calculates the mean square across the entire intermediate dimension, while the upstream fused training path normalizes each SSM group separately. The reporter inspected a specified development commit and ran a minimal reproduction with Transformers 5.5.0. The issue remains open, with no related PR so far. Public report
The discrepancy lies in the normalization scope. The current MambaRMSNormGated computes the mean square over the last dimension and does not receive group information at construction; the fused path passes ngroups. If different groups have different value scales, the overall mean allows one group’s inputs to affect another group’s outputs. Using a small FP32 tensor, the report compared the implementation with the official reference and found a maximum output difference of about 0.802. It also measured nonzero cross-group derivatives in the current implementation where the grouped reference gives zero. These are the reporter’s measurements and have not been independently reproduced. Reproduction code and results
The check has practical deployment implications: the official documentation lists 8 as the default value for n_groups, states that Codestral Mamba uses eight groups, and describes both PyTorch and CUDA kernel forward paths. The documentation cautions that the two paths may have small numerical differences; based on the computational analysis in this report, a difference in grouping scope concerns the function’s definition. Engineers need to distinguish it from errors due to operation ordering or low precision. Official Mamba2 documentation
The reporter proposed reshaping the tensor by group and computing RMSNorm separately for each group, preserving the existing weight shape, and said local output and gradient comparisons passed. However, the full fused SSM operation, pretrained models, and generation quality have not been evaluated. Codestral’s gating order is also outside the scope of this reproduction, so the report does not establish that its quality has been affected. Follow-up should focus on maintainer review and consistency tests across fused training, ordinary forward passes, and cache-based decoding in a full model. Validation scope and proposed fix