From acbb839945e280685f1bfe229fcce448a36d7281 Mon Sep 17 00:00:00 2001 From: dafu-wu Date: Tue, 28 Jul 2026 04:44:58 +0000 Subject: [PATCH 1/2] fix: prevent PP hang in MoE aux-loss metric reduction When moe_router_load_balancing_type=aux_loss is enabled with pipeline_model_parallel_size > 1, training could silently hang inside get_moe_metrics(). Root cause: reduce_aux_losses_tracker_across_ranks() runs torch.distributed.all_reduce over the pipeline-parallel group for each name present in the *local* MoE logging tracker. Megatron creates the tracker entry lazily (save_to_aux_losses_tracker only allocates torch.zeros(num_layers) the first time a rank saves a loss). If some PP rank did not save 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), that rank skips the all_reduce for the name while other PP ranks perform it, so the collective mismatches participants and hangs with no traceback. Fix: mirror Megatron's own track_moe_metrics(force_initialize=True) guard. Before the reduction, pre-initialize the tracker on every rank so each tracked name exists with an equally-sized zero tensor of length (num_layers + mtp_num_layers), matching the size the router uses in save_to_aux_losses_tracker. num_layers/mtp_num_layers are threaded from model_config in both the policy and value Megatron workers. When num_layers is None the behaviour is unchanged, so this is a no-op for all existing call sites. Signed-off-by: dafu-wu --- nemo_rl/models/megatron/common.py | 36 +++++++++++++++++++ .../policy/workers/megatron_policy_worker.py | 5 +++ .../value/workers/megatron_value_worker.py | 5 +++ 3 files changed, 46 insertions(+) diff --git a/nemo_rl/models/megatron/common.py b/nemo_rl/models/megatron/common.py index eb43051b3b5..cce65e2d421 100644 --- a/nemo_rl/models/megatron/common.py +++ b/nemo_rl/models/megatron/common.py @@ -125,6 +125,9 @@ 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. @@ -141,7 +144,40 @@ def get_moe_metrics( 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 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 save an + aux loss this step (e.g. 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() runs torch.distributed.all_reduce over + # the pipeline-parallel group for each name present in the *local* tracker. The + # tracker entry is created lazily (Megatron save_to_aux_losses_tracker only + # allocates torch.zeros(num_layers) the first time a rank saves a loss). If any PP + # rank did not save 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. + # + # Mirror Megatron's own track_moe_metrics(force_initialize=True) guard: allocate a + # zero tensor of size (num_layers + mtp_num_layers) for each tracked name on every + # rank, matching the size the router uses in save_to_aux_losses_tracker. + if num_layers is not None: + if track_names is None: + track_names = ["load_balancing_loss"] + tracker_num_layers = num_layers + (mtp_num_layers or 0) + tracker = get_moe_layer_wise_logging_tracker() + for name in track_names: + if name not in tracker: + tracker[name] = { + "values": torch.zeros(tracker_num_layers, device="cuda"), + "reduce_group": None, + "avg_group": None, + "reduce_group_has_dp": False, + } + reduce_aux_losses_tracker_across_ranks() tracker = get_moe_layer_wise_logging_tracker() diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index b6ae9a23a95..1d0b62e2526 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -1052,6 +1052,11 @@ 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 saved no aux loss this step (e.g. MTP MoE on the last stage). + num_layers=getattr(model_config, "num_layers", None), + mtp_num_layers=getattr(model_config, "mtp_num_layers", None), ) if moe_metrics: metrics["moe_metrics"] = moe_metrics diff --git a/nemo_rl/models/value/workers/megatron_value_worker.py b/nemo_rl/models/value/workers/megatron_value_worker.py index a4bf63b41d1..12308a82c29 100644 --- a/nemo_rl/models/value/workers/megatron_value_worker.py +++ b/nemo_rl/models/value/workers/megatron_value_worker.py @@ -635,6 +635,11 @@ 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 saved no aux loss this step (e.g. MTP MoE on the last stage). + num_layers=getattr(model_config, "num_layers", None), + mtp_num_layers=getattr(model_config, "mtp_num_layers", None), ) if moe_metrics: metrics["moe_metrics"] = moe_metrics From 18572530cc8495bbb61c2e9df470d59e49106eef Mon Sep 17 00:00:00 2001 From: dafu-wu Date: Tue, 11 Aug 2026 01:14:50 +0000 Subject: [PATCH 2/2] fix: port MoE aux-loss PP-hang guard to the current mcore tracker API Addresses review feedback on the aux-loss pre-initialization added to get_moe_metrics(). The previous approach did not take effect on the pinned megatron-core, and would have introduced a zero-valued metric once it did. Reach the live tracker. get_moe_layer_wise_logging_tracker() is a deprecated shim that rebuilds a dict copy on every call, so writing the pre-initialized entry through it was discarded and reduce_aux_losses_tracker_across_ranks() still skipped the all_reduce -- the hang was not actually prevented. Use MoEMetricsTracker.ensure_initialized() instead, which is what Megatron's own report(force_initialize=True) calls. It also picks the device itself, removing the hardcoded device="cuda", and drops the stale reduce_group_has_dp field (renamed to needs_dp_avg in the MoE logging refactor). Derive track_names from the model config. The router records a distinct name per balancing type (load_balancing_loss, seq_load_balancing_loss, global_load_balancing_loss, z_loss), each driving its own all_reduce, and moe_router_load_balancing_type may be a list. Hardcoding "load_balancing_loss" left seq_aux_loss and global_aux_loss users hanging. get_aux_loss_track_names() matches the balancing type exactly and requires a non-zero coefficient, which mirrors MoETopKRouter.get_aux_loss_coeff/is_aux_loss_enabled: a substring test would treat seq_aux_loss as aux_loss, and configs that name a balancing type without setting moe_aux_loss_coeff leave the router recording nothing. Gate pre-initialization on track_names being non-empty. Workers only check num_moe_experts > 1, not whether load balancing is on. With the default moe_router_load_balancing_type: "none" the tracker stays empty today and the metric is dropped by the falsy guard; an effective pre-init would have started reporting a permanently-zero load_balancing_loss for every MoE config in the repo. Deriving track_names yields an empty list under "none", so nothing is pre-initialized. Also pass the pre-init arguments at the third call site, _finish_train_step_body (single-controller / split-API path), which was missed and has no MoE coverage in the nightly suite; document the new parameters in the Args block; and rename the loop-local num_layers to num_tracked_layers so it no longer shadows the parameter. Tests assert against the real get_moe_metrics_tracker().metrics rather than stubbing the deprecated accessor -- a stubbed plain dict is live, so such a test would pass while production stayed a no-op. Signed-off-by: dafu-wu --- nemo_rl/models/megatron/common.py | 132 ++++++++++++--- .../policy/workers/megatron_policy_worker.py | 16 +- .../value/workers/megatron_value_worker.py | 5 +- .../unit/models/megatron/test_moe_metrics.py | 156 ++++++++++++++++++ 4 files changed, 280 insertions(+), 29 deletions(-) diff --git a/nemo_rl/models/megatron/common.py b/nemo_rl/models/megatron/common.py index cce65e2d421..26991b2aeb2 100644 --- a/nemo_rl/models/megatron/common.py +++ b/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,6 +122,78 @@ 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, @@ -138,6 +211,17 @@ 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, @@ -146,37 +230,33 @@ def get_moe_metrics( form "moe/{name}_layer_{i}". Note: - num_layers/mtp_num_layers 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 save an - aux loss this step (e.g. an MTP MoE layer that lives only on the last stage). + 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() runs torch.distributed.all_reduce over - # the pipeline-parallel group for each name present in the *local* tracker. The - # tracker entry is created lazily (Megatron save_to_aux_losses_tracker only - # allocates torch.zeros(num_layers) the first time a rank saves a loss). If any PP - # rank did not save 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. + # 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 track_moe_metrics(force_initialize=True) guard: allocate a - # zero tensor of size (num_layers + mtp_num_layers) for each tracked name on every - # rank, matching the size the router uses in save_to_aux_losses_tracker. - if num_layers is not None: - if track_names is None: - track_names = ["load_balancing_loss"] + # 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) - tracker = get_moe_layer_wise_logging_tracker() + # 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: - if name not in tracker: - tracker[name] = { - "values": torch.zeros(tracker_num_layers, device="cuda"), - "reduce_group": None, - "avg_group": None, - "reduce_group_has_dp": False, - } + mcore_tracker.ensure_initialized(name, tracker_num_layers) reduce_aux_losses_tracker_across_ranks() tracker = get_moe_layer_wise_logging_tracker() @@ -186,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: diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 1d0b62e2526..98e55dad215 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -59,7 +59,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, @@ -1054,9 +1057,11 @@ def train( 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 saved no aux loss this step (e.g. MTP MoE on the last stage). + # 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 @@ -1602,6 +1607,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 diff --git a/nemo_rl/models/value/workers/megatron_value_worker.py b/nemo_rl/models/value/workers/megatron_value_worker.py index 12308a82c29..637a7c9f8be 100644 --- a/nemo_rl/models/value/workers/megatron_value_worker.py +++ b/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 ( @@ -637,9 +638,11 @@ def train( ), # 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 saved no aux loss this step (e.g. MTP MoE on the last stage). + # 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 diff --git a/tests/unit/models/megatron/test_moe_metrics.py b/tests/unit/models/megatron/test_moe_metrics.py index 6a8c2ea7fab..0d76cde0f47 100644 --- a/tests/unit/models/megatron/test_moe_metrics.py +++ b/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()