Chapter 04
NeMo Megatron 与 Megatron-Core 的关系
NeMo 系列里最容易让人迷糊的是"NeMo / NeMo Megatron / Megatron-LM / Megatron-Core"这一串名字。本章按时间线和分层架构把关系讲清楚。
4.1三层栈
flowchart TB
App[NeMo collections.llm
用户写的 recipe] Mid[NeMo lightning
PT-Lightning + 模型抽象] Core[Megatron-Core
TP/PP/DP 通信原语 + Optimizer] TE[TransformerEngine
FP8 / Flash kernel] PT[PyTorch + CUDA] App --> Mid --> Core --> TE --> PT
用户写的 recipe] Mid[NeMo lightning
PT-Lightning + 模型抽象] Core[Megatron-Core
TP/PP/DP 通信原语 + Optimizer] TE[TransformerEngine
FP8 / Flash kernel] PT[PyTorch + CUDA] App --> Mid --> Core --> TE --> PT
4.2各层职责
| 层 | 仓库 | 角色 |
|---|---|---|
| NeMo | NVIDIA-NeMo/NeMo | 顶层框架：recipe、CLI、aligner、ASR/TTS 等多模态 |
| Megatron-Core (MCore) | NVIDIA/Megatron-LM | 分布式底座：ParallelLinear / DistributedDataParallel / DistOptim / PP schedule |
| TransformerEngine | NVIDIA/TransformerEngine | kernel 库：FP8 GEMM / fused attention / RMSNorm |
4.3"NeMo Megatron" 这个名字
历史名词：早期 NeMo 内部叫 "NeMo Megatron Collection"，指的是 NeMo 里用 Megatron 风格做 LLM 训练的子模块（nemo.collections.nlp.models.language_modeling.megatron_gpt_model 这类）。2024+ NeMo 2.x 后改名 nemo.collections.llm，更直白。
所以你可能在不同文档里看到这几个名字：
- NeMo Megatron = NeMo 里跑 Megatron 风格 LLM 训练的子模块（旧名）；
- Megatron-LM = NVIDIA 的原版分布式训练仓库；
- Megatron-Core = Megatron-LM 提取出来的、能被 NeMo / 其他框架引用的 Python 包；
- NeMo collections.llm = 2.x 时代的 NeMo Megatron。
4.4NeMo 给 Megatron-Core 增加了什么
- PyTorch Lightning 集成：Lightning 的 Trainer / Callback 生态；
- recipe 范式：把 Megatron 命令行 hundred-flag 风格收敛成 Python 函数；
- Hydra 配置：YAML + override 的另一条路；
- multi-modal collections：ASR / TTS / Vision；
- checkpoint 通用层：NeMo ↔ HF 双向转换；
- Aligner：SFT/DPO/PPO/SteerLM 整体 RLHF 工具集；
- NIM / Triton 部署集成。
4.5选用建议
| 需求 | 用什么 |
|---|---|
| 极致性能 + 不要 PT-Lightning | Megatron-LM 原版 |
| 同等性能 + PT-Lightning + 高层 API | NeMo |
| 只关心 Megatron-Core 的并行原语 | 直接 import megatron-core |
| 多模态 (ASR / TTS / Vision) | NeMo（独家） |
| RLHF 流程 | NeMo Aligner（独家） |
| 转换 HF checkpoint | NeMo（独家） |
4.6性能对照
Llama-3 70B、64×H100、bf16、seqlen=8K：
| 栈 | MFU | tokens/s/GPU |
|---|---|---|
| Megatron-LM 原版 | ~55% | ~3.6k |
| NeMo 2.x（底层走 MCore） | ~54% | ~3.5k |
| NeMo + TE FP8 | ~63% | ~4.2k |
结论：NeMo 几乎没有性能损失，开了 FP8 还能更快。
4.7这章你需要带走的
- 三层栈：NeMo → Megatron-Core → TransformerEngine；
- "NeMo Megatron" = NeMo 里跑 Megatron-Core 的子模块（2.x 改名 collections.llm）；
- NeMo 不重新发明分布式，全用 MCore；
- NeMo 独有：Lightning + recipe + 多模态 + Aligner + checkpoint 转换。