`bridge.training.gtp`#

Generalized Tensor Parallelism helpers for the standard Bridge runtime.

Module Contents#

Functions#

`get_transformer_config`	Return the MCore transformer config nested in a Bridge model config.
`is_gtp_remat_active`	Return whether dense or expert GTP weight rematerialization is enabled.
`configure_gtp_remat`	Configure process-global GTP state before constructing model modules.
`classify_gtp_remat_chains`	Classify all model chunks after distributed wrapping and before first forward.
`get_data_distribution_group`	Return the group spanning every rank that consumes distinct input data.

API#

bridge.training.gtp.get_transformer_config(model_config: Any) → Any#: Return the MCore transformer config nested in a Bridge model config.

bridge.training.gtp.is_gtp_remat_active(model_config: Any) → bool#: Return whether dense or expert GTP weight rematerialization is enabled.

bridge.training.gtp.configure_gtp_remat(model_config: Any) → None#: Configure process-global GTP state before constructing model modules.

bridge.training.gtp.classify_gtp_remat_chains( model: list[torch.nn.Module], model_config: Any, ) → None#: Classify all model chunks after distributed wrapping and before first forward.

bridge.training.gtp.get_data_distribution_group( pg_collection: megatron.core.process_groups_config.ProcessGroupCollection, model_config: Any, *, with_context_parallel: bool = False, ) → torch.distributed.ProcessGroup#: Return the group spanning every rank that consumes distinct input data.