bridge.training.flex_dispatcher_backend

Module Contents

Functions

apply_flex_dispatcher_backend	Configure the requested dispatcher without inspecting the local GPU.
validate_flex_dispatcher_backend	Validate the final backend on the training GPU without changing the config.

Data

logger

API

bridge.training.flex_dispatcher_backend.logger: logging.Logger

‘getLogger(…)’

bridge.training.flex_dispatcher_backend.apply_flex_dispatcher_backend(

model_config: megatron.core.transformer.TransformerConfig,

moe_flex_dispatcher_backend: str | None = None,

) → None

Configure the requested dispatcher without inspecting the local GPU.

Recipe construction may run on a different host from training, and child recipes may replace a parent’s backend. Validate hardware support with validate_flex_dispatcher_backend after all recipe and user overrides. An unsupported backend is never replaced with another dispatcher.

Parameters:

• model_config – Model configuration to update in place.

• moe_flex_dispatcher_backend – deepep, hybridep, or ncclep to select flex dispatch; None explicitly selects alltoall.

Raises:

ValueError – If the backend name is not recognized.

bridge.training.flex_dispatcher_backend.validate_flex_dispatcher_backend(

model_config: megatron.core.transformer.TransformerConfig,

) → None

Validate the final backend on the training GPU without changing the config.

Raises:

ValueError – If flex has no recognized backend or the GPU is unsupported.