bridge.training.flex_dispatcher_backend

Module Contents

Functions

_fallback_to_alltoall	Clear flex dispatcher state when the requested backend cannot be used.
apply_flex_dispatcher_backend	Apply a supported flex dispatcher backend to the model config.
validate_flex_dispatcher_backend	Validate the selected flex dispatcher backend for the current GPU architecture.

Data

logger

API

bridge.training.flex_dispatcher_backend.logger: logging.Logger

‘getLogger(…)’

bridge.training.flex_dispatcher_backend._fallback_to_alltoall(

model_config: megatron.core.transformer.TransformerConfig,

) → None

Clear flex dispatcher state when the requested backend cannot be used.

bridge.training.flex_dispatcher_backend.apply_flex_dispatcher_backend(

model_config: megatron.core.transformer.TransformerConfig,

moe_flex_dispatcher_backend: str | None = None,

) → None

Apply a supported flex dispatcher backend to the model config.

DeepEP is applicable only for MoE models on Ampere, Hopper, B200 and B300 GPUs. HybridEP is applicable only for MoE models on GB200, GB300 with NVL72 and on Ampere, Hopper, B200 and B300 GPUs. NCCL EP is applicable only for MoE models on Hopper and Blackwell GPUs.

bridge.training.flex_dispatcher_backend.validate_flex_dispatcher_backend(

model_config: megatron.core.transformer.TransformerConfig,

) → None

Validate the selected flex dispatcher backend for the current GPU architecture.