Enable TP for Mamba-based models (#14811)

Signed-off-by: Roi Koren <roik@nvidia.com>
2025-12-10 18:03:50 -08:00
parent bd7824b24d
commit 7c985331dd
1 changed files with 0 additions and 1 deletions
@@ -339,7 +339,6 @@ class MambaMixer2(torch.nn.Module):
            input_is_parallel=True,
            quant_config=quant_config,
            prefix=f"{prefix}.out_proj",
            reduce_results=False,
        )
        self.norm = Mixer2RMSNormGated(