`core.models.common.model_chunk_schedule_plan`#

Module Contents#

Classes#

`ModelChunkState`	State shared across a model chunk.
`RecomputeSegment`	A contiguous group of layers recomputed as one unit.
`TransformerLayerSchedulePlan`	Schedule the executing plan of the nodes in a transformer/mtp layer.
`TransformerModelChunkSchedulePlan`	Schedule the executing plan of the sub-modules in a model chunk sub-modules.

Functions#

_copy_chunk_state_value

Copy list-valued chunk state so a later in-place append cannot corrupt a snapshot.

Data#

_MUTABLE_CHUNK_STATE_FIELDS

API#

class core.models.common.model_chunk_schedule_plan.ModelChunkState#

State shared across a model chunk.

This class holds state that is shared between different components of a model chunk, such as input tensors, parameters, and configuration.

core.models.common.model_chunk_schedule_plan._MUTABLE_CHUNK_STATE_FIELDS#: (‘input_ids’, ‘position_ids’, ‘padding_mask’, ‘mtp_input_mask’, ‘mtp_hidden_states’)

core.models.common.model_chunk_schedule_plan._copy_chunk_state_value(value: Any) → Any#: Copy list-valued chunk state so a later in-place append cannot corrupt a snapshot.

class core.models.common.model_chunk_schedule_plan.RecomputeSegment( layers: List[core.models.common.model_chunk_schedule_plan.TransformerLayerSchedulePlan], chunk_state: core.models.common.model_chunk_schedule_plan.ModelChunkState, start_index: int, )#

A contiguous group of layers recomputed as one unit.

The initial forward runs under no_grad and is replayed with grad enabled at the start of the segment’s own backward. A segment drops its layers’ activations as soon as its own forward is done (release_activations) and the replayed ones as each node backwards (ScheduleNode._release_state), so only one segment’s activations are live at a time and what crosses the forward->backward gap is the segment’s input tensor plus the _MUTABLE_CHUNK_STATE_FIELDS snapshot. Layers left eager under block keep their activations throughout. The A2A issued by a replay is exposed by design; normal fwd/bwd A2A overlap is preserved.

start_index is the segment’s first layer index within the chunk, for NVTX labels.

Initialization

capture( layer: core.models.common.model_chunk_schedule_plan.TransformerLayerSchedulePlan, f_input: torch.Tensor, ) → None#: When the head layer starts its forward, snapshot what the replay needs: input tensor, mutable chunk state, RNG.

release_activations( layer: core.models.common.model_chunk_schedule_plan.TransformerLayerSchedulePlan, ) → None#

Drop this segment’s forward activations once its tail layer has forwarded.

Runs while the chunk forward is still going, not at the end of it: the nodes hold their boundary tensors in inputs / output / detached until something unsets them, so deferring this to the end of the chunk would keep one such set alive per recomputed layer for the rest of the forward. The tail layer’s output is the caller’s f_input and is what the next segment captures, so it survives.

Only the Python references are dropped. Resizing the storages (what ScheduleNode’s free_input path and tensor_parallel.random.checkpoint do) is not usable here: the replay rebuilds the graph on top of these very tensors, and freeing their storage corrupts it.

release_input( layer: core.models.common.model_chunk_schedule_plan.TransformerLayerSchedulePlan, ) → None#: Drop the retained input once the head layer has finished its backward.

recompute( layer: core.models.common.model_chunk_schedule_plan.TransformerLayerSchedulePlan, ) → None#

Replay the segment with grad enabled when its tail layer enters the backward.

Must run before any of that layer’s nodes backwards: mtp_post_process is in the recompute scope, so its graph does not exist until the replay has run. Re-running the same nodes in the same order repopulates each node’s inputs / output / detached / layer_state as in the initial forward.

class core.models.common.model_chunk_schedule_plan.TransformerLayerSchedulePlan( layer, event, chunk_state, comp_stream, comm_stream, extra_args={}, )#

Schedule the executing plan of the nodes in a transformer/mtp layer.

This class organizes the sub-modules of a transformer/mtp layer, including attention, post attention, MLP, dispatch, combine and mtp post process nodes.

layer (TransformerLayerSchedulePlan) ├── pre_dispatch_computation (TransformerLayerNode): │ attention -> layernorm -> router -> dispatch preprocess ├── moe_dispatch (TransformerLayerNode): dispatch All2All ├── mlp (TransformerLayerNode): mlp module ├── moe_combine (TransformerLayerNode): combine All2All └── mtp_post_process (PostProcessNode): mtp post process

Note that MTP layer has the same operation and execution order with TransformerLayer regarding moe_dispatch, mlp, moe_combine, but contains extra operations in pre_dispatch_computation and mtp_post_process:

mtp.pre_dispatch_computation wraps around transformer_layer.pre_dispatch_computation with extra norm, proj and embedding operations.
mtp.mtp_post_process contains output_layer, mtp loss operations, whereas transformer_layer.mtp_post_process is empty.

Initialization

Initializes a transformer layer schedule plan.

Parameters:

layer (TransformerLayer) – split a transformer layer into multiple nodes for fine-grained scheduling.
event (torch.cuda.Event) – record CUDA event across multiple nodes on different streams for synchronization.
chunk_state (ModelChunkState) – model state shared in the model chunk.
comp_stream (Callable) – Func that returns CUDA stream for computation.
comm_stream (Callable) – Func that returns CUDA stream for communication.
extra_args (dict) – extra arguments for the layer.

The event and chunk_state are binded to the TransformerModelChunkSchedulePlan and shared across all layers in the model chunk.

pre_dispatch_computation#: None

moe_dispatch#: None

mlp#: None

moe_combine#: None

mtp_post_process#: None

release_state()#: Release reference, this helps avoid memory leak.

_build_callable_nodes(event, comp_stream, comm_stream, extra_args)#: Builds the callable nodes for the transformer/mtp layer: pre_dispatch_computation, moe_dispatch, mlp, moe_combine, and mtp_post_process.

set_fsdp_reshard_hooks(post_forward_hook, post_backward_hook)#

Wire FSDP parameter release callbacks for the fine-grained overlap schedule.

The EP overlap schedule bypasses the normal FSDP forward/backward hooks (registered on the FSDP unit module) because it calls sub-modules directly instead of going through TransformerLayer.forward(). This method attaches explicit release hooks to individual schedule nodes so that all-gathered parameters are freed at the right time.

Parameters:

post_forward_hook – Callable(module) that releases forward-pass params (bwd=False). Typically fsdp_wrapper.post_forward_release_module.
post_backward_hook – Callable(module) that releases backward-pass params (bwd=True). Typically fsdp_wrapper.post_backward_release_module.

_iter_layer_nodes()#

Yield this layer’s nodes in forward order, NoopScheduleNodes included.

Single definition of the sequence for recompute_forward/set_forward_no_grad/ reset_for_recompute; run() issues the same order inline (it interleaves each node with the backward layer’s), so a new node must be added in both places.

This is also the full-recompute scope. mtp_post_process must stay in it: it builds a torch.cat over the MTP pre_dispatch node’s detached input, so leaving it outside would anchor that cat on a no_grad leaf and drop the decoder’s gradient via MTP.

_iter_recomputed_nodes()#

Yield the real ScheduleNodes (skips NoopScheduleNode, which holds no state).

Only for state toggling; the forward helpers walk _iter_layer_nodes, since a Noop still passes the tensor through and still consumes a low-precision context in run().

set_forward_no_grad(no_grad: bool)#

Toggle no-grad forward for this layer’s recomputed nodes.

The initial forward runs with no_grad=True (no autograd graph retained); the backward-time recompute runs with no_grad=False to rebuild the graph.

recompute_forward(f_input: torch.Tensor) → torch.Tensor#

Re-run this layer with grad enabled and return its output.

One low-precision context per node, matching run()’s forward half node for node.

reset_for_recompute()#: Free the retained forward activations of this layer, keeping the nodes reusable for a recompute forward. Also clears the per-layer shared state.

get_low_precision_context()#: Get the low-precision context for the transformer layer.

static run( f_layer, b_layer, f_input=None, b_grad=None, is_last_layer_in_bwd=False, )#

Schedule one-forward-one-backward operations for a single transformer layer.

This function interleaves forward and backward operations, overlapping the communications (dispatch or combine) of one with the computations (pre_dispatch or mlp) of the other to maximize parallelism and efficiency.

When f_layer and b_layer are not None, forward and backward pass are overlapped as follows: comm_stream: combine_bwd | dispatch_fwd->dispatch_bwd | combine_fwd comp_stream: pre_dispatch_fwd | mlp_bwd->mlp_bwd_dw->mlp_fwd| pre_dispatch_bwd For MTP, mtp_post_process_fwd is executed after the combine_fwd in the comp_stream, and mtp_post_process_bwd is executed before the combine_bwd in the comp_stream.

Parameters:

f_layer (TransformerLayerSchedulePlan) – Forward layer (for current microbatch)
b_layer (TransformerLayerSchedulePlan) – Backward layer (for previous microbatch)
f_input (Tensor) – Input for forward computation
b_grad (Tensor) – Gradient for backward computation
is_last_layer_in_bwd (bool) – Whether the current layer is the last layer in the backward pass.

Returns:

Functions or values for next iteration’s computation

class core.models.common.model_chunk_schedule_plan.TransformerModelChunkSchedulePlan( model, input_ids: torch.Tensor, position_ids: torch.Tensor, attention_mask: torch.Tensor, decoder_input: torch.Tensor = None, labels: torch.Tensor = None, packed_seq_params=None, extra_block_kwargs=None, runtime_gather_output: Optional[bool] = None, loss_mask: Optional[torch.Tensor] = None, padding_mask=None, *, mtp_input_mask: Optional[torch.Tensor] = None, output_processor: Optional[Callable[..., Any]] = None, output_processor_context: Optional[Any] = None, )#

Bases: megatron.core.pipeline_parallel.utils.AbstractSchedulePlan

Schedule the executing plan of the sub-modules in a model chunk sub-modules.

This class organizes the computation nodes for a model chunk, including preprocessing, transformer layers, and postprocessing.

TransformerModelChunkSchedulePlan ├── pre_process: PreProcessNode ├── layers: List[TransformerLayerSchedulePlan] │ ├── layer[0]: TransformerLayerSchedulePlan │ ├── layer[1]: TransformerLayerSchedulePlan │ └── … └── post_process: PostProcessNode

Subclasses can swap the per-layer schedule plan by overriding the LAYER_SCHEDULE_PLAN_CLASS class attribute (e.g. HybridStack uses a layer plan that understands grouped/inferred layer types). They can also swap the pre/post-process node classes via PRE_PROCESS_NODE_CLASS / POST_PROCESS_NODE_CLASS so each model owns its own embedding / output layer node implementations.

Initialization

Initialize the schedule plan of all Transformer layers’ sub-modules.

This function creates a schedule plan for a model chunk, including preprocessing, transformer layers, and postprocessing.

Parameters:

model – The model to build a schedule plan for.
input_ids – Input token IDs.
position_ids – Position IDs.
attention_mask – Attention mask.
decoder_input – Decoder input tensor.
labels – Labels for loss computation.
packed_seq_params – Parameters for packed sequences.
extra_block_kwargs – Additional keyword arguments for blocks.
runtime_gather_output – Whether to gather output at runtime.
loss_mask (torch.Tensor) – Used to mask out some portions of the loss
mtp_input_mask (Optional[torch.Tensor]) – Tensor of shape [batch, sequence] whose values are converted to booleans. Nonzero/True marks a valid MTP conditioning token; zero/False invalidates that token and any deeper MTP prediction path that crosses it. None applies no additional conditioning mask.
output_processor (Callable) – Custom postprocess hook to run instead of the default logits/loss path.
output_processor_context (Any) – User-defined context object forwarded to output_processor.

Returns:

The model chunk schedule plan.

LAYER_SCHEDULE_PLAN_CLASS#: None

PRE_PROCESS_NODE_CLASS#: None

POST_PROCESS_NODE_CLASS#: None

_build_recompute_segments(config)#

Split this chunk’s layers into full-recompute segments.

Decoder segmentation mirrors megatron.core.recompute.checkpointed_forward, so recompute_method / recompute_num_layers mean the same thing with and without –overlap-moe-expert-parallel-comm: ‘uniform’ groups every decoder layer by recompute_num_layers, ‘block’ recomputes the first recompute_num_layers decoder layers one per segment. MTP gets one segment per depth under both methods.

Unlike the non-overlap ‘block’ branch there is no recompute_skip_num_layers window under fp8/fp4: no checkpoint primitive is involved, since the replay marks the retained segment input grad-tracking itself.

_build_layer_schedule_plan(module, comp_stream, comm_stream)#

_extra_args_for_layer(module, layer_idx, num_layers)#

Per-layer extra_args dict passed to the layer plan constructor.

Subclasses extend this hook to thread additional metadata (e.g. hybrid layer-type symbols) without overriding _build_layer_schedule_plan.

property event#: Gets the CUDA event for synchronization.

record_current_stream()#: Records the current CUDA stream in the event.

wait_current_stream()#: Waits for the event to complete on the current CUDA stream.

get_layer(i)#: Gets the transformer layer at the specified index.

pop_layer()#: Pops the transformer layer in FILO order.

num_layers()#: Gets the number of transformer layers.

property state#: Gets the model chunk state.

release_state()#: Release reference, this helps avoid memory leak.

static run( f_schedule_plan, b_schedule_plan, b_grad=None, pre_forward=None, pre_backward=None, post_forward=None, post_backward=None, )#

Model Chunk level 1f1b fine-grained scheduler.

This function schedules the forward and backward passes for a model chunk, which interleaves forward and backward function of multiple Transformer layers within a model chunk, and this is needed to overlap the submodules between the individual forward and backward functions.

Assume there are 4 layers in the given model chunk: Phase 0: p2p_comm_sync -> forward_preprocess -> p2p_comm_sync -> backward_postprocess Phase 1: forward_layer[0] + backward_layer[3], overlapped execution by schedule_layer_1f1b Phase 2: forward_layer[1] + backward_layer[2], overlapped execution by schedule_layer_1f1b Phase 3: forward_layer[2] + backward_layer[1], overlapped execution by schedule_layer_1f1b Phase 4: forward_layer[3] + backward_layer[0], overlapped execution by schedule_layer_1f1b Phase 5: send_forward_recv_backward -> send_backward_recv_forward Phase 6: backward_dw of the first layer -> forward_postprocess -> backward_preprocess

Parameters:

f_schedule_plan (TransformerModelChunkSchedulePlan) – The forward schedule plan
b_schedule_plan (TransformerModelChunkSchedulePlan) – The backward schedule plan
b_grad (Tensor or None) – The gradient of the loss function
pre_forward (callable or None) – The function to call before the forward pass
pre_backward (callable or None) – The function to call before the backward pass
post_forward (callable or None) – The function to call after the forward pass
post_backward (callable or None) – The function to call after the backward pass

Returns:

The output of the forward pass.

core.models.common.model_chunk_schedule_plan#

Module Contents#

Classes#

Functions#

Data#

API#

`core.models.common.model_chunk_schedule_plan`#