nemo_rl.models.megatron.common#

Module Contents#

Functions#

_round_up_to_multiple

broadcast_tensor

Broadcasts a tensor from src_rank to all ranks in the group using broadcast_object_list for metadata.

get_aux_loss_track_names

Returns the aux-loss tracker names the router records for a model config.

get_moe_metrics

Returns Mixture of Experts (MoE) auxiliary-loss metrics.

get_mtp_metrics

Returns Multi-Token Prediction (MTP) loss and acceptance rate metrics.

Data#

_AUX_LOSS_TRACK_NAMES

API#

nemo_rl.models.megatron.common._round_up_to_multiple(value: int, multiple: int) → int#
nemo_rl.models.megatron.common.broadcast_tensor(
tensor: torch.Tensor | None,
src_rank: int,
group: torch.distributed.ProcessGroup,
) → torch.Tensor#

Broadcasts a tensor from src_rank to all ranks in the group using broadcast_object_list for metadata.

Handles the case where the input tensor might be None on non-source ranks. If the input tensor is provided on non-source ranks, it must have the correct shape and dtype matching the tensor on the source rank.

Parameters:

  • tensor – The tensor to broadcast on the source rank. Can be None on non-source ranks (will be created with correct shape/dtype). If not None on non-source ranks, it’s used as the buffer for the broadcast and must match the source tensor’s metadata.

  • src_rank (int) – The global rank of the source process.

  • group – The process group for communication.

Returns:

The broadcasted tensor. On non-source ranks, this will be the tensor received from the source.

Return type:

torch.Tensor

Raises:

  • ValueError – If the tensor is None on the source rank, or if a tensor provided on a non-source rank has mismatched shape/dtype/device.

  • TypeError – If broadcasting metadata fails (e.g., due to pickling issues).

nemo_rl.models.megatron.common._AUX_LOSS_TRACK_NAMES: dict[str, str]#

None

nemo_rl.models.megatron.common.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.

Parameters:

model_config – Megatron TransformerConfig (or any object exposing moe_router_load_balancing_type / moe_aux_loss_coeff).

Returns:

Aux-loss tracker names to pre-initialize, in the order Megatron records them. Empty when no aux loss is enabled.

Return type:

list[str]

nemo_rl.models.megatron.common.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.

This function reduces MoE auxiliary losses across ranks, aggregates them, and returns a dictionary of metrics.

Parameters:

  • 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:

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}”.

Return type:

dict[str, Any]

.. 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).

nemo_rl.models.megatron.common.get_mtp_metrics(loss_scale: float = 1.0) → dict[str, Any]#

Returns Multi-Token Prediction (MTP) loss and acceptance rate metrics.

This function reduces MTP metrics across ranks and returns a dictionary of metrics.

Parameters:

loss_scale – Scale factor applied to each MTP layer’s loss (e.g., 1/num_microbatches). MTPLossLoggingHelper accumulates the per-microbatch loss across microbatches without dividing, so callers must pass 1/num_microbatches to recover the mean (mirroring get_moe_metrics). Acceptance rate is a ratio of counts and is not scaled. Defaults to 1.0.

Returns:

A flat dict of metrics. Each MTP layer’s loss is returned under the key “mtp_{i}loss” and acceptance rate under “mtp{i}_acceptance_rate” where i is 1-indexed (matching Megatron-LM).

Return type:

dict[str, Any]