bridge.training.utils.pg_utils

Module Contents

Functions

get_pg_collection

Return the ProcessGroupCollection from a model or list of model chunks.

API

bridge.training.utils.pg_utils.get_pg_collection(

model: Union[megatron.core.transformer.MegatronModule, list[megatron.core.transformer.MegatronModule]],

) → megatron.core.process_groups_config.ProcessGroupCollection

Return the ProcessGroupCollection from a model or list of model chunks.

This mirrors the style of utility accessors like get_model_config, but for retrieving the communication process group collection from the model wrapper.

Parameters:

model – A MegatronModule or a list of MegatronModule chunks.

Returns:

The model’s process group collection.

Return type:

ProcessGroupCollection