bridge.training.utils.pg_utils

Module Contents

Classes

DistTrainProcessGroupCollection

Process group collection for dist train.

Functions

get_pg_collection	Return the ProcessGroupCollection from a model or list of model chunks.
pipeline_stage_roles	Return (is_first, is_last) for this rank’s model chunk.

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

bridge.training.utils.pg_utils.pipeline_stage_roles(

model: megatron.core.transformer.MegatronModule,

) → tuple[bool, bool]

Return (is_first, is_last) for this rank’s model chunk.

Both are True without pipeline parallelism, so callers can fall through to the ordinary single-stage path. Virtual pipeline chunks count as first or last only when they are the edge chunk of an edge stage.

Parameters:

model – One model chunk, possibly wrapped by DDP or a precision wrapper.

Returns:

Whether the chunk is the first and whether it is the last stage of the pipeline.

class bridge.training.utils.pg_utils.DistTrainProcessGroupCollection(

pg_collection: megatron.core.process_groups_config.ProcessGroupCollection,

language_model_module_name: Optional[str] = None,

)

Bases: megatron.core.process_groups_config.ProcessGroupCollection

Process group collection for dist train.

Initialization

Initialize the dist train process group collection.

Parameters:

• pg_collection – The process group collection.

• language_model_module_name – The name of the language model module.

get_language_model_cp_size() → int

Get context parallel size for the language model.

Returns:

Context parallel size for the language model.

Raises:

ValueError – If no language model is specified for this collection.

has_language_model() → bool

Check if this rank has a language model.

Returns:

True if this rank has a language model, False otherwise.