solving PP=1 + EP>1 issue, ignoring moe_freq and using only pattern for moe layers count

Author: chochowski

src/megatron/bridge/models/conversion/model_bridge.py

@@ -266,6 +266,11 @@ def _megatron_local_name_to_global(
     vp_stage: Optional[int] = None,
 ) -> str:
     """Adjust layer number and expert number from local to global numbering."""
+    # ``layer_module`` is resolved lazily and shared between the PP and EP
+    # branches. It stays ``None`` on the EP-only path (PP=1, EP>1) until the EP
+    # branch resolves it, so it must never be read before being assigned.
+    layer_module = None
+
     # PP
     pp_group = _get_pp_group(models)
     if "layers." in param_name and get_pg_size(pp_group) > 1:
@@ -290,7 +295,25 @@ def _megatron_local_name_to_global(
     is_expert_param = (is_grouped_expert_param or is_local_expert_param) and ".adapter." not in param_name
     ep_group = _get_ep_group(models) if is_expert_param else None
     if is_expert_param and ep_group is not None and get_pg_size(ep_group) > 1:
-        num_experts_per_rank = layer_module.mlp.num_local_experts   # per-layer, heterogeneous-safe
+        # Resolve the layer module independently of the PP branch above: on the
+        # EP-only path (PP=1, EP>1) that branch never runs, so ``layer_module``
+        # is still ``None`` here. Prefer the per-layer module count
+        # (heterogeneous-safe); fall back to the config-derived count when the
+        # model isn't available (e.g. name-only conversions in tests).
+        if layer_module is None:
+            layer_match = re.match(r"^(.+?\.layers\.\d+)", param_name)
+            if models is not None and layer_match is not None:
+                try:
+                    _, layer_module = get_module_and_param_from_name(
+                        models=models, param_name=layer_match.group(1), vp_stage=vp_stage
+                    )
+                except (ValueError, AttributeError):
+                    layer_module = None
+
+        if isinstance(layer_module, MegatronModule):
+            num_experts_per_rank = layer_module.mlp.num_local_experts  # per-layer, heterogeneous-safe
+        else:
+            num_experts_per_rank = config.num_moe_experts // get_pg_size(ep_group)
 
         def _update_grouped_expert_number(param_name: str, param_type: str) -> str:
             """Update expert number from local to global for weight or bias parameters."""

src/megatron/bridge/training/utils/train_utils.py

@@ -1044,7 +1044,15 @@ def training_log(
         num_layers = getattr(config.model, "num_layers", None)
         moe_layer_freq = getattr(config.model, "moe_layer_freq", None)
         pattern = getattr(config.model, "hybrid_layer_pattern", "") or ""
-        if moe_layer_freq is None and pattern:                 # per-layer averaging denominator
+        # For hybrid models the MoE layers are defined by the hybrid pattern
+        # ("E" == MoE layer), not by ``moe_layer_freq`` — which defaults to 1 on
+        # the inherited TransformerConfig/MambaModelProvider path and would
+        # otherwise make ``_count_moe_layers`` average aux/z-loss over *every*
+        # layer instead of only the MoE ones. The pattern is authoritative
+        # whenever present, so always derive the per-layer averaging denominator
+        # from it. ``num_layers`` stays at full depth for the tracker storage
+        # size (the NCCL-deadlock fix).
+        if pattern:
             moe_layer_freq = [1 if c == "E" else 0 for c in pattern]
 
         # Wrap the TB writer so MoE/MTP metrics also reach MLFlow / Comet (issue #2989).

tests/unit_tests/training/utils/test_train_utils.py

@@ -1134,6 +1134,85 @@ def test_moe_logging_without_z_loss(
         assert "z_loss" not in track_names
         assert len(track_names) == 1
 
+    @mock.patch("megatron.bridge.training.utils.train_utils.get_num_microbatches")
+    @mock.patch("megatron.bridge.training.utils.train_utils.reduce_max_stat_across_model_parallel_group")
+    @mock.patch("megatron.bridge.training.utils.train_utils.get_world_size_safe")
+    @mock.patch("megatron.bridge.training.utils.train_utils.print_rank_last")
+    @mock.patch("megatron.bridge.training.utils.train_utils.track_moe_metrics")
+    @mock.patch("megatron.bridge.training.utils.train_utils.report_runtime")
+    @mock.patch("megatron.bridge.training.utils.train_utils.report_throughput")
+    @mock.patch("megatron.bridge.training.utils.train_utils.report_l2_norm_grad")
+    def test_moe_logging_hybrid_pattern_default_freq(
+        self,
+        mock_report_l2_norm_grad,
+        mock_report_throughput,
+        mock_report_runtime,
+        mock_track_moe,
+        mock_print_rank_last,
+        mock_get_world_size,
+        mock_reduce_lr,
+        mock_get_microbatches,
+        mock_config,
+        mock_global_state,
+        loss_dict,
+    ):
+        """Hybrid-MoE models must derive the averaging denominator from the pattern.
+
+        Regression: recipes such as ``nemotron_3_nano`` set a ``hybrid_layer_pattern``
+        with a handful of ``E`` (MoE) layers but leave ``moe_layer_freq`` at the
+        inherited default (``1``). ``num_layers`` must stay at full depth for the
+        tracker storage size (the NCCL-deadlock fix), while ``moe_layer_freq`` is
+        passed as the per-layer binary MoE mask derived from the pattern so aux/z-loss
+        are averaged over only the MoE layers.
+        """
+        total_loss_dict = self.get_fresh_total_loss_dict()
+
+        mock_report_l2_norm_grad.return_value = {}
+        mock_report_throughput.return_value = {}
+        mock_report_runtime.return_value = {}
+        mock_get_microbatches.return_value = 8
+        mock_reduce_lr.return_value = 1e-4
+        mock_get_world_size.return_value = 32
+
+        # Hybrid Mamba+MoE config: pattern places MoE ("E") on 2 of 8 layers,
+        # while moe_layer_freq is the inherited default (1, not None).
+        pattern = "M-M*E-M*E-"
+        mock_config.model.num_moe_experts = 8
+        mock_config.model.moe_router_load_balancing_type = "aux_loss"
+        mock_config.model.moe_z_loss_coeff = None
+        mock_config.model.moe_per_layer_logging = True
+        mock_config.model.num_layers = len(pattern)
+        mock_config.model.moe_layer_freq = 1
+        mock_config.model.mtp_num_layers = None
+        mock_config.model.is_hybrid_model = True
+        mock_config.model.hybrid_layer_pattern = pattern
+
+        training_log(
+            loss_dict=loss_dict,
+            total_loss_dict=total_loss_dict,
+            learning_rate=1e-4,
+            decoupled_learning_rate=None,
+            loss_scale=1024.0,
+            report_memory_flag=False,
+            skipped_iter=0,
+            grad_norm=2.5,
+            params_norm=15.2,
+            num_zeros_in_grad=0,
+            config=mock_config,
+            global_state=mock_global_state,
+            history_wct=None,
+            model=None,
+        )
+
+        mock_track_moe.assert_called_once()
+        call_args = mock_track_moe.call_args
+        # Storage size stays at full depth so PP tensor sizes match (deadlock fix).
+        assert call_args.kwargs["num_layers"] == len(pattern)
+        # Averaging denominator is the per-layer MoE mask, summing to only the "E"s.
+        moe_layer_freq = call_args.kwargs["moe_layer_freq"]
+        assert moe_layer_freq == [1 if c == "E" else 0 for c in pattern]
+        assert sum(moe_layer_freq) == pattern.count("E") == 2
+
     @mock.patch("megatron.bridge.training.utils.train_utils.get_num_microbatches")
     @mock.patch("megatron.bridge.training.utils.train_utils.reduce_max_stat_across_model_parallel_group")
     @mock.patch("megatron.bridge.training.utils.train_utils.get_world_size_safe")