Pai-Megatron-Patch：围绕Megatron-Core打造大模型训练加速生态

作者：李鹏，胡凯文，王明，黄俊

单位：阿里云智能集团人工智能平台PAI算法团队

引言

Pai-Megatron-Patch（https://github.com/alibaba/Pai-Megatron-Patch）是阿里云人工智能平台PAI研发的围绕Nvidia MegatronLM的大模型开发配套工具，旨在帮助开发者快速上手大模型，完成大模型（LLM）相关的高效分布式训练，有监督指令微调，下游任务评估等大模型开发链路。最近一年来，我们持续打磨Pai-Megatron-Patch的性能和扩展功能，围绕Megatron-Core（以下简称MCore）进一步打造大模型训练加速技术生态，推出更多的的训练加速、显存优化特性，主要包括：

• 更新多款热门LLM大模型的Dense&MoE版本的预训练&微调最佳实践：llama3系列，qwen1.5，qwen2系列，deepseek-v2系列以及mistral系列等等。
• 支持模型并行条件下的高精度低损耗的Huggingface和MCore模型权重互转转换。
• 支持Transformer Engine训练技术引擎加持下的FlashAttention-3，FP8低精度训练加速且确保收敛正常。
• 支持基于torchrun启动的Communication Overlapping通信优化，能提供10%以上的吞吐性能提升。
• 支持Distributed Optimizer CPU Offloading训练显存优化技术，相对开源Megatron-Core在同等硬件资源条件下能训练更长的序列且吞吐速度相比关闭Offloading不会显著降低。

NVIDIA Megatron-Core在设计上更加解耦和模块化，不仅提高了LLM训练的效率和稳定性，也为二次开发和探索新的LLM架构提供了更大的灵活性。我们在Megatron-Core上面搭建了多种热门大模型比如LLama3，Qwen2等等，并为Mcore自带的Distributed Optmizer植入了CPU Offloading特性。CPU Offloading降显存策略可以和Transformer Engine，MoE以及流水并行等模块协同使用。整体技术栈如下图所示：

本文以开源LLama3.1为例，首先描述在Huggingface模型继续预训练/微调时权重转换过程中遇到的精度对齐问题，接着介绍PAI算法团队自研的带cpu offloading功能的分布式优化器以及Transformer Engine中推出的带cpu offloading功能的激活检查点技术以及相关实验效果，然后介绍最新的FlashAttention-3以及分布式通信优化相关的技术以及吞吐&收敛实验效果，最后给出在预训练&微调场景下综合运用这些加速开关的使用示例。

HF和Mcore模型双向转换精度对齐

概述

由于MCore在Transformer模型关键组件实现方面与Huggingface存在一定差异，导致在继续预训练或者微调时将Huggingface模型权重转换Mcore模型权重后执行前向推理会产生误差，该误差主要来自以下六个方面：

其中RMS Norm的精度计算差异如下所示：

消融实验

我们在A100/H100机型上比较了Mcore版的LLaMA3.1与HuggingFace版的LLaMA3.1的误差，实验数据如下表所示：

主要结论如下：

1. RoPE实现差异导致的误差非常小。
2. Bias SwiGLU Fusion及RMSNorm相比RoPE会造成更大的误差，并且这些误差逐层累积会对输出logits造成明显的影响。
3. 对于不同设备，由于底层架构&算子实现等原因，导致实际误差有所不同。
4. 即使不修改Attention和SwiGLU，模型在多个下游任务评估指标上的平均相对误差仅0.15%。（来自Nvidia[1]）

Mcore训练降显存&加速技术

本章节我们围绕Offloading，FlashAttention-3以及Overlapping三项技术展开讨论。卸载(Offloading)技术是将部分模型参数、优化器状态或者中间计算结果等从显存移动到其他存储设备（如RAM，NvME SSD等）来节省显存空间的使用，从而允许在少量机器上进行更大规模的模型训练。目前卸载技术主要有针对优化器参数更新过程的卸载（Optmizer CPU Offloading）以及对中间激活及对应模型权重的卸载（Activation/Weight CPU Offloading）。

Optimizer CPUOffloading

技术原理

目前大模型几乎全部采用Adam优化器进行混合精度训练/半精度训练。对于一个参数量为M的模型，以Megatron为例，当采用bfloat16格式进行训练时，显存占用能达到 [参数显存占用 + 参数梯度占用 + Adam状态占用] = [(2 + 4) + 4 + (4 + 4)] = 18M Bytes的大小。为了对这一部分进行优化，DeepSpeed提出了ZeRO-Offload[5]，也就是本节的优化器CPU卸载技术，旨在利用CPU的计算资源，将计算量较小而显存占用较多的参数更新部分(至少需要占用12M Bytes，包括全精度参数及优化器状态部分)放在CPU上进行，在参数更新后复制回GPU。尽管能有效解决显存占用问题，但我们观察到ZeRO-Offload技术仍存在以下不足：

1. 在CPU上进行参数更新影响了吞吐速度，而GPU在这一过程中几乎是idle的。
2. 由于3D并行的存在，各GPU负载不均。

为此，在ZeRO-Offload的基础上，我们实现了更细粒度的Optimizer Offloading，允许用户手动/半自动地控制优化器卸载比例，来达到吞吐及显存占用的最佳tradeoff。模块整体架构如图所示

Chunk Manager

尽管在优化器卸载中，优化器状态所在的设备会在初始化(静态策略)或前几个迭代(动态策略)确定，为了降低D2H/H2D的拷贝次数以及减少显存碎片，将多个参数/优化器状态/梯度打包成chunk是必要的。同时，为了能更好确定卸载比例，实时追踪当前显存分配量，Chunk Manager也是不可或缺的模块。在最新的Megatron-patch中，我们针对参数更新过程的实际需要，设计了Chunk及ChunkManager类。其中，Chunk通过提供AllocTensor、AppendTensor等API，实现张量内存的低层级统一分配管理。

在初始化Chunk对象后，通过AllocTensor/AppendTensor在chunk上新建或分配一个tensor，后续通过ResetDevice函数即可将该chunk上的所有Tensor移动至其他设备，为了更好地将通信与计算重叠，Chunk的关键API均支持异步特性。由于Optimizer Offload的需求，我们还设计了ChunkClone函数来方便创建对应的梯度/优化器状态的Chunk。

随后在Manager中，通过实现类似的接口来让将外部Tensor注册到Chunk上，并分配对应的梯度及优化器状态。由于涉及Chunk的拷贝以及内存回收，Chunk的大小选择及注册Tensor时的Chunk选择策略会影响实际运行时的显存利用率。针对上述两个问题，我们分别设计了Chunk Size预搜索以及Chunk Select函数为后续拓展保留可能。在当前实现中，我们针对Chunk Size采取了最少内部碎片的策略，而针对Chunk Select，由于其事实上是一个装箱问题，我们采用了较常见的Next-Fit策略。

HybridAdam

当参数、梯度及优化器状态分布在CPU RAM及GPU上时，需要对Optimizer进行对应修改。参考ColossalAI，我们在Megatron-Patch内实现了HybridAdam，来允许不同设备上的参数更新。在CPU及GPU参数更新上，HybridAdam采用了SIMD/CUDA实现来进行加速。此外，还需要在ChunkManger内提前分配优化器状态来进行内存管理。

Static Optimizer Offload Policy

当采用静态优化器卸载策略时，根据用户输入的卸载比例，选择部分参数在CPU上进行更新。由于ChunkManager统计了各Chunk的占用大小，我们能快速确定需要Offload的Chunk并将它们提前移动至GPU，Offload流程如下所示。

Auto Optimizer Offload Policy

通过仔细调整卸载比例，使用静态策略已能获得不错的效果。然而，由于各GPU负载不均，使用静态策略会将相同比例的参数量移动至GPU，部分降低了训练吞吐上限。为了减少调参，同时使各GPU的显存负载更加均衡，更好的解决方案是允许各卡自动采用不同的卸载比例。 [6]提出了通过实时追踪当前显存占用，并将显存占用分解为可计算的模型数据（如前文提到的 18M Bytes）及不可计算的非模型数据（主要为临时张量），在训练过程中保证当前可分配显存不小于历史最大非模型数据，即可确定各GPU上的卸载比例。

受Megatron框架限制，Megatron-Patch对这一feature采取简化实现。我们首先实现MemStatsCollector来统计追踪当前显存占用以及历史非模型数据的大小。在模型训练过程中，一次迭代可分为前向-后向以及参数更新两部分，分别具有不同的非模型数据大小，通常前向-反向的非模型数据会多于参数更新步。因此在每次调用optimizer.step()时，我们都调用一次optimizer.update_layout()，根据当前显存使用情况更新各个chunk的设备。如下图，在MCore中，模型数据包括不可移动的Fixed MD（即DDP内置的参数/梯度buffer）以及可移动的Chunk（C1～5），当观察到最大非模型数据加上已使用的模型数据大于安全阈值时，就讲显存上的部分Chunk移动到内存中更新，反之就移动更多Chunk到CUDA设备上。我们采用贪心策略，尽可能让GPU拥有更多的Chunk。

但在实际训练中，由于显存碎片的存在，第二个迭代的非模型数据往往大于第一个迭代的值，但小于第一次迭代前向-反向步及参数更新步的值之和。在参数更新结束后，我们采取warmup，当本次迭代非模型数据增加时，使用当前迭代各步骤的非模型数据之和作为下一个迭代非模型数据的估计值。

吞吐&收敛实验

相对于其他offloading技术，Optimizer CPU offloading能在少量性能损耗的情况下节省显存占用。在本节中，我们首先测试了在4K上下文长度且DP=8时LLaMA 3.1 8B的吞吐受optimizer CPU offload配置变化的影响情况，可以看到，原本无法在A100上训练的LLaMA3.1-8B的设置，通过offloading能实现运行。在static策略下，通过调整offload比例到65%，模型训练吞吐能达到188.0，同时显存仅占用约77GiB。通过采用Auto策略，Megatron-Patch能进一步分配显存占用，获得约4.5%的吞吐提升。这一数据，相对于目前我们使用A100测得的LLaMA3.1-8B最佳吞吐212.1相比，仅损失了7.5%。

为了探索Optimizer CPU offloading技术在实际场景中的可用性，我们对LLaMA 3.1 70B展开实验。LLaMA3.1的一个重要更新是将支持的上下文长度从8K拓展到128K。这一变化显著增加了通常情况下用户的微调成本。为了解决这一超长上下文场景中的训练问题，采用Activation Checkpointing以及optimizer offloading技术是十分必要的。因此，我们还测试了在4机32卡环境内使用长上下文训练LLaMA3.1-70B的吞吐及显存占用情况，实验结果如表所示，主要结论如下：

1. 对于16K及以上4机32卡无法训练的上下文配置，采用Activation Checkpointing以及optimizer offloading技术使训练成为可能
2. 相对于Activation Checkpointing，使用optimizer offloading损耗的性能更少（16K序列长度， 376 vs 289）
3. 由于两种降显存技术正交，同时使用两个技术能将4机32卡上可训练的上下文长度从32K进一步提升至64K

另外我们也进行了Qwen2-57B-A14B使用Optimizer CPU Offloading的吞吐实验。训练Qwen2-57B-A14B至少需要两台机器才能拉起，且达到最优吞吐的模型并行配置是tp2/pp2/cp1/ep4。从下表的实验数据可以看出关闭Offloading的局限性在于只能处理2048的序列长度，而打开Offloading可以处理4096的序列长度。

下图展示了打开&关闭Optimizer CPU Offloading的从头预训练收敛实验，可以看出收敛曲线基本是重合的，这进一步验证了我们自研的Optimizer CPU Offloading的可靠性。

Activation/WeightCPUOffloading

技术原理

相对于优化器卸载技术，直接在前向-反向过程中对不活跃的中间激活及权重进行卸载是一种更方便的手段。Transformer Engine 1.3实现了这一技术，用户通过控制Megatron中的--cpu-offloading开关即可方便地调用这一功能。

在Transformer Engine中，这一功能通过context实现，当用户开启offloading时，初始化自动生成一个offload上下文，这个上下文会在进入时注册两个hook：on_save_for_backward 以及 on_get_saved_tensor，用于异步拷贝待卸载的weights/activation，这两个hook记录了当前需要进行offload的tensor，随后在forward过程中，通过调用group_prefetch_offload_commit_async，最终触发当前layer的异步offload。

吞吐实验

我们在LLaMA3.1-8B上使用A100机器进行了测试，结果如下

结论：

1. 对于activation offloading技术，尽管增加offload层数能降低大量显存占用，但增加的通信量对于吞吐的影响较大，建议仅启用1～2层。
2. 对于activation recomputation技术，设置每隔几层重新计算对于最终占用及吞吐的影响均不大。
3. 在最优参数设置下activation recomputation相对于Optimizer Offloading损失了约12%的性能，而Activation Offloading则损失了近40%。

FlashAttention-3

技术原理

在长序列场景下Tansformer模型的复杂度主要来自于计算Attention的开销。FlashAttention（及后续的FlashAttention-2）开创了一种通过最小化内存读写来降低存储访问开销（Memory Acceess Cost）进而加速GPU上注意力计算的方法。这项技术促进了大模型上下文长度的大幅增长。虽然在A100上，FlashAttention-2最大能达到230TFLOPs/s（最大FLOPs利用率约等于73%），比标准的self-attention快3至10倍，但是在H100 GPU上的最大FLOPs利用率仅为35%。FlashAttention-3介绍了在NVIDIA前沿GPU架构上加速注意力计算的三种主要技术：（1）通过warp-specialization来overlap计算与数据移动（2）交错块级矩阵乘法和softmax操作（3）利用硬件对FP8低精度格式的支持进行非一致性处理。

FlashAttention-3带来的显著性能提升主要有以下几个方面[2]：

1. 更高效的GPU利用率：可将H100上最大FLOPs利用率从原先的35%提升至75%，最高达到740 TFLOPS。这使得当采用FP16计算精度时它比FlashAttention-2快1.5到2.0倍。如果是采用FP8计算精度，FlashAttention-3接近达到1.2PFLOPS的性能。
2. 低精度下的高性能计算：FlashAttention-3能够在保持准确性的前提下进行更低精度的FP8训练。这不仅允许更快的数据处理，还能减少内存使用，从而为运行大模型节约成本节约和提升效率。
3. 在LLMs中使用更长的上下文信息：通过加速注意力机制，FlashAttention-3使得大模型能更高效地处理更长的文本片段。这可能使得应用能够理解并生成更长、更复杂的内容，而不会减慢处理速度。

我们先简单回顾下经典的FlashAttention算法。它通过重新排序注意力计算过程，并利用平铺（tiling）及重新计算（recomputation）技术，极大地加速了计算过程并减少了内存使用，将原本随序列长度呈二次平方增长的内存需求降低至线性。通过利用平铺技术，将输入数据块从HBM（GPU内存）加载到SRAM（快速缓存）中，针对该数据块执行注意力计算，并将输出结果更新回HBM中。通过不将庞大的中间注意力矩阵写入HBM，从而减少了内存读写量，这带来了2到4倍的运行加速。下图展示的是FlashAttention前向传播过程：通过采用平铺和softmax重缩放技术，FlashAttention以数据块为单位进行操作，避免了从HBM读写数据的需要，同时在没有任何近似的情况下得到了正确的输出结果。

NVIDIA前沿架构引入了三种新特性分别是WGMMA (Warpgroup Matrix Multiply-Accumulate), TMA (Tensor Memory Accelerator) 和低精度的FP8支持。FlashAttention-3通过重写FlashAttention内核以利用这些新特性来显著提升吞吐速度。新特性中的WGMMA和TMA的异步性质为QKV矩阵运算和Softmax的非矩阵运算的Overlapping提供了算法性能提升机会。因为GPU上的非矩阵乘法操作比矩阵乘法操作要慢得多，所以我们希望矩阵乘法和softmax能够并行执行。当TensorCore忙于矩阵乘法时Multi-Function单元应该同时计算指数，这样可以最大化利用GPU资源，避免因非矩阵运算的瓶颈而造成性能损失。

FlashAttention-3采用一种叫做Intra-warpgroup的overlapping技术。这种pipeline技术通过增加寄存器压力的代价，将FP16注意力前向传播的吞吐量从大约620 TFLOPS提升到了大约640-660 TFLOPS。这种技术需要更多的寄存器来同时存储GEMM（通用矩阵乘法）的两个累加器以及softmax的操作数和结果。

吞吐&收敛实验

我们在单机八卡H100上利用llama3.1-8B模型对FlashAttention-3以及FP8进行性能测试。测试过程中选择的seqlen是4096，并固定TP（tensor model parallel size）为2，PP(pipeline model parallel size)为1，DP(data parallel size)为4。

同时我们也在NVIDIA前沿GPU卡上对FP8以及FlashAttention-3的收敛稳定性进行了评测。下图中绿线表示FP8和FlashAttention-3开关同时打开，蓝线表示仅开启FlashAttention-3开关，红线表示仅开启FlashAttention-2开关当作基准Baseline。

结论：

1. FlashAttention-3相比Fused Attention的优势是GPU利用率高些，同时消耗的显存更少。
2. 叠加FP8开关能带来GPU利用率的显著提升，FP8的hybird格式比e4m3的利用率稍微高些。
3. 将MBS提升到最大值3以后，显存占用也接近极限的80G，这时GPU利用率可以进一步提升到72.7%。
4. FlashAttention-3和FlashAttention-2的收敛曲线基本是重合的，但是开启FP8开关后导致训练精度损失。

Communication Overlapping

技术原理

在LLM训练中，为了提升吞吐，将通信与计算重叠十分重要。目前在Megatron中，Distributed Optimizer将优化器状态和高精度主参数分布在多个GPU上，采取增加通信的方式来显著降低显存占用。为了减少额外通信带来的性能损失，Distributed Optimizer同时也实现了多种不同的通信计算重叠技术[3]。

1. 数据并行（Data-Parallel, DP）中通信与计算的重叠

在使用Distributed Optimizer进行分布式训练时，各DP Rank间需要进行梯度同步，引入了额外的梯度的reduce-scatter和更新参数的all-gather通信。在Megatron中，通过设置overlap_grad_sync=true和overlap_param_sync=true启用重叠技术后，这些DP通信会按单个TransformerLayer或单个Virtual Pipeline的粒度进行分块并与计算进行重叠，使通信仅占用一个块的时间。

2. 张量并行（Tensor-Parallel，TP）中通信与计算的重叠

如上图，当使用序列并行激活切分（sequence-parallel activation sharding）时，张量并行需要引入额外的Reduce-Scatter（图中绿色RS部分）以及All-Gather（图中绿色AG部分）通信。对于那些无计算依赖的TP通信（即黄色框内的部分），Megatron默认采用批式方法对其进行重叠，而对于有计算依赖的其它TP通信（红色框内的线性层与TP通信对），通信与计算会被分块，然后以流水线方式进行重叠。在这一过程中，张量的All-Gather会采用多步的输入P2P环交换代替，而Reduce-Scatter则采用多步的GEMM输出P2P环交换以及对输出的reduction代替。

在Megatron中，流式TP通信重叠通过TE后端实现，并通过设置ub_tp_comm_overlap=true启用。

3. 管道并行（Pipeline-Parallel，PP）中通信与计算的重叠

管道并行引入了各PP Rank间GPU的P2P激活及梯度的通信，随着虚拟管道并行大小的增加，由于每个micro batch中执行的Transformer层数减少，PP通信频率和开销也随之增加，这可能会抵消虚拟管道化减少的bubble带来的吞吐提升。在在1F1B阶段（即管道化的主体部分，其中前向和后向微批次执行交错），Megatron默认启用当前无数据依赖的通信与计算重叠。

吞吐实验

为了说明Communication Overlapping的优势，我们在不同Overlap情况下对LLaMA3.1-70B的训练吞吐进行测量。实验配置采用4机32卡A100

从表中可以明显看到：

1. 单独开启TP Comm Overlap能为LLaMA3.1-70B的训练带来约10%的提升
2. 单独开启Grad Reduce以及Param Gather的overlap效果不明显，但通过结合TP Comm Overlap，能将训练吞吐进一步提升至201.3TFLOP/s/GPU，相对于基线有17%的提升

LLM训练加速应用指南

预训练&微调命令统一描述

ENV=$1                          # 运行环境配置开关: dsw单机训练训练，dlc表示多机训练环境
MODEL_SIZE=$2                   # 模型结构参数量级: 8B, 70B
BATCH_SIZE=$3                   # 一次迭代一个数据并行内的样本数
GLOBAL_BATCH_SIZE=$4            # 一次迭代多个数据并行的总样本数
LR=$5                           # 学习率
MIN_LR=$6                       # 最小学习率
SEQ_LEN=$7                      # 序列长度
PAD_LEN=$8                      # Padding长度
PR=${9}                          # 训练精度: fp16, bf16, fp8
TP=${10}                        # 模型并行度
PP=${11}                        # 流水并行度
CP=${12}                        # 上下文并行度
SP=${13}                        # 是否使用序列并行: true, false
DO=${14}                        # 是否使用Megatron版Zero-1降显存优化器: true, false
FL=${15}                        # 是否优先使用Flash Attention: true, false
SFT=${16}                       # 是否执行微调训练: true, false
AC=${17}                        # 激活检查点模式: sel, full, offload, none
OPTIMIZER_OFFLOAD=${18}          # 是否启用Offload optimizer: false, static, auto
SAVE_INTERVAL=${19}             # 保存ckpt的间隔
DATASET_PATH=${20}              # 训练数据集路径
VALID_DATASET_PATH=${21}        # 验证数据集路径
PRETRAIN_CHECKPOINT_PATH=${22}  # 预训练模型路径
TRAIN_TOKENS_OR_ITERS=${23}     # 训练TOKEN或者Iter数
WARMUP_TOKENS_OR_ITERS=${24}    # 预热TOKEN或者Iter数        
OUTPUT_BASEPATH=${25}           # 训练输出日志文件路径

预训练&微调示例

如下展示的对llama3.1的8B量级模型进行TP4PP2切分后的继续预训练示例

cd /workspace/Pai-Megatron-Patch/examples/llama3_1
sh run_mcore_llama3_1.sh  \
dsw  \
8B   \
1    \
8 \
1e-5   \
1e-6   \
128  \
128  \
bf16  \
4   \
2  \
1 \
true \
true   \
true \
false \
false   \
false \
100000  \
/mnt/llama3-datasets/wudao_llama3bpe_content_document   \
/mnt/llama3-datasets/wudao_llama3bpe_content_document   \
/mnt/llama3-ckpts/Meta-Llama-3.1-8B/mcore-tp4-pp2  \
10000  \
100   \
/workspace/output_mcore_llama3_1

通过设置MP_DATASET_TYPE环境变量实现使用原始json格式的数据集进行指令微调

export MP_DATASET_TYPE="raw"
cd /workspace/Pai-Megatron-Patch/examples/llama3_1
sh run_mcore_llama3_1.sh  \
dsw  \
8B   \
1    \
8 \
1e-5   \
1e-6   \
128  \
128  \
bf16  \
4   \
2  \
1 \
true \
true   \
true \
true \
false   \
false \
100000  \
/mnt/llama3-datasets/alpaca_zh-llama3-train.json   \
/mnt/llama3-datasets/alpaca_zh-llama3-valid.json   \
/mnt/llama3-ckpts/Meta-Llama-3.1-8B/mcore-tp4-pp2  \
10000  \
100   \
/workspace/output_mcore_llama3_1

总结

在基于Megatron-Core的LLama3.1大模型最佳实践开发过程中，我们围绕大模型训练测试了以下核心技术的性能：

• 自研带CPU Offloading功能的分布式优化器的鲁棒性。
• 基于Megatron-Core的多重训练加速技术的可靠性。
• 运用综合加速技术来训练LLama3.1过程中的易用性和稳定性。

后续在Pai-Megatron-Patch中还会陆续放出更多高质量的大模型最佳实践以及最新的训练加速技术应用示例，敬请期待。

参考文献

1. https://github.com/NVIDIA/Megatron-LM/blob/main/docs/llama_mistral.md
2. https://www.together.ai/blog/flashattention-3
3. https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/features/optimizations/communication_overlap.html
4. Reducing Activation Recomputation in Large Transformer Models
5. ZeRO-Offload: Democratizing Billion-Scale Model Training
6. https://github.com/Tencent/PatrickStar