`core.models.common.utils`#

Schedule-plan helpers shared by GPTModel and HybridModel.

These pieces used to live in core/models/gpt/fine_grained_callables.py and were imported by core/models/common/model_chunk_schedule_plan.py and the hybrid schedule plan via that path. They are model-agnostic in practice — the Pre/PostProcessNode classes call the model’s _preprocess / _postprocess methods and don’t otherwise care which model implements them — so they live here now.

Module Contents#

Classes#

`LayerState`	State shared between the schedule nodes that come from one logical layer.
`PreProcessNode`	Run the model’s `_preprocess` (embedding + rotary + padding mask).
`PostProcessNode`	Run the model’s `_postprocess` (final norm, output layer, loss).
`TransformerLayerNode`	Schedule node for one slot of a fine-grained transformer layer plan.
`_BackwardDWWrapper`	Backward weight-gradient wrapper for a transformer pre-dispatch slot.

Functions#

`weak_method`	Wrap `method` in a weakref-keyed dispatcher to break refcycles.
`should_free_input`	Whether the schedule node named `name` can free its input after forward.

API#

core.models.common.utils.weak_method(method)#

Wrap method in a weakref-keyed dispatcher to break refcycles.

ScheduleNode keeps a reference to the bound forward / backward functions of every node in the plan; using a strong reference would keep the layer plan (and the model chunk through it) alive after the iteration completes. The weakref.WeakMethod lets the schedule plan be torn down between iterations without manual del chains.

core.models.common.utils.should_free_input(name, is_moe, config, num_local_experts)#

Whether the schedule node named name can free its input after forward.

The schedule decomposes a transformer layer into pre_dispatch_computation, moe_dispatch, mlp, and moe_combine nodes; the inputs to some of those nodes are not needed in backward and can be released early to lower peak activation memory. Dense layers and the pre_dispatch_computation node always need their input retained (the attention residual flows through the post-MLP BDA).

Parameters:

name – Schedule node name.
is_moe – True for MoE layers; dense layers always retain inputs.
config – TransformerConfig for the layer.
num_local_experts – Local expert count on this rank (None for dense).

Returns:

True iff the named node may free its input after forward.

class core.models.common.utils.LayerState#

State shared between the schedule nodes that come from one logical layer.

Empty placeholder; nodes attach their own attributes (residual, dispatched probs, shared-expert outputs) for downstream nodes in the same layer to consume. Kept as a real class so weakrefs work uniformly.

class core.models.common.utils.PreProcessNode(model, chunk_state, event, stream)#

Bases: megatron.core.pipeline_parallel.utils.ScheduleNode

Run the model’s _preprocess (embedding + rotary + padding mask).

The schedule plan wraps a model that exposes a _preprocess method returning the canonical 6-tuple (decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset, padding_mask) (slots a given model doesn’t use are returned as None). The chunk state is mutated in-place so layer nodes can read the same fields by name.

Initialization

Initialize a schedule node.

Parameters:

forward_func (callable) – Function to execute during the forward pass.
stream (Callable) –
Func that returns CUDA stream for computation. This can be either a ‘compute’ stream or a ‘communicate’ stream.
- ’compute’ stream: Used for computational nodes like attention and experts.
- ’communicate’ stream: Used for nodes that handle token communication, such as token dispatch and combine operations in MoE layers.
event (torch.cuda.Event) – The CUDA event used for synchronization. Each microbatch within a model chunk shares the same event, which is used to manage dependencies between nodes operating on different streams.
backward_func (callable, optional) – Function for the backward pass.
free_input (bool) – Flag to indicate if the input should be freed after the forward pass.
name (str) – Name of the node for debugging purposes.

forward_impl()#: Run model preprocessing and store chunk-level inputs for layer nodes.

class core.models.common.utils.PostProcessNode(model, chunk_state, event, stream)#

Bases: megatron.core.pipeline_parallel.utils.ScheduleNode

Run the model’s _postprocess (final norm, output layer, loss).

Calls _postprocess with mtp_in_postprocess=False because the schedule plan handles MTP layers as sibling layer nodes inside the same chunk; the model’s MTP block is not invoked here. The optional final layernorm — applied only when this rank holds an empty decoder shard (early stage of pipeline parallel) — is handled here so the chunk plan does not need a separate node for it.

Initialization

Initialize a schedule node.

Parameters:

forward_func (callable) – Function to execute during the forward pass.
stream (Callable) –
Func that returns CUDA stream for computation. This can be either a ‘compute’ stream or a ‘communicate’ stream.
- ’compute’ stream: Used for computational nodes like attention and experts.
- ’communicate’ stream: Used for nodes that handle token communication, such as token dispatch and combine operations in MoE layers.
event (torch.cuda.Event) – The CUDA event used for synchronization. Each microbatch within a model chunk shares the same event, which is used to manage dependencies between nodes operating on different streams.
backward_func (callable, optional) – Function for the backward pass.
free_input (bool) – Flag to indicate if the input should be freed after the forward pass.
name (str) – Name of the node for debugging purposes.

forward_impl(hidden_states)#: Run model postprocessing for the chunk’s final hidden states.

class core.models.common.utils.TransformerLayerNode( stream, event, layer_state, chunk_state, submodule, name='default', bwd_dw_callables=None, extra_args={}, )#

Bases: megatron.core.pipeline_parallel.utils.ScheduleNode

Schedule node for one slot of a fine-grained transformer layer plan.

Each transformer layer is decomposed into pre_dispatch_computation, moe_dispatch, mlp, and moe_combine slots; this class is the scheduler-side handle for one slot. It owns the slot’s stream / event, the per-slot free_input policy, and the optional delayed weight-gradient hook. Subclasses override _resolve_free_input to specialize the policy (HybridStackNode does this for grouped layers).

Initialization

Initialize a schedule node.

Parameters:

forward_func (callable) – Function to execute during the forward pass.
stream (Callable) –
Func that returns CUDA stream for computation. This can be either a ‘compute’ stream or a ‘communicate’ stream.
- ’compute’ stream: Used for computational nodes like attention and experts.
- ’communicate’ stream: Used for nodes that handle token communication, such as token dispatch and combine operations in MoE layers.
event (torch.cuda.Event) – The CUDA event used for synchronization. Each microbatch within a model chunk shares the same event, which is used to manage dependencies between nodes operating on different streams.
backward_func (callable, optional) – Function for the backward pass.
free_input (bool) – Flag to indicate if the input should be freed after the forward pass.
name (str) – Name of the node for debugging purposes.

static _resolve_free_input(name, is_moe, config, num_local_experts)#: Free-input policy hook. Subclasses override to specialize.

detach(t)#: Detach a tensor and remember it for backward through the schedule node.

forward_impl(*args)#: Invoke the slot’s submodule forward.

backward_impl(outputs, output_grad)#: Run the slot’s backward and return the input grads.

forward(*inputs)#: Execute forward and fire the per-layer post-forward hook on the last slot.

backward(*output_grad)#

Execute backward and fire the per-layer post-backward hook on the first slot.

When delay_wgrad_compute is set, the hook fires after backward_dw instead, because the wgrad work has not yet run when backward returns.

backward_dw()#: Run the slot’s delayed weight-gradient callables on the slot’s stream.

set_post_forward_hook(hook)#: Mark this slot as the layer’s last fwd node and register the hook.

set_post_backward_hook(hook)#: Mark this slot as the layer’s first bwd node and register the hook.

__del__()#

class core.models.common.utils._BackwardDWWrapper(layer)#

Backward weight-gradient wrapper for a transformer pre-dispatch slot.

Runs the layer’s self_attention.backward_dw plus, on MoE layers, the shared-expert backward_dw; coordinates with the cuda-graph wgrad capture (set_graphed_backward_dw_callable) so that scopes covered by the graph are not re-run eagerly. Used when overlap_moe_expert_parallel_comm and delay_wgrad_compute are both enabled.

Initialization

backward_dw()#: Run eager or graphed backward wgrad callables for the wrapped layer.

set_graphed_backward_dw_callable(graphed_backward_dw_callable)#: Plug the cuda-graph backward wgrad replay callable.

parameters()#

Yield parameters from the wrapped layer’s wgrad submodules.

Mirrors torch.nn.Module.parameters so callers (notably TransformerLayerNode.backward_dw) can collect post_wgrad_grad_acc_hook without knowing the concrete layer layout.

core.models.common.utils#

Module Contents#

Classes#

Functions#

API#

`core.models.common.utils`#