大模型训练框架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 符号说明

一些用于描述并行