`nemo_rl.models.megatron.pipeline_parallel`#

Pipeline parallel utilities for Megatron models.

Module Contents#

Functions#

`broadcast_obj_from_pp_rank`	Broadcast an object across pipeline parallel ranks.
`_materialize_loss_metrics_for_object_broadcast`	Convert tensor metrics to host values before object transport.
`broadcast_loss_metrics_from_last_stage`	Broadcast loss metrics from the last pipeline stage to all stages.
`broadcast_tensors_from_last_stage`	Broadcast multiple tensors from the last pipeline stage to all stages.

Data#

`LossMetricPayloadT`
`LossMetricPayload`

API#

nemo_rl.models.megatron.pipeline_parallel.LossMetricPayloadT#: ‘TypeVar(…)’

nemo_rl.models.megatron.pipeline_parallel.LossMetricPayload: TypeAlias#: None

nemo_rl.models.megatron.pipeline_parallel.broadcast_obj_from_pp_rank(obj: Any) → Any#

Broadcast an object across pipeline parallel ranks.

This utility function handles broadcasting an object from the rank that owns it to all other pipeline parallel ranks. If only one rank has the object (non-None), it will be broadcast to all other ranks.

Parameters:: obj – The object to broadcast. Can be None on ranks that don’t own it.
Returns:: The object on all ranks (either the original or the broadcast copy).
Raises:: ValueError – If the object doesn’t exist on any pipeline parallel rank.

nemo_rl.models.megatron.pipeline_parallel._materialize_loss_metrics_for_object_broadcast( loss_metrics: nemo_rl.models.megatron.pipeline_parallel.LossMetricPayload, ) → nemo_rl.models.megatron.pipeline_parallel.LossMetricPayload#: Convert tensor metrics to host values before object transport.

nemo_rl.models.megatron.pipeline_parallel.broadcast_loss_metrics_from_last_stage( loss_metrics: Optional[nemo_rl.models.megatron.pipeline_parallel.LossMetricPayloadT] = None, ) → nemo_rl.models.megatron.pipeline_parallel.LossMetricPayloadT#

Broadcast loss metrics from the last pipeline stage to all stages.

Loss is computed only on the last pipeline stage, so every other stage needs the result handed to it before it can report metrics.

Parameters:

loss_metrics – The payload on the last pipeline stage, None on every other stage. Either a dict of metrics or a list of such dicts; the same shape comes back on every rank.

Returns:

The payload, on every rank. On the last stage this is a materialized copy (tensors moved to host) rather than the object that was passed in, because it is what actually went over the wire.

Raises:

ValueError – The last pipeline stage passed loss_metrics=None.
RuntimeError – The broadcast completed but delivered no payload.

nemo_rl.models.megatron.pipeline_parallel.broadcast_tensors_from_last_stage( tensors: dict[str, Optional[torch.Tensor]], ) → dict[str, torch.Tensor]#

Broadcast multiple tensors from the last pipeline stage to all stages.