Rate this Page

★ ★ ★ ★ ★

Pipeline Parallelism#

Created On: Jun 16, 2025 | Last Updated On: Oct 15, 2026

Note

torch.distributed.pipelining is currently in alpha state and under development. API changes may be possible. It was migrated from the PiPPy project.

Why Pipeline Parallel?#

Pipeline Parallelism is one of the primitive forms of parallelism for deep learning. It allows the execution of a model to be partitioned such that multiple micro-batches can execute different parts of the model code concurrently. Pipeline parallelism can be an effective technique for:

large-scale training
bandwidth-limited clusters
large model inference

The above scenarios share a commonality that the computation per device cannot hide the communication of conventional parallelism, for example, the weight all-gather of FSDP.

What is `torch.distributed.pipelining`?#

While promising for scaling, pipelining is often difficult to implement because it needs to partition the execution of a model in addition to model weights. The partitioning of execution often requires intrusive code changes to your model. Another aspect of complexity comes from scheduling micro-batches in a distributed environment, with data flow dependency considered.

The pipelining package provides a toolkit that does said things automatically which allows easy implementation of pipeline parallelism on general models.

It consists of two parts: a splitting frontend and a distributed runtime. The splitting frontend takes your model code as-is, splits it up into “model partitions”, and captures the data-flow relationship. The distributed runtime executes the pipeline stages on different devices in parallel, handling things like micro-batch splitting, scheduling, communication, and gradient propagation, etc.

Overall, the pipelining package provides the following features:

Splitting of model code based on simple specification.
Rich support for pipeline schedules, including GPipe, 1F1B, Interleaved 1F1B and Looped BFS, and providing the infrastructure for writing customized schedules.
First-class support for cross-host pipeline parallelism, as this is where PP is typically used (over slower interconnects).
Composability with other PyTorch parallel techniques such as data parallel (DDP, FSDP) or tensor parallel. The TorchTitan project demonstrates a “3D parallel” application on the Llama model.

Step 1: build `PipelineStage`#

Before we can use a PipelineSchedule, we need to create PipelineStage objects that wrap the part of the model running in that stage. The PipelineStage is responsible for allocating communication buffers and creating send/recv ops to communicate with its peers. It manages intermediate buffers e.g. for the outputs of forward that have not been consumed yet, and it provides a utility for running the backwards for the stage model.

A PipelineStage needs to know the input and output shapes for the stage model, so that it can correctly allocate communication buffers. The shapes must be static, e.g. at runtime the shapes can not change from step to step. A class PipeliningShapeError will be raised if runtime shapes do not match the expected shapes. When composing with other parallelisms or applying mixed precision, these techniques must be taken into account so the PipelineStage knows the correct shape (and dtype) for the output of the stage module at runtime.

Users may construct a PipelineStage instance directly, by passing in an nn.Module representing the portion of the model that should run on the stage. This may require changes to the original model code. See the example in Option 1: splitting a model manually.

Alternatively, the splitting frontend can use graph partitioning to split your model into a series of nn.Module automatically. This technique requires the model is traceable with torch.Export. Composability of the resulting nn.Module with other parallelism techniques is experimental, and may require some workarounds. Usage of this frontend may be more appealing if the user cannot easily change the model code. See Option 2: splitting a model automatically for more information.

Step 2: use `PipelineSchedule` for execution#

We can now attach the PipelineStage to a pipeline schedule, and run the schedule with input data. Here is a GPipe example:

from torch.distributed.pipelining import ScheduleGPipe

# Create a schedule
schedule = ScheduleGPipe(stage, n_microbatches)

# Input data (whole batch)
x = torch.randn(batch_size, in_dim, device=device)

# Run the pipeline with input `x`
# `x` will be divided into microbatches automatically
if rank == 0:
    schedule.step(x)
else:
    output = schedule.step()

If your input pipeline already produces microbatches, pass them directly:

arg_mbs = [(x0,), (x1,)]
kwarg_mbs = [{"mask": mask0}, {"mask": mask1}]
target_mbs = [target0, target1]

schedule.step(
    arg_mbs=arg_mbs,
    kwarg_mbs=kwarg_mbs,
    target_mbs=target_mbs,
)

If your loss takes extra arguments, pass the ones shared by every microbatch in loss_kwargs, and the ones that differ per microbatch, such as per-token weights, in loss_kwarg_mbs, one dict per microbatch. Like the other pre-split inputs, loss_kwarg_mbs goes with target_mbs. Microbatch i computes loss_fn(output, target_mbs[i], **loss_kwargs, **loss_kwarg_mbs[i]), so a key must not appear in both:

def loss_fn(output, target, *, num_tokens, token_weights):
    per_token_loss = torch.nn.functional.cross_entropy(
        output.transpose(1, 2), target, reduction="none"
    )
    return (per_token_loss * token_weights).sum() / num_tokens

schedule = ScheduleGPipe(stage, n_microbatches, loss_fn=loss_fn)
schedule.step(
    arg_mbs=arg_mbs,
    target_mbs=target_mbs,
    loss_kwargs={"num_tokens": num_tokens},  # same for every microbatch
    loss_kwarg_mbs=[{"token_weights": w0}, {"token_weights": w1}],
    losses=losses,
)

While inferring stage metadata (on the first step by default), the schedule also calls loss_fn once with microbatch 0’s inputs, so loss_fn must not accumulate into objects passed through loss_kwarg_mbs.

Note that the above code needs to be launched for each worker, thus we use a launcher service to launch multiple processes:

torchrun --nproc_per_node=2 example.py

Options for Splitting a Model#

Option 1: splitting a model manually#

To directly construct a PipelineStage, the user is responsible for providing a single nn.Module instance that owns the relevant nn.Parameters and nn.Buffers, and defines a forward() method that executes the operations relevant for that stage. For example, a condensed version of the Transformer class defined in Torchtitan shows a pattern of building an easily partitionable model.

class Transformer(nn.Module):
    def __init__(self, model_args: ModelArgs):
        super().__init__()

        self.tok_embeddings = nn.Embedding(...)

        # Using a ModuleDict lets us delete layers without affecting names,
        # ensuring checkpoints will correctly save and load.
        self.layers = torch.nn.ModuleDict()
        for layer_id in range(model_args.n_layers):
            self.layers[str(layer_id)] = TransformerBlock(...)

        self.output = nn.Linear(...)

    def forward(self, tokens: torch.Tensor):
        # Handling layers being 'None' at runtime enables easy pipeline splitting
        h = self.tok_embeddings(tokens) if self.tok_embeddings else tokens

        for layer in self.layers.values():
            h = layer(h, self.freqs_cis)

        h = self.norm(h) if self.norm else h
        output = self.output(h).float() if self.output else h
        return output

A model defined in this manner can be easily configured per stage by first initializing the whole model (using meta-device to avoid OOM errors), deleting undesired layers for that stage, and then creating a PipelineStage that wraps the model. For example:

with torch.device("meta"):
    assert num_stages == 2, "This is a simple 2-stage example"

    # we construct the entire model, then delete the parts we do not need for this stage
    # in practice, this can be done using a helper function that automatically divides up layers across stages.
    model = Transformer()

    if stage_index == 0:
        # prepare the first stage model
        del model.layers["1"]
        model.norm = None
        model.output = None

    elif stage_index == 1:
        # prepare the second stage model
        model.tok_embeddings = None
        del model.layers["0"]

    from torch.distributed.pipelining import PipelineStage
    stage = PipelineStage(
        model,
        stage_index,
        num_stages,
        device,
    )

When composing with other Data or Model parallelism techniques, output_args may also be required, if the output shape/dtype of the model chunk will be affected.

Option 2: splitting a model automatically#

If you have a full model and do not want to spend time on modifying it into a sequence of “model partitions”, the pipeline API is here to help. Here is a brief example:

class Model(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.emb = torch.nn.Embedding(10, 3)
        self.layers = torch.nn.ModuleList(
            Layer() for _ in range(2)
        )
        self.lm = LMHead()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.emb(x)
        for layer in self.layers:
            x = layer(x)
        x = self.lm(x)
        return x

If we print the model, we can see multiple hierarchies, which makes it hard to split by hand:

Model(
  (emb): Embedding(10, 3)
  (layers): ModuleList(
    (0-1): 2 x Layer(
      (lin): Linear(in_features=3, out_features=3, bias=True)
    )
  )
  (lm): LMHead(
    (proj): Linear(in_features=3, out_features=3, bias=True)
  )
)

Let us see how the pipeline API works:

from torch.distributed.pipelining import pipeline, SplitPoint

# An example micro-batch input
x = torch.LongTensor([1, 2, 4, 5])

pipe = pipeline(
    module=mod,
    mb_args=(x,),
    split_spec={
        "layers.1": SplitPoint.BEGINNING,
    }
)

The pipeline API splits your model given a split_spec, where SplitPoint.BEGINNING stands for adding a split point before execution of certain submodule in the forward function, and similarly, SplitPoint.END for split point after such.

If we print(pipe), we can see:

GraphModule(
  (submod_0): GraphModule(
    (emb): InterpreterModule()
    (layers): Module(
      (0): InterpreterModule(
        (lin): InterpreterModule()
      )
    )
  )
  (submod_1): GraphModule(
    (layers): Module(
      (1): InterpreterModule(
        (lin): InterpreterModule()
      )
    )
    (lm): InterpreterModule(
      (proj): InterpreterModule()
    )
  )
)

def forward(self, x):
    submod_0 = self.submod_0(x);  x = None
    submod_1 = self.submod_1(submod_0);  submod_0 = None
    return (submod_1,)

The “model partitions” are represented by submodules (submod_0, submod_1), each of which is reconstructed with original model operations, weights and hierarchies. In addition, a “root-level” forward function is reconstructed to capture the data flow between those partitions. Such data flow will be replayed by the pipeline runtime later, in a distributed fashion.

The Pipe object provides a method for retrieving the “model partitions”:

stage_mod : nn.Module = pipe.get_stage_module(stage_idx)

The returned stage_mod is a nn.Module, with which you can create an optimizer, save or load checkpoints, or apply other parallelisms.

Pipe also allows you to create a distributed stage runtime on a device given a ProcessGroup:

stage = pipe.build_stage(stage_idx, device, group)

Alternatively, if you would like to build the stage runtime later after some modification to the stage_mod, you can use a functional version of the build_stage API. For example:

from torch.distributed.pipelining import build_stage
from torch.nn.parallel import DistributedDataParallel

dp_mod = DistributedDataParallel(stage_mod)
info = pipe.info()
stage = build_stage(dp_mod, stage_idx, info, device, group)

Note

The pipeline frontend uses a tracer (torch.export) to capture your model into a single graph. If your model is not full-graph’able, you can use our manual frontend below.

Hugging Face Examples#

In the PiPPy repo where this package was original created, we kept examples based on unmodified Hugging Face models. See the examples/huggingface directory.

Examples include:

GPT2
Llama

Technical Deep Dive#

How does the `pipeline` API split a model?#

First, the pipeline API turns our model into a directed acyclic graph (DAG) by tracing the model. It traces the model using torch.export – a PyTorch 2 full-graph capturing tool.

Then, it groups together the operations and parameters needed by a stage into a reconstructed submodule: submod_0, submod_1, …

Different from conventional submodule access methods like Module.children(), the pipeline API does not only cut the module structure of your model, but also the forward function of your model.

This is necessary because model structure like Module.children() merely captures information during Module.__init__(), and does not capture any information about Module.forward(). Said differently, Module.children() lacks information about the following aspects key to pipelininig:

Execution order of child modules in forward
Activation flows between child modules
Whether there are any functional operators between child modules (for example, relu or add operations will not be captured by Module.children()).

The pipeline API, on the contrary, makes sure that the forward behavior is truly preserved. It also captures the activation flow between the partitions, helping the distributed runtime to make correct send/receive calls without human intervention.

Another flexibility of the pipeline API is that split points can be at arbitrary levels within your model hierarchy. In the split partitions, the original model hierarchy related to that partition will be reconstructed at no cost to you. At a result, fully-qualified names (FQNs) pointing to a submodule or parameter would be still valid, and services that relies on FQNs (such as FSDP, TP or checkpointing) can still run with your partitioned modules with almost zero code change.

Implementing Your Own Schedule#

You can implement your own pipeline schedule by extending one of the following two classes:

PipelineScheduleSingle
PipelineScheduleMulti

PipelineScheduleSingle is for schedules that assigns only one stage per rank. PipelineScheduleMulti is for schedules that assigns multiple stages per rank. All stages assigned to one rank must execute on one device; adjacent same-rank stages pass activations and gradients directly without device copies.

For example, ScheduleGPipe and Schedule1F1B are subclasses of PipelineScheduleSingle. Whereas, ScheduleInterleaved1F1B, ScheduleLoopedBFS, ScheduleInterleavedZeroBubble, and ScheduleZBVZeroBubble are subclasses of PipelineScheduleMulti.

Accessing Stage Forward Information#

A runtime surrounding a pipeline stage may need to select resources by the global logical stage and current microbatch. Physical pipeline rank is not a substitute for logical stage identity because one rank may own several stages in an interleaved schedule.

Register a context factory on the stage without changing the wrapped module’s forward signature:

from contextlib import contextmanager

from torch.distributed.pipelining import PipelineStageInfo


@contextmanager
def stage_forward_context(info: PipelineStageInfo):
    planner.enter(
        stage_index=info.stage_index,
        microbatch_index=info.microbatch_index,
        is_metadata_inference=info.is_metadata_inference,
    )
    try:
        yield
    finally:
        planner.exit()


handle = stage.register_forward_context(stage_forward_context)

The factory is called around each built-in stage forward. Dynamic metadata inference uses microbatch zero and sets is_metadata_inference=True, allowing a consumer to distinguish the representative probe from real microbatch-zero execution. Static metadata setup does not execute the module and therefore does not enter the context.

Only one context can be registered on a stage at a time. Remove its handle before registering another context. Register before the first schedule step if the consumer must observe dynamic metadata inference.

The context executes outside a compiled or exported stage module, so it does not add graph inputs or change the module signature. CUDA graph capture runs the Python context while recording the stage computation, but replay does not re-enter Python; consumers must bind replay-stable state during capture. A custom schedule action receives this context only when it delegates execution to stage.forward_one_chunk().

class torch.distributed.pipelining.PipelineStageInfo(*, stage_index, microbatch_index, is_metadata_inference=False)[source]#

Identify one pipeline-stage forward invocation.

This immutable value owns no tensors, process groups, schedule state, or storage. Future fields must have defaults so existing keyword-only construction remains source compatible.

Variables:

stage_index (int) – Global logical pipeline-stage index. A physical pipeline rank may own several logical stages.
microbatch_index (int) – Microbatch index within the current pipeline step.
is_metadata_inference (bool) – Whether this invocation is the representative dynamic metadata-inference probe rather than real execution.

PipelineStage.register_forward_context(context_factory)[source]#

Register a context manager around this stage’s forward computation.

The factory receives a PipelineStageInfo for every built-in runtime forward and dynamic metadata-inference probe. Static metadata setup does not execute the stage module and therefore does not invoke the factory. The returned handle removes the registration.

Only one context may be registered at a time. A custom schedule action receives the context only when it delegates execution to forward_one_chunk(). CUDA graph replay does not execute this Python callback; consumers must bind replay-stable state during capture.

Parameters:: context_factory (Callable[[PipelineStageInfo], AbstractContextManager[None]]) – Callable that returns a fresh context manager for each stage invocation. The context must not suppress exceptions raised by the stage module.
Returns:: A handle whose remove() method removes the registration.
Raises:: RuntimeError – If a forward context is already registered.
Return type:: RemovableHandle

Controlling FSDP Unshard Lookahead#

Multi-stage runtime schedules lower a compute-only schedule into explicit FSDP UNSHARD and RESHARD actions. Two parameters control different parts of that lowering:

max_active_stages is the target parameter-residency window. It determines which stages remain unsharded and where RESHARD actions are inserted.
unshard_lookahead is the issue-distance window. It determines how many upcoming distinct logical stages may begin unsharding.

Separating these windows allows a schedule to issue fewer all-gathers early without evicting parameters sooner or adding another unshard/reshard cycle. An asynchronous unshard is still real GPU work:

pre-all-gather cast or quantization
  -> copy-in and packing
  -> all-gather
  -> copy-out
  -> post-all-gather quantization or layout preparation
  -> parameter ready

async_op=True avoids a host-side wait, but these kernels, copies, collective traffic, allocations, and stream dependencies can still contend with the forward. Issuing the entire residency window at once can therefore put non-critical parameter preparation ahead of useful compute.

The "auto" policy estimates how much of this work fits into each rank’s pipeline startup bubble. Let a balanced stage forward take F, and let the composite critical-path cost of preparing one stage’s unsharded parameters be U. Ignoring pipeline transfer latency, rank r waits approximately U + rF before its first useful forward. This can complete approximately floor((U + rF) / U) unshards; issuing one more allows the next unshard to overlap that first forward. With the simplifying assumption F = U = T, the lookahead is r + 2, capped by max_active_stages.

For PP4 with max_active_stages=4, "auto" resolves to (2, 3, 4, 4):

interval           | 0..T    | T..2T   | 2T..3T  | 3T..4T  | 4T..5T
-------------------+---------+----------+----------+----------+---------
rank 0 forward     | blocked | F(first) |          |          |
rank 0 preparation | U0      | U1       |          |          |  => 2
-------------------+---------+----------+----------+----------+---------
rank 1 forward     | blocked | blocked  | F(first) |          |
rank 1 preparation | U0      | U1       | U2       |          |  => 3
-------------------+---------+----------+----------+----------+---------
rank 2 forward     | blocked | blocked  | blocked  | F(first) |
rank 2 preparation | U0      | U1       | U2       | U3       |  => 4
-------------------+---------+----------+----------+----------+---------
rank 3 forward     | blocked | blocked  | blocked  | blocked  | F(first)
rank 3 preparation | U0      | U1       | U2       | U3       |  => 4

The diagram is an analytical starting point, not an exact CUDA-stream model. Uk denotes preparation of the kth upcoming rank-local stage, not a global stage index. Vertically aligned preparation and forward cells are intended to overlap. Real stages may be unbalanced, unshard phases may overlap only partially, and network or memory-bandwidth contention may change the best distance. Choose a policy accordingly:

Policy	Per-rank issue distance	Intended use
`"full"`	`max_active_stages`	Compatibility default matching the original full-window behavior.
`"auto"`	`min(pp_rank + 2, max_active_stages)`	Deterministic startup-bubble estimate that avoids recipe-level tuning; not a universal optimum.
Tuple	The corresponding positive integer for each PP rank	Expert tuning for measured model, topology, and fabric behavior.

A custom compute-only schedule may be combined with a tuple before lowering. For exact UNSHARD and communication placement, supply an already lowered compute_comms schedule; unshard_lookahead cannot retune actions that are already present.

An atomic compound action remains indivisible, so it may extend either window by up to the action’s number of stages minus one. Non-full policies can also move absolute P2P action positions because unshards consume lowering rounds. They do not change compute order, RESHARD placement, residency episodes, or collective counts. Applications should benchmark an explicit tuple when stage costs differ materially from the balanced model.

Logging#

You can turn on additional logging using the TORCH_LOGS environment variable from torch._logging:

TORCH_LOGS=+pp will display logging.DEBUG messages and all levels above it.
TORCH_LOGS=pp will display logging.INFO messages and above.
TORCH_LOGS=-pp will display logging.WARNING messages and above.

API Reference#

Model Split APIs#

The following set of APIs transform your model into a pipeline representation.

class torch.distributed.pipelining.SplitPoint(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]#: Enum representing the points at which a split can occur in the execution of a submodule. :ivar BEGINNING: Represents adding a split point before the execution of a certain submodule in the forward function. :ivar END: Represents adding a split point after the execution of a certain submodule in the forward function.

torch.distributed.pipelining.pipeline(module, mb_args, mb_kwargs=None, split_spec=None, split_policy=None)[source]#

Split a module based on a specification.

See Pipe for more details.

Parameters:

module (Module) – The module to be split.
mb_args (tuple[Any, ...]) – Example positional inputs, in micro-batch form.
mb_kwargs (dict[str, Any] | None) – Example keyword inputs, in micro-batch form. (default: None)
split_spec (dict[str, SplitPoint] | None) – A dictionary using submodule names as split marker. (default: None)
split_policy (Callable[[GraphModule], GraphModule] | None) – The policy to use for splitting the module. (default: None)

Return type:

A pipeline representation of class Pipe.

class torch.distributed.pipelining.Pipe(split_gm, num_stages, has_loss_and_backward, loss_spec)[source]#

torch.distributed.pipelining.pipe_split()[source]#

pipe_split is a special operator that is used to mark the boundary between stages in a module. It is used to split the module into stages. It is a no-op if your annotated module is run eagerly.

Example

>>> def forward(self, x):
>>>     x = torch.mm(x, self.mm_param)
>>>     x = torch.relu(x)
>>>     pipe_split()
>>>     x = self.lin(x)
>>>     return x

The above example will be split into two stages.

Microbatch Utilities#

class torch.distributed.pipelining.microbatch.TensorChunkSpec(split_dim)[source]#

Class used to specify chunking of inputs

torch.distributed.pipelining.microbatch.split_args_kwargs_into_chunks(args, kwargs, chunks, args_chunk_spec=None, kwargs_chunk_spec=None)[source]#

Given a sequence of args and kwargs, split them into a number of chunks according to their respective chunking specs.

Parameters:

args (tuple[Any, ...]) – Tuple of args
kwargs (dict[str, Any] | None) – Dict of kwargs
chunks (int) – Number of chunks to split the args and kwargs into
args_chunk_spec (tuple[TensorChunkSpec, ...] | None) – chunking specs for args, in same shape as args
kwargs_chunk_spec (dict[str, TensorChunkSpec] | None) – chunking specs for kwargs, in same shape as kwargs

Returns:

List of sharded args kwargs_split: List of sharded kwargs

Return type:

args_split

torch.distributed.pipelining.microbatch.merge_chunks(chunks, chunk_spec)[source]#

Given a list of chunks, merge them into a single value according to the chunk spec.

Parameters:

chunks (list[Any]) – list of chunks
chunk_spec – Chunking spec for the chunks

Returns:

Merged value

Return type:

value

Pipeline Stages#

class torch.distributed.pipelining.stage.PipelineStage(submodule, stage_index, num_stages, device, input_args=None, output_args=None, output_grads=None, input_grads=None, group=None, dw_builder=None, get_mesh=None)[source]#

A pipeline stage for pipeline parallelism with sequential model partitioning.

Supports both static and dynamic metadata inference:

Static mode:: All of input_args, output_args (and input_grads/output_grads when DTensors are present) are provided at construction time.
Dynamic mode:: Metadata is inferred from the first microbatch at runtime; any statically provided args are used for validation only.

Parameters:

submodule (Module) – The nn.Module wrapped by this stage.
stage_index (int) – Zero-based stage ID.
num_stages (int) – Total number of stages in the pipeline.
device (device) – Device this stage runs on.
input_args (Tensor | tuple[Tensor, ...] | None) – Example input tensors (single tensor or tuple). Optional.
output_args (Tensor | tuple[Tensor, ...] | None) – Example output tensors. Optional.
output_grads (Tensor | tuple[Tensor | None, ...] | None) – Example output gradients (received from next stage). Optional.
input_grads (Tensor | tuple[Tensor | None, ...] | None) – Example input gradients (sent to previous stage). Optional.
group (ProcessGroup | None) – Process group for P2P communication. Defaults to the world process group.
dw_builder (Callable[[], Callable[[...], None]] | None) – Builder for deferred weight-update runners used by zero-bubble (F/I/W) schedules.
get_mesh (GetMeshCallback | None) – GetMeshCallback used during dynamic DTensor inference. Ignored in fully static DTensor mode.

torch.distributed.pipelining.stage.build_stage(stage_module, stage_index, pipe_info, device, group=None)[source]#

Create a pipeline stage given a stage_module to be wrapped by this stage and pipeline information.

Parameters:

stage_module (torch.nn.Module) – the module to be wrapped by this stage
stage_index (int) – the index of this stage in the pipeline
pipe_info (PipeInfo) – information about the pipeline, can be retrieved by pipe.info()
device (torch.device) – the device to be used by this stage
group (Optional[dist.ProcessGroup]) – the process group to be used by this stage

Returns:

a pipeline stage that can run with PipelineSchedules.

Return type:

_PipelineStage

Pipeline Schedules#

Activation-liveness analysis#

Pipeline runtimes and memory planners can use analyze_pipeline_activation_liveness to determine how many reusable logical slots are needed for activations retained from forward through backward. The analysis does not allocate tensors. It returns a PipelineActivationLiveness plan whose slot_by_stage_and_microbatch[(stage_index, microbatch_index)] values are slot IDs that a caller may map to buffers or arena regions.

An activation becomes live at its forward (F) action. Full backward (B) releases it. For schedules that separate input backward (I) from weight backward (W), I does not release the activation because W may still need the saved forward state; W releases it. Lifetimes include both endpoint positions, so actions grouped into the same compound schedule position overlap.

For example, consider two stages and two microbatches on one pipeline rank:

position:  0     1     2     3     4     5     6     7
action:   F0,0  F1,0  F0,1  B1,0  B0,0  F1,1  B1,1  B0,1

With granularity="stage_microbatch", the four activation lifetimes are (0, 4), (1, 3), (2, 7), and (5, 6); the final lifetime may reuse the first slot. With granularity="microbatch", the selected stages for each microbatch share one conservative lifetime: (0, 4) for microbatch 0 and (2, 7) for microbatch 1. The latter mode is useful when a consumer manages all selected stages for one microbatch as one storage unit. Stage indices are global logical indices and commonly identify virtual stages hosted by the same pipeline rank.

torch.distributed.pipelining.schedules.analyze_pipeline_activation_liveness(schedule, *, pp_rank, stage_indices, granularity)[source]#

Analyze activation lifetimes and assign reusable pipeline slots.

The analysis follows the finalized action sequence for pp_rank. A forward action starts the lifetime of activation state retained for backward. Full backward ends that lifetime; when input and weight backward are split, weight backward ends it. Input backward does not release the activation because weight backward may still consume it. Both interval endpoints are inclusive, including sub-actions in one compound schedule position.

This utility only computes liveness and logical slot IDs. It does not allocate tensors or prescribe a storage layout. Consumers can use the plan to reason about peak activation memory or map slots to stable buffers and arena regions.

Parameters:

schedule (PipelineScheduleMulti) – Constructed multi-stage schedule whose finalized order will execute the pipeline step.
pp_rank (int) – Rank within the pipeline process group to analyze. This may differ from the rank constructing schedule because every rank holds the complete finalized schedule.
stage_indices (Sequence[int]) – Global logical stages whose activations should be analyzed. One pipeline rank commonly hosts multiple such virtual stages. Every selected stage must execute once per microbatch on pp_rank.
granularity (Literal['microbatch', 'stage_microbatch']) – "stage_microbatch" assigns independent lifetimes to every selected (stage, microbatch) pair. "microbatch" uses one lifetime per microbatch spanning the earliest selected-stage forward through the latest selected-stage release.

Returns:

An immutable activation-liveness plan. Its PipelineActivationLiveness.slot_by_stage_and_microbatch mapping contains a logical slot ID for every selected stage and microbatch.

Raises:

ValueError – If the granularity or stage selection is invalid, the rank is absent, or the finalized schedule does not contain one valid forward-to-backward activation lifetime for every selected pair.

Return type:

PipelineActivationLiveness

class torch.distributed.pipelining.schedules.PipelineActivationLiveness(pp_rank, granularity, stage_indices, num_microbatches, slot_by_stage_and_microbatch, num_slots, _activation_lifetime_by_stage_and_microbatch)[source]#

Reusable activation slots derived from a pipeline schedule.

This result describes logical storage slots; it does not allocate or own tensors. Each key in slot_by_stage_and_microbatch is a global (stage_index, microbatch_index) pair, and each value is the integer ID of the slot assigned to that activation. A caller may map those IDs to its own buffers or arena regions. If activations have different sizes, the caller must size a shared slot for every activation assigned to it or group compatible activations before allocating storage.

An activation becomes live when its forward action executes and remains live through its full-backward or weight-backward action. Input-backward alone does not end the lifetime because a later weight-backward action may still need the forward activation. Schedule positions are inclusive, so two lifetimes that end and begin in the same compound action cannot share a slot.

With granularity="stage_microbatch", each stage and microbatch pair has its own lifetime. With granularity="microbatch", all selected stages for a microbatch share one lifetime from the earliest forward through the latest release, and therefore share one slot ID.

Instances are returned by analyze_pipeline_activation_liveness().

Variables:

pp_rank (int) – Rank in the pipeline process group whose schedule was analyzed. This is not necessarily the global distributed rank.
granularity (Literal['microbatch', 'stage_microbatch']) – Whether slots represent individual stage/microbatch pairs or whole microbatches across all selected stages.
stage_indices (tuple[int, ...]) – Global logical pipeline-stage indices included in the analysis. These commonly identify virtual stages hosted by pp_rank.
num_microbatches (int) – Number of microbatches in the analyzed schedule.
slot_by_stage_and_microbatch (collections.abc.Mapping[tuple[int, int], int]) – Immutable mapping from (stage_index, microbatch_index) to logical activation-slot ID.
num_slots (int) – Number of logical slots required by the assignment.

get_activation_lifetime(stage_index, microbatch_index)[source]#

Return the inclusive schedule interval for an activation.

The returned (forward_position, release_position) indexes the finalized action sequence analyzed for pp_rank. Runtime schedules use their compute-and-communication sequence; other multi-stage schedules use their compute sequence.

At stage_microbatch granularity, the interval belongs to the exact stage and microbatch pair. At microbatch granularity, every selected stage for the microbatch returns the same interval: the convex hull from the earliest selected-stage forward to the latest selected-stage full- or weight-backward action.

Parameters:

stage_index (int) – Global logical pipeline-stage index.
microbatch_index (int) – Microbatch index within the pipeline step.

Returns:

The inclusive (forward_position, release_position) interval.

Raises:

ValueError – If the pair was not included in the analysis.

Return type:

tuple[int, int]

slot_for(stage_index, microbatch_index)[source]#

Return the activation slot for a stage and microbatch.

Parameters:

stage_index (int) – Global logical pipeline-stage index.
microbatch_index (int) – Microbatch index within the pipeline step.

Returns:

The logical activation-slot ID assigned to the pair.

Raises:

ValueError – If the pair was not included in the analysis.

Return type:

int

class torch.distributed.pipelining.schedules.ScheduleGPipe(stage, n_microbatches, loss_fn=None, args_chunk_spec=None, kwargs_chunk_spec=None, output_merge_spec=None, scale_grads=True)[source]#

The GPipe schedule. Will go through all the microbatches in a fill-drain manner.

class torch.distributed.pipelining.schedules.Schedule1F1B(stage, n_microbatches, loss_fn=None, args_chunk_spec=None, kwargs_chunk_spec=None, output_merge_spec=None, scale_grads=True)[source]#

The 1F1B schedule. Will perform one forward and one backward on the microbatches in steady state.

class torch.distributed.pipelining.schedules.ScheduleInterleaved1F1B(stages, n_microbatches, loss_fn=None, args_chunk_spec=None, kwargs_chunk_spec=None, output_merge_spec=None, scale_grads=True, backward_requires_autograd=True, defer_pp_recv=False, max_active_stages=3, unshard_lookahead='full', *, defer_reduce_grad_wait=False, max_outstanding_sends=None)[source]#

The Interleaved 1F1B schedule. See https://arxiv.org/pdf/2104.04473 for details. Will perform one forward and one backward on the microbatches in steady state and supports multiple stages per rank. When microbatches are ready for multiple local stages, Interleaved 1F1B prioritizes the earlier microbatch (also called “depth first”).

This schedule is mostly similar to the original paper. It differs by being relaxing the requirement of num_microbatch % pp_size == 0. Using the flex_pp schedule, we will have num_rounds = max(1, n_microbatches // pp_group_size) and it works as long as n_microbatches % num_rounds is 0. As a few examples, support

pp_group_size = 4, n_microbatches = 10. We will have num_rounds = 2 and n_microbatches % 2 is 0.
pp_group_size = 4, n_microbatches = 3. We will have num_rounds = 1 and n_microbatches % 1 is 0.

Parameters:

max_active_stages (int) – Positive target number of local FSDP stages whose unsharded parameters may remain resident. An atomic compound action may transiently require up to its size minus one additional stages.
unshard_lookahead (Literal['full', 'auto'] | tuple[python:int, ...]) – Number of upcoming distinct stages whose asynchronous unshards may be issued. "full" matches max_active_stages; "auto" uses min(rank + 2, max_active_stages); a tuple supplies one positive integer per pipeline rank.
defer_reduce_grad_wait (bool) – If True, delay each FSDP WAIT_REDUCE_GRAD until before the next reduction or schedule end. Otherwise, each wait follows its REDUCE_GRAD. At most one reduction may be pending per pipeline rank.
max_outstanding_sends (int | None) – Maximum number of pending forward and backward send batches on each pipeline rank. None applies no hard limit. A smaller limit may move a wait before its receive is guaranteed to be posted and may stall. Schedule construction raises ValueError when no causally safe wait placement can meet the limit.

class torch.distributed.pipelining.schedules.ScheduleLoopedBFS(stages, n_microbatches, loss_fn=None, output_merge_spec=None, scale_grads=True, backward_requires_autograd=True, defer_pp_recv=False, max_active_stages=3, unshard_lookahead='full', *, defer_reduce_grad_wait=False, max_outstanding_sends=None)[source]#

Breadth-First Pipeline Parallelism. See https://arxiv.org/abs/2211.05953 for details. Similar to Interleaved 1F1B, Looped BFS supports multiple stages per rank. What is different is that when microbatches are ready for multiple local stages, Loops BFS will prioritizes the earlier stage, running all available microbatches at once.

Parameters:

max_active_stages (int) – Positive target number of local FSDP stages whose unsharded parameters may remain resident. An atomic compound action may transiently require up to its size minus one additional stages.
unshard_lookahead (Literal['full', 'auto'] | tuple[python:int, ...]) – Number of upcoming distinct stages whose asynchronous unshards may be issued. "full" matches max_active_stages; "auto" uses min(rank + 2, max_active_stages); a tuple supplies one positive integer per pipeline rank.
defer_reduce_grad_wait (bool) – If True, delay each FSDP WAIT_REDUCE_GRAD until before the next reduction or schedule end. Otherwise, each wait follows its REDUCE_GRAD. At most one reduction may be pending per pipeline rank.
max_outstanding_sends (int | None) – Maximum number of pending forward and backward send batches on each pipeline rank. None applies no hard limit.

class torch.distributed.pipelining.schedules.ScheduleInterleavedZeroBubble(stages, n_microbatches, loss_fn=None, args_chunk_spec=None, kwargs_chunk_spec=None, output_merge_spec=None, scale_grads=True, backward_requires_autograd=True, defer_pp_recv=False, max_active_stages=3, unshard_lookahead='full', *, defer_reduce_grad_wait=False, max_outstanding_sends=None)[source]#

The Interleaved Zero Bubble schedule. See https://arxiv.org/pdf/2401.10241 for details. Will perform one forward and one backward on inputs for the microbatches in steady state and supports multiple stages per rank. Uses the backward for weights to fill in the pipeline bubble.

In particular this is implementing the ZB1P schedule in the paper.

Parameters:

max_active_stages (int) – Positive target number of local FSDP stages whose unsharded parameters may remain resident. An atomic compound action may transiently require up to its size minus one additional stages.
unshard_lookahead (Literal['full', 'auto'] | tuple[python:int, ...]) – Number of upcoming distinct stages whose asynchronous unshards may be issued. "full" matches max_active_stages; "auto" uses min(rank + 2, max_active_stages); a tuple supplies one positive integer per pipeline rank.
defer_reduce_grad_wait (bool) – If True, delay each FSDP WAIT_REDUCE_GRAD until before the next reduction or schedule end. Otherwise, each wait follows its REDUCE_GRAD. At most one reduction may be pending per pipeline rank.
max_outstanding_sends (int | None) – Maximum number of pending forward and backward send batches on each pipeline rank. None applies no hard limit.

class torch.distributed.pipelining.schedules.ScheduleZBVZeroBubble(stages, n_microbatches, loss_fn=None, args_chunk_spec=None, kwargs_chunk_spec=None, output_merge_spec=None, scale_grads=True, backward_requires_autograd=True, defer_pp_recv=False, max_active_stages=3, unshard_lookahead='full', *, defer_reduce_grad_wait=False, max_outstanding_sends=None)[source]#

The Zero Bubble schedule (ZBV variant). See https://arxiv.org/pdf/2401.10241 Section 6 for details.

This schedules requires exactly two stages per rank.

This schedule will perform one forward and one backward on inputs for the microbatches in steady state and supports multiple stages per rank. Uses backward with respect to weights to fill in the pipeline bubble.

This ZB-V schedule would have the “zero bubble” property only if time forward == time backward input == time backward weights. In practice, this is not likely true for real models so alternatively a greedy scheduler could be implemented for unequal/unbalanced time.

Parameters:

max_active_stages (int) – Positive target number of local FSDP stages whose unsharded parameters may remain resident. An atomic compound action may transiently require up to its size minus one additional stages.
unshard_lookahead (Literal['full', 'auto'] | tuple[python:int, ...]) – Number of upcoming distinct stages whose asynchronous unshards may be issued. "full" matches max_active_stages; "auto" uses min(rank + 2, max_active_stages); a tuple supplies one positive integer per pipeline rank.
defer_reduce_grad_wait (bool) – If True, delay each FSDP WAIT_REDUCE_GRAD until before the next reduction or schedule end. Otherwise, each wait follows its REDUCE_GRAD. At most one reduction may be pending per pipeline rank.
max_outstanding_sends (int | None) – Maximum number of pending forward and backward send batches on each pipeline rank. None applies no hard limit.

class torch.distributed.pipelining.schedules.ScheduleDualPipeV(stages, n_microbatches, loss_fn=None, args_chunk_spec=None, kwargs_chunk_spec=None, output_merge_spec=None, scale_grads=True, backward_requires_autograd=True, defer_pp_recv=False, max_active_stages=3, unshard_lookahead='full', *, defer_reduce_grad_wait=False, max_outstanding_sends=None)[source]#

The DualPipeV schedule. A more efficient schedule variant based on the DualPipe schedule introduced by DeepSeek in https://arxiv.org/pdf/2412.19437

Based on the open sourced code from deepseek-ai/DualPipe

Parameters:

max_active_stages (int) – Positive target number of local FSDP stages whose unsharded parameters may remain resident. An atomic compound action may transiently require up to its size minus one additional stages.
unshard_lookahead (Literal['full', 'auto'] | tuple[python:int, ...]) – Number of upcoming distinct stages whose asynchronous unshards may be issued. "full" matches max_active_stages; "auto" uses min(rank + 2, max_active_stages); a tuple supplies one positive integer per pipeline rank.
defer_reduce_grad_wait (bool) – If True, delay each FSDP WAIT_REDUCE_GRAD until before the next reduction or schedule end. Otherwise, each wait follows its REDUCE_GRAD. At most one reduction may be pending per pipeline rank.
max_outstanding_sends (int | None) – Maximum number of pending forward and backward send batches on each pipeline rank. None applies no hard limit.

class torch.distributed.pipelining.schedules.PipelineScheduleSingle(stage, n_microbatches, loss_fn=None, args_chunk_spec=None, kwargs_chunk_spec=None, output_merge_spec=None, scale_grads=True)[source]#

Base class for single-stage schedules. Implements the step method. Derived classes should implement _step_microbatches.

Gradients are scaled by num_microbatches depending on the scale_grads argument, defaulting to True. This setting should match the configuration of your loss_fn, which may either average losses (scale_grads=True) or sum losses (scale_grads=False).

step(*args, target=None, losses=None, return_outputs=True, loss_kwargs=None, arg_mbs=None, kwarg_mbs=None, target_mbs=None, loss_kwarg_mbs=None, finalize_gradients=True, **kwargs)[source]#

Run one training iteration of a single-stage pipeline schedule.

By default, args, kwargs, and target are full-batch values that the schedule splits into microbatches. When the caller has already split inputs into microbatches, pass them through arg_mbs, kwarg_mbs, and target_mbs instead.

Parameters:

*args (Any) – Whole-batch positional inputs for the first pipeline stage. Do not pass positional inputs with pre-split inputs.
target (Any, optional) – Whole-batch target for loss computation. When passing pre-split inputs, pass targets through target_mbs instead. Default: None.
losses (list, optional) – Mutable list populated with one loss per microbatch when this schedule owns the last stage and a loss_fn was configured. Default: None.
return_outputs (bool, optional) – Whether to merge and return output chunks on the last stage. Default: True.
loss_kwargs (dict, optional) – Keyword arguments passed unchanged to the configured loss_fn for every microbatch. Unlike kwargs, they are not split. Default: None.
arg_mbs (list[tuple], optional) – Pre-split positional inputs, one tuple per microbatch. Default: None.
kwarg_mbs (list[dict], optional) – Pre-split keyword inputs, one dict per microbatch. Default: None.
target_mbs (list, optional) – Pre-split targets, one entry per microbatch. Default: None.
loss_kwarg_mbs (list[dict], optional) – Pre-split loss keyword arguments, one dict per microbatch, such as per-token tensors. Like the other pre-split inputs, pass targets through target_mbs with it. Microbatch i computes loss_fn(output, target_mbs[i], **loss_kwargs, **loss_kwarg_mbs[i]), so a key must not appear in both. While inferring stage metadata (on the first step by default), the schedule also calls loss_fn once with microbatch 0’s inputs, so loss_fn must not accumulate into objects passed this way. Default: None.
finalize_gradients (bool, optional) – Whether to reduce accumulated FSDP gradients and reshard FSDP parameters at the end of this call. If False, FSDP parameters stay unsharded, and a later call with True is required before reading .grad. With scale_grads=True, the final call divides by n_microbatches only. Divide by the number of accumulated calls separately if needed. This has no effect on non-FSDP stages. Default: True.
**kwargs (Any) – Whole-batch keyword inputs for the first pipeline stage. Do not pass keyword inputs with pre-split inputs.

Returns:

The merged output from the last stage when this rank: owns the last stage and return_outputs=True; otherwise None.

Return type:

Any or None

Raises:

RuntimeError – If backward computation is enabled and step is called under torch.no_grad(). Use eval() for forward-only execution.
TypeError – If a pre-split microbatch container has the wrong type.
ValueError – If whole-batch and pre-split inputs are mixed, or if a pre-split container does not have one entry per microbatch.

Examples:

>>> output = schedule.step(x, mask=mask, target=target)
>>> arg_mbs = [(x0,), (x1,)]
>>> kwarg_mbs = [{"mask": mask0}, {"mask": mask1}]
>>> target_mbs = [target0, target1]
>>> output = schedule.step(
...     arg_mbs=arg_mbs,
...     kwarg_mbs=kwarg_mbs,
...     target_mbs=target_mbs,
... )

class torch.distributed.pipelining.schedules.PipelineScheduleMulti(stages, n_microbatches, loss_fn=None, args_chunk_spec=None, kwargs_chunk_spec=None, output_merge_spec=None, use_full_backward=None, scale_grads=True, backward_requires_autograd=True)[source]#

Base class for multi-stage schedules. Implements the step method.

Gradients are scaled by num_microbatches depending on the scale_grads argument, defaulting to True. This setting should match the configuration of your loss_fn, which may either average losses (scale_grads=True) or sum losses (scale_grads=False).

step(*args, target=None, losses=None, return_outputs=True, loss_kwargs=None, arg_mbs=None, kwarg_mbs=None, target_mbs=None, loss_kwarg_mbs=None, finalize_gradients=True, **kwargs)[source]#

Run one training iteration of a multi-stage pipeline schedule.

By default, args, kwargs, and target are full-batch values that the schedule splits into microbatches. When the caller has already split inputs into microbatches, pass them through arg_mbs, kwarg_mbs, and target_mbs instead.

Parameters:

*args (Any) – Whole-batch positional root inputs when this rank owns the first pipeline stage. Do not pass positional inputs with pre-split inputs.
target (Any, optional) – Whole-batch target for loss computation. When passing pre-split inputs, pass targets through target_mbs instead. Default: None.
losses (list, optional) – Mutable list populated with one loss per microbatch when this schedule owns the last stage and a loss_fn was configured. Default: None.
return_outputs (bool, optional) – Whether to merge and return output chunks on the last stage. Default: True.
loss_kwargs (dict, optional) – Keyword arguments passed unchanged to the configured loss_fn for every microbatch. Unlike kwargs, they are not split. Default: None.
arg_mbs (list[tuple], optional) – Pre-split positional inputs, one tuple per microbatch. Default: None.
kwarg_mbs (list[dict], optional) – Pre-split keyword inputs, one dict per microbatch. Default: None.
target_mbs (list, optional) – Pre-split targets, one entry per microbatch. Default: None.
loss_kwarg_mbs (list[dict], optional) – Pre-split loss keyword arguments, one dict per microbatch, such as per-token tensors. Like the other pre-split inputs, pass targets through target_mbs with it. Microbatch i computes loss_fn(output, target_mbs[i], **loss_kwargs, **loss_kwarg_mbs[i]), so a key must not appear in both. While inferring stage metadata (on the first step by default), the schedule also calls loss_fn once with microbatch 0’s inputs, so loss_fn must not accumulate into objects passed this way. Default: None.
finalize_gradients (bool, optional) – Whether to reduce accumulated FSDP gradients and reshard FSDP parameters at the end of this call. If False, FSDP parameters stay unsharded, and a later call with True is required before reading .grad. With scale_grads=True, the final call divides by n_microbatches only. Divide by the number of accumulated calls separately if needed. This has no effect on non-FSDP stages. Default: True.
**kwargs (Any) – Whole-batch keyword root inputs when this rank owns the first pipeline stage. Do not pass keyword inputs with pre-split inputs.

Returns:

The merged output from the last stage when this rank: owns the last stage and return_outputs=True; otherwise None.

Return type:

Any or None

Raises:

RuntimeError – If backward computation is enabled and step is called under torch.no_grad(). Use eval() for forward-only execution.
TypeError – If a pre-split microbatch container has the wrong type.
ValueError – If whole-batch and pre-split inputs are mixed, or if a pre-split container does not have one entry per microbatch.

Examples:

>>> output = schedule.step(x, mask=mask, target=target)
>>> arg_mbs = [(x0,), (x1,)]
>>> kwarg_mbs = [{"mask": mask0}, {"mask": mask1}]
>>> target_mbs = [target0, target1]
>>> output = schedule.step(
...     arg_mbs=arg_mbs,
...     kwarg_mbs=kwarg_mbs,
...     target_mbs=target_mbs,
... )

torch.distributed.pipelining.schedules.get_schedule_class(schedule_name)[source]#

Maps a schedule name (case insensitive) to its corresponding class object.

Parameters:: schedule_name (str) – The name of the schedule.

Pipeline Parallelism#

Why Pipeline Parallel?#

What is `torch.distributed.pipelining`?#

Step 1: build `PipelineStage`#

Step 2: use `PipelineSchedule` for execution#

Options for Splitting a Model#

Option 1: splitting a model manually#

Option 2: splitting a model automatically#

Hugging Face Examples#

Technical Deep Dive#

How does the `pipeline` API split a model?#

Implementing Your Own Schedule#

Accessing Stage Forward Information#

Controlling FSDP Unshard Lookahead#

Logging#

API Reference#

Model Split APIs#

Microbatch Utilities#

Pipeline Stages#

Pipeline Schedules#

Activation-liveness analysis#

Docs

Tutorials

Resources

Pipeline Parallelism#

Why Pipeline Parallel?#

What is torch.distributed.pipelining?#

Step 1: build PipelineStage#

Step 2: use PipelineSchedule for execution#

Options for Splitting a Model#

Option 1: splitting a model manually#

Option 2: splitting a model automatically#

Hugging Face Examples#

Technical Deep Dive#

How does the pipeline API split a model?#

Implementing Your Own Schedule#

Accessing Stage Forward Information#

Controlling FSDP Unshard Lookahead#

Logging#

API Reference#

Model Split APIs#

Microbatch Utilities#

Pipeline Stages#

Pipeline Schedules#

Activation-liveness analysis#

Docs

Tutorials

Resources

What is `torch.distributed.pipelining`?#

Step 1: build `PipelineStage`#

Step 2: use `PipelineSchedule` for execution#

How does the `pipeline` API split a model?#