bridge.utils.fusions#

Fusion capability checks for Megatron models.

This module provides functions to check if various fusion optimizations can be enabled based on the current environment and dependencies.

Module Contents#

Functions#

can_enable_gradient_accumulation_fusion

Check if gradient accumulation fusion can be enabled.

validate_rope_fusion_compatibility

Validate if RoPE fusion is compatible with the current model configuration.

Data#

logger

LOG_FUSION_DISABLE

API#

bridge.utils.fusions.logger#

‘getLogger(…)’

bridge.utils.fusions.LOG_FUSION_DISABLE#

None

bridge.utils.fusions.can_enable_gradient_accumulation_fusion() → bool#

Check if gradient accumulation fusion can be enabled.

There are two independent backends that support wgrad fusion:

  • TransformerEngine (TE): handles fuse_wgrad_accumulation internally, no extra CUDA extension needed.

  • APEX fused_weight_gradient_mlp_cuda: required for the non-TE (ColumnParallelLinear) path.

In TE-based training containers (the common case for high-perf models), fused_weight_gradient_mlp_cuda is not installed, so the APEX check alone incorrectly returns False even though TE can perform the fusion.

Returns:

True if gradient accumulation fusion is available via either backend.

Return type:

bool

bridge.utils.fusions.validate_rope_fusion_compatibility(
config: megatron.core.transformer.transformer_config.TransformerConfig,
) → bool#

Validate if RoPE fusion is compatible with the current model configuration.

Parameters:

config – The TransformerConfig instance to validate.

Returns:

True if RoPE fusion is compatible, False otherwise.

Return type:

bool