PyTorch 深度学习——在多块 GPU 上训练模型

本章涵盖以下内容：

• 分布式训练的基本概念
• PyTorch 的分布式包（torch.distributed）
• 不同形式的并行策略

在前面的章节中，我们大多一直专注于在单块 GPU 上训练模型。但随着模型越来越大、数据集越来越庞大，只依赖单块 GPU 训练就会变得不可行。尤其是近几年，大语言模型的流行让模型规模出现了爆炸式增长。其实从名字就能看出来：large language model，确实就是“很大”的语言模型。比如，Meta 开源的 LLaMA 3.1 就有 70 亿参数、800 亿参数和 4050 亿参数 三种版本。其中，4050 亿参数版本仅仅在推理阶段就大约需要 800 GB 内存（这里还没有考虑量化等优化手段）。

注意：读者应根据自己的使用方式，查阅相关模型的官方许可证条款。

为了解决这一问题，本章将探讨如何使用 PyTorch 的分布式子包，在多块 GPU 上训练模型，并覆盖分布式训练的核心概念以及多种并行形式。

16.1 并行编程简介

并行编程，就是把一个问题拆分成多个更小的任务，并让这些任务同时执行。这么做的目的，是提升性能与效率。

可以把它想象成在拼一幅很大的拼图。如果你一个人拼，可能要花很久才能完成；但如果有几个朋友一起帮忙，每个人就可以负责拼图的不同区域：一个人拼天空，一个人拼树木，另一个人拼建筑。因为这些部分是同时进行的，所以整幅拼图会更快完成。在机器学习中，并行编程允许同一个模型同时处理数据的不同部分，从而加速训练过程。

来看下面这个简单示例（1_motivator.py）：

代码清单 16.1 顺序执行的计时示例

# In[]
def some_work(x):
    time.sleep(0.5)
    return x

start_time = time.time()
results_sequential = sum([some_work(i) for i in range(10)])
sequential_time = time.time() - start_time

print(results_sequential)
print(f"Sequential Time: {sequential_time:.2f} seconds")

# Out[]
45
Sequential Time: 5.04 seconds

在这个例子中，我们有一个函数，它会执行一些工作，并耗时 0.5 秒。然后我们在循环中调用它 10 次。整个循环执行完成大约需要 5 秒。

接下来，我们使用 concurrent.futures 模块来把这项工作并行化。

代码清单 16.2 多线程执行计时

# In[]
start_time = time.time()
with concurrent.futures.ThreadPoolExecutor() as executor:
    results_multithreaded = sum(list(executor.map(some_work, range(10))))
multithreaded_time = time.time() - start_time

print(results_multithreaded)
print(f"Multithreaded Time: {multithreaded_time:.2f} seconds")

# Out[]
45
Multithreaded Time: 0.51 seconds

这里我们创建了多个可以并发运行的线程。整个循环完成只需要大约 0.5 秒。相比顺序执行，这就是一次非常显著的加速。

关于线程与进程

对于不熟悉线程与进程区别的读者来说：线程（thread） 是一种轻量级执行单元，可以与其他线程并发运行。多个线程共享同一块内存空间，因此它们的创建与切换开销更小。

而进程（process） 则拥有各自独立的内存空间，因此更加“重量级”。进程的创建和切换通常更慢，但由于它们彼此隔离，在处理 GPU 这类底层加速器时复杂度更低，因此通常更适合用于分布式训练。

16.1.1 分布式计算术语

在分布式训练中，我们要面对的是多块 GPU，这些 GPU 可能位于同一台机器上，也可能分布在多台机器上。因此，这里会涉及一套不同的术语与概念。torch.distributed 所采用的术语体系，来自 MPI（Message Passing Interface，消息传递接口）标准。图 16.1 展示了这些概念。

注意：MPI 标准定义了一套库例程的语法与语义，可用于在 C 和 C++ 中实现并行程序。

图 16.1 单 GPU 训练（左） vs. 多 GPU 训练（中） vs. 多机器多 GPU 训练（右）

接下来本章会频繁用到一些关键术语，包括：

• World size（世界大小） ——整个分布式系统被称为“world”。world size 指的是分布式系统中的总进程数。通常我们会限制为“每个进程对应一块 GPU”，因此 world size 通常就等于所使用的 GPU 数量。每个进程在这个 world 中都有一个唯一的 rank。图 16.1 中，单 GPU 训练的 world size 是 1，单机多 GPU 训练的 world size 是 4，多机训练的 world size 是 8。
• Rank（秩） ——rank 是分布式系统中分配给每个进程的唯一标识符。它是一个整数，范围从 0 到 world size - 1。rank 用于标识某个参与分布式系统的具体进程。
• Process group（进程组） ——进程组是一组能够彼此通信的进程集合。在 PyTorch 中，当我们调用 torch.distributed.init_process_group 时，就会创建一个进程组。这个进程组用于管理各个进程之间的通信。
• Node（节点） ——node 指的是分布式系统中的一台机器。它可以是真实的物理机，也可以是虚拟机。在深度学习语境中，一个 node 通常指的是一台拥有一块或多块 GPU 的机器。
• Global rank 与 local rank（全局秩与本地秩） ——当涉及多台机器时，我们需要区分 global rank 和 local rank。global rank 是某个进程在整个分布式系统中的 rank；local rank 是它在单台节点内部的 rank。图 16.1 中，多机训练里每台机器有 4 块 GPU，因此每台机器的 local rank 最多到 4，而 global rank 则覆盖到 8。
• SPMD（Single Program Multiple Data，单程序多数据） ——SPMD 是一种编程模型：同一个程序由多个进程执行，但每个进程处理自己的那份数据。在深度学习里，这意味着每个进程运行相同的模型，但处理的是不同的数据分片。当我们考虑执行清单 16.1 里的脚本时，每个进程运行的是相同的逻辑和操作。当然，也有可能不同进程执行不同操作。

在本章中，我们将把重点放在单机多 GPU 训练上。由于这种情况下 local 和 global 语境差别不大，我们主要会使用 rank 这个术语。虽然跨机器训练需要不同的初始化方式与网络通信设置，但底层概念和训练脚本的整体思路其实是非常相似的。

16.1.2 硬件要求

通常来说，要使用分布式训练，你需要多块 GPU。我们知道不是每个人手头都有这样的硬件，所以本章也会提供一些可以在单机 CPU 上运行的示例。

当然，只用 CPU 基本看不到什么明显加速，但代码依然能跑，你也依然可以借此理解这些概念。如果你能接触到多块 GPU，那么就可以在那些 GPU 上运行这些代码，亲眼看到训练加速的效果。

另外，PyTorch 的分布式 API 在 Windows 上支持比较有限。我们建议使用 Linux 或 macOS 来做分布式训练。如果你使用 Windows，可以考虑通过 Windows Subsystem for Linux（WSL） 来运行。

注意：使用 VS Code 配合 WSL 运行这些代码非常直接，具体可参考这里：mng.bz/dWqX。

16.1.3 初始化一个分布式程序

把一个单 GPU 程序改造成分布式程序，需要做几项改动。在真正开始使用多块 GPU 之前，我们必须先初始化分布式程序。正如图 16.2 所示，这包括：

• 创建多个进程
• 建立通信引导（bootstrap），并确保所有进程之间都能通信
• 修改原始程序中的训练逻辑

图 16.2 把单 GPU 程序改造成分布式程序所需的步骤

对于第 1 步，我们可以使用 torch.multiprocessing 模块来创建多个进程。这个模块提供了一种简单方法，用于创建和管理多个进程。torch.multiprocessing.spawn 函数可以创建多个进程，并在每个进程中运行同一个函数。spawn 方法需要传入要运行的函数、传给它的参数，以及要创建的进程数量（2_initialization.py）。

代码清单 16.3 使用 torch.multiprocessing 启动多个进程

import torch.multiprocessing as mp

if __name__ == "__main__":
    num_processes = 4
    mp.spawn(init_process_with_store, args=(num_processes,), nprocs=num_processes)

每个被创建出来的进程都会执行 init_process_with_store 函数，其中第一个参数是该进程的索引（0、1、2 或 3），第二个参数则是我们传入的 num_processes。

接下来，在第 2 步中，我们需要完成通信引导，并确保所有进程之间都能互相通信。PyTorch 提供了 torch.distributed 包，专门用于初始化和管理分布式程序（2_initialization.py）。

代码清单 16.4 初始化进程组

import torch.distributed as dist
import os
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "12355"

def init_process_with_store(rank, world_size, backend="gloo"):
    store = dist.TCPStore(
        host_name=os.environ["MASTER_ADDR"],
        port=int(os.environ["MASTER_PORT"]),
        world_size=world_size,
        is_master=rank == 0
    )
    dist.init_process_group(backend, store=store, rank=rank, world_size=world_size)

第一步是借助一个 TCPStore 来初始化分布式程序，用它完成通信引导。TCPStore 是一个简单的键值存储，它允许多个进程之间共享信息。我们需要设置 MASTER_ADDR 和 MASTER_PORT 环境变量，用来指定 master 进程（持有 TCPStore 的那个进程）的地址与端口。这里我们用 os 模块设置这些环境变量。

还需要说明的是：把 store 显式传给 init_process_group 其实是可选的。如果不传，PyTorch 会为我们自动创建一个默认 store。比如在同一个文件里的 init_process 示例函数中，我们没有显式创建 store，初始化照样能够正常工作。

除此之外，torch 包里还带了一个很方便的命令行工具，随 PyTorch 一起安装，它叫 torchrun（pytorch.org/docs/stable…）。这个工具可以帮助我们直接启动分布式训练任务。它会替我们处理初始化与通信设置。借助torchrun，我们就不必自己写 mp.spawn 去创建多个进程，也不必手动设置那些用于 bootstrap 的环境变量。此时，初始化逻辑会变得非常简洁：我们只需读取 torchrun 自动写入的环境变量，再直接调用 init_process_group 即可（3_torchrun.py）。

代码清单 16.5 使用 torchrun

# torchrun --nproc-per-node=4 3_torchrun.py    #1
import torch.distributed as dist
import os

def main():
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    print(f"{rank=} initializing")
    dist.init_process_group(backend="gloo", rank=rank, world_size=world_size)

if __name__ == "__main__":
    main()

#1 使用 torchrun 的命令，它会创建 4 个进程，并在每个进程中运行同一个脚本。同时还会自动填充 "RANK"、"WORLD_SIZE" 等环境变量。

图 16.2 的第 3 步——也就是修改训练程序——我们稍后会更详细地讨论。现在既然初始化已经搭好了，就先利用它，顺势进入一个关于**集体通信（collective communication）**的小插曲。

16.2 集体通信

在上一节中，我们花了不少工夫才走到可以调用 torch.distributed.init_process_group() 这一步。但这个函数到底在做什么？它的作用是初始化分布式进程组。它会把各个进程之间的通信通道建立起来，使它们可以彼此通信。这正是分布式训练的核心，因为在分布式训练中，各个进程必须能够高效地互相传递 tensor 数据。

在分布式训练中，我们经常需要在进程之间传递模型梯度、模型参数以及其他 tensor 数据。因此，保证这些通信尽可能高效，对整体训练效率至关重要。通信花的时间越多，留给真正计算与训练的时间就越少。

我们可以选择不同的 backend。最常见的 backend 是 gloo 和 nccl。gloo 是基于 CPU 的，适用于在 CPU 上进行分布式训练，所以前面的示例用了它。nccl 是面向 GPU 的 backend，用于 GPU 分布式训练。mpi 则是一个标准 backend，适用于多机分布式训练。

注意：gloo 对某些 collective 操作支持 GPU tensor，但它通常需要先把 tensor 从 GPU 拷到 CPU，再完成通信，最后再拷回 GPU，这会引入额外开销，因此性能不如原生 GPU 通信。

在这些 backend 之上，我们可以使用 torch.distributed 内置的各种集体通信操作（collective communication operations） 。先来看几个重要的操作：

• Broadcast（广播） ——广播就像广播电台向所有听众同时发消息。它把某个进程上的数据发送给组里的所有其他进程。源进程中的数据会被复制到所有其他进程上。图 16.3 左侧展示了这一过程：数据（tensor_0）由 rank 0 发出，结果所有其他进程（包括 rank 0 自己）最终都拥有了这份数据。
• Allreduce（全归约） ——allreduce 可以想象成一群人一起做一道菜：每个人（进程）都带来自己的食材（数据），最后把所有食材合成一道共享菜品，再把这道菜分给每个人。这个操作会把所有进程的数据收集起来，执行某种归约运算（例如求和、取最大值），然后再把结果发送回所有进程。图 16.3 右侧中，每个 rank 各自持有一份数据（t0、t1、t2 或 t3），allreduce 之后，每个进程都拿到了所有数据的和。

图 16.3 broadcast（左）与 allreduce（右）两种集体通信操作示意图

这里只是举了几个 collectives 的例子。实际上还有很多其他 collective，分别描述了 tensor 数据如何在多个进程之间交换。PyTorch 官网有一个很好的可视化教程资源：mng.bz/rZqZ。理解了这些collectives，我们就可以开始设想：不同进程之间到底能以哪些方式协作，从而实现不同类型的并行。

把这些内容串起来，我们可以运行一个示例脚本，它展示了 broadcast 和 allreduce 在实践中的效果，以及这些操作会如何改变 tensor（4_collectives.py）。

代码清单 16.6 PyTorch 的 broadcast 与 all_reduce 操作

import torch.distributed as dist
...
def perform_broadcast(rank):
    if rank == 0:
        tensor = torch.tensor([123.0])
    else:
        tensor = torch.zeros(1)
    dist.broadcast(tensor, src=0)    #1

def perform_all_reduce(rank):
    tensor = torch.tensor([float(rank)])
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)    #2
...

# Out[]
Process 1 performing broadcast
Process 2 performing broadcast
Process 0 performing broadcast
Process 2 has tensor tensor([0.])
Process 1 has tensor tensor([0.])
Process 0 has tensor tensor([123.])
Process 0 received tensor: tensor([123.])
Process 1 received tensor: tensor([123.])
Process 2 received tensor: tensor([123.])
Process 0 performing all_reduce
Process 2 performing all_reduce
Process 1 performing all_reduce
Process 0 has tensor tensor([0.])
Process 1 has tensor tensor([1.])
Process 2 has tensor tensor([2.])
Process 0 received tensor: tensor([3.])
Process 2 received tensor: tensor([3.])
Process 1 received tensor: tensor([3.])

#1 rank 0 把 tensor 123 广播给 rank 1 和 rank 2。
#2 每个 rank 都广播自己的 rank 值，因此最终求和结果是 0+1+2。

可以看到，broadcast 操作会把 rank 0 上的 tensor 发送给所有其他进程；而 allreduce 会把所有进程上的 tensor 求和，并把结果（3.0）返回给所有进程。需要注意的是，由于多进程本身具有非确定性，你在自己机器上看到的打印顺序可能与这里不同。为了让输出更容易读，我们通常会加同步点（通过 barrier 这类 collective），不过前面的代码为简洁起见省略了这些。

有了 collective communication 这套工具，我们就可以开始实现深度学习训练中常见的各种分布式并行方式了。

16.3 并行方式简介

神经网络的前向传播，本质上是一个“天然适合并行（embarrassingly parallel）”的问题，也就是说，它非常容易被分摊到多块 GPU 上执行。每个输入样本都用同样的模型权重进行处理，但彼此之间并不需要对方的计算结果，因此不同 GPU 可以独立地处理各自那一份数据，而无需互相通信。于是，只要把模型复制到多个设备上，每块 GPU 就能独立处理自己那部分数据。

比如说，如果我们有两块 GPU，而原本在单块 GPU 上训练时 batch size 是 20，那么就可以把数据平均分成两份。每块 GPU 各自处理 batch size 为 10 的那一半数据，并独立完成对应的前向传播。这样一来，每个设备都能独立工作，从而最大化计算效率。图 16.4 展示了这一点：前向传播在每个设备上独立执行。

图 16.4 分布式训练中的步骤：前向传播（1）、反向传播（2）以及参数更新（3）；但这里其实还缺少一个关键步骤，因为每个模型上得到的梯度并不相同

在初始前向传播之后，每个模型都会得到一个 loss 值。由于每个模型处理的是不同的数据分片，所以这些 loss 并不相同。（在分布式训练里，shard 指的是整个数据集被切分之后，某个计算单元所负责的那一部分。）

类似地，反向传播也可以独立执行。每个模型都会根据自己先前算出的 loss，独立计算自己的梯度。虽然模型参数本身是相同的，但由于它们对应的数据分片不同，因此得到的梯度 tensor 也不同。

这时就会出现一个很自然的问题：参数该怎么更新？ 如果我们像单设备训练那样，直接在每个模型上调用 optimizer.step()，那么每个模型就会用不同的梯度更新自己的参数。结果就是，这些模型会逐渐分叉，最后变成多个彼此不同的模型，而不是一个被多个设备共同训练出来的统一模型。

为了解决这个问题，我们必须在训练流程中增加额外一步：在反向传播之后、optimizer.step() 之前，把所有设备上的梯度先聚合起来。这一步在分布式训练中通常被称为 gradient reduction（梯度归约） 。它的作用是，在参数更新之前，先保证所有设备上的梯度是一致的。只有这样，我们训练的才是一个真正统一的模型。

而这正是我们前面学过的 collective communication 操作派上用场的地方。我们可以用 allreduce 来把所有设备上的梯度做平均。一旦梯度同步完成，就可以像平常一样执行 optimizer.step() 来更新参数。这样就能保证所有模型副本始终保持一致，而不会互相漂移。下面来看看，在 PyTorch 里该如何实现这种数据并行（data parallel） 。

16.4 数据并行

正如前面提到的，数据并行的基本做法是：把模型复制到多个设备上，同时把数据划分开，让每个设备处理其中的一部分。为了在参数更新之前保持一致性，我们需要通过 collective communication 来同步所有设备上的模型梯度，如图 16.5 所示。

图 16.5 加入了“同步步骤”后的训练流程

先看下面这段代码，看看实际实现起来大概是什么样子（5_ddp_from_scratch.py）。

代码清单 16.7 一个简单的梯度同步示例

input_data = torch.randn(3) + rank     #1

output = model(input_data)        #2
loss = output.mean()              #2

opt.zero_grad()                    #3
loss.backward()         #4

for param in model.parameters():                                #5
    dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM)      #5
    param.grad.data /= world_size  #6

opt.step()   #7

#1 做模型前向传播。
#2 这里使用的是伪造数据（rand），并加上 rank，从而确保每个进程拿到的数据不同。
#3 这里开始反向传播。
#4 loss.backward() 与单进程训练中完全一样。
#5 对模型梯度做 all_reduce。
#6 因为 all_reduce 做的是求和，所以我们再除以 world_size，得到平均梯度。
#7 更新参数；opt.step() 与单进程训练中完全一样。

在上面的代码里，我们先完成前向传播并计算 loss。接着调用 loss.backward() 来计算梯度。之后，使用 dist.all_reduce 在所有设备之间同步这些梯度。最后，再调用 opt.step() 来更新参数。

不过，虽然这个实现很好地说明了数据并行的基本思想，它本身在性能上并不高效。实际中，PyTorch 已经通过 torch.nn.parallel.DistributedDataParallel（简称 DDP）模块，提供了一个更加高效而稳健的实现。DDP 会使用优化过的通信模式来尽量降低梯度同步开销，因此它是一种可扩展性更强的分布式训练方案。只要把模型包上一层 DDP，我们就能轻松启用数据并行，而 PyTorch 会自动处理底层复杂性，包括梯度同步（6_ddp.py）。

代码清单 16.8 使用 DistributedDataParallel 包装模型

from torch.nn.parallel import DistributedDataParallel as DDP
model = SimpleModel().to(device)
model = DDP(model)         #1

#1 只需要把模型用 DDP 包起来，之后就可以像单进程训练时那样正常使用。

在前面的基础实现里，我们其实还忽略了一个重要方面：数据加载。在真实场景下，必须正确地把数据集在各个进程之间切分开来。PyTorch 为此提供了 DistributedSampler，它会把数据加载限制在某个特定子集上。这样做可以保证 DataLoader 不会重复读取数据集中的同一条样本，也能保证一个 epoch 是由所有设备共同处理完整个数据集（6_ddp.py）。

代码清单 16.9 使用 DistributedSampler 创建 DataLoader

def get_data_loader(batch_size, rank, world_size):
    transform = transforms.Compose([transforms.ToTensor()])
    dataset = datasets.MNIST(root='../data/p2ch16', train=True, download=True, transform=transform)
    sampler = DistributedSampler(dataset,
        num_replicas=world_size, rank=rank)        #1
    loader = DataLoader(dataset, batch_size=batch_size,
        sampler=sampler)      #2
    return loader

#1 DistributedSampler 需要知道 rank 和 world_size，才能判断当前 rank 应该处理数据集中的哪一部分。
#2 这个 sampler 会被传给 DataLoader。

使用 DDP 来训练一个 MNIST 分类模型的完整示例，可以在 6_ddp.py 中看到。

16.5 模型并行

数据并行很好用，但它有一个前提限制：模型必须能够完整装进单块 GPU。因为数据并行依赖于把整个模型复制到多个设备上，然后让每个设备处理一部分数据。从本质上说，数据并行是在用更多显存去换取更多计算并行度，以便加快训练。如果模型本身已经大到单块 GPU 根本装不下，那么单靠数据并行就没法解决问题。

这时候，就需要使用 模型并行（model parallelism） 。模型并行的思想是：把模型本身拆分到多个设备上。每个设备只保存模型的一部分，并负责这一部分对应的前向和反向计算。

对于一个被拆分到多个设备上的模型，其前向传播过程和单 GPU 上的前向传播很相似，只是会多出额外的通信步骤。当模型的某一部分在当前设备上处理完一个 batch 后，它产生的中间输出——通常叫做 activation（激活值） ——就必须被发送到下一个设备，按顺序继续处理。你可以把这个模型想象成被拆成了三个小模型，它们相互协作，最终产生与原始完整模型相同的输出。正如图 16.6 所示，batch 先在 device 0 上处理，得到输出，再把这个输出交给 device 1；device 1 继续处理后，再把结果传给 device 2，如此继续。

图 16.6 一个被改造成可运行在多个设备上的模型。箭头方向表示一个 batch 在模型中的流动方向。

在代码示例中，我们用几个简单的线性层来表示一个被拆分在三台设备上的模型。每台设备负责处理模型的一段，并通过 点对点通信（peer-to-peer, P2P） 把激活值传给下一个设备。P2P 操作与前面讲过的 collective communication 有点类似，但它只发生在两个进程之间。在 torch.distributed 中，P2P 操作主要包括 send() 和 recv()。这两个函数都需要一个 tensor，以及源 rank 和目标 rank，来完成通信。下面的示例中，我们正是通过这些操作来传递激活值（7_model_parallel.py）。

代码清单 16.10 使用 send 和 recv 实现模型并行

...                 #1
    if rank == 0:                                                      #2
        model_part1 = nn.Linear(1, 2)
        input_batch = torch.tensor([[1.0]])
        part1_activations = model_part1(input_batch)
        print(f"Process {rank} activations from part 1: {part1_activations}")
        dist.send(part1_activations, dst=1)                           
    elif rank == 1:
        model_part2 = nn.Linear(2, 3)
        part1_activations = torch.zeros(2)
        dist.recv(part1_activations, src=0)       #3
        part2_activations = model_part2(part1_activations)
        dist.send(part2_activations, dst=2)
        print(f"Process {rank} activations from part 2: {part2_activations}")
    elif rank == 2:
        model_part3 = nn.Linear(3, 1)
        part2_activations = torch.zeros(3)
        dist.recv(part2_activations, src=1)
        part3_output = model_part3(part2_activations)
        print(f"Process {rank} output from part 3: {part3_output}")

#1 这里省略了初始化部分。
#2 每个 rank 持有模型的一部分。
#3 使用点对点操作（P2P）在不同模型分区之间发送和接收激活 tensor。

可以看出，模型并行使我们能够把一个模型分布在多个设备上，同时依然得到与单设备完整模型相同的输出。这样一来，那些原本单块 GPU 装不下的大模型也就有办法训练了。训练时还必须包含反向传播，其过程与前向传播类似：每个设备先对自己那一段 loss 调用 .backward() 来计算梯度，再把这些梯度传递给其他设备，以确保整个模型都能被正确更新。

16.5.1 流水线并行

前面讲的模型并行虽然可行，但还比较原始。问题在于：当某一台设备在做前向传播时，其他设备往往是闲着的，这就导致资源利用率不高。想进一步优化，就可以引入 流水线并行（pipeline parallelism） 。

流水线并行的做法是：把模型切成多个部分，也就是多个 stage（阶段） ，分布在不同设备上；与此同时，还会把输入数据切成更小的单元，称为 microbatch（微批次） 。这样，一个 stage 在处理某个 microbatch 的同时，另一个 stage 可以并行处理另一个 microbatch，从而让计算更充分地重叠起来。

如图 16.7 所示，假设我们有三台设备。在流水线并行方案中，我们会把模型拆成三个 stage，分别放在三台设备上。输入数据则被切成多个 microbatch，使得每台设备都能并行处理模型中属于自己的那一段。

图 16.7 三个 microbatch 的流水线并行示意。每个 microbatch 在某一设备处理后，会被通信到下一台设备，如箭头所示。这样可以让多台设备并发工作。

流水线并行如今被广泛用于大语言模型的训练与推理。它特别适合那种由大量重复模块构成的模型，而大语言模型正好符合这一特征，因为它们内部会反复堆叠 transformer 层。

当然，流水线并行也带来一些新的挑战，比如如何管理设备间通信开销，以及如何确保各个 stage 的计算负载足够均衡，避免出现瓶颈或“拖后腿”的 stage。流水线并行通常需要一定程度的手动调优，并要求你对模型架构本身有比较好的理解，才能真正发挥它的优势并尽量减少其缺点。

PyTorch 在 torch.distributed.pipelining 子包中提供了流水线并行的支持，里面包含了模型切分、调度策略定义与执行等工具。

下面的代码展示了如何使用 PyTorch 内置工具来设置流水线并行。首先，我们定义一个示例模型，其中有一个 split_spec 属性。这个属性用于指定模型应该在什么地方切开，其中 SplitPoint.BEGINNING 表示应在指定层之前进行切分。

代码清单 16.11 带有切分点的模型（8_pp.py）

class SimpleMLP(nn.Module):
    def __init__(self, hidden_size=512, n_layers=4):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Linear(hidden_size, hidden_size) for _ in range(n_layers)
        ])
        self.split_spec = {
            f"layers.{i}": SplitPoint.BEGINNING for i in range(1, n_layers)
        }

    def forward(self, x):
        for layer in self.layers:
            x = torch.relu(layer(x))
        return x

模型定义好之后，我们就可以为每个 rank 构建对应的 pipeline stage；每个 stage 只包含模型的一部分。因为已经定义好了 split_spec，所以可以直接使用 pipeline API，根据这些切分点自动完成模型分区。

代码清单 16.12 构建一个 pipeline stage

pipe = pipeline(model, mb_args=(x_mb,), split_spec=model.split_spec)
stage = pipe.build_stage(rank, device)

最后，我们还需要定义一个 schedule（调度策略） 。调度策略决定了这些 microbatch 如何在各个 stage 之间流动与执行。在这个例子里，我们使用的是简单的 GPipe 调度方式。GPipe 会先把所有 microbatch 的前向传播全部执行完，再统一执行所有反向传播。

代码清单 16.13 执行流水线调度

if rank == 0:
    schedule.step(x)
elif rank == world_size - 1:
    output = schedule.step(target=target, losses=losses)
else:
    # Middle ranks just forward/backward
    schedule.step()

这个示例可以像前面讲过的那样，通过 torchrun 在多个进程上运行。完整代码可见 code/p2ch16/8_pp.py。

16.5.2 张量并行

张量并行（tensor parallelism） 是另一种模型并行方式。它通过把模型参数（也就是 tensor）拆分到多个设备上，来分摊神经网络的计算负载。和流水线并行是在“模块级别”切模型不同，张量并行更进一步，它是在模块内部进行切分。

如图 16.8 所示，张量并行会把模型参数切成更小的片段。例如，一个线性层的权重和偏置可以拆到两台设备上。这样，每台设备都可以独立地拿输入 tensor 做部分矩阵乘法。不过，这样做也意味着必须增加额外通信，用来交换中间结果。只有这样，每台设备才能拿到自己完成计算所需要的全部信息，最终拼出与单设备运行时相同形状的正确输出。在张量并行中，模型权重既可以按行方向（dimension 0） 切分，也可以按列方向（dimension 1） 切分，从而减少通信开销并加速训练。至于为什么会采用这些具体策略，本书就不展开了；如果想深入了解，可以查看最早提出张量并行的论文（arxiv.org/pdf/1909.08…）。

图 16.8 流水线并行（上）与张量并行（下）的对比。流水线并行是在模块级别切分，而张量并行是在模块参数（tensor）级别切分。

16.5.3 如何在流水线并行与张量并行之间做选择

那么，该如何在流水线并行和张量并行之间做选择呢？很遗憾，答案通常是：这取决于模型的架构。

对于那些更瘦、更深的模型——也就是层数很多、但每层参数相对较少——通常更适合使用流水线并行。这种方式会把模型按顺序切成多个 stage，让不同设备分别处理不同阶段，从而让所有设备都能保持忙碌。

相反，对于那些更宽、更浅的模型——也就是层数较少、但单层参数量非常大——则往往更适合张量并行。因为在这种情况下，每一层的计算负担本身就很重，把参数拆到多个设备上，更容易有效分担这部分工作量。

另外，如果模型中的某一层已经大到单块设备显存根本装不下，那张量并行就变成了必需品。通过把这一层的参数切分到多个设备上，张量并行可以让模型在可用显存范围内运行下去，从而避免因为显存不足而根本无法训练或推理。

16.6 n 维并行

前面我们分别介绍了数据并行和模型并行，但它们并不是互斥的。事实上，在现实系统中，往往会把它们组合起来，以便训练更大的模型并提升整体效率。这就是所谓的 n-dimensional parallelism（n 维并行） 。

之所以称为 “n 维”，是因为实践中通常会把多种并行方式叠加在一起，从而得到更高效率。比如，张量并行会引入大量通信开销，因此通常倾向于在单机内部使用，以便利用高速 GPU 互联；而流水线并行的通信需求相对更少，因此更适合跨多台主机分配工作。

当我们训练真正大规模模型时，往往需要几百甚至几千块 GPU，才能把训练时间控制在合理范围内。如果只依赖某一种并行方式，例如数据并行，那么随着规模扩大，像 allreduce 这类 collective 操作很快就会变成瓶颈。通过组合不同的并行策略，我们就可以像调 learning rate 或 batch size 一样，对整个系统进行“性能调参”。例如，LLaMA 3 在训练时就使用了 4D parallelism，把数据并行、张量并行、流水线并行以及上下文并行（context parallelism）结合在了一起。

注意：分布式训练对大语言模型非常重要，而 LLaMA 论文对此有非常详细的描述（arxiv.org/pdf/2407.21…）。

要描述 n 维并行其实并不容易。单单管理一个“所有设备都在同一个进程组里”的通信模式就已经不简单了；而当我们还需要只让设备的某些子集彼此通信时，事情会更加复杂。好在，torch.distributed 提供了一个叫做 DeviceMesh 的抽象，它允许我们定义一个设备网格（mesh），并为不同维度定制通信模式。它在底层其实会创建多个进程组，因此对实现 n 维并行特别有用。

假设我们要设计一种分布式架构：有两个完整模型副本，而每个副本内部又通过模型并行拆分到多块 GPU 上。在这种设置下，输入数据沿着“行”方向穿过模型，而模型副本则沿着“列”方向同步。图 16.9 展示了如何用一个 device mesh 表示这种结构。

图 16.9 一个 2 × 4 的 device mesh 可视化，其中行表示模型副本，列表示模型分区。虚线箭头表示不同设备之间的通信方式。

在代码中，我们可以这样初始化这个 device mesh（device_mesh.py）：

from torch.distributed.device_mesh import init_device_mesh

mesh_2d = init_device_mesh(                                          #1
    "cpu", (2, 4), mesh_dim_names=("replicate", "model_parallel")    #1
) #1

#1 创建一个 2 维设备网格：2 行 4 列。第一个参数指定设备类型，第二个参数定义 mesh 的形状，第三个参数为各维度命名（这里的名字可以是任意字符串）。

当创建一个 2 × 4 的 device mesh 时，底层实际上会生成 6 个不同的进程组。在“行”这个维度上，会有两个进程组，分别包含 ranks [0, 1, 2, 3] 和 [4, 5, 6, 7]；在“列”这个维度上，会有四个进程组，分别是 [0, 4]、[1, 5]、[2, 6]、[3, 7]。如果我们想同步模型权重，就会使用“列”方向上的那些进程组。要做到这一点，只需要按名称 "replicate" 从 mesh 中取出对应维度即可（device_mesh.py）。

代码清单 16.14 同步模型参数

replica_mesh = mesh_2d["replicate"]
replica_group = replica_mesh.get_group()   #1
def sync_model(replica_mesh):
    print(f"Mesh {replica_mesh} averaging model parameters")
    for p in model_part.parameters():
        dist.all_reduce(p, group=replica_group)    #2
        p.data /= replica_mesh.size()

#1 从 mesh 中选取对应的 group。
#2 使用这个 group 进行通信。

同步完成后，所有副本上的模型参数就会保持一致。类似地，当输入通过每个副本时，我们也可以借助 device mesh 中“列”方向的进程组来协调通信（device_mesh.py）：

model_parallel_mesh = mesh_2d["model_parallel"]
forward_pass(model_part, model_parallel_mesh)    #1

#1 forward_pass 中实现了跨模型分区执行前向传播的逻辑。

有了 device mesh，表达各种并行方式、管理多个进程组，就会容易得多。

16.7 完全分片数据并行

最后，必须重点介绍一下最重要的分布式 API 之一：Fully Sharded Data Parallel（FSDP） 。FSDP 结合了数据并行与模型并行的优势，允许你同时对模型和数据进行分片，并分布到多个设备上。这种混合方式让训练大模型时的扩展性和资源利用率都更高。

FSDP 最早于 2021 年由 Facebook AI Research 在一篇博客文章中提出（mng.bz/xZqY）。它的核心思想是做 full parameter sharding（完整参数分片） ：每个设备只存储并参与计算本地所需的那一小部分模型参数、梯度与优化器状态。与这一思路非常接近的一个著名实现是微软推广开来的 ZeRO-3（arxiv.org/pdf/2101.06…），它证明了完整参数分片在减少冗余复制和提升训练效率方面的有效性。

和张量并行类似，FSDP 也会把模块参数切分到多个设备上，如图 16.9 所示。当某次计算需要处理某个输入时，会执行一次 all_gather，临时把前向传播所需的完整权重收集到每个设备上。计算完成后，这些额外权重就可以释放掉，参数重新回到分片状态。反向传播时也有类似模式：梯度会通过 all_gather 暂时聚合，然后再通过 reduce_scatter 分发回各设备。这些通信通常会尽量与计算重叠，从而提升整体性能并减少通信开销。

图 16.10 FSDP 实现的伪代码示意

要用 FSDP 初始化一个模型，你需要指定：模型中哪些部分需要被分片，以及这些分片分布在哪些设备上。torch.distributed.fsdp 包提供了一个 fully_shard 方法，接收你要分片的模块，以及一个可选的 device mesh 参数，用于定义这些分片应该落在哪些设备上（fsdp_example.py）。

代码清单 16.15 使用 fully sharded data parallel

from torch.distributed.fsdp import fully_shard

model = SimpleModel()

mesh = init_device_mesh(device, (2,))
print(mesh)
for module in model.modules():
    if isinstance(module, nn.Linear):
        fully_shard(module, mesh=mesh)
fully_shard(model, mesh=mesh)

为了使用 FSDP，你通常需要多次调用 fully_shard：先对线性层等需要分片的子模块调用一次，再对整个根模型调用一次。

当模型大到单块 GPU 放不下，而 DDP 又不再适用时，FSDP 通常是更推荐的方案。FSDP 已经展示出了很强的可扩展性，成功运行到 512 块 GPU 的规模，这对大多数生产场景都已经足够。不过，如果在更大规模下，all_gather 或 reduce_scatter 的通信开销开始成为瓶颈，那就需要进一步叠加其他并行方式。

16.8 面向大语言模型的特殊并行方式

大语言模型（LLM）有一些自身特征，使得我们可以利用它们来更高效地使用 GPU。既然 transformer 和 LLM 已经成为许多自然语言处理任务中的主流架构，就值得专门提一下几种对它们尤其相关的并行策略：上下文并行（context parallelism） 和 专家并行（expert parallelism） 。

16.8.1 上下文并行

上下文并行主要针对长序列训练。它的做法是：沿着序列长度维度，把整个序列拆分给不同 rank，让每个 rank 负责一段连续 token。每个 rank 只处理自己那一段 token 对应的 query，但在注意力计算时，仍需与其他 rank 协调，以保证跨分段边界的因果掩码（causal masking）是正确的。它本质上是在增加少量同步轮次的代价下，换取更低的单 rank 显存占用和更高吞吐量，尤其适合节点内部高速互联的场景。

在长序列训练中，注意力计算量和激活值显存占用通常会随着序列长度 (L) 的平方增长，也就是大致 (O(L^2))。按序列维度做切分，可以降低每个 rank 的激活占用和临时 tensor 的峰值显存压力，同时又能通过跨分区边界的 masking 保持自回归约束不被破坏。它的代价是额外同步，但在节点内部有高速互联时，这种代价通常是值得的。

因此，上下文并行特别适合那些“长上下文训练、且激活值成为主要瓶颈”的场景，尤其适合单机多 GPU 且设备之间互联很快的系统。

16.8.2 专家并行

在大语言模型中，一种常见架构是 Mixture of Experts（MoE，专家混合） 。在这种架构里，传统的前馈网络会被一组“专家网络”替换，不同专家会针对不同类型数据或不同任务逐渐形成 specialization。专家并行（expert parallelism） 的做法，就是把这些专家分布到不同设备上，让每个设备只承载其中一部分专家。

MoE 的一个关键特性是：它采用稀疏计算。也就是说，对于每个 token，通常只会真正激活少量几个专家（比如 top-1 或 top-2）。一个轻量级的 router 会先给各个专家打分，然后为每个 token 选出 top-k 专家，同时还会配合一个额外损失项，以保证各个专家的负载大体均衡。正因为如此，专家并行通常依赖 all-to-all 通信：先把需要送去某个专家的 token 收集到对应设备，再在专家处理完成后把输出结果散回去。对于拥有很多专家的大模型来说，这种模式非常适合，因为它能在不按比例增加每个设备计算和显存负担的前提下，把模型规模继续做大。

16.9 把所有并行方式串起来

在本章中，我们依次考察了多种并行策略，以及如何用 PyTorch API 来实现它们。虽然每一种方式单独用都能发挥作用，但在真实训练系统里，把多种并行方式组合起来之后，又会引入新的挑战，例如：如何做 checkpoint 管理、不同并行方式之间如何组合、以及资源该如何高效利用。这些复杂性通常都需要比较精细的工程设计，尤其是在保存与加载模型 checkpoint，或者协调多设备通信时。现实系统必须正面解决这些问题，才能真正吃到分布式训练的全部收益。

为了给这些问题提供一个更实用的参考实现，PyTorch 团队推出了 TorchTitan。这是一个演示如何只使用原生 PyTorch API 来完成分布式训练的仓库。TorchTitan 提供了一套干净而简洁的代码，用于训练大语言模型，其中包含了 LLaMA、DeepSeekV3 以及部分 Hugging Face 模型的现成配置。它还集成了 tokenizer 等必要组件，因此只需一条命令，就可以启动一个从零开始训练大语言模型的分布式任务。

TorchTitan 仓库地址为：github.com/pytorch/tor…。若想更深入了解 TorchTitan 的设计细节与实验结果，可以参考它配套的论文（arxiv.org/pdf/2410.06…）。

16.10 结论

在本章中，我们系统梳理了 PyTorch 分布式 API 的整体图景。我们讲解了分布式训练的基本原理，并展示了如何使用 PyTorch 的分布式包在多块 GPU 上训练模型。在这一过程中，我们也讨论了并行编程中的关键概念，以及如何更高效地利用多种硬件资源。与此同时，我们还介绍了分布式训练脚本的初始化过程。

接着，我们进一步讨论了集体通信（collective communication） ，以及它在多设备间同步梯度中的作用，并说明了为什么分布式训练对扩展到更大的模型和数据集至关重要。整章中，我们陆续介绍了多种并行形式，包括数据并行、模型并行，以及更高级的流水线并行、张量并行和完全分片数据并行。我们还展示了如何借助 DeviceMesh 与 n 维并行 来实现这些策略。最后，我们通过 TorchTitan 这个端到端的实践示例，把这些概念串联了起来。

16.11 练习

• 你能如何改进“从零实现”的数据并行版本？去查一下 bucketing，它在 DDP 内部设计文档中有解释（mng.bz/AGn7）。

• 实现一个不同的模型并行示例，其中输入与输出 tensor 的形状各不相同：

  • 在数据传输和同步方面，你会遇到哪些挑战？

• 阅读流水线并行文档（mng.bz/Z9Ba），并实现 GPipe。

• 自己创建一个 3D device mesh。

• 调研 PyTorch 中的 fully sharded data parallel（FSDP） ：

  • 它相比传统数据并行有哪些优势？
  • 它在什么场景下最有价值？

• 使用 TorchTitan，运行一次 LLaMA3 8B 的分布式训练任务：

  • 需要做哪些配置步骤？
  • 与标准单 GPU 训练代码相比，有哪些差异？

小结

• 分布式训练使得深度学习模型可以扩展到多块 GPU 或多台机器上，从而更高效地处理大型数据集和复杂模型。
• 并行编程思想有助于把计算拆分为可以同时执行的多个部分，从而提升训练速度和硬件利用率。
• rank 是分布式环境中每个进程的唯一标识，而 world size 是参与训练的总进程数。
• 搭建分布式训练通常需要初始化进程组并启动多个进程，常用工具包括 torchrun。
• 像 broadcast 和 allreduce 这样的集体通信操作，对于在多个设备之间同步数据（例如梯度）至关重要，从而保证模型更新保持一致。
• 数据并行会在每个设备上复制一份模型，并把数据拆分开，因此必须同步梯度，以保证所有副本一致。
• 模型并行则是把模型本身拆到多个设备上，适用于单块 GPU 放不下的大模型，因为计算与参数都被拆分了。
• 流水线并行会把模型切成多个顺序 stage，让不同的 microbatch 可以并发流动，从而提高吞吐量。
• 张量并行会把单层参数或 tensor 拆分到多个设备上，适合训练特别宽、特别大的模型。
• n 维并行会组合多种并行策略，例如数据并行、模型并行、流水线并行和张量并行，以便在大规模训练中最大化扩展性与效率。
• PyTorch 中的 DeviceMesh 抽象可以帮助管理复杂的通信模式和多个进程组，从而更容易实现高级并行策略。
• Fully sharded data parallel 是一种混合式方案，它同时对模型和数据做分片，从而优化显存与计算资源的使用。
• 上下文并行 和 专家并行 是专门面向大语言模型扩展的两类特殊并行技术。
• 像 TorchTitan 这样的仓库，为实际分布式训练提供了很好的实践入口，它们展示了现代训练系统中的最佳实践与技术组合。
• 如何选择并组合合适的并行策略，是把深度学习扩展到现代模型与数据规模的关键；而 PyTorch 已经为这些需求提供了相当完整而强大的支持。