大模型训练框架Megatron原理

这篇文章是关于NVIDIA Megatron框架的深入分析，该框架用于训练超大的Transformer语言模型。

摘要

Megatron：基于PyTorch的分布式训练框架，用于训练超大的Transformer语言模型。
目标：通过综合应用数据并行、Tensor并行和Pipeline并行来复现如GPT-3这样的大型模型。

引言

大模型的重要性：在NLP领域，大模型能带来更精准的语义理解和推理能力。
挑战：显存限制、计算需求巨大，以及不同并行策略的通信模式和通信量问题。

1.数据并行

概念：在每个worker上复制一份模型，输入数据集被分片，梯度定期汇总。
限制：GPU的batch_size过小，降低GPU利用率，增加通信成本；最大设备数受限于batch_size。

2.模型并行

概念：将模型的内存和计算分布在多个worker之间，解决单个设备无法容纳的问题。
类型：
流水线并行：将模型的不同层放到不同设备上。
张量并行：在层内进行分割，将矩阵运算分配到不同设备上。

技术组合

并行技术的组合使用：不同的并行技术可以结合使用，但会导致复杂的相互作用。
Megatron-LM：展示了如何结合流水线、张量和数据并行技术，在1000个GPU上训练大型语言模型。

张量模型并行（Tensor Model Parallelism）

原理：通过将矩阵乘法分割到多个GPU上，减少通信和控制GPU计算范围。
Transformer模型并行：特指层内切分，即Tensor Model Parallel。
切分Transformer：Megatron将Masked Multi Self Attention和Feed Forward都进行切分以并行化。
切分MLP：MLP块的第一部分是GEMM，后面是GeLU，Megatron提供了两种并行化GEMM的选项。
切分self attention：利用多头注意力操作中固有的并行性，以列并行方式对GEMM进行分区。

3.并行配置

3.1 符号说明

一些用于描述并行配置的符号，这些符号包括：

t：张量模型并行度（Tensor Model Parallelism degree）
p：流水线模型并行度（Pipeline Model Parallelism degree）
d：数据并行度（Data Parallelism degree）
m：微批次数量（microbatch count）
B：总批次大小（total batch size）
b：微批次大小（microbatch size）
s：序列长度（sequence length）
h：隐藏层大小（hidden size）

3.2 Tensor and Pipeline Model Parallelism

张量模型并行和流水线模型并行的配置：

张量模型并行：在同一服务器内的不同GPU之间分割模型的参数和计算。
流水线模型并行：跨服务器的不同GPU之间分割模型的层，每个阶段执行模型的一部分计算。

当使用流水线并行时，会引入流水线气泡（pipeline bubbles），这是由于流水线中的微批次需要在设备间传递，导致某些设备在等待数据时出现空闲。气泡的大小与流水线深度和微批次数量有关。

3.3 Data and Model Parallelism

数据并行和模型并行的结合使用：

流水线模型并行：当 t=1t=1（即不考虑张量模型并行）时，流水线的微批次数量是 m=B/(d⋅b)m=B/(d⋅b)，流水线气泡大小是 (n/d−1)⋅(B/(d⋅b))(n/d−1)⋅(B/(d⋅b))。
数据和张量模型并行：当使用张量模型并行时，每个微批次都需要执行all-reduce通信，这在多GPU服务器之间可能非常昂贵。而数据并行对于每个批次只需执行一次all-reduce。

结论：

结论#1：在使用多GPU服务器时，通常应该将张量模型并行度控制在单个服务器内的GPU数量以内，然后使用流水线模型并行来跨服务器扩展到更大的模型。
结论#2：当使用数据和模型并行时，总的模型并行大小应该是 M=t⋅pM=t⋅p，这样模型参数和中间元数据可以放入GPU内存。数据并行性可用于将训练扩展到更多GPU。

3.4 Microbatch Size

微批次大小 bb 对模型训练吞吐量的影响。微批次大小既影响操作的算术强度，也影响流水线气泡大小。最佳微批次大小取决于模型的吞吐量和内存占用特性，以及流水线深度 pp、数据并行尺寸 dd 和批尺寸 BB。

经验之谈#3：

经验之谈#3：最佳微批次尺寸 bb 取决于模型的吞吐量和内存占用特性，以及流水线深度 pp、数据并行尺寸 dd 和批尺寸 BB。

3.5 对比

对比了不同的并行机制：

张量与流水线并行：张量模型并行在节点内最佳，因为它减少了通信量。流水线模型并行使用更便宜的点对点通信，可以跨节点执行。
流水线与数据并行：流水线模型并行主要用于支持不适合单个worker的大型模型训练，数据并行用于扩大训练规模。
张量与数据并行：数据并行通信不频繁，而张量模型并行需要对每个微批进行all-to-all通信，这在多GPU节点上可能非常昂贵。

结论

Megatron的PTD-P技术：结合了流水线并行、张量并行和数据并行，在训练万亿参数的大型模型时达到了高聚合吞吐量。

张量模型并行（Tensor Model Parallelism）

在张量模型并行中，模型的参数和计算被分割到不同的GPU上。如果 tt 是张量模型并行度，那么每个GPU上只需要存储和计算模型参数的一部分。这允许模型的总大小超过单个GPU的内存限制。

数据并行（Data Parallelism）

数据并行涉及复制模型到多个GPU上，并在每个GPU上处理数据的一个子集。然后，所有GPU上的梯度被汇总（通过all-reduce操作）以更新模型参数。如果有 dd 个数据并行副本，那么总的模型副本数就是 dd。

微批次大小（Microbatch Size）

微批次大小 bb 是指在流水线模型并行中，每个GPU在单个阶段处理的数据量。选择合适的微批次大小对于优化训练吞吐量和减少流水线气泡至关重要。微批次大小影响计算时间，因为它决定了每个GPU上计算的算术强度。

通信量

在张量模型并行中，每个层的前向传播和后向传播都需要进行all-reduce操作，这增加了通信量。如果 tt 是张量模型并行度，那么每个微批次需要 2⋅t2⋅t 次all-reduce操作（前向和后向各一次）。因此，总的通信量与微批次数量 mm 和张量模型并行度 tt 成正比。

例如，有2台机器，每台机器8卡，一共16张卡；在tensor-parallel-size=2, pipeline-parallel-size=4的时候，每个模型的“子块”是如何被分配到16张卡上的？

_TENSOR_MODEL_PARALLEL_GROUP：

取值示例为[g0, g1], [g2, g3], ..., [g14, g15]

_PIPELINE_MODEL_PARALLEL_GROUP：

[g0, g4, g8, g12]，[g1, g5, g9, g13]，[g2, g6, g10, g14]，[g3, g7, g11, g15]

_DATA_PARALLEL_GROUP：

[g0, g2],[g1, g3], [g4, g6], [g5, g7], [g8, g10], [g9, g11], [g12, g14], [g13, g15]

world_size/(tensor_model_parallel_size* pipeline_model_parallel_size) = 16/(2*4)=2。

参考

[细读经典]Megatron论文和代码详细分析(2)

[细读经典]Megatron论文和代码详细分析(1)

Megatron-LM源码阅读（一）

Megatron-LM源码阅读（二）

megatron学习总结

GTC 2020: Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

www.DeepL.com/Translator

NVIDIA Megatron：超大Transformer语言模型的分布式训练框架（一）

NVIDIA Megatron：超大Transformer语言模型的分布式训练框架 (二)