作者： 李鹏，王明，施晨，黄俊
导读
随着深度学习大语言模型的不断发展，其模型结构和量级在快速演化，依托大模型技术的应用更是层出不穷。对于广大开发者来说不仅要考虑如何在复杂多变的场景下有效的将大模型消耗的算力发挥出来，还要应对大模型的持续迭代。开发简单易用的大模型训练工具就成了应对以上问题广受关注的技术方向，让开发者专注于大模型解决方案的开发，降低大模型训练加速性能优化和训练/推理全流程搭建的人力开发成本。阿里云机器学习平台 PAI 开源了业内较早投入业务应用的大模型训练工具 Pai-Megatron-Patch，本文将详解 Pai-Megatron-Patch 的设计原理和应用。
Pai-Megatron-Patch 是什么
Pai-Megatron-Patch 工具是阿里云机器学习平台 PAI 算法团队研发，基于阿里云智算服务 PAI-灵骏平台的大模型最佳实践解决方案配套工具，旨在帮助大模型开发者快速上手灵骏产品，完成大语言模型（LLM）的高效分布式训练，有监督指令微调，模型离线推理验证等完整大模型开发链路。该项目提供了业界主流开源大模型基于 Megatron-LM 的训练 &离线推理验证流程，方便用户快速上手大模型训练。
主要特性
多款热门大模型支持：llama，llama-2，codellama, 百川，通义千问，Falcon，GLM，Starcoder，Bloom，chatglm 等
支持模型权重互转转换：在 Huggingface，Megatron 和 Transformer Engine 之间进行算子命名空间映射
支持 Flash Attention 2.0 和 Transformer Engine 模式下的 FP8 训练加速且确保收敛
丰富且简单易用的使用示例，支持大模型预训练，微调，评估和推理，强化学习全流程最佳实践
开源地址
https://github.com/alibaba/Pai-Megatron-Patch
技术架构
Pai-Megatron-Patch 的设计理念是不对 Megatron-LM 的源码进行侵入式修改，即不在 Megatron-LM 里面添加新的功能特性，将需要扩充完善的部分以 patch 补丁的方式呈现。在 patch 中构建 LLM 训练链路通过依赖 Megatron-LM 核心库的方法实现和 Megatron-LM 的解耦合。这样解耦合的好处就是 Megatron-LM 的升级不会影响用户的 LLM 最佳实践体验。
Pai-Megatron-Patch 中包含模型库，分词器，模型转换，强化学习，离线文本生成以及使用示例和工具集等用于构建 LLM 训练的关键要素。在模型库中包含热门大模型的 Megatron 版本实现，例如 baichuan，bloom，chatglm，falcon，galactica，glm，llama，qwen 和 starcoder，后续还会根据需要及时添加新的 Megatron 版大模型实现。同时 patch 还提供了 huggingface 模型权重和 Megatron 模型权重之间的双向转换。一方面是方便用户加载 huggingface 的权重在 Megatron 中继续预训练或者微调，另一方面是方便用户对训练好的 Megatron 模型使用 huggingface 的评估/推理流程对模型质量进行客观评估。在强化学习部分，patch 提供了 PPO 训练流程等，方便用户使用 SFT 模型和 RM 模型进行强化学习。最后 patch 提供了大量的使用示例帮助用户快速开始大模型训练 &离线推理。具体请参考阿里云灵骏产品的使用流程: 智算服务 PAI 灵骏大模型分布式训练方案。
关键技术
1. 模型权重转换
研发 Megatron-Patch 的初衷之一就是能将世界各地研发机构在 Huggingface 上放出的热门大模型使用 Megatron 引擎进行继续预训练或者继续微调。这就需要首先将 Huggingface 模型格式的 ckpt 转换成 Megatron 模型格式，才能正确加载进来，否则会出现 pytorch 加载模型失败。Megatron-Patch 的一个核心可靠性保障特征就是在采用算子拆分，流水并行，序列并行，Zero 显存优化，BF16 混合精度，梯度检查点等训练加速技术确保模型训练吞吐速度平均提升 1.5 倍以上的同时，在评估任务模式下的单一样本前向 loss 值，预训练/微调任务模式下的 loss 曲线，离线文本生成任务模式下的生成效果这三个方面和 Huggingface 是对齐的，从而确保 Megatron 版模型的可靠性。
另一方面，Megatron 版的 transformer 实现方式提供了一种让用户仅仅通过设置开关就能实现不同种类 GPT 模式的能力。比如 llama 模型打开如下开关即可。
如果想将 llama 模式变成 baichuan 模型，那么仅仅需要添加采用--use-alibi-mask 开关，同时关闭 Rotary Embeeding 开关即可，具体配置如下所示：
下面我们以 llama-2 为例，详解从 huggingface 到 megatron 的模型权重转换技术。下表总结了两者在不同 module 上的命名对应关系。在 patch 实现过程中，我们首先将 HF 格式的 ckpt 转换到一种内部格式，然后再把这种内部格式转换成对应的外部格式。这样做可以最大程度复用已有的转换逻辑来处理新模型。在转换为内部格式的过程中，q_proj, k_proj, v_proj 需要沿着第 0 维拼接在一起后赋值给内部变量 query_key_value。
当用户在资源受限情况下需要按照 TP>1 来拆分权重的时候，这里需要注意的是针对 MLP 层的 gate_proj 和 up_proj 的操作。不能像 qkv 那样在转换成内部格式的时候进行 merge 再执行算子拆分。需要在拆分前加入如下针对 MLP 层的权重合并的代码逻辑才能确保正确收敛。
2. 基于 TE 的 FP8 训练收敛
Transformer Engine(TE)是一个在英伟达 GPUS 上运行的针对 Transformer 模型的加速库，其中包括针对 Hopper GPU 的 FP8 混合精度，该精度可以在较低的显存利用率下提供更好的训练 &推理速度。在 TE 内部封装了 Flash Attention 实现，同时 TE 还提供了一组高度优化后的算子用来构建 Transformer 模型。比如 LayerNormLinear 就是将 LayerNorm 和 QKV-Proojection 进行算子融合，LayerNormMLP 就是将 layernorm 和 mlp 进行算子融合。如下图所示：
从 Huggingface 到 TE 模型的权重转换技术和之前是类似的，也需要事先找到两者之间的映射关系。从下表可以看出，TE 中多了_extra_state 是用来存 fp8 训练的 scale 和 history 的，这些在加载的时候会出现冲突，这时只要将 load_state_dict 函数的 strict 设置成 False 就可以了，比如 load_state_dict(state_dict_, strict=False)。
在 Megatron-Patch 中使用示例中打开 FP8 混合精度训练开关也很容易，如下所示：
我们可以使用如下训练脚本 run_pretrain_megatron_llama_enwiki.sh 来测试打开 FP8 开关后的预训练收敛性。下图展示了 llama-7B 和 llama-2-70B 模型在打开和关闭 FP8 时的 loss 曲线对比，可以看出基本是重合的。
LLama-7B
LLama2-70B
3. 大模型训练 &推理
从 github 上获取 Megatron 模型训练工具 PAI-Megatron-Patch（https://github.com/alibaba/Pai-Megatron-Patch）源代码并拷贝到工作目录/mnt/workspace/下。
1）模型格式转换
使用我们提供的模型转换脚本，将 huggingface 格式的模型文件转换为 megatron 格式：
2）继续预训练
中文继续预训练汉化指引
Step1: 获取需要扩充词表的模型（如 llama-13b-hf）
Step2: 获取需要扩充的词表
使用 sentence-piece 代码库从自有文本语料中学习词表，得到 randeng-sp.model 文件
Step3: 词表扩充
扩充模型 tokenizer：将 randeng-sp.model 中的词表添加到 llama-13b-hf 文件夹下 tokenizer.model 中
扩充模型词表对应的参数矩阵
word_embedding、lm_head
新词向量可以使用原词向量均值作为初始化，比如“天气”=mean([“天”，“气”])
修改与词表大小相关的文件并保存，如 config.json
运行继续预训练脚本 run_pretrain_megatron_llama.sh，需要传入的参数列表如下：
注意设置正确的数据集挂载路径 WORK_DIR 以及运行环境 ENV，运行示例如下所示：
3）有监督微调
在微调开始之前，请先进入
获取 json 文件。运行 run_finetune_megatron_llama.sh 脚本，需要传入的参数列表如下：
多节点运行示例如下所示：
4）离线推理
模型训练完成后，可以进行离线推理，评估模型效果。根据上面的训练流程不同，我们提供了 Megatron 格式的推理链路。对于 Megatron 训练的模型，可以直接用 Megatron 框架进行推理。
此处提供一个离线推理输出的文件，推理的数据组织形式需要与微调时的保持一致。
测试样本：
https://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/
pai-megatron-patch/llama2-datasets/pred_input.jsonl
注意：
模型保存的路径下缺少 tokenizer 依赖的文件，需要将微调前模型路径下所有 json 文件及 tokenizer.model 拷贝至保存模型的路径下（位于{OUTPUT_BASEPATH }/checkpoint），与 latest_checkpointed_iteration.txt 同级。
以下有监督微调过程保存模型的推理代码，需要将 run_text_generation_megatron_llama.sh 脚本中 CUDA_VISIBLE_DEVICES 参数设置为 0；GPUS_PER_NODE 参数设置为 1；同时使用下列代码进行推理。此时使用单卡进行推理。注意：此处模型 tp 为 1，可使用单卡推理；如果 tp>1，则需使用相应卡数进行推理。
4. 大模型强化学习
一般来说，SFT 微调过的模型在对话场景已经会有不错的表现了。如果想进一步提升模型效果，可以再加上 RLHF 训练。包括奖励模型（Reward Model）的训练和强化学习（PPO）的训练。这里展示了如何使用当前最常用的 RLHF 开源代码框架，DeepSpeed-Chat 和 trlx，来进行奖励函数训练（RM），以及强化学习优化（PPO）。
1）模型格式转换
如果基于 huggingface 格式的模型直接进行奖励模型训练（RM）和强化学习优化（PPO），可以跳过此步骤。如果基于 Megatron 格式的模型，如 PAI-Megatron-Patch 训练好的 SFT 模型，进行 RM 和 PPO 训练，需要使用我们提供的模型转换脚本，先将 Megatron 格式的模型文件转换为 huggingface 格式。
LLaMA2 模型转换：
BLOOM 模型转换：
2）DeepSpeed-Chat
下载安装开源社区 DeepSpeed-Chat 源代码：
基于 LLaMA2 模型训练奖励模型（RM）：
基于 LLaMA2 进行强化学习优化训练（PPO）：
3）trlx
下载安装开源社区 trlx 源代码：
基于 BLOOM 模型训练奖励模型（RM）：
基于 GPT-J 模型训练奖励模型（RM）：
基于 BLOOM 模型进行强化学习优化训练（PPO）：
基于 GPT-J 模型进行强化学习优化训练（PPO）：
PPO 单测
如果您想跳过 有监督微调（SFT）与 奖励模型训练（RM）两个步骤，只单独测试 PPO 模块的性能，可以运行如下指令单测 PPO：
开源生态——构想和未来
在 PAI-Megatron-Patch 的开发过程中，我们围绕中文大模型训练加速落地沉淀了以下几个方面的内容：
Huggingface 的模型权重无损转换成 Megatron 或者 Transformer Engine 可读的模型权重。
H800 集群开启 FP8 混合精度训练确保收敛。
LLM 大模型在 PAI 灵骏智算平台上的最佳实践。
强化学习技术在 PAI 灵骏智算平台上的最佳实践。
后续在 PAI-Megatron-Patch 中还会陆续放出更多高质量的大模型和最佳实践。此外，在中长期，我们在 Megatron 版的 Lora 流程以及 Transformer Engine 方向上会持续投入精力，也欢迎各种维度的反馈和改进建议以及技术讨论，同时我们十分欢迎和期待对开源社区建设感兴趣的同行一起参与共建，钉钉群号是 29605038042。
参考文献
[1]. Attention Is All You Need
[2]. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
[3]. Reducing Activation Recomputation in Large Transformer Models
[4]. FP8 Formats for Deep Learning
[5]. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
[6]. LLaMA: Open and Efficient Foundation Language Models
[7]. Llama 2: Open Foundation and Fine-Tuned Chat Models
[8]. Benchmarking Large Language Models on NVIDIA H100 GPUs with CoreWeave
