2、分布式训练基础：数据并行、模型并行、流水线并行的基本概念与区别

当单块 GPU 的显存容量（如 80GB H100）和算力无法满足大语言模型（LLM）的训练需求时，分布式训练成为必然选择。以 175B 为例，仅模型参数就需要约 350GB 的 FP16 显存，加上优化器状态（Adam 需额外 2 倍参数量的 FP32 副本）和梯度，总显存需求超过 1.2TB，远超任何单卡能力。分布式训练的核心思想是将计算和显存压力分摊到多块 GPU 上，而数据并行、模型并行、流水线并行正是三种最基本的并行策略。理解它们的原理、适用场景与组合方式，是掌握 Megatron-LM 等框架的基石。

上图展示了三种并行策略的核心思想与组合关系。数据并行关注“数据切分”，模型并行关注“参数切分”，流水线并行关注“层间切分”。在实际工程中，它们往往不是互斥的，而是像乐高积木一样组合使用，形成三维并行。下面逐一深入剖析。

2.1 数据并行（Data Parallelism, DP）

核心原理

数据并行是最直观、最常用的并行方式。其核心思想是：每张 GPU 都保存一份完整的模型参数副本，将训练数据按 batch 维度切分成多份，分别送入不同的 GPU 进行前向和反向计算。每个 GPU 独立计算得到梯度后，通过 AllReduce 通信操作将所有 GPU 的梯度进行求和平均，然后每个 GPU 使用相同的平均梯度更新自己的模型副本。这样，所有 GPU 上的模型始终保持一致，但处理的数据量是原来的 N 倍（N 为 GPU 数量）。

在 Megatron-LM 中，数据并行通常通过 PyTorch 的 DistributedDataParallel (DDP) 实现。DDP 在反向传播时，会自动对每个参数的梯度执行 AllReduce。关键参数是 world_size（总 GPU 数）和 rank（当前 GPU 编号）。

判断方法

判断一个训练任务是否适合数据并行，主要看两点：

• 模型能否放入单卡显存：如果模型参数 + 梯度 + 优化器状态的总和小于单卡显存的 70%，数据并行是首选。
• 训练数据量是否足够大：数据并行要求每个 GPU 分到的 batch size 不能太小（通常 ≥ 8），否则梯度噪声大，影响收敛。

案例拆解：BERT-Large 的数据并行训练

假设我们使用 8 张 A100 80GB 训练 BERT-Large（3.4 亿参数）。模型参数 FP16 约 680MB，梯度 680MB，Adam 优化器状态（FP32 的 momentum 和 variance）约 2.7GB，总计约 4GB。单卡显存完全放得下。

初始配置：全局 batch size = 256，单卡 batch size = 32。8 张卡并行，每张卡处理 32 个样本。

关键操作：每张卡前向计算 loss → 反向计算梯度 → 执行 AllReduce 对梯度求平均 → 每张卡用平均梯度更新本地模型。

通信开销分析：每步训练需要通信的梯度总量为 680MB × 8 = 5.44GB（AllReduce 的通信量约为 2×(N-1)/N × 数据量）。在 NVLink 600GB/s 的带宽下，通信耗时约 9ms，而计算耗时约 200ms，通信占比不到 5%，效率很高。

结果：训练吞吐量接近线性扩展，8 卡相比单卡加速比约 7.6 倍（受通信和负载不均影响）。

2.2 模型并行（Model Parallelism, MP）—— 张量并行

核心原理

模型并行（在 Megatron-LM 中特指张量并行，Tensor Parallelism）解决的是“单层太大放不进单卡”的问题。其核心思想是：将 Transformer 层中的权重矩阵按行或列切分到多张 GPU 上，每张 GPU 只保存一部分权重，计算时通过通信聚合结果。

以 Transformer 中的线性层 Y = XW 为例（X 为输入，W 为权重矩阵），Megatron-LM 采用两种切分方式：

• 列并行（Column Parallel）：将 W 按列切分为 [W1, W2]，分别放在 GPU0 和 GPU1 上。输入 X 需要广播到两张卡，每张卡计算 XW_i，最后通过 AllGather 拼接得到完整输出 Y。
• 行并行（Row Parallel）：将 W 按行切分为 [W1; W2]，输入 X 也按列切分为 [X1, X2]，每张卡计算 X_i W_i，最后通过 AllReduce 求和得到 Y。

在 Transformer 中，通常对 attention 的 QKV 投影使用列并行，对输出投影使用行并行，对 MLP 的第一层使用列并行、第二层使用行并行。这样可以在不增加额外通信的前提下，将计算和显存均匀分摊。

判断方法

模型并行的适用场景非常明确：

• 单个 Transformer 层的权重矩阵过大（如 hidden size 超过 8192），无法放入单卡。
• 卡间通信带宽极高（如 NVLink 或高速 InfiniBand），因为每个 Transformer 层的前向和反向都需要多次 AllGather/AllReduce。

案例拆解： 175B 的张量并行

175B 的 hidden size 为 12288，单个 attention 的 QKV 投影矩阵大小为 12288×36864，约 1.8GB（FP16）。单卡 80GB 虽然能放下，但计算量巨大，单卡算力不足。

配置：使用 8 卡张量并行（MP=8）。QKV 投影矩阵被切分为 8 份，每份 12288×4608，约 226MB。每张卡只计算自己负责的部分。

通信模式：前向传播中，输入 X 需要广播到 8 张卡（AllGather），输出需要 AllReduce 求和。每个 Transformer 层有 2 次 AllGather 和 2 次 AllReduce。

性能数据：在 A100 集群上，MP=8 时，单层计算时间约 15ms，通信时间约 5ms，通信占比 25%。如果使用 MP=16，通信占比会上升到 40% 以上，效率明显下降。

2.3 流水线并行（Pipeline Parallelism, PP）

核心原理

流水线并行是将模型按层（Layer）切分到不同的 GPU 上。例如，一个 40 层的 Transformer 模型，可以切分为 4 段，每段 10 层，分别放在 4 张 GPU 上。数据像流水线一样，依次经过 GPU0 → GPU1 → GPU2 → GPU3 完成前向传播，反向传播则按相反顺序进行。

朴素的流水线并行存在严重的“气泡”（Bubble）问题：在第一个 micro-batch 完成前向之前，后面的 GPU 处于空闲状态。为了解决这个问题，Megatron-LM 引入了 micro-batch 流水线调度：将一个大 batch 切分为多个 micro-batch（如 16 个），依次送入流水线。当 GPU0 处理完 micro-batch 1 并传给 GPU1 后，GPU0 立即开始处理 micro-batch 2，从而让所有 GPU 尽量保持忙碌。

Megatron-LM 默认使用 1F1B（One-Forward-One-Backward）调度策略：每个 GPU 交替执行一个 micro-batch 的前向和一个 micro-batch 的反向，以平衡显存占用和气泡率。

判断方法

流水线并行适合以下场景：

• 模型层数非常多（如 100+ 层），但单层大小适中。
• 卡间通信带宽相对有限（如跨节点），因为流水线并行只需要在相邻 GPU 之间传输激活值和梯度，通信量远小于张量并行。
• 需要跨节点扩展：通常与张量并行结合，节点内用 MP，节点间用 PP。

案例拆解： 175B 的流水线并行

有 96 层 Transformer。假设我们使用 4 节点（每节点 8 卡），共 32 卡。采用 PP=4，MP=8 的组合：每个节点内 8 卡做张量并行，4 个节点之间做流水线并行。

配置：每个节点负责 24 层（96/4）。每个 micro-batch 的激活值在节点间传输，通信量约为每层激活值大小 × 层数边界。

调度参数：global batch = 2048，micro-batch = 16，则每个 micro-batch 大小为 128。流水线深度为 4，气泡率约为 (PP-1)/(PP-1+micro-batch) = 3/19 ≈ 15.8%。

优化手段：增大 micro-batch 数量可以降低气泡率，但会增加显存占用（需要保存更多中间激活）。Megatron-LM 还支持虚拟流水线（Virtual Pipeline），将每个 GPU 上的层再细分为更小的 chunk，进一步降低气泡。

结果对比：PP=4 时，训练吞吐量约为单节点 MP=8 的 3.2 倍（理论 4 倍，气泡损失 15.8%，跨节点通信损失约 5%）。

2.4 三种并行策略的对比与组合

在实际的 Megatron-LM 训练中，三种并行策略通常组合使用，形成三维并行：

• 节点内：使用张量并行（MP=8），利用 NVLink 高速带宽。
• 节点间：使用流水线并行（PP=节点数），减少跨节点通信量。
• 全局：使用数据并行（DP=总卡数/(MP×PP)），扩大整体 batch size。

例如，使用 64 卡训练 175B，可以配置为 DP=2, MP=8, PP=4。总卡数 = 2×8×4 = 64。每张卡只保存 1/8 的参数（张量并行），负责 1/4 的层（流水线并行），处理 1/2 的数据（数据并行）。

常见误区与总结