Is this page helpful?

core.models.gpt.fine_grained_callables#

Module Contents#

Functions#

build_transformer_layer_callables

Create callables for transformer layer nodes. Divides the transformer layer’s operations into a sequence of smaller, independent functions. This decomposition separates computation-heavy tasks (e.g., self-attention, MLP) from communication-heavy tasks (e.g., MoE’s All-to-All).

API#

core.models.gpt.fine_grained_callables.build_transformer_layer_callables(
layer: megatron.core.transformer.transformer_layer.TransformerLayer,
)#

Create callables for transformer layer nodes. Divides the transformer layer’s operations into a sequence of smaller, independent functions. This decomposition separates computation-heavy tasks (e.g., self-attention, MLP) from communication-heavy tasks (e.g., MoE’s All-to-All).

The five callables align with the schedule plan’s slot order:

  1. pre_dispatch_computation (computation): attention -> pre-MLP layernorm -> router -> dispatch preprocess. For dense layers this is just the attention pass.

  2. moe_dispatch (communication): MoE dispatch All-to-All.

  3. mlp / moe_experts (computation): dense MLP or routed-experts compute.

  4. moe_combine (communication): MoE combine All-to-All + post-MLP residual.

  5. mtp_post_process (computation): always None here; only the MTP wrapper in common/fine_grained_callables.py fills this slot.

By assigning these functions to different CUDA streams (e.g., a compute stream and a communication stream), the scheduler can overlap their execution, preventing tasks from competing for resources and hiding communication latency by running them in parallel with functions from other micro-batches.

Parameters:

layer – The transformer layer to build callables for.

Returns:

  • forward_funcs: List of 5 callables, one per slot in the schedule plan (pre_dispatch_computation, moe_dispatch, mlp, moe_combine, mtp_post_process=None).

  • backward_dw: Dict mapping slot name to the delayed-wgrad callable (keys: “pre_dispatch_computation”, “mlp”).

Return type:

A tuple containing