`core.pipeline_parallel.utils`#

Module Contents#

Classes#

`StageDispatchBwdGrad`	1F1B + NCCL-EP zero-copy only: redirect the dispatch-backward grad into the persistent symm buffer so the one-sided `dispatch_bwd` can consume it.
`NoopScheduleNode`	A placeholder node in the computation graph that simply passes through inputs and outputs.
`ScheduleNode`	Base node for fine-grained scheduling.
`AbstractSchedulePlan`	To use combined 1f1b, model must implement build_schedule_plan while take the same signature as model forward but return an instance of AbstractSchedulePlan

Functions#

`is_pp_first_stage`	Return True if in the first pipeline model-parallel stage, False otherwise.
`is_pp_last_stage`	Return True if in the last pipeline-model-parallel stage, False otherwise.
`is_vp_first_stage`	Return True if in the first virtual pipeline model-parallel stage, False otherwise.
`is_vp_last_stage`	Return True if in the last virtual pipeline model-parallel stage, False otherwise.
`get_pp_first_rank`	Return the global rank of the first rank in the pipeline parallel group.
`get_pp_last_rank`	Return the global rank of the last rank in the pipeline parallel group.
`get_pp_next_rank`	Return the global rank of the next rank in the pipeline parallel group, or None if last stage.
`get_pp_prev_rank`	Return the global rank of the previous rank in the pipeline parallel group, or None if first stage.
`make_viewless`	Make_viewless util func
`set_ideal_affinity_for_current_gpu`	Set CPU affinity for the current GPU to optimize host-device transfers.
`set_streams`	Set the stream for communication operations.
`get_comp_stream`	Get the stream for computation
`get_comm_stream`	Get the stream for communication

Data#

`logger`
`_USE_DYNAMIC_COMP_STREAM`
`_COMP_STREAM`
`_COMM_STREAM`

API#

core.pipeline_parallel.utils.logger#: ‘getLogger(…)’

class core.pipeline_parallel.utils.StageDispatchBwdGrad#

Bases: torch.autograd.Function

1F1B + NCCL-EP zero-copy only: redirect the dispatch-backward grad into the persistent symm buffer so the one-sided dispatch_bwd can consume it.

Under the 1F1B overlap schedule the dispatch output is consumed by the next node, which detaches it into a leaf; autograd therefore hands dispatch_bwd a non-symm AccumulateGrad clone. Applying this identity node to the dispatch output — while it is still inside the dispatch node’s own graph segment — makes it the sole consumer, moving that accumulation to our output; the backward then does a single plain->symm copy into the dispatcher’s _zc_bwd_token_buf. That buffer is free to stage into precisely because get_expert_zero_copy_buffers withholds it from the op-fuser under overlap. Forward is identity (no numeric effect).

static forward(ctx, dispatched_tokens, token_dispatcher)#: Identity forward; stashes the dispatcher so backward can reach its symm buffer.

static backward(ctx, grad)#: Stage the incoming gradient into the symm dispatch-backward buffer.

core.pipeline_parallel.utils.is_pp_first_stage(pp_group: torch.distributed.ProcessGroup)#: Return True if in the first pipeline model-parallel stage, False otherwise.

core.pipeline_parallel.utils.is_pp_last_stage(pp_group: torch.distributed.ProcessGroup)#: Return True if in the last pipeline-model-parallel stage, False otherwise.

core.pipeline_parallel.utils.is_vp_first_stage(vp_stage: int, vp_size: int | None)#: Return True if in the first virtual pipeline model-parallel stage, False otherwise.

core.pipeline_parallel.utils.is_vp_last_stage(vp_stage: int, vp_size: int | None)#: Return True if in the last virtual pipeline model-parallel stage, False otherwise.

core.pipeline_parallel.utils.get_pp_first_rank(pp_group: torch.distributed.ProcessGroup)#: Return the global rank of the first rank in the pipeline parallel group.

core.pipeline_parallel.utils.get_pp_last_rank(pp_group: torch.distributed.ProcessGroup)#: Return the global rank of the last rank in the pipeline parallel group.

core.pipeline_parallel.utils.get_pp_next_rank(pp_group: torch.distributed.ProcessGroup)#: Return the global rank of the next rank in the pipeline parallel group, or None if last stage.

core.pipeline_parallel.utils.get_pp_prev_rank(pp_group: torch.distributed.ProcessGroup)#: Return the global rank of the previous rank in the pipeline parallel group, or None if first stage.

core.pipeline_parallel.utils.make_viewless(e)#: Make_viewless util func

core.pipeline_parallel.utils.set_ideal_affinity_for_current_gpu()#: Set CPU affinity for the current GPU to optimize host-device transfers.

class core.pipeline_parallel.utils.NoopScheduleNode#

A placeholder node in the computation graph that simply passes through inputs and outputs.

This class is used as a no-op node in the scheduling system when a real computation node is not needed but the interface must be maintained (e.g., dense layer doesn’t need moe_dispatch and moe_combine). It simply returns its inputs unchanged in both forward and backward passes.

forward(inputs)#: Passes through inputs unchanged in the forward pass.

backward(outgrads)#: Passes through gradients unchanged in the backward pass.

class core.pipeline_parallel.utils.ScheduleNode( forward_func: Callable, stream: torch.cuda.Stream, event: torch.cuda.Event, backward_func: Optional[Callable] = None, free_input: bool = False, name: str = 'schedule_node', ncclep_zero_copy: bool = False, )#

Base node for fine-grained scheduling.

This class represents a computational node in the pipeline schedule. It handles the execution of forward and backward operations on a stream.

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.

default_backward_func(outputs, output_grad)#: Default backward function

forward(inputs=())#: Schedule node forward

_forward(*inputs)#

get_output()#: Get the forward output

backward(output_grad)#: Schedule node backward

_backward(*output_grad)#

get_grad()#: Get the grad of inputs

stream_acquire_context(name=None)#

Stream acquire context that handles event synchronization, NVTX profiling, and stream context.

This context manager consolidates:

Event wait/record for synchronization between streams
NVTX range for profiling (if name is provided)
torch.cuda.stream context for execution on the specified stream

Parameters:: name – Optional name for NVTX range profiling

_release_state()#: Clear the state of the node

class core.pipeline_parallel.utils.AbstractSchedulePlan#

Bases: abc.ABC

To use combined 1f1b, model must implement build_schedule_plan while take the same signature as model forward but return an instance of AbstractSchedulePlan

abstractmethod static run( f_schedule_plan, b_schedule_plan, grad=None, pre_forward=None, pre_backward=None, post_forward=None, post_backward=None, )#: run() is the protocol between our schedule logic and model, which is used to schedule the forward and backward schedule plans for the models.

core.pipeline_parallel.utils._USE_DYNAMIC_COMP_STREAM#: None

core.pipeline_parallel.utils._COMP_STREAM#: None

core.pipeline_parallel.utils._COMM_STREAM#: None

core.pipeline_parallel.utils.set_streams(comm_stream=None, high_priority=False)#: Set the stream for communication operations.

core.pipeline_parallel.utils.get_comp_stream()#: Get the stream for computation

core.pipeline_parallel.utils.get_comm_stream()#: Get the stream for communication

core.pipeline_parallel.utils#