Megatron-LM

NVIDIA 开源的大规模 Transformer 训练参考实现，以及可组合的 Megatron Core 库，面向多 GPU / 多机。

它是做什么的

以官方仓库说明为准，包含两部分：

• Megatron-LM：Megatron Core + 预配置训练脚本与示例，适合研究团队、学分布式、快速实验
• Megatron Core：GPU 优化的 Transformer 积木（并行、混合精度、模型结构等），给框架开发者与自建训练管线用

目标是在大量 GPU 上高效训练从数十亿参数到更大的 Transformer。文档中的并行指南说明：多种并行维度可以组合，覆盖从数十亿到万亿参数量级的扩展思路——具体能训多大取决于硬件、序列长度与配置，勿把文档里的量级描述当成你机器上的保证。相对「只会用 Trainer 开 ZeRO」的路径，Megatron 更偏底层：你要理解 TP/PP 等如何切层、切矩阵，以及数据预处理成二进制索引格式。适合当参考实现与 Core 库，而不是第一堂微调课。

要点

• 并行可组合：Data Parallel（DP）、Tensor Parallel（TP）、Pipeline Parallel（PP）、Context Parallel（CP，长序列）、Expert Parallel（EP，MoE）；另有 Megatron-FSDP 等切分策略
• 混合精度：文档覆盖 FP16 / BF16 / FP8 / FP4 等路径——按 GPU 代际与当前文档选用，不要凭印象硬开
• 数据管线：自有语料需预处理成 .bin / .idx；tools/preprocess_data.py 支持 JSONL（text 字段）与多种 tokenizer
• 仓库结构：megatron/core（内核与并行）、training 脚本、以及推理/导出、后训练、RL 等相关目录（以当前 main 为准）
• 入门路径：官方 Quickstart 提供 mock 数据的最小训练环，以及 LLaMA-3 示例脚本

适合谁

• 做或读大规模分布式训练代码，需要 TP/PP/CP/EP 组合
• 已有单机微调与基础 DDP 经验，要把「参考实现」落到多机集群
• 框架/平台开发者，想嵌入 Megatron Core 积木而非只跑脚本

官方入口

• GitHub: NVIDIA/Megatron-LM
• Megatron Core 文档
• Your First Training Run
• Parallelism Strategies

建议怎么用

先完成官方安装文档中的环境准备，再按 Quickstart 跑最小环：

torchrun --nproc_per_node=2 examples/run_simple_mcore_train_loop.py

跑通后再读 Parallelism Guide：按层宽/激活显存选 TP、按深度选 PP、按超长上下文选 CP、MoE 再开 EP；各维度乘积需与 world_size 对齐。自有数据：JSONL（text）→ preprocess_data.py 得到 .bin/.idx → 再挂训练脚本。不要把它当「第一堂微调课」；先有单机训练与 DDP 基础，再上 Megatron。