Fix Mamba conv params under fine-grained FSDP gather

Mamba's fused path reads conv1d weights directly instead of calling Conv1d.forward(), so fine-grained Megatron-FSDP never gathered those child parameters before the second forward. Register the conv module as an extra forward-gather source and resolve context-parallel Mamba params from the live mixer object.

Author: Tom Long

megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py

@@ -759,6 +759,17 @@ def _pre_forward_param_unshard(module: nn.Module, *unused):
             if self.enable_fine_grained_param_gather_hook:
                 param_list = list(module.parameters(recurse=False))
 
+            extra_forward_param_modules = getattr(module, "_fsdp_extra_forward_param_modules", ())
+            if isinstance(extra_forward_param_modules, nn.Module):
+                extra_forward_param_modules = (extra_forward_param_modules,)
+            if extra_forward_param_modules:
+                seen_param_ids = {id(param) for param in param_list}
+                for extra_module in extra_forward_param_modules:
+                    for extra_param in extra_module.parameters():
+                        if id(extra_param) not in seen_param_ids:
+                            param_list.append(extra_param)
+                            seen_param_ids.add(id(extra_param))
+
             # All-gather the parameters before the forward pass.
             self.all_gather_and_wait_parameters_ready(
                 params=param_list,

megatron/core/ssm/mamba_context_parallel.py

@@ -70,6 +70,7 @@ def __init__(
         A_log_cp1: torch.Tensor,
         D_cp1: torch.Tensor,
         D_has_hdim: bool,
+        mixer=None,
     ) -> None:
         if not HAVE_EINOPS:
             raise ImportError("einops is required by the Mamba model but cannot be imported")
@@ -84,6 +85,7 @@ def __init__(
         self.A_log_cp1 = A_log_cp1
         self.D_cp1 = D_cp1
         self.D_has_hdim = D_has_hdim
+        self._mixer = mixer
 
         self.cp_size = self.cp_group.size()
 
@@ -231,24 +233,29 @@ def conv1d_channels(self):
     def get_conv1d_weight(self) -> torch.Tensor:
         """Returns a slice of the conv1d weight relevant to the current context parallel rank"""
         # weight shape: [conv_dim, 1, d_conv]
-        return self._slice_conv_param(self.conv1d_cp1.weight)
+        conv1d = self._mixer.conv1d if self._mixer is not None else self.conv1d_cp1
+        return self._slice_conv_param(conv1d.weight)
 
     def get_conv1d_bias(self) -> torch.Tensor:
         """Returns a slice of the conv1d bias relevant to the current context parallel rank"""
         # bias shape: [conv_dim]
-        return self._slice_conv_param(self.conv1d_cp1.bias)
+        conv1d = self._mixer.conv1d if self._mixer is not None else self.conv1d_cp1
+        return self._slice_conv_param(conv1d.bias)
 
     def get_dt_bias(self) -> torch.Tensor:
         """Returns a slice of dt_bias relevant to the current context parallel rank"""
-        return self._slice_vector_param(self.dt_bias_cp1)
+        param = self._mixer.dt_bias if self._mixer is not None else self.dt_bias_cp1
+        return self._slice_vector_param(param)
 
     def get_A_log(self) -> torch.Tensor:
         """Returns a slice of A_log relevant to the current context parallel rank"""
-        return self._slice_vector_param(self.A_log_cp1)
+        param = self._mixer.A_log if self._mixer is not None else self.A_log_cp1
+        return self._slice_vector_param(param)
 
     def get_D(self) -> torch.Tensor:
         """Returns a slice of D relevant to the current context parallel rank"""
-        return self._slice_vector_param(self.D_cp1, has_hdim=self.D_has_hdim)
+        param = self._mixer.D if self._mixer is not None else self.D_cp1
+        return self._slice_vector_param(param, has_hdim=self.D_has_hdim)
 
     def _slice_conv_param(self, param: torch.Tensor) -> torch.Tensor:
         """

megatron/core/ssm/mamba_mixer.py

@@ -316,6 +316,9 @@ def __init__(
                 else:
                     nn.init.kaiming_uniform_(self.conv1d.weight, a=math.sqrt(5))
 
+        # The fused Mamba path reads conv1d weights directly instead of calling conv1d.forward().
+        self._fsdp_extra_forward_param_modules = (self.conv1d,)
+
         self.activation = "silu"
         self.act = nn.SiLU()
 
@@ -408,6 +411,7 @@ def __init__(
             A_log_cp1=self.A_log,
             D_cp1=self.D,
             D_has_hdim=self.D_has_hdim,
+            mixer=self,
         )
         self.tp_group = pg_collection.tp
 
@@ -700,17 +704,22 @@ def _ssm_training(
             assert sequence_packing_available, reason_for_no_sequence_packing
             seq_idx = packed_seq_params.seq_idx
 
+        conv1d_weight = rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w")
+        conv1d_bias = self.cp.get_conv1d_bias()
+        dt_bias = self.cp.get_dt_bias().float()
+        D = (
+            rearrange(self.cp.get_D().float(), "(h p) -> h p", p=self.headdim)
+            if self.D_has_hdim
+            else self.cp.get_D()
+        )
+
         y = mamba_split_conv1d_scan_combined(
             zxBCdt,
-            rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w"),
-            self.cp.get_conv1d_bias(),
-            self.cp.get_dt_bias().float(),
+            conv1d_weight,
+            conv1d_bias,
+            dt_bias,
             A,
-            D=(
-                rearrange(self.cp.get_D().float(), "(h p) -> h p", p=self.headdim)
-                if self.D_has_hdim
-                else self.cp.get_D()
-            ),
+            D=D,
             chunk_size=self.chunk_size,
             activation=self.activation,
             headdim=None if self.D_has_hdim else self.headdim,