AI流水线并行折腾手记

GPT-3的论文里就提了这种思路，Megatron-LM也把它用得挺狠。

用一个简单例子说明：假设模型有8层，流水线并行切成4个stage，每个stage负责2层。

起因：GPU没剩几块，模型却越来越大

手里的GPU资源有限，但模型规模没停下来。数据并行是常规做法，但到了单张GPU装不下模型的时候，要么换更大显存的GPU，要么把模型切片扔到多张卡上。

模型并行说起来很简单：把模型按层切开，每张卡负责一部分。GPT-3的论文里就提了这种思路，Megatron-LM也把它用得挺狠。但真正动手搞的时候，发现事情没那么干净。

基本原理：把模型切成段

先说清楚几个概念，后面才好讲踩坑。

数据并行：多张卡各自跑一份完整模型，梯度同步。问题在于每张卡都得装得下整个模型。

模型并行：把模型切分到多张卡上，每张卡只存一部分。又分两种：张量并行把层内切开，流水线并行把层间切开。

流水线并行：把模型按层分到多个stage，每个stage在一张或一组卡上。数据在stage之间流过，前一个stage算完，结果传给下一个。

用一个简单例子说明：假设模型有8层，流水线并行切成4个stage，每个stage负责2层。

理想情况下，4个stage能同时工作：stage1在处理下一个batch，stage2在处理当前batch，stage3在处理上一个batch，stage4在处理上上个batch。但这有个前提：stage之间不能互相等待。

实践：用Megatron-LM搞一把

环境和准备工作

硬件：4个节点，每节点8张A100（40GB），共32张卡。

软件环境：

# CUDA版本
nvcc --version
# nvcc: NVIDIA (R) Cuda compiler driver
# Copyright (c) 2005-2023 NVIDIA Corporation
# Built on Fri_Jan__6_16:45:21_PST_2023
# Cuda compilation tools, release 12.0, V12.0.140

# 检查NCCL版本
python -c "import torch; print(torch.cuda.nccl.version())"
# 21903123

Megatron-LM版本：v3.0.2

git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
git checkout v3.0.2

基础配置

先用一个相对小的模型验证流程，175B那种等调通再说。

配置文件 examples/megatron_gpt_config.yaml：

# 模型配置
tensor-model-parallel-size: 1
pipeline-model-parallel-size: 4
num-layers: 32
hidden-size: 4096
num-attention-heads: 32
seq-length: 2048
max-position-embeddings: 2048

# 训练配置
micro-batch-size: 4
global-batch-size: 256
train-iters: 100000
lr: 0.00015
lr-decay-style: cosine
weight-decay: 0.1

# 优化器配置
optimizer: adam
clip-grad: 1.0

第一次运行：连代码都跑不起来

bash examples/pretrain_gpt.sh \
    --tensor-model-parallel-size 1 \
    --pipeline-model-parallel-size 4 \
    --num-layers 32 \
    --hidden-size 4096 \
    --num-attention-heads 32 \
    --micro-batch-size 4 \
    --global-batch-size 256 \
    --seq-length 2048 \
    --max-position-embeddings 2048 \
    --train-iters 100000 \
    --lr 0.00015 \
    --lr-decay-style cosine \
    --weight-decay 0.1 \
    --optimizer adam \
    --clip-grad 1.0

直接报错：

RuntimeError: NCCL error in: /workspace/pytorch/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:925, unhandled system error, NCCL version 2.17.1
ncclSystemError: System call (e.g. socket, malloc) or external library call failed or device error.

第一个坑：NCCL通信问题。多节点环境对网络要求很高，尤其是跨节点流水线并行。

检查网络配置：

# 检查网卡绑定
ibstat | grep -A 5 "State:"
# 确保每个节点上IB卡都在Active状态

# 测试节点间带宽
ib_write_bw -d mlx5_0 -s 4096 -n 1000000 <对端IP>
# 发现跨节点带宽只有正常值的一半

问题定位到：不同节点的网卡没有正确绑定到同一个高速互联网络，部分流量走了较慢的以太网。

调整网络拓扑，确保所有流水线通信都走InfiniBand，再试：

# 强制使用指定网卡
export NCCL_SOCKET_IFNAME=ib0
export NCCL_IB_DISABLE=0

这次能跑起来了，但速度很慢。

第二个问题：GPU利用率太低

观察GPU利用率：

nvidia-smi dmon -s u -c 100

发现GPU利用率波动很大，平均只有40%左右。

问题在于流水线填充。流水线并行不是一开始就能并行，需要先把流水线填满。用最naive的GPipe方式，每个micro batch要串行流过所有stage，大部分时间stage都在等数据。

Megatron-LM提供了1F1B（One-Forward-One-Backward）调度来缓解这个问题，但配置上需要仔细调整。

修改配置启用1F1B：

--pipeline-model-parallel-split-rank 1 \
--num-layers-per-virtual-pipeline-stage 4

这里把32层切成4个物理stage，每个物理stage再内部切分成虚拟pipeline，每个虚拟stage负责4层。

调整后GPU利用率提升到70%左右，但还是不理想。

第三个坑：micro batch大小的选择

流水线并行里，micro batch size是个很敏感的参数。太小了，通信开销占大头；太大了，显存不够。

尝试不同配置：

micro batch从2提到4，速度提升明显；但再提到8就OOM了。增加虚拟pipeline阶段数（从4到8）能进一步流水化，带来一些提升。

但还有一个问题：节点间通信开销太大。

通信优化：多网卡绑定和数据并行混用

Megatron-LM支持混合并行：流水线并行 + 数据并行。这样可以在节点内部用数据并行，减少跨节点通信。

调整配置：

tensor-model-parallel-size: 1
pipeline-model-parallel-size: 2
data-parallel-size: 2

这里把流水线并行从4减到2，每个节点内部用2个数据并行worker。这样跨节点通信减少了一半。

bash examples/pretrain_gpt.sh \
    --tensor-model-parallel-size 1 \
    --pipeline-model-parallel-size 2 \
    --data-parallel-size 2 \
    --num-layers 32 \
    --hidden-size 4096 \
    --num-attention-heads 32 \
    --micro-batch-size 4 \
    --global-batch-size 256 \
    --seq-length 2048 \
    --num-layers-per-virtual-pipeline-stage 8

训练速度提升到1850 tokens/s。

第四个问题：checkpoint恢复

长时间训练不可避免会遇到中断。流水线并行的checkpoint恢复比数据并行复杂，因为每个stage的状态不一样。

默认checkpoint配置：

--save /path/to/checkpoints \
    --save-interval 2000 \
    --load /path/to/checkpoints

但第一次尝试恢复时遇到问题：

RuntimeError: mismatch in tensor parallel size

原因：恢复时没有正确指定parallel size，导致恢复的checkpoint和当前配置不匹配。

正确做法：恢复时必须和保存时的并行配置完全一致。

# 保存时记录配置
cat > checkpoint_config.txt << EOF
tensor-model-parallel-size: 1
pipeline-model-parallel-size: 2
data-parallel-size: 2
num-layers: 32
hidden-size: 4096
micro-batch-size: 4
global-batch-size: 256
EOF

# 恢复时使用相同配置
bash examples/pretrain_gpt.sh \
    --tensor-model-parallel-size 1 \
    --pipeline-model-parallel-size 2 \
    --data-parallel-size 2 \
    --load /path/to/checkpoints

调优后的最终配置

经过一轮调优，最终配置：

# 并行策略
tensor-model-parallel-size: 1
pipeline-model-parallel-size: 2
data-parallel-size: 2
num-layers-per-virtual-pipeline-stage: 8

# 模型
num-layers: 32
hidden-size: 4096
num-attention-heads: 32
seq-length: 2048

# 批次
micro-batch-size: 4
global-batch-size: 256

# 训练
train-iters: 100000
lr: 0.00015
lr-decay-style: cosine
weight-decay: 0.1
clip-grad: 1.0

# 优化
optimizer: adam
bf16: true

训练速度：1950 tokens/s，GPU利用率82%。

实际踩过的坑总结

1. NCCL通信问题

现象：多节点环境下NCCL报错，训练无法启动。

原因：网卡配置不当，部分流量走了低速网络。

解决：确保所有通信走高速互联网络，必要时强制指定网卡：

export NCCL_SOCKET_IFNAME=ib0
export NCCL_IB_DISABLE=0

2. GPU利用率低

现象：GPU利用率平均只有40%，训练速度慢。

原因：流水线填充开销，大部分时间stage在等数据。

解决：启用1F1B调度，增加虚拟pipeline阶段数：

--num-layers-per-virtual-pipeline-stage 8

3. 跨节点通信开销

现象：多节点训练比单节点训练还慢。

原因：流水线并行导致跨节点通信频繁。

解决：混合流水线并行和数据并行，减少跨节点通信：

pipeline-model-parallel-size: 2
data-parallel-size: 2

4. Micro batch大小选择

现象：太小速度慢，太大OOM。

原因：micro batch太小通信开销大，太大显存不够。

解决：通过实验找到平衡点，一般4-8之间。

5. Checkpoint恢复失败

现象：checkpoint恢复时报错。

原因：恢复时并行配置和保存时不一致。

解决：恢复时使用完全相同的并行配置。

经验判断

折腾一圈下来，有几个实际感受。

流水线并行不是银弹：它解决的是单卡显存不够的问题，不是提升训练速度的。如果你的单卡能装下模型，数据并行通常更快。

混合并行是现实选择：纯流水线并行通信开销太大，混合数据并行和流水线并行更实用。

虚拟pipeline有边际效应：从4个虚拟stage提到8个有提升，再继续提升收益递减，但调参成本增加。

网络比算力更关键：多节点环境下，网络配置和拓扑设计直接影响训练速度。再强的算子，都在等网络。

模型设计要考虑并行：从一开始就把模型设计得更容易并行化，比后期再改更容易。

写在最后

流水线并行说到底是把问题拆解：模型太大，就切成段；算力不够，就分工做。但分工不等于简单切分，还要考虑协作效率。

工业流水线之所以高效，是因为每个工序都高度标准化，工位之间传输也高度优化。深度学习的流水线现在还在发展阶段，很多细节需要手工调优。

也许未来会有更自动化的并行策略，但现在还是得靠经验、靠试错、靠对系统行为的理解。

折腾了一轮，对分布式训练的理解更深了一些，但也意识到还有很多东西要学。模型的规模还在增长，硬件和软件也都在演进，这条路还很长。

版权声明: 本文首发于 指尖魔法屋-AI流水线并行折腾手记（https://blog.thinkmoon.cn/post/234-deep-dive-ai-pipeline-parallelism-model-pipeline-practice/） 转载或引用必须申明原指尖魔法屋来源及源地址！