Pipeline Parallelism

Created On: Jun 16, 2025 | Last Updated On: Aug 13, 2025

注意

torch.distributed.pipelining 目前处于 Alpha 阶段，正在开发中。API 可能会有变更。它从 PiPPy 项目迁移而来。

Why Pipeline Parallel?

Pipeline Parallelism 是深度学习中一种 **基础** 的并行方式。它允许将模型的 **执行** 分割，从而使多个 **微批次 (micro-batches)** 可以并发地执行模型的不同部分的代码。Pipeline parallelism 可以是有效的技术，用于

• 大规模训练

• 带宽受限的集群

• 大型模型推理

上述场景的共同点是，每台设备的计算量无法掩盖传统并行方式（例如 FSDP 的权重 all-gather）的通信开销。

What is torch.distributed.pipelining?

虽然在扩展性方面前景广阔，但流水线并行化 (pipelining) 通常难以实现，因为它除了分割模型权重外，还需要 **分割模型的执行**。分割执行通常需要对模型进行侵入式代码更改。另一个复杂性来自于 **在分布式环境中调度微批次**，并考虑 **数据流依赖**。

该 pipelining 包提供了一个 **自动** 完成这些事情的工具集，可以轻松地在 **通用** 模型上实现流水线并行。

它由两部分组成：一个 **分割前端** 和一个 **分布式运行时**。分割前端接收您的模型代码，将其分割成“模型分区”，并捕获数据流关系。分布式运行时在不同设备上并行执行流水线阶段，处理微批次分割、调度、通信和梯度传播等事宜。

总的来说，pipelining 包提供了以下功能：

• 基于简单规范的}(\text{model})代码分割。

• 丰富的流水线调度支持，包括 GPipe、1F1B、Interleaved 1F1B 和 Looped BFS，并提供编写自定义调度的基础设施。

• 对跨主机流水线并行的一流支持，因为 PP 通常在此场景下使用（跨越较慢的互连）。

• 与其他 PyTorch 并行技术（如数据并行（DDP、FSDP）或张量并行）的可组合性。 TorchTitan 项目展示了在 Llama 模型上的“3D 并行”应用。

Step 1: build PipelineStage

在使用 PipelineSchedule 之前，我们需要创建 PipelineStage 对象，它们包装了在该阶段运行的模型部分。 PipelineStage 负责分配通信缓冲区并创建发送/接收操作与对等节点通信。它管理中间缓冲区（例如，尚未使用的前向传播的输出），并提供一个用于运行该阶段模型后向传播的实用程序。

PipelineStage 需要知道阶段模型的输入和输出形状，以便正确分配通信缓冲区。形状必须是静态的，例如，在运行时，形状不能从一个步骤变到另一个步骤。如果运行时形状与预期形状不匹配，将引发 PipeliningShapeError 类。在与其他并行技术组合或应用混合精度时，必须考虑这些技术，以便 PipelineStage 在运行时了解阶段模块输出的正确形状（和数据类型）。

用户可以直接构造 PipelineStage 实例，传入一个代表应在该阶段运行的模型部分的 nn.Module。这可能需要更改原始模型代码。请参阅 选项 1：手动分割模型 中的示例。

或者，分割前端可以使用图分区技术自动将您的模型分割成一系列 nn.Module。此技术要求模型可以使用 torch.Export 进行跟踪。生成的 nn.Module 与其他并行技术的组合是实验性的，可能需要一些变通方法。如果用户无法轻松更改模型代码，使用此前端可能更具吸引力。有关更多信息，请参阅 选项 2：自动分割模型。

Step 2: use PipelineSchedule for execution

现在我们可以将 PipelineStage 附加到流水线调度中，并使用输入数据运行调度。以下是一个 GPipe 示例：

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()

请注意，上面的代码需要为每个工作进程启动，因此我们使用一个启动器服务来启动多个进程。

torchrun --nproc_per_node=2 example.py

Options for Splitting a Model

Option 1: splitting a model manually

要直接构造 PipelineStage，用户负责提供一个单独的 nn.Module 实例，该实例拥有相关的 nn.Parameters 和 nn.Buffers，并定义一个 forward() 方法来执行与该阶段相关的操作。例如，TorchTitan 中定义的 Transformer 类的精简版本展示了一种构建易于分割的模型的方法。

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

以这种方式定义的模型可以通过先使用 meta-device 初始化整个模型（以避免 OOM 错误）、删除该阶段不需要的层，然后创建包装模型的 PipelineStage 来为每个阶段进行配置。例如：

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,
    )

当与其他数据并行或模型并行技术组合时，如果模型块的输出形状/数据类型受到影响，可能还需要 output_args。

Option 2: splitting a model automatically

如果您有一个完整的模型，并且不想花时间将其修改为一系列“模型分区”，那么 pipeline API 可以提供帮助。这是一个简化的示例：

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

如果我们打印模型，可以看到多个层级，这使得手动分割变得困难。

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)
  )
)

让我们看看 pipeline API 的工作原理：

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,
    }
)

该 pipeline API 根据 split_spec 分割您的模型，其中 SplitPoint.BEGINNING 表示在 forward 函数中 **某个子模块执行之前** 添加一个分割点，类似地，SplitPoint.END 表示 **在此之后** 添加一个分割点。

如果我们 print(pipe)，我们可以看到：

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,)

“模型分区”由子模块（submod_0、submod_1）表示，每个子模块都使用原始模型的操作、权重和层级重建。此外，还会重建一个“根级” forward 函数来捕获这些分区之间的数据流。流水线运行时稍后将以分布式方式重放这种数据流。

该 Pipe 对象提供了一个检索“模型分区”的方法：

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

返回的 stage_mod 是一个 nn.Module，您可以使用它来创建优化器、保存或加载检查点，或应用其他并行技术。

Pipe 还允许您在给定 ProcessGroup 的设备上创建分布式阶段运行时：

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

或者，如果您想在修改 stage_mod 后再构建阶段运行时，可以使用 build_stage API 的函数式版本。例如：

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)

注意

该 pipeline 前端使用跟踪器 (torch.export) 将您的模型捕获为单个图。如果您的模型无法完全图化，可以使用下面的手动前端。

Hugging Face Examples

在此包最初创建的 PiPPy 仓库中，我们保留了基于未修改的 Hugging Face 模型的示例。请参阅 examples/huggingface 目录。

示例包括：

• GPT2

• Llama

Technical Deep Dive

How does the pipeline API split a model?

首先，pipeline API 通过跟踪模型将其转换为有向无环图 (DAG)。它使用 torch.export（一个 PyTorch 2 的全图捕获工具）来跟踪模型。

然后，它将一个阶段所需的 **操作和参数** 组合到一个重建的子模块中：submod_0、submod_1、……

与传统的子模块访问方法（如 Module.children()）不同，pipeline API 不仅切割了您模型的模块结构，还切割了您模型的 **前向传播** 函数。

这是必要的，因为像 Module.children() 这样的模型结构仅在 Module.__init__() 期间捕获信息，而不捕获关于 Module.forward() 的任何信息。换句话说，Module.children() 缺少关于以下对流水线至关重要方面的信息：

• forward 中子模块的执行顺序

• 子模块之间的激活流

• 子模块之间是否存在任何函数式运算符（例如，relu 或 add 操作不会被 Module.children() 捕获）。

相反，pipeline API 确保 **前向传播** 的行为得到真正保留。它还捕获分区之间的激活流，帮助分布式运行时在无需人工干预的情况下做出正确的发送/接收调用。

该 pipeline API 的另一个灵活性在于，分割点可以位于模型层级内的任意级别。在分割的分区中，与该分区相关的原始模型层级将免费重建。因此，指向子模块或参数的完全限定名 (FQNs) 仍然有效，并且依赖于 FQNs 的服务（如 FSDP、TP 或 checkpointing）仍然可以几乎零代码更改地与您的分区模块一起运行。

Implementing Your Own Schedule

您可以通过扩展以下两个类之一来实现自己的流水线调度：

• PipelineScheduleSingle

• PipelineScheduleMulti

PipelineScheduleSingle 适用于 **只分配一个** 阶段到每个 rank 的调度。 PipelineScheduleMulti 适用于将多个阶段分配到每个 rank 的调度。

例如，ScheduleGPipe 和 Schedule1F1B 是 PipelineScheduleSingle 的子类。而 ScheduleInterleaved1F1B、ScheduleLoopedBFS、ScheduleInterleavedZeroBubble 和 ScheduleZBVZeroBubble 是 PipelineScheduleMulti 的子类。

日志记录

您可以使用 torch._logging 的 TORCH_LOGS 环境变量来开启额外的日志记录。

• TORCH_LOGS=+pp 将显示 logging.DEBUG 消息及其以上所有级别。

• TORCH_LOGS=pp 将显示 logging.INFO 消息及其以上级别。

• TORCH_LOGS=-pp 将显示 logging.WARNING 消息及其以上级别。

API Reference

Model Split APIs

以下 API 集会将您的模型转换为流水线表示。

class torch.distributed.pipelining.SplitPoint(value)

枚举，表示在子模块执行中可以发生分割的点。:ivar BEGINNING: 表示在 forward 函数中 **某个子模块执行之前** 添加一个分割点。:ivar END: 表示在 forward 函数中 **某个子模块执行之后** 添加一个分割点。

torch.distributed.pipelining.pipeline(module, mb_args, mb_kwargs=None, split_spec=None, split_policy=None)

根据规范分割模块。

更多详情请参见 Pipe。

参数:

• module (Module) – 要分割的模块。

• mb_args (tuple[Any, ...]) – 示例位置输入，以微批次形式。

• mb_kwargs (dict[str, Any] | None) – 示例关键字输入，以微批次形式。（默认为 None）

• split_spec (dict[str, SplitPoint] | None) – 一个使用子模块名称作为分割标记的字典。（默认为 None）

• split_policy (Callable[[GraphModule], GraphModule] | None) – 用于分割模块的策略。（默认为 None）

返回类型:

类 Pipe 的流水线表示。

class torch.distributed.pipelining.Pipe(split_gm, num_stages, has_loss_and_backward, loss_spec)

torch.distributed.pipelining.pipe_split()

pipe_split 是一个特殊的操作符，用于标记模块中阶段之间的边界。它用于将模块分割成阶段。如果以急切方式运行带注释的模块，则它是一个 no-op。

示例

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

上面的示例将被分割成两个阶段。

Microbatch Utilities

class torch.distributed.pipelining.microbatch.TensorChunkSpec(split_dim)

用于指定输入分块的类。

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

给定一系列 args 和 kwargs，根据它们各自的分块规范将它们分割成指定数量的块。

参数:

• args (tuple[Any, ...]) – Args 元组。

• kwargs (dict[str, Any] | None) – Kwargs 字典。

• chunks (int) – 要将 args 和 kwargs 分割成的块数。

• args_chunk_spec (tuple[TensorChunkSpec, ...] | None) – args 的分块规范，形状与 args 相同。

• kwargs_chunk_spec (dict[str, TensorChunkSpec] | None) – kwargs 的分块规范，形状与 kwargs 相同。

返回:

分片 args 的列表

返回类型:

分片 kwargs 的列表

torch.distributed.pipelining.microbatch.merge_chunks(chunks, chunk_spec)

给定一个块列表，根据分块规范将它们合并成一个单一值。

参数:

• chunks (list[Any]) – 块的列表。

• chunk_spec – 块的分块规范。

返回:

合并后的值。

返回类型:

value

Pipeline Stages

class torch.distributed.pipelining.stage.PipelineStage(submodule, stage_index, num_stages, device, input_args=None, output_args=None, group=None, dw_builder=None)

表示流水线并行设置中流水线阶段的类。

PipelineStage 假定模型的顺序分割，即模型被分割成块，其中一个块的输出馈送到下一个块的输入，没有跳接连接。

PipelineStage 通过在线性顺序中传播从 stage0 到 stage1 等的输出，自动执行运行时形状/数据类型推断。要绕过形状推断，请将 input_args 和 output_args 传递给每个 PipelineStage 实例。

参数:

• submodule (nn.Module) – 由此阶段包装的 PyTorch 模块。

• stage_index (int) – 此阶段的 ID。

• num_stages (int) – 总阶段数。

• device (torch.device) – 此阶段所在设备。

• input_args (Union[torch.Tensor, Tuple[torch.tensor]], optional) – 子模块的输入参数。

• output_args (Union[torch.Tensor, Tuple[torch.tensor]], optional) – 子模块的输出参数。

• group (dist.ProcessGroup, optional) – 分布式训练的进程组。如果为 None，则使用默认组。

• dw_builder (Optional[Callable[[], Callable[..., None]]) – 如果提供，dw_builder 将构建一个新的 dw_runner 函数，该函数将为 F、I、W（前向、输入、权重）零气泡调度执行 W 操作（输入权重）。

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

给定要由该阶段包装的 stage_module 和流水线信息，创建一个流水线阶段。

参数:

• stage_module (torch.nn.Module) – 要由此阶段包装的模块。

• stage_index (int) – 管道中此阶段的索引

• pipe_info (PipeInfo) – 有关管道的信息，可以通过 pipe.info() 检索

• device (torch.device) – 此阶段要使用的设备

• group (Optional[dist.ProcessGroup]) – 此阶段要使用的进程组

返回:

一个流水线阶段，可以与 PipelineSchedules 一起运行。

返回类型:

_PipelineStage

Pipeline Schedules

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)

GPipe 调度。将以填充-排空（fill-drain）的方式遍历所有微批次。

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)

1F1B 调度。在稳态下将对微批次执行一次前向和一次后向。

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)

交错 1F1B 调度。详情请参阅 https://arxiv.org/pdf/2104.04473。在稳态下将对微批次执行一次前向和一次后向，并支持每 rank 多个阶段。当微批次可用于多个本地阶段时，交错 1F1B 会优先处理更早的微批次（也称为“深度优先”）。

此调度与原始论文非常相似。它通过放宽 num_microbatch % pp_size == 0 的要求来区分。使用 flex_pp 调度，我们将有 num_rounds = max(1, n_microbatches // pp_group_size)，并且只要 n_microbatches % num_rounds == 0 即可正常工作。例如，支持

1. pp_group_size = 4, n_microbatches = 10。我们将有 num_rounds = 2 且 n_microbatches % 2 == 0。

2. pp_group_size = 4, n_microbatches = 3。我们将有 num_rounds = 1 且 n_microbatches % 1 == 0。

class torch.distributed.pipelining.schedules.ScheduleLoopedBFS(stages, n_microbatches, loss_fn=None, output_merge_spec=None, scale_grads=True, backward_requires_autograd=True)

广度优先流水线并行。详情请参阅 https://arxiv.org/abs/2211.05953。与交错 1F1B 类似，Looped BFS 支持每 rank 多个阶段。不同之处在于，当微批次可用于多个本地阶段时，Looped BFS 会优先处理更早的阶段，一次性运行所有可用的微批次。

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)

交错零气泡调度。详情请参阅 https://arxiv.org/pdf/2401.10241。在稳态下将对微批次的输入执行一次前向和一次后向，并支持每 rank 多个阶段。利用关于权重的后向来填充流水线气泡。

特别是，这实现了论文中的 ZB1P 调度。

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)

零气泡调度（ZBV 变体）。详情请参阅 https://arxiv.org/pdf/2401.10241 第 6 节。

此调度要求每个 rank 恰好有两个阶段。

此调度在稳态下将对微批次的输入执行一次前向和一次后向，并支持每 rank 多个阶段。利用关于权重的后向来填充流水线气泡。

此 ZB-V 调度只有在 time forward == time backward input == time backward weights 时才具有“零气泡”属性。在实践中，对于实际模型来说，这种情况不太可能发生，因此可以实现一个贪婪调度器来处理不相等/不平衡的时间。

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)

DualPipeV 调度。这是基于 DeepSeek 在 https://arxiv.org/pdf/2412.19437 中引入的 DualPipe 调度的一个更高效的调度变体。

基于 deepseek-ai/DualPipe 的开源代码。

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)

单阶段调度的基类。实现了 step 方法。派生类应实现 _step_microbatches。

根据 scale_grads 参数（默认为 True），梯度会按 num_microbatches 进行缩放。此设置应与您的 loss_fn 的配置相匹配，loss_fn 可以平均损失（scale_grads=True）或求和损失（scale_grads=False）。

step(*args, target=None, losses=None, return_outputs=True, **kwargs)

使用整批输入运行流水线调度的一个迭代。它将自动将输入分块为微批次，并根据调度实现遍历微批次。

args：模型的位置参数（与非流水线情况相同）。kwargs：模型的关键字参数（与非流水线情况相同）。target：损失函数的 target。losses：一个列表，用于存储每个微批次的损失。return_outputs：是否返回最后一个阶段的输出。

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)

多阶段调度的基类。实现了 step 方法。

根据 scale_grads 参数（默认为 True），梯度会按 num_microbatches 进行缩放。此设置应与您的 loss_fn 的配置相匹配，loss_fn 可以平均损失（scale_grads=True）或求和损失（scale_grads=False）。

step(*args, target=None, losses=None, return_outputs=True, **kwargs)

使用整批输入运行流水线调度的一个迭代。它将自动将输入分块为微批次，并根据调度实现遍历微批次。

args：模型的位置参数（与非流水线情况相同）。kwargs：模型的关键字参数（与非流水线情况相同）。target：损失函数的 target。losses：一个列表，用于存储每个微批次的损失。return_outputs：是否返回最后一个阶段的输出。