阿里云 PAI- 灵骏大模型训练工具 Pai-Megatron-Patch 正式开源！

作者： 李鹏，王明，施晨，黄俊

导读

 随着深度学习大语言模型的不断发展，其模型结构和量级在快速演化，依托大模型技术的应用更是层出不穷。对于广大开发者来说不仅要考虑如何在复杂多变的场景下有效的将大模型消耗的算力发挥出来，还要应对大模型的持续迭代。开发简单易用的大模型训练工具就成了应对以上问题广受关注的技术方向，让开发者专注于大模型解决方案的开发，降低大模型训练加速性能优化和训练/推理全流程搭建的人力开发成本。阿里云机器学习平台 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 模型打开如下开关即可。

  --swiglu     --use-rotary-position-embeddings     --no-position-embedding     --untie-embeddings-and-output-weights     --disable-bias-linear