core.models.common.fine_grained_callables

Layer-callable builders for the combined-1F1B fine-grained schedule plan.

These build_* functions assemble the per-layer (forward_funcs, backward_dw) tuple that the schedule plan plugs into TransformerLayerNode.

The TransformerLayer-specific builder lives in gpt/fine_grained_callables.py because it depends on GPT’s MoE wiring; the MTP builder and the dispatcher build_layer_callables are model-agnostic — both GPTModel and HybridModel schedule MTP layers identically — so they live here.

Module Contents

Functions

build_mtp_layer_callables	Callables for multi-token prediction layer nodes.
get_layer_moe_metadata	Return (is_moe, num_local_experts) for schedule-node construction.
build_layer_callables	Dispatch to the appropriate layer-callable builder.

API

core.models.common.fine_grained_callables.build_mtp_layer_callables(layer)

Callables for multi-token prediction layer nodes.

Wraps the inner layer.mtp_model_layer’s callables with MTP-specific pre-process (chunk and concat embeddings) and post-process (gather across depths) steps.

core.models.common.fine_grained_callables.get_layer_moe_metadata(layer)

Return (is_moe, num_local_experts) for schedule-node construction.

core.models.common.fine_grained_callables.build_layer_callables(layer)

Dispatch to the appropriate layer-callable builder.

Returns (forward_funcs, backward_dw).