大模型部署必备知识：分布式训练一次性讲清

原创已于 2025-12-05 17:24:55 修改 · 768 阅读

28 ·

CC 4.0 BY-SA版权

文章标签：

#分布式 #人工智能 #缓存 #神经网络 #深度学习 #语言模型

于 2025-12-05 17:24:32 首次发布

从零开始一步步构建大模型专栏收录该内容

12 篇文章

订阅专栏

1.宏观架构与核心需求

所谓分布式训练，本质上是为了解决两个核心矛盾：

算力矛盾（时间不够）： 数据量太大，单卡跑完需要几年，必须多卡并行加速。
存储矛盾（空间不够）： 模型参数量太大，单张显卡的显存（VRAM）根本装不下模型。

我们通过两个具体的实例来看看这两种模式是如何诞生的。

1. 场景一：数据并行 (Data Parallelism, DP)

核心目的： 加速训练（解决“算力矛盾”）。 适用场景： 模型能塞进单卡，但数据量巨大。

实例推演：ResNet-50 on ImageNet

假设我们要训练一个经典的计算机视觉模型 ResNet-50。

模型大小： 约 98MB（2500万参数），单张显卡完全放得下。
数据量： ImageNet 数据集约 128 万张图片。
单卡瓶颈： 假设使用一张 NVIDIA V100，设置 Batch Size 为 256。完成一次前向+反向传播需要 0.5 秒。
单卡耗时： 跑完一个 Epoch（所有图片过一遍）需要 $\frac{1,280,000}{256} \times 0.5s \approx 2500s \approx 42 \text{min}$ 。如果要跑 100 个 Epoch，需要 70 小时。

分布式方案

如果我们有 8 张 V100 显卡：

复制模型： 我们在 8 张卡上，每张都完整复制一份 ResNet-50 模型（Model Replica）。
切分数据： 我们将一个大 Batch（比如 $256 \times 8 = 2048$ ）切分成 8 个小份（Micro-Batch），每张卡拿到 256 张图片。
并行计算： 8 张卡同时计算各自数据的梯度（Gradients）。
同步（关键步骤）： 此时每张卡的梯度是不一样的（因为吃的数据不一样）。我们需要把 8 张卡的梯度求平均，然后更新所有的模型。

结果：理想情况下，时间缩短为原来的 1/8（约 9 小时）。

代价：需要在每一步计算后进行通信（All-Reduce），通信会消耗时间，所以实际加速比通常小于 8。

2. 场景二：模型并行 (Model Parallelism, MP)

核心目的： 突破显存极限（解决“存储矛盾”）。 适用场景： 单张显卡根本装不下模型。

实例推演：GPT-3 (175B)

假设我们要训练 GPT-3。

参数量： 1750 亿 (175 Billion)。

显存硬性计算：

如果是 FP16（半精度）存储，每个参数占 2 Bytes。

仅模型权重就需要： $175 \times 10^9 \times 2 \text{ Bytes} \approx 350 \text{ GB}$ 。

训练时还需要存储梯度（Gradients）和优化器状态（Optimizer States，如 Adam 通常需要存 2 份状态）。

实际上，训练 GPT-3 至少需要 TB 级别 的显存空间。

硬件限制： 目前最强的 NVIDIA A100 80GB 或 H100 80GB，单卡显存仅 80GB。

结论： 物理上，你无法把 GPT-3 放进任何一张单卡里。

分布式方案

必须把模型“切开”。切法有两种：

A. 流水线并行 (Pipeline Parallelism - PP) -> 纵向切 想象模型是 96 层的 Transformer Layer 堆叠。

做法： 我们把前 12 层放在 GPU 0，第 13-24 层放在 GPU 1，以此类推...直到 GPU 7 放最后 12 层。
数据流向： 数据进入 GPU 0，算出中间结果，传给 GPU 1... 像工厂流水线一样。

B. 张量并行 (Tensor Parallelism - TP) -> 横向切 想象模型内部的一个巨大的矩阵乘法 $Y = X \times W$ 。假设 $W$ 的形状是 $[4096, 4096]$ 。

做法： 我们不切分层数，而是切分这个矩阵 W。
GPU 0 负责维护 W的左半部分 [4096, 2048]。
GPU 1 负责维护 W 的右半部分 [4096, 2048]。
计算时，两张卡分别计算一部分结果，然后拼接起来。

3. 分布式训练的基本“物理单位”

在进入下一部分的具体算法（如 DDP, Ring-AllReduce）之前，我们需要统一几个基础术语，否则看日志会很晕：

假设我们有2台服务器，里面插了 8 张 GPU：

Node (节点)： 物理服务器的概念。如果是单机多卡，Node=1。如果是多机多卡（比如 10 台服务器），Node=10。

World Size (世界大小)： 全局进程总数（通常等于 GPU 总数）。

实例： 2 台服务器，每台 8 卡。World Size = 16。

Rank (全局排名)： 每一个进程（GPU）的全局唯一 ID。

实例： 上述环境中，ID 范围是 0 到 15。Rank 0 通常是主节点（Master），负责协调。

Local Rank (本地排名)： 在当前服务器（Node）内的 ID。

实例： 第 2 台服务器的第一张卡。Rank可能是 8，但 Local Rank 是 0。

2.数据并行与 Ring All-Reduce

在单机单卡训练中，流程很简单：前向 -> 反向 -> 梯度更新。在多卡数据并行中，流程变成了：前向 -> 反向 -> [梯度同步] -> 梯度更新。

这个 [梯度同步] 是性能的生死线。如果同步太慢，显卡计算完就在空转（Idle），造成算力浪费。

1. 通信拓扑的演进：从中心化到去中心化

假设我们有 4 张 GPU (Rank 0 ~ 3)，每张卡计算出了一个梯度向量 $g$ 。我们的目标是让每张卡都得到 $\sum g = g_0 + g_1 + g_2 + g_3$ 。

方案 A：Parameter Server (PS, 参数服务器模式) —— 旧时代的“中心化”

逻辑： 选 Rank 0 做班长。Rank 1, 2, 3 把自己的梯度发给 Rank 0。Rank 0 算出总和，再把结果发回给 Rank 1, 2, 3。
瓶颈： 这构成了 N-to-1 的通信。Rank 0 的网络带宽会被瞬间打满，而其他卡在等待。随着卡数增加，Rank 0 必死无疑。这在早期的 TensorFlow 分布式中很常见。

方案 B：Ring All-Reduce (环状全归约) —— 现代的“去中心化”

逻辑： 没有班长。所有 GPU 排成一个逻辑上的圆环（0->1->2->3->0）。大家同时收发数据。
优势： 带宽利用率最优。无论有多少张卡，每张卡的通信量是恒定的，不会随卡数增加而显著变慢。这是百度 Silicon Valley AI Lab (SVAIL) 在 2017 年引入深度学习领域的关键算法，也是 NVIDIA NCCL 库的基础。

2. 实例拆解：Ring All-Reduce 是如何工作的？

为了讲清楚这个过程，我们不使用比喻，而是通过数据块的移动来演示。

假设有 3 张 GPU (GPU 0, 1, 2)。每张卡上有一个梯度向量，长度为 N。我们将这个向量切分为 3 个块 (Chunk)。

GPU 0 的梯度数据： $[a_0, b_0, c_0]$
GPU 1 的梯度数据： $[a_1, b_1, c_1]$
GPU 2 的梯度数据： $[a_2, b_2, c_2]$

目标： 每张卡最后都要拥有 $[a_{sum}, b_{sum}, c_{sum}]$ 。

Ring All-Reduce 分为两个阶段：Scatter-Reduce (散播-归约) 和 All-Gather (全收集)。

第一阶段：Scatter-Reduce (大家各自算一部分和)

在这个阶段，显卡之间互传数据，同时进行累加。

Step 1:
- GPU 0 发送 $b_0$ 给 GPU 1，接收 GPU 2 的 $a_2$ 。
- GPU 1 发送 $c_1$ 给 GPU 2，接收 GPU 0 的 $b_0$ 。
- GPU 2 发送 $a_2$ 给 GPU 0，接收 GPU 1 的 $c_1$ 。
- 计算： GPU 1 收到 $b_0$ 后，计算 $b_0+b_1$ 。其他卡同理。

Step 2 (继续传):

GPU 1 把刚才算好的 $(b_0+b_1)$ 发给 GPU 2。
GPU 2 计算： $(b_0+b_1) + b_2$ -> 得到了完整的 $b_{sum}$ 。
同时，GPU 0 得到了完整的 $a_{sum}$ 。
同时，GPU 1 得到了完整的 $c_{sum}$ 。

阶段结论： 此时，每张卡只拥有完整结果的三分之一。

GPU 0: $[a_{sum}, \dots, \dots]$
GPU 1: $[\dots, b_{sum}, \dots]$
GPU 2: $[\dots, \dots, c_{sum}]$

第二阶段：All-Gather (互通有无)

现在大家把各自算好的完整部分，沿着环传给下一家，不再计算，只是复制。

经过两轮传输，所有卡都凑齐了 $[a_{sum}, b_{sum}, c_{sum}]$ 。

原理总结：

Ring All-Reduce 将大数据的通信压力，打散到了每一条 GPU-to-GPU 的链路上（如 NVLink），让所有链路在同一时刻都在全速工作，没有闲置。

3. PyTorch DDP 的工程优化：Bucketing (分桶)

理解了算法，我们看实际落地。在 PyTorch 的 DistributedDataParallel (DDP) 中，如果每算出一个参数的梯度（比如一个 Bias 的梯度只有 1KB）就跑一次 All-Reduce，那通信启动的延迟 (Latency) 会远大于实际传输时间。

为了解决这个问题，DDP 引入了 Bucket (桶) 的概念。

实例讲解

假设你的模型是 ResNet-50，有 160 个参数层（Tensors）。

倒序排列： 反向传播是从最后一层（输出层）往前算的。DDP 会按照反向传播的顺序（Layer 160 -> Layer 159 -> ... -> Layer 1）来管理梯度。

装桶：

DDP 默认开辟一个 25MB 的 Bucket（由 bucket_cap_mb 参数控制）。
当 Layer 160 的梯度算出来，放入 Bucket。
Layer 159 算出来，放入 Bucket。
...
当 Bucket 满了（或者积攒了一定数量），DDP 就会立刻触发一次 Ring All-Reduce 异步传输。

计算与通信重叠 (Computation-Communication Overlap)：

这是 DDP 快的核心原因。
当 GPU 还在计算 Layer 100 的梯度（计算任务）时，GPU 后台已经在通过网络传输 Layer 160-150 的梯度（通信任务）。
理想状态下，除了第一个 Bucket，后续的通信时间都被计算时间“掩盖”了。

4. 关键代码对应的物理意义

当你写下这段代码时：

model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

底层发生了以下事情：

Model Broadcast: Rank 0 将自己的模型权重参数复制一份，强制覆盖到所有其他 Rank，确保初始状态数学上完全一致。
Reducer 初始化: 初始化 Bucket 机制，准备拦截反向传播产生的梯度。
Hook 注册: 在每个 Parameter 上注册 Autograd Hook。一旦梯度计算完成，Hook 就会触发，把梯度搬运到 Bucket 里，准备 All-Reduce。

3.流水线并行与张量并行

在前两部分，我们解决了“数据量大”的问题（DDP）。但在面对千亿参数的大模型（LLM）时，单张显卡的 80GB 显存根本装不下模型权重。这时，我们必须把模型切开，分配到多张卡上。这就是 模型并行 (Model Parallelism, MP)。

切分模型主要有两个维度：纵向切（流水线并行）和横向切（张量并行）。本部分将深入这两个硬核技术。

1. 纵向切分：流水线并行 (Pipeline Parallelism, PP)

核心思想： 将模型的层 (Layers) 分组，分配给不同的 GPU。

实例推演：Naive Pipeline（朴素流水线）的问题

假设一个 Transformer 模型有 4 层（L1, L2, L3, L4），我们有 4 张 GPU（G0, G1, G2, G3）。

分配： G0 负责 L1，G1 负责 L2，... G3 负责 L4。
数据流（前向）：
1. Batch 1 进入 G0。G0 计算 L1，把激活值（Activations）发给 G1。此时 G1, G2, G3 闲置。
2. G1 收到数据，计算 L2，发给 G2。此时 G0, G2, G3 闲置。 ...
痛点：气泡 (Bubble) 在任意时刻，只有 1 张 GPU 在工作，其他 3 张都在干等数据。这种现象被称为“气泡”。在这个例子中，硬件利用率极低（接近 1/4）。

解决方案：GPipe 与 Micro-Batch

为了填补这些气泡，工业界（如 Google 的 GPipe）引入了 Micro-Batch（微批次） 的概念。

操作： 假设原本的一个 Batch Size 是 1024。我们不一次性把 1024 个数据灌进去，而是切分成 4 个 Micro-Batch（每个 256）。
流水线填充：
1. (Time 1) G0 处理 Micro-Batch 1。
2. (Time 2) G0 处理 Micro-Batch 2；同时 G1 处理 G0 刚刚传过来的 Micro-Batch 1。
3. (Time 3) G0 处理 MB 3；G1 处理 MB 2；G2 处理 MB 1。
效果： 通过切碎数据，让后续的 GPU 尽早开始工作。虽然在开头（填充阶段）和结尾（排空阶段）依然有气泡，但中间大部分时间，4 张 GPU 是同时工作的。

代价： 显存占用增加（需要缓存所有 Micro-Batch 的中间激活值用于反向传播）。

2. 横向切分：张量并行 (Tensor Parallelism, TP)

PP 解决了层数太多的问题，但对于某些超大的层（比如词表层或超宽的 MLP 层），可能单层参数就撑爆显存。这时需要把张量（矩阵）本身切开。这是 NVIDIA Megatron-LM 的核心技术。

实例讲解：MLP 层的切分

在 Transformer 的 MLP 块中，主要包含两个线性层（Linear Layer）：

第一个层 A： 把维度从 $h$ 放大到 $4h$ (例如 $1024 \to 4096$ )。公式： $Y = X A$ 。
非线性激活： GeLU。
第二个层 B： 把维度从 $4h$ 缩小回 $h$ (例如 $4096 \to 1024$ )。公式： $Z = Y B$ 。

假设我们有 2 张 GPU。我们要切分这两个大矩阵 $A$ 和 $B$ 。

步骤一：列并行 (Column Parallelism) —— 切分矩阵 A

矩阵 $A$ 的形状是 $[1024, 4096]$ 。我们按列把它一分为二。

G0 持有： $A_1$ (形状 $[1024, 2048]$ ，即前一半列)。
G1 持有： $A_2$ (形状 $[1024, 2048]$ ，即后一半列)。

输入 $X$ (形状 $[B, 1024]$ ) 复制到两张卡上。

G0 计算： $Y_1 = X \times A_1 \rightarrow$ 得到 $[B, 2048]$ 的结果。
G1 计算： $Y_2 = X \times A_2 \rightarrow$ 得到 $[B, 2048]$ 的结果。

注意： 此时 Y 被劈开了，G0 只有左半部分特征，G1 只有右半部分特征。

步骤二：行并行 (Row Parallelism) —— 切分矩阵 B

矩阵 B 的形状是 [4096, 1024]。为了配合上一步的输出，我们按行把它一分为二。

G0 持有： $B_1$ (形状 $[2048, 1024]$ ，即上半部分行)。
G1 持有： $B_2$ (形状 $[2048, 1024]$ ，即下半部分行)。

步骤三：巧妙的结合 (Megatron-LM 的精髓)

现在 G0 手里有 $Y_1$ (左半特征) 和 $B_1$ (上半权重)。

G1 手里有 $Y_2$ (右半特征) 和 $B_2$ (下半权重)。

根据矩阵乘法原理：

$Z = Y \times B = [Y_1, Y_2] \times \begin{bmatrix} B_1 \\ B_2 \end{bmatrix} = Y_1 B_1 + Y_2 B_2$

G0 计算 $Z_1 = Y_1 B_1$ 。

G1 计算 $Z_2 = Y_2 B_2$ 。

关键点： 在这两个线性层中间（即 $Y$ 产生后），G0 和 G1 不需要通信。因为 G0 下一步正好只需要 $Y_1$ ，G1 正好只需要 $Y_2$ 。只有在最后计算 $Z$ 时，才需要一次通信。

步骤四：All-Reduce

G0 算出 $Z_1$ ，G1 算出 $Z_2$ 。

此时两张卡执行一次 All-Reduce (Sum)：

G0 得到 $Z_1 + Z_2 = Z$ 。
G1 得到 $Z_1 + Z_2 = Z$ 。

总结： 通过“列并行+行并行”的组合，我们在一个 MLP 块内部，只发生了一次 All-Reduce 通信（即在 Output 处），极大地降低了通信成本。

3. 3D 并行：这一部分的总结

在实际训练万亿参数模型（如 GPT-4）时，通常是三种并行混合使用，称为 3D Parallelism：

Data Parallel (DP): 比如有 1000 张卡。我们分成 100 个组，每组 10 张卡。这 100 个组之间做数据并行（复制模型）。
Pipeline Parallel (PP): 在这 10 张卡的一组内，把 96 层模型切分成 5 个阶段 (Stage)，每个阶段 2 张卡。
Tensor Parallel (TP): 在这 2 张卡内部，把具体的矩阵切开计算。

并行方式	切分对象	核心通信操作	解决痛点	瓶颈
DP	数据 (Batch)	All-Reduce (Gradients)	训练太慢	显存容量
PP	模型层 (Layers)	P2P (Send/Recv)	层数太多，单卡放不下	气泡 (Idle time)
TP	矩阵 (Tensors)	All-Reduce (Activations)	单层参数太大	通信带宽要求极高 (通常需要 NVLink)

无论是 DP, PP 还是 TP，只要是模型并行，都非常复杂且依赖高速网络。如果你只有几张显卡（比如 4张 3090），想微调一个 LLaMA-7B，用 TP/PP 太重了，用 DDP 显存又不够（因为 DDP 要复制模型）。

在 Part 2 中，我们介绍了 DDP（速度快，但显存浪费严重，每张卡都要存一份完整的模型和优化器状态）。在 Part 3 中，我们介绍了 MP（能跑大模型，但工程实现极难，且对显卡间的网络带宽要求极高）。

现在我们要讲的是目前微调大模型（如 LLaMA, ChatGLM）时最常用的技术：ZeRO (Zero Redundancy Optimizer)，也就是大家熟知的 DeepSpeed 的核心技术。

它的核心理念是：能不能既保持数据并行（DDP）的简单逻辑，又要把显存占用降下来？

ZeRO 与显存革命

1. 显存到底被谁吃掉了？（硬核算账）

要理解 ZeRO，必须先算一笔账。很多人以为显存不够是因为模型太大（参数多），但这只是冰山一角。

假设我们要训练一个 15 亿参数 (1.5B) 的 GPT-2 模型，使用 混合精度 (FP16) 训练。

模型权重 (Parameters): 1.5B $\times$ 2 Bytes (FP16) = 3 GB。

现在的显卡动不动就 24GB、80GB，为什么 3GB 的模型经常跑不起来？因为训练状态才是显存杀手。使用最常用的 Adam 优化器时，除了模型权重，还需要存储：

梯度 (Gradients): FP16，大小同权重。 = 3 GB。

优化器状态 (Optimizer States): Adam 需要维护参数的 Momentum (动量) 和 Variance (方差)，且为了精度通常用 FP32 存储。

Momentum (FP32): 1.5B x 4 Bytes = 6 GB
Variance (FP32): 1.5B x 4 Bytes = 6 GB
Master Weights (FP32备份): 1.5B x 4 Bytes = 6 GB
优化器总计: 18 GB。

总结： 一个本体只有 3GB 的模型，训练时仅仅是“静态”占用的显存就高达 3GB (Param) + 3GB (Grad) + 18GB (Opt) = 24GB。这还没算中间的激活值（Activation）和临时显存。这就是为什么单卡 24GB 的 3090/4090 很难跑大模型全量微调的原因。

2. DDP 的痛点：冗余

在标准的 DDP（数据并行）中，如果我们有 8 张显卡：

每张卡都必须保留这 24GB 的完整数据。
但实际上，每张卡只处理 1/8 的数据，更新梯度时也是同步一样的数值。
痛点： 既然大家最后的参数都一样，为什么每张卡都要完整存一份 Adam 的状态？这就是巨大的显存冗余。

3. ZeRO 的核心：切分 (Sharding)

ZeRO 的思想非常直观：既然存不下，那就切碎了大家分着存。 它不切分计算（像模型并行那样），而是切分存储。

ZeRO 定义了三个阶段（Stages），切分程度逐级递增：

ZeRO Stage 1: 切分优化器状态 ( $P_{os}$ )

原理： 把那最占内存的 18GB 优化器状态切成 8 份。每张卡只负责更新 1/8 的参数。
流程：
1. 每张卡正常算梯度。
2. 进行 All-Reduce 同步梯度。
3. 关键点： 更新参数时，GPU 0 只负责更新模型的前 1/8 参数，GPU 1 更新接下来的 1/8...
4. 更新完后，大家再通过 All-Gather 把最新的参数同步回来。
收益： 显存显著下降（减少了约 75% 的静态显存需求）。

ZeRO Stage 2: 切分梯度 ( $P_{os+g}$ )

原理： 既然优化器状态切了，梯度也没必要大家都存完整的。
流程：
1. 每张卡算出梯度后，不再进行 All-Reduce（即不再让大家都获得完整梯度）。
2. 而是进行 Reduce-Scatter：GPU 0 只收聚合后的前 1/8 梯度，GPU 1 只收第 2/8...
3. GPU 0 用这 1/8 的梯度更新自己维护的那 1/8 优化器状态。
收益： 进一步节省了 3GB 的梯度显存。

ZeRO Stage 3: 全切分 ( $P_{os+g+p}$ )

原理： 最激进的一步。连模型权重（那 3GB）也切了。
现状： 在 ZeRO-3 下，每张显卡里没有完整的模型。
- GPU 0 只存 Layer 1-2。
- GPU 1 只存 Layer 3-4。
关键问题： 没完整模型怎么做前向传播（计算）？
动态抓取机制 (On-the-fly Fetching):
1. 前向传播开始： 到了 Layer 1。
2. Broadcast: GPU 0 把 Layer 1 的权重广播给所有其他 GPU。
3. Compute: 大家拿到 Layer 1，计算，得出结果。
4. Discard (丢弃): 除了 GPU 0（持有者），其他 GPU 立即删除 Layer 1 的权重。
5. 进入 Layer 2... 重复上述过程。
收益： 显存占用降到最低。理论上，显存占用与 GPU 数量成反比。卡越多，单卡占用越小。
代价： 通信量剧增（以时间换空间）。

4. 进阶：ZeRO-Offload

即便有了 ZeRO-3，显存可能还是不够（比如想用单卡跑 70B 模型）。 DeepSpeed 于是推出了 ZeRO-Offload。

原理： 显卡显存（VRAM）很贵且小（24GB），但内存（RAM）很便宜且大（可以插到 512GB）不过最近确实涨价很厉害，但是还是比显存便宜。

做法： 把 优化器状态 和 梯度更新的计算 全部搬到 CPU 和内存 里去执行。

流程：
1. GPU 疯狂计算前向+反向，算出梯度。
2. GPU 把梯度扔给 CPU。
3. CPU 在内存里慢慢更新参数（用 Adam）。
4. CPU 把更新好的参数写回 GPU。
结果： 只要显卡能塞下一层模型的参数，就能跑起来。这就是为什么现在很多人可以用消费级显卡微调大模型的原因。