开源大模型训练框架Megatron-LLaMA来了 32卡训练可加速176%

要点:

淘天集团联合爱橙科技正式对外开源大模型训练框架Megatron-LLaMA，以提高大语言模型训练性能，降低训练成本。

测试显示，相比HuggingFace代码，Megatron-LLaMA在32卡训练上获得176%的加速;在大规模训练上表现出几乎线性的扩展性。

Megatron-LLaMA已在GitHub开源，将持续关注社区发展，与开发者共同推进自适应配置、更多模型支持等方向。

9月12日，淘天集团联合爱橙科技正式对外开源大模型训练框架Megatron-LLaMA，以让技术开发者们更方便地提升大语言模型训练性能，降低训练成本。

据悉，Megatron-LLaMA是一个基于Megatron-LM的分布式深度学习训练框架，专门用于大规模语言模型LLaMA的训练。LLaMA已经成为开源社区中最杰出的大规模语言模型之一，它集成了BPE词元化、预规范化、旋转嵌入、SwiGLU激活函数、RMSNorm和非绑定嵌入等多项优化技术，在客观和主观评价中都展现出卓越的结果。

项目地址:https://github.com/alibaba/Megatron-LLaMA

LLaMA开发了7B、13B、30B和65B/70B多个模型规模的版本。在开源社区中，也出现了许多基于LLaMA的成功变体，无论是通过连续训练/监督微调还是从零开始训练，都进一步证明了LLaMA在长上下文理解、长上下文生成、代码编写、数学问题求解、工具使用等任务上的卓越能力。

然而，由于大规模语言模型的训练或微调都需要强大的计算资源，开发者自己尝试在LLaMA上实现新的设计往往非常昂贵。Megatron-LM是一个集成了张量并行、流水线并行和序列并行的分布式训练解决方案，可以在训练具有数百亿参数的模型时充分利用硬件资源，使资源利用率远高于基于Huggingface和DeepSpeed实现的公开版LLaMA。但是，原生的Megatron-LM在极大规模下会遭受分布式优化器的通信瓶颈。

测试显示，相比HuggingFace直接获得的代码，Megatron-LLaMA在32卡训练上获得176%的加速效果。在大规模训练上，Megatron-LLaMA相对原生Megatron-LM有几乎线性的扩展性，且对网络稳定性表现出高容忍度。

Megatron-LLaMA改进了DistributedOptimizer的梯度聚合机制，实现梯度通信与计算的并行，从而优化了反向传播流程。

当前Megatron-LLaMA已在GitHub上开源，淘天集团和爱橙科技将共同维护，并积极关注社区发展，与开发者们在自适应配置选择、模型支持等方面进行合作，以推动Megatron和LLaMA生态建设。

Megatron-LLaMA的主要特点如下:

在Megatron-LM中实现了标准的LLaMA，可以按需配置其中的优化技术。未来还将支持Alibi和FlashAttention2等特性。

通过实现高度重叠的通信和计算，提升了通信计算并行性，与DeepSpeed ZeRO Stage2类似，大大减少了GPU内存占用，提高了吞吐量。

提供了分布式检查点保存/恢复等实用工具，加速了训练过程，支持与HDFS等分布式文件系统集成;支持与HuggingFace transformers库的tokenizer集成，方便下游任务迁移。

Megatron-LLaMA使LLaMA模型的大规模训练变得快速、经济高效且具备可扩展性。它降低了训练大模型的门槛，是开源社区一个重要的贡献。