diff --git a/nemo_rl/models/megatron/common.py b/nemo_rl/models/megatron/common.py index eb43051b3b5..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,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: diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index fc913dbb8ce..032b9b32893 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/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 diff --git a/nemo_rl/models/value/workers/megatron_value_worker.py b/nemo_rl/models/value/workers/megatron_value_worker.py index a4bf63b41d1..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 ( @@ -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 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()