core.transformer.residual_recompute_plan

Residual recompute planning functions.

Only immutable boundary plans may be cached, such as a plan that keeps a ShortcutMoE producer/consumer pair in the same block. Managers own checkpoints from one forward invocation; callers choose which operations to register.

Module Contents

Functions

build_recompute_block_end_plan	Mark block ends without splitting disjoint adjacent atomic layer pairs.
build_recompute_layer_managers	Allocate fresh managers for an active plan, sharing one within each block.
finalize_recompute_block	Discard registered outputs and attach replay at the block’s live boundary.

API

core.transformer.residual_recompute_plan.build_recompute_block_end_plan(

num_layers: int,

block_size: int | None,

*,

atomic_layer_pairs: collections.abc.Sequence[tuple[int, int]] = (),

) → tuple[bool, ...]

Mark block ends without splitting disjoint adjacent atomic layer pairs.

Indices use the caller’s local layer order. A pair may exceed a block size of one; otherwise blocks contain at most block_size layers. None uses the entire local stack. Empty stacks return an empty plan.

core.transformer.residual_recompute_plan.build_recompute_layer_managers(

block_ends: collections.abc.Sequence[bool],

) → list[megatron.core.tensor_parallel.random.CheckpointWithoutOutputManager]

Allocate fresh managers for an active plan, sharing one within each block.

core.transformer.residual_recompute_plan.finalize_recompute_block(

manager: megatron.core.tensor_parallel.random.CheckpointWithoutOutputManager | None,

hidden_states: torch.Tensor,

is_block_end: bool,

) → None

Discard registered outputs and attach replay at the block’s live boundary.