NVIDIA/Megatron-LM

解决的问题

Megatron-LM 提供了一个 GPU 优化的框架，可在数千个 GPU 上训练大规模 Transformer 模型（参数量从 20 亿到 4620 亿），通过实现高级分布式训练策略，解决了大规模语言模型训练中的内存和计算瓶颈问题。

如何工作

该项目分为两个主要组件：

• Megatron Core：一个可组合的 GPU 优化构建模块库。提供基本的 Transformer 组件，并支持多种并行策略，包括张量并行（TP）、流水线并行（PP）、数据并行（DP）、专家并行（EP）和上下文并行（CP）。还支持混合精度训练（FP16、BF16、FP8 和 FP4）。
• Megatron-LM：一个参考实现，将 Megatron Core 与预配置的训练脚本结合，支持快速实验和研究。

适用人群

• 机器学习工程师：构建大规模模型自定义训练流水线的人。
• 研究团队：从事分布式训练和 Transformer 架构研究的团队。
• 框架开发者：使用优化的 GPU 内核开发自有训练框架的开发者。

主要亮点

• 极致可扩展性：可在最多 6,144 个 H100 GPU 集群上训练高达 4620 亿参数的模型。
• 细粒度并行：全面支持 TP、PP、DP、EP 和 CP，以优化 GPU 利用率。
• 高效率：在 H100 集群上实现高达 47% 的模型 FLOP 利用率（MFU）。
• 跨格式支持：包含 Megatron Bridge，支持 Hugging Face 与 Megatron 格式之间的双向检查点转换。
• 现代架构支持：初步支持 DeepSeek-V4 和 Falcon-H1 混合 Transformer-Mamba 架构。