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