fix: prevent PP hang in MoE aux-loss metric reduction

nemo_rl/models/megatron/common.py

@@ -16,6 +16,7 @@
 
 import torch
 import torch.distributed as dist
+from megatron.core.transformer.moe.moe_logging import get_moe_metrics_tracker
 from megatron.core.transformer.moe.moe_utils import (
     clear_aux_losses_tracker,
     get_moe_layer_wise_logging_tracker,
@@ -121,10 +122,85 @@ def broadcast_tensor(
     return tensor
 
 
+#: Mapping from a Megatron ``moe_router_load_balancing_type`` value to the aux-loss
+#: name the router records for it. Mirrors ``MoETopKRouter.is_aux_loss_enabled``, which
+#: treats these three types as first-class and drives one all_reduce per recorded name.
+_AUX_LOSS_TRACK_NAMES: dict[str, str] = {
+    "aux_loss": "load_balancing_loss",
+    "seq_aux_loss": "seq_load_balancing_loss",
+    "global_aux_loss": "global_load_balancing_loss",
+}
+
+
+def get_aux_loss_track_names(model_config: Any) -> list[str]:
+    """Returns the aux-loss tracker names the router records for a model config.
+
+    Megatron's router only records an aux loss when its balancing type is configured
+    *and* the matching coefficient is non-zero (``MoETopKRouter.get_aux_loss_coeff``
+    returns 0.0 otherwise, and ``_apply_aux_loss`` returns early). Deriving the names
+    the same way keeps the pre-initialization in ``get_moe_metrics`` aligned with what
+    the router actually tracks, so no permanently-zero metric is reported for models
+    that have load balancing disabled (e.g. ``moe_router_load_balancing_type: "none"``).
+
+    ``moe_router_load_balancing_type`` may be a single string or a list, in which case
+    ``moe_aux_loss_coeff`` is a list of the same length (validated by Megatron's
+    ``TransformerConfig``), so more than one aux loss can be live at once.
+
+    Args:
+        model_config: Megatron ``TransformerConfig`` (or any object exposing
+            ``moe_router_load_balancing_type`` / ``moe_aux_loss_coeff``).
+
+    Returns:
+        list[str]: Aux-loss tracker names to pre-initialize, in the order Megatron
+        records them. Empty when no aux loss is enabled.
+    """
+    routing_type = getattr(model_config, "moe_router_load_balancing_type", None)
+    aux_loss_coeff = getattr(model_config, "moe_aux_loss_coeff", 0.0)
+
+    routing_types: list[Any] = []
+    coeffs: list[Any] = []
+    if isinstance(routing_type, str):
+        routing_types = [routing_type]
+        # A single balancing type pairs with a scalar coefficient.
+        coeffs = [aux_loss_coeff]
+    elif isinstance(routing_type, (list, tuple)):
+        routing_types = list(routing_type)
+        if isinstance(aux_loss_coeff, (list, tuple)):
+            coeffs = list(aux_loss_coeff)
+        else:
+            # Defensive: Megatron validates that the lists have matching lengths, but
+            # tolerate a scalar coefficient rather than raising while collecting metrics.
+            coeffs = [aux_loss_coeff] * len(routing_types)
+
+    track_names: list[str] = []
+    for index, single_routing_type in enumerate(routing_types):
+        if not isinstance(single_routing_type, str):
+            continue
+        name = _AUX_LOSS_TRACK_NAMES.get(single_routing_type)
+        if name is None:
+            continue
+        coeff = coeffs[index] if index < len(coeffs) else 0.0
+        # Only a real number can enable the loss; anything else (None, or a stand-in
+        # object from a partially-populated config) means "not configured".
+        if isinstance(coeff, (int, float)) and coeff > 0 and name not in track_names:
+            track_names.append(name)
+
+    # z_loss is recorded independently of the load balancing type, gated only on its
+    # own coefficient being set (see MoETopKRouter.apply_z_loss).
+    z_loss_coeff = getattr(model_config, "moe_z_loss_coeff", None)
+    if isinstance(z_loss_coeff, (int, float)):
+        track_names.append("z_loss")
+
+    return track_names
+
+
 def get_moe_metrics(
     loss_scale: float,
     total_loss_dict: Optional[dict] = None,
     per_layer_logging: bool = False,
+    num_layers: Optional[int] = None,
+    mtp_num_layers: Optional[int] = None,
+    track_names: Optional[list[str]] = None,
 ) -> dict[str, Any]:
     """Returns Mixture of Experts (MoE) auxiliary-loss metrics.
 
@@ -135,13 +211,53 @@ def get_moe_metrics(
         loss_scale: Scale factor to apply to each auxiliary loss (e.g., 1/num_microbatches).
         total_loss_dict: If provided, accumulate means into this dict (by name).
         per_layer_logging: If True, include per-layer values in the returned dict.
+        num_layers: Total number of transformer layers. When provided together with a
+            non-empty ``track_names``, the aux-loss tracker is pre-initialized on every
+            rank before the reduction (see Note). Defaults to None, which disables
+            pre-initialization.
+        mtp_num_layers: Extra layers contributed by Multi-Token Prediction, added to
+            ``num_layers`` to size the pre-initialized tensor, matching the size the
+            router uses when recording. Defaults to None (treated as 0).
+        track_names: Aux-loss names to pre-initialize; must mirror what the router
+            records for the configured ``moe_router_load_balancing_type``, so callers
+            should derive it via ``get_aux_loss_track_names(model_config)``. Defaults to
+            None, which disables pre-initialization.
 
     Returns:
         dict[str, Any]: A flat dict of aggregated metrics. For each aux loss name,
         the mean value is returned under the same key (e.g., "load_balancing_loss").
         If per_layer_logging is True, per-layer values are returned under keys of the
         form "moe/{name}_layer_{i}".
+
+    Note:
+        num_layers/mtp_num_layers/track_names pre-initialize the aux-loss tracker so
+        every pipeline-parallel rank participates in the collective all_reduce below
+        with an equally-sized tensor, preventing a hang when some PP rank did not
+        record an aux loss this step (e.g. a stage with no MoE layer, or an MTP MoE
+        layer that lives only on the last stage).
     """
+    # Pre-initialize the aux-loss tracker so every PP rank has the same set of
+    # named, equally-sized tensors BEFORE the collective all_reduce below.
+    #
+    # reduce_aux_losses_tracker_across_ranks() all_reduces over the pipeline-parallel
+    # group for each name present in the *local* tracker. Tracker entries are created
+    # lazily (MoEMetricsTracker.record only allocates torch.zeros(num_layers) the first
+    # time a rank records a loss), and _sync_metrics skips names it does not have. If
+    # any PP rank did not record an aux loss this step, it skips the all_reduce for that
+    # name while other PP ranks perform it -> the collective mismatches participants and
+    # hangs with no traceback.
+    #
+    # Mirror Megatron's own report(force_initialize=True) guard, which calls
+    # ensure_initialized() for the same reason. Sizing must be
+    # (num_layers + mtp_num_layers) to match what the router passes to record().
+    if num_layers is not None and track_names:
+        tracker_num_layers = num_layers + (mtp_num_layers or 0)
+        # Use the live tracker rather than get_moe_layer_wise_logging_tracker(), which is
+        # a deprecated shim returning a fresh dict copy -- writes to it are discarded.
+        mcore_tracker = get_moe_metrics_tracker()
+        for name in track_names:
+            mcore_tracker.ensure_initialized(name, tracker_num_layers)
+
     reduce_aux_losses_tracker_across_ranks()
     tracker = get_moe_layer_wise_logging_tracker()
 
@@ -150,8 +266,8 @@ def get_moe_metrics(
         aux_losses = {k: v["values"].float() * loss_scale for k, v in tracker.items()}
         for name, loss_list in aux_losses.items():
             # Megatron-LM aggregates aux losses across layers and normalizes by number of MoE layers
-            num_layers = int(loss_list.numel()) if loss_list.numel() > 0 else 1
-            aggregated_value = loss_list.sum() / num_layers
+            num_tracked_layers = int(loss_list.numel()) if loss_list.numel() > 0 else 1
+            aggregated_value = loss_list.sum() / num_tracked_layers
             metrics[name] = float(aggregated_value.item())
             if total_loss_dict is not None:
                 if name not in total_loss_dict:

nemo_rl/models/policy/workers/megatron_policy_worker.py

@@ -64,7 +64,10 @@
     MegatronGenerationRefitMixin,
 )
 from nemo_rl.models.generation.vllm.config import VllmConfig
-from nemo_rl.models.megatron.common import get_moe_metrics
+from nemo_rl.models.megatron.common import (
+    get_aux_loss_track_names,
+    get_moe_metrics,
+)
 from nemo_rl.models.megatron.data import (
     get_microbatch_iterator,
     process_global_batch,
@@ -1077,6 +1080,13 @@ def train(
             moe_metrics = get_moe_metrics(
                 loss_scale=moe_loss_scale,
                 per_layer_logging=self.cfg["megatron_cfg"]["moe_per_layer_logging"],
+                # Pre-initialize the aux-loss tracker on every PP rank so the
+                # cross-PP all_reduce inside get_moe_metrics does not hang when a
+                # rank recorded no aux loss this step (e.g. a stage with no MoE
+                # layer, or MTP MoE on the last stage).
+                num_layers=getattr(model_config, "num_layers", None),
+                mtp_num_layers=getattr(model_config, "mtp_num_layers", None),
+                track_names=get_aux_loss_track_names(model_config),
             )
             if moe_metrics:
                 metrics["moe_metrics"] = moe_metrics
@@ -1772,6 +1782,13 @@ def _scale_metric(name: str, value: Any) -> Any:
             moe_metrics = get_moe_metrics(
                 loss_scale=moe_loss_scale,
                 per_layer_logging=self.cfg["megatron_cfg"]["moe_per_layer_logging"],
+                # Pre-initialize the aux-loss tracker on every PP rank so the
+                # cross-PP all_reduce inside get_moe_metrics does not hang when a
+                # rank recorded no aux loss this step (e.g. a stage with no MoE
+                # layer, or MTP MoE on the last stage).
+                num_layers=getattr(model_config, "num_layers", None),
+                mtp_num_layers=getattr(model_config, "mtp_num_layers", None),
+                track_names=get_aux_loss_track_names(model_config),
             )
             if moe_metrics:
                 metrics["moe_metrics"] = moe_metrics

nemo_rl/models/value/workers/megatron_value_worker.py

@@ -57,6 +57,7 @@
 from nemo_rl.distributed.named_sharding import NamedSharding
 from nemo_rl.models.megatron.common import (
     broadcast_tensor,
+    get_aux_loss_track_names,
     get_moe_metrics,
 )
 from nemo_rl.models.megatron.data import (
@@ -635,6 +636,13 @@ def train(
                 per_layer_logging=self.cfg["megatron_cfg"].get(
                     "moe_per_layer_logging", False
                 ),
+                # Pre-initialize the aux-loss tracker on every PP rank so the
+                # cross-PP all_reduce inside get_moe_metrics does not hang when a
+                # rank recorded no aux loss this step (e.g. a stage with no MoE
+                # layer, or MTP MoE on the last stage).
+                num_layers=getattr(model_config, "num_layers", None),
+                mtp_num_layers=getattr(model_config, "mtp_num_layers", None),
+                track_names=get_aux_loss_track_names(model_config),
             )
             if moe_metrics:
                 metrics["moe_metrics"] = moe_metrics

tests/unit/models/megatron/test_moe_metrics.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from types import SimpleNamespace
 from typing import Any, Dict
 
 import pytest
@@ -117,3 +118,158 @@ def _clear():
     assert metrics["moe/z_loss_layer_1"] == pytest.approx(2.0)
 
     assert cleared["called"], "clear_aux_losses_tracker should be called"
+
+
+@pytest.mark.mcore
+@pytest.mark.parametrize(
+    "routing_type,aux_loss_coeff,z_loss_coeff,expected",
+    [
+        # Load balancing disabled: nothing to pre-initialize, so no permanently-zero
+        # metric is reported for the many MoE configs that use "none".
+        ("none", 0.0, None, []),
+        # Configured type but zero coefficient: the router returns before recording.
+        ("aux_loss", 0.0, None, []),
+        ("aux_loss", 1e-2, None, ["load_balancing_loss"]),
+        # seq_aux_loss must not be reported as "load_balancing_loss": the router
+        # records a distinct name, each driving its own all_reduce.
+        ("seq_aux_loss", 1e-2, None, ["seq_load_balancing_loss"]),
+        ("global_aux_loss", 1e-2, None, ["global_load_balancing_loss"]),
+        # A list of balancing types pairs with a list of coefficients; only the
+        # entries with a non-zero coefficient are live.
+        (
+            ["aux_loss", "seq_aux_loss"],
+            [1e-2, 1e-2],
+            None,
+            ["load_balancing_loss", "seq_load_balancing_loss"],
+        ),
+        (["aux_loss", "seq_aux_loss"], [0.0, 1e-2], None, ["seq_load_balancing_loss"]),
+        # z_loss is tracked independently of the balancing type.
+        ("none", 0.0, 1e-3, ["z_loss"]),
+        ("aux_loss", 1e-2, 1e-3, ["load_balancing_loss", "z_loss"]),
+        # Unknown/unsupported balancing types record no aux loss.
+        ("sinkhorn", 1e-2, None, []),
+    ],
+)
+def test_get_aux_loss_track_names(routing_type, aux_loss_coeff, z_loss_coeff, expected):
+    """track_names must mirror exactly the aux losses the router records."""
+    from nemo_rl.models.megatron.common import get_aux_loss_track_names
+
+    model_config = SimpleNamespace(
+        moe_router_load_balancing_type=routing_type,
+        moe_aux_loss_coeff=aux_loss_coeff,
+        moe_z_loss_coeff=z_loss_coeff,
+    )
+
+    assert get_aux_loss_track_names(model_config) == expected
+
+
+@pytest.mark.mcore
+def test_get_aux_loss_track_names_tolerates_missing_attrs():
+    """A config without MoE attributes must yield no names rather than raise."""
+    from nemo_rl.models.megatron.common import get_aux_loss_track_names
+
+    assert get_aux_loss_track_names(SimpleNamespace()) == []
+    # z_loss is gated only on its own coefficient, independent of the balancing type.
+    assert get_aux_loss_track_names(SimpleNamespace(moe_z_loss_coeff=1e-3)) == [
+        "z_loss"
+    ]
+
+
+@pytest.mark.mcore
+def test_get_aux_loss_track_names_ignores_non_numeric_coeffs():
+    """Non-numeric placeholders must not enable a loss (and must not raise).
+
+    Worker tests build ``model.config`` as a ``MagicMock``, whose attribute access
+    returns a truthy mock rather than a number. Comparing that against 0 raises
+    ``TypeError``, so the coefficient checks are type-guarded.
+    """
+    from unittest.mock import MagicMock
+
+    from nemo_rl.models.megatron.common import get_aux_loss_track_names
+
+    mock_config = MagicMock()
+    mock_config.num_moe_experts = 4
+    assert get_aux_loss_track_names(mock_config) == []
+
+    # A real balancing type paired with a non-numeric coefficient is still "off".
+    partial_config = MagicMock()
+    partial_config.moe_router_load_balancing_type = "aux_loss"
+    assert get_aux_loss_track_names(partial_config) == []
+
+
+@pytest.mark.mcore
+def test_preinit_registers_entry_on_live_tracker(monkeypatch):
+    """A PP rank that recorded no aux loss must still enter the collective.
+
+    This asserts against Megatron's real tracker rather than a stubbed dict:
+    ``get_moe_layer_wise_logging_tracker()`` is a deprecated shim that returns a fresh
+    dict copy, so writing the pre-initialized entry through it would be silently lost.
+    """
+    from megatron.core.transformer.moe.moe_logging import (
+        destroy_moe_metrics_tracker,
+        get_moe_metrics_tracker,
+    )
+
+    from nemo_rl.models import megatron as megatron_module
+    from nemo_rl.models.megatron.common import get_moe_metrics
+
+    destroy_moe_metrics_tracker()
+    try:
+        monkeypatch.setattr(
+            megatron_module.common,  # type: ignore[attr-defined]
+            "reduce_aux_losses_tracker_across_ranks",
+            lambda *args, **kwargs: None,
+        )
+
+        get_moe_metrics(
+            loss_scale=1.0,
+            num_layers=4,
+            mtp_num_layers=1,
+            track_names=["load_balancing_loss"],
+        )
+
+        live = get_moe_metrics_tracker().metrics
+        assert set(live) == {"load_balancing_loss"}
+        # Size must be num_layers + mtp_num_layers to match what the router records.
+        assert live["load_balancing_loss"].values.numel() == 5
+    finally:
+        destroy_moe_metrics_tracker()
+
+
+@pytest.mark.mcore
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        # No pre-initialization requested at all (existing callers / dense models).
+        {},
+        # num_layers known but no aux loss is live (e.g. balancing type "none"):
+        # pre-initializing would report a permanently-zero metric.
+        {"num_layers": 4, "track_names": []},
+        # track_names known but num_layers unavailable: cannot size the tensor.
+        {"track_names": ["load_balancing_loss"]},
+    ],
+)
+def test_no_preinit_leaves_tracker_empty(monkeypatch, kwargs):
+    """Without both num_layers and track_names, no tracker entry may be created."""
+    from megatron.core.transformer.moe.moe_logging import (
+        destroy_moe_metrics_tracker,
+        get_moe_metrics_tracker,
+    )
+
+    from nemo_rl.models import megatron as megatron_module
+    from nemo_rl.models.megatron.common import get_moe_metrics
+
+    destroy_moe_metrics_tracker()
+    try:
+        monkeypatch.setattr(
+            megatron_module.common,  # type: ignore[attr-defined]
+            "reduce_aux_losses_tracker_across_ranks",
+            lambda *args, **kwargs: None,
+        )
+
+        metrics = get_moe_metrics(loss_scale=1.0, **kwargs)
+
+        assert metrics == {}
+        assert get_moe_metrics_tracker().metrics == {}
+    finally:
+        destroy_moe_metrics_tracker()