make FP4 param gather work with the mixed precisions in NVFP4 recipe

megatron/core/models/hybrid/hybrid_block.py

@@ -86,7 +86,12 @@ def __init__(
         dtype=None,
         pg_collection: ProcessGroupCollection = None,
         is_mtp_layer: bool = False,
+        name: str | None = None,
     ) -> None:
+        """
+        Args:
+            name (str | None): module instance name passed top-down from its paranet module
+        """
         super().__init__(config=config)
         self.pre_process = pre_process
         self.post_layer_norm = post_layer_norm
@@ -126,6 +131,7 @@ def __init__(
                         layer_number=layer_number,
                         pp_layer_offset=pp_layer_offset,
                         pg_collection=pg_collection,
+                        name=(name + f".layers.{i}") if name is not None else None,
                     )
                 elif layer_type == LayerSymbols.ATTENTION:
                     layer = build_module(
@@ -136,6 +142,7 @@ def __init__(
                         is_mtp_layer=is_mtp_layer,
                         add_layer_offset=False,
                         pp_layer_offset=pp_layer_offset,
+                        name=(name + f".layers.{i}") if name is not None else None,
                     )
                 elif layer_type == LayerSymbols.DS_ATTENTION:
                     layer = build_module(
@@ -146,6 +153,7 @@ def __init__(
                         is_mtp_layer=is_mtp_layer,
                         add_layer_offset=False,
                         pp_layer_offset=pp_layer_offset,
+                        name=(name + f".layers.{i}") if name is not None else None,
                     )
                 elif layer_type == LayerSymbols.MLP:
                     layer = build_module(
@@ -154,6 +162,7 @@ def __init__(
                         layer_number=layer_number,
                         pg_collection=pg_collection,
                         add_layer_offset=False,
+                        name=(name + f".layers.{i}") if name is not None else None,
                     )
                 elif layer_type == LayerSymbols.MOE:
                     layer = build_module(
@@ -162,6 +171,7 @@ def __init__(
                         layer_number=layer_number,
                         pg_collection=pg_collection,
                         add_layer_offset=False,
+                        name=(name + f".layers.{i}") if name is not None else None,
                     )
                 elif layer_type == LayerSymbols.GDN:
                     layer = build_module(
@@ -171,6 +181,7 @@ def __init__(
                         pg_collection=pg_collection,
                         # Set to False as we do not want to change offset.
                         add_layer_offset=False,
+                        name=(name + f".layers.{i}") if name is not None else None,
                     )
                 else:
                     raise ValueError("unexpected layer_type")

megatron/core/models/hybrid/hybrid_model.py

@@ -260,6 +260,7 @@ def __init__(
             post_process=self.post_process,
             dtype=config.params_dtype,
             pg_collection=self.pg_collection,
+            name="decoder",
         )
 
         # MTP block - uses mtp_block_spec from hybrid_stack_spec.submodules
@@ -279,6 +280,7 @@ def __init__(
                 mtp_layer_pattern=self.mtp_pattern,
                 mtp_num_depths=self.mtp_num_depths,
                 hybrid_submodules=hybrid_submodules,
+                name="mtp",
             )
             self._setup_mtp_cuda_graphs()
 

megatron/core/post_training/modelopt/layers.py

@@ -124,6 +124,7 @@ def __init__(
         tp_comm_buffer_name: str = None,  # Not used
         disable_grad_reduce: bool = False,
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
+        name: str | None = None,  # Not used
     ):
         self.config = config
         self.tp_group = tp_group

megatron/core/quantization/utils.py

@@ -7,10 +7,10 @@
 
 
 def get_quant_config_or_none(
-    module_path: str, recipe: Optional[RecipeConfig] = None
+    module_path: Optional[str], recipe: Optional[RecipeConfig] = None
 ) -> Union[QuantizationConfig, None]:
     """Resolve quantization config for a layer."""
-    if recipe is None:
+    if recipe is None or module_path is None:
         return None
     re_match = re.search(r'layers\.(\d+)', module_path)
     if re_match:

megatron/core/ssm/gated_delta_net.py

@@ -85,6 +85,7 @@ def __init__(
         use_qk_l2norm: bool = True,
         A_init_range: Tuple[float, float] = (1, 16),
         pg_collection: ProcessGroupCollection = None,
+        name: str | None = None,
     ):
         """
         Args:
@@ -98,6 +99,7 @@ def __init__(
             A_init_range: The initialization range for the attention weights.
             pg_collection: The required process groups to use for tensor model parallel and context
                 parallel.
+            name (str | None): module instance name passed top-down from its paranet module
         """
 
         if not HAVE_FLA:
@@ -158,6 +160,7 @@ def __init__(
             is_expert=False,
             tp_comm_buffer_name="fc1",
             tp_group=self.pg_collection.tp,
+            name=(name + ".in_proj") if name is not None else None,
         )
 
         # Conv1d for QKV
@@ -230,6 +233,7 @@ def __init__(
             is_expert=False,
             tp_comm_buffer_name="fc2",
             tp_group=self.pg_collection.tp,
+            name=(name + ".out_proj") if name is not None else None,
         )
 
         self.reset_parameters()

megatron/core/ssm/mamba_layer.py

@@ -72,8 +72,13 @@ def __init__(
         layer_number: int = 1,
         pg_collection: ProcessGroupCollection = None,
         pp_layer_offset: int = 0,
+        name: str | None = None,
     ):
-        """Initialize Mamba Layer."""
+        """Initialize Mamba Layer.
+
+        Args:
+            name (str | None): module instance name passed top-down from its paranet module
+        """
         super().__init__(config)
         assert pg_collection is not None, "pg_collection must be provided for MambaLayer"
 
@@ -88,6 +93,7 @@ def __init__(
             layer_number=layer_number,
             pg_collection=pg_collection,
             pp_layer_offset=pp_layer_offset,
+            name=(name + f".mixer") if name is not None else None,
         )
         self.norm = submodules.norm(self.config, self.config.hidden_size)
         self.mamba_bda = build_module(submodules.mamba_bda)

megatron/core/ssm/mamba_mixer.py

@@ -177,7 +177,12 @@ def __init__(
         layer_number=None,
         pg_collection: ProcessGroupCollection = None,
         pp_layer_offset: int = 0,
+        name: str | None = None,
     ):
+        """
+        Args:
+            name (str | None): module instance name passed top-down from its paranet module
+        """
         if not HAVE_MAMBA_SSM:
             raise ImportError(
                 "MambaSSM is not installed. Please install it with `pip install mamba-ssm`."
@@ -258,6 +263,7 @@ def __init__(
             is_expert=False,
             tp_comm_buffer_name="fc1",
             tp_group=self.pg_collection.tp,
+            name=(name + f".in_proj") if name is not None else None,
         )
         # in_proj packs [z, x, B, C, dt] into one ColumnParallelLinear.  Each
         # component is independently TP-sharded but with different sizes.  When
@@ -392,6 +398,7 @@ def __init__(
             is_expert=False,
             tp_comm_buffer_name="fc2",
             tp_group=self.pg_collection.tp,
+            name=(name + f".out_proj") if name is not None else None,
         )
 
         # Regarding `conv1d`.{`weight`, `bias`}, `dt_bias`, `A_log`, and `D`: these are the

megatron/core/ssm/mlp_layer.py

@@ -21,6 +21,7 @@ def __init__(
         hidden_dropout: float = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
         add_layer_offset: bool = True,
+        name: str | None = None,
     ):
         super().__init__(
             config=config,
@@ -29,4 +30,5 @@ def __init__(
             hidden_dropout=hidden_dropout,
             pg_collection=pg_collection,
             add_layer_offset=add_layer_offset,
+            name=name,
         )

megatron/core/tensor_parallel/inference_layers.py

@@ -83,6 +83,7 @@ def __init__(
         is_expert: bool = False,
         symmetric_ar_type: Optional[str] = None,
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
+        name: str | None = None,
     ):
         assert HAVE_TE, "--transformer-impl=inference_optimized requires transformer engine"
         super().__init__(
@@ -98,6 +99,7 @@ def __init__(
             is_expert=is_expert,
             symmetric_ar_type=symmetric_ar_type,
             tp_group=tp_group,
+            name=name,
         )
 
     def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, None]:
@@ -129,6 +131,7 @@ def __init__(
         skip_weight_param_allocation: bool = False,
         tp_comm_buffer_name: Optional[str] = None,
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
+        name: str | None = None,
     ):
         assert HAVE_TE, "--transformer-impl=inference_optimized requires transformer engine"
         super().__init__(
@@ -144,6 +147,7 @@ def __init__(
             skip_weight_param_allocation=skip_weight_param_allocation,
             tp_comm_buffer_name=tp_comm_buffer_name,
             tp_group=tp_group,
+            name=name,
         )
         self.tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert)
         self.tp_size = dist.get_world_size(self.tp_group)
@@ -256,6 +260,7 @@ def __init__(
         skip_weight_param_allocation: bool = False,
         tp_comm_buffer_name: Optional[str] = None,
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
+        name: str | None = None,
     ):
         assert HAVE_TE, "--transformer-impl=inference_optimized requires transformer engine"
         super().__init__(
@@ -271,6 +276,7 @@ def __init__(
             skip_weight_param_allocation=skip_weight_param_allocation,
             tp_comm_buffer_name=tp_comm_buffer_name,
             tp_group=tp_group,
+            name=name,
         )
         self.tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert)
         self.tp_size = dist.get_world_size(self.tp_group)
@@ -352,6 +358,7 @@ def __init__(
         is_expert: bool,
         tp_comm_buffer_name: Optional[str] = None,
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
+        name: str | None = None,
     ):
         assert HAVE_TE, "--transformer-impl=inference_optimized requires transformer engine"
         super().__init__(
@@ -365,6 +372,7 @@ def __init__(
             is_expert=is_expert,
             tp_comm_buffer_name=tp_comm_buffer_name,
             tp_group=tp_group,
+            name=name,
         )
         self.tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert)
         self.tp_size = dist.get_world_size(self.tp_group)

megatron/core/tensor_parallel/layers.py

@@ -835,6 +835,7 @@ def __init__(
         tp_comm_buffer_name: Optional[str] = None,  # Not used
         disable_grad_reduce: bool = False,
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
+        name: str | None = None,
     ):
         super(ColumnParallelLinear, self).__init__()
 
@@ -1181,6 +1182,7 @@ def __init__(
         is_expert: bool = False,
         tp_comm_buffer_name: str | None = None,  # Not used
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
+        name: str | None = None,
     ):
         super(RowParallelLinear, self).__init__()
 

megatron/core/transformer/attention.py

@@ -290,7 +290,12 @@ def __init__(
         cp_comm_type: str | None = None,
         pg_collection: ProcessGroupCollection | None = None,
         pp_layer_offset: Optional[int] = None,
+        name: str | None = None,
     ):
+        """
+        Args:
+            name (str | None): module instance name passed top-down from its paranet module
+        """
         super().__init__(config=config)
 
         self.config = config
@@ -397,6 +402,7 @@ def __init__(
             is_expert=False,
             tp_comm_buffer_name='proj',
             tp_group=self.pg_collection.tp,
+            name=(name + ".linear_proj") if name is not None else None,
         )
 
         if (
@@ -1379,7 +1385,12 @@ def __init__(
         cp_comm_type: str | None = None,
         pg_collection: ProcessGroupCollection | None = None,
         pp_layer_offset: Optional[int] = None,
+        name: str | None = None,
     ):
+        """
+        Args:
+            name (str | None): module instance name passed top-down from its paranet module
+        """
         super().__init__(
             config=config,
             submodules=submodules,
@@ -1389,6 +1400,7 @@ def __init__(
             cp_comm_type=cp_comm_type,
             pg_collection=pg_collection,
             pp_layer_offset=pp_layer_offset,
+            name=name,
         )
 
         self.linear_qkv_out_dim = self.query_projection_size + 2 * self.kv_projection_size
@@ -1405,6 +1417,7 @@ def __init__(
             is_expert=False,
             tp_comm_buffer_name='qkv',
             tp_group=self.pg_collection.tp,
+            name=(name + ".linear_qkv") if name is not None else None,
         )
 
         # Resolve which norm class to use for Q and K.
@@ -1788,7 +1801,12 @@ def __init__(
         attn_mask_type: AttnMaskType = AttnMaskType.padding,
         cp_comm_type: str | None = None,
         pg_collection: ProcessGroupCollection | None = None,
+        name: str | None = None,
     ):
+        """
+        Args:
+            name (str | None): module instance name passed top-down from its paranet module
+        """
         super().__init__(
             config=config,
             submodules=submodules,
@@ -1797,6 +1815,7 @@ def __init__(
             attention_type="cross",
             cp_comm_type=cp_comm_type,
             pg_collection=pg_collection,
+            name=name,
         )
 
         if self.config.num_query_groups != self.config.num_attention_heads:
@@ -1812,6 +1831,7 @@ def __init__(
             bias=self.config.add_bias_linear,
             skip_bias_add=False,
             is_expert=False,
+            name=(name + ".linear_q") if name is not None else None,
         )
 
         self.linear_kv = submodules.linear_kv(
@@ -1823,6 +1843,7 @@ def __init__(
             bias=self.config.add_bias_linear,
             skip_bias_add=False,
             is_expert=False,
+            name=(name + ".linear_kv") if name is not None else None,
         )
 
     def get_query_key_value_tensors(