Per-layer expert count + MoE aux-loss tracker sizing (the NCCL-deadlock fix)

Author: chochowski

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

@@ -290,8 +290,7 @@ 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 = config.num_moe_experts
-        num_experts_per_rank = num_experts // ep_group.size()
+        num_experts_per_rank = layer_module.mlp.num_local_experts   # per-layer, heterogeneous-safe
 
         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

@@ -1041,10 +1041,11 @@ def training_log(
         if getattr(config.model, "moe_z_loss_coeff", None) is not None:
             track_names.append("z_loss")
 
-        if getattr(config.model, "is_hybrid_model", False):
-            layers = getattr(config.model, "hybrid_layer_pattern", "").count("E")
-        else:
-            layers = getattr(config.model, "num_layers", None)
+        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
+            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).
         # No-op when neither logger is configured: the original writer is returned as-is.
@@ -1058,8 +1059,8 @@ def training_log(
             per_layer_logging=getattr(config.model, "moe_per_layer_logging", False),
             force_initialize=True,
             track_names=track_names,
-            num_layers=layers,
-            moe_layer_freq=getattr(config.model, "moe_layer_freq", None),
+            num_layers=num_layers,
+            moe_layer_freq=moe_layer_freq,
             mtp_num_layers=getattr(config.model, "mtp_num_layers", None),
             pg_collection=pg_collection,
         )