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.
_get_checkpoint_weight_topology	Read dense/expert TP and GTP sizes from runtime or serialized configs.
_validate_checkpoint_weight_topology	Reject native resharding when either side uses the GTP SwiGLU layout.
_get_dataloader_process_group	Return DP x dense GTP, excluding CP ranks that repeat the same samples.
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._get_checkpoint_weight_topology(

model_config: Any,

) → tuple[int, int, int, int]

Read dense/expert TP and GTP sizes from runtime or serialized configs.

bridge.training.gtp._validate_checkpoint_weight_topology(

*,

saved: tuple[int, int, int, int],

requested: tuple[int, int, int, int],

) → None

Reject native resharding when either side uses the GTP SwiGLU layout.

bridge.training.gtp._get_dataloader_process_group(

pg_collection: megatron.core.process_groups_config.ProcessGroupCollection,

) → torch.distributed.ProcessGroup

Return DP x dense GTP, excluding CP ranks that repeat the same samples.

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,

*,

reduce_scatter_with_fp32_accumulation: bool = False,

nccl_ub: bool = False,

pg_collection: megatron.core.process_groups_config.ProcessGroupCollection | None = None,

) → None

Configure process-global GTP state before constructing model modules.

Parameters:

• model_config – Model provider or builder config containing the GTP sizes.

• reduce_scatter_with_fp32_accumulation – Accumulate GTP reduce-scatter results locally in FP32.

• nccl_ub – Register an NCCL symmetric-memory pool for the dense GTP group.

• pg_collection – Initialized process groups, required when nccl_ub is enabled.

Raises:

• RuntimeError – GTP is active but the required Transformer Engine support is missing.

• ValueError – GTP is active and nccl_ub is enabled without process groups.

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.