Megatron-LM 源码精读

从入口到并行策略，逐文件解读 NVIDIA Megatron-LM 的核心实现

入口
pretrain_gpt.py
pretrain_gpt.py 源码精读
用户代码入口：get_batch、loss_func、forward_step、model_provider 四大回调的定义与数据流全景
推荐首读 数据流
基础设施 · training.py
training.py · 初始化篇
initialize_megatron() → get_model() → setup_model_and_optimizer()
分布式环境初始化、SPMD 模型构建（PP/VPP）、DDP/FSDP 包装、优化器创建、Checkpoint 加载
SPMD DDP VPP
training.py · 训练循环篇
train() → train_step() → forward_backward_func()
训练主循环、前向反向策略选择、梯度累积、评估与检查点、CUDA Graph 优化
train_step forward_backward
显存计算器
GPU 显存计算逻辑详解
静态显存（权重/梯度/优化器状态/FP8缓存）与动态显存（Activation）的逐模块计算公式，PP Layout/VPP 影响分析
显存 Activation MoE
并行策略
parallel_state.py
Rank 与并行组
GPU 编号如何映射到 TP/PP/DP/CP/EP 五维并行网格、RankGenerator 通信组生成算法
TP PP DP CP / EP
集合通信操作详解
AllReduce / AllGather / ReduceScatter / P2P / AllToAll
六种集合通信原语在 Megatron 中的具体用法、分布式优化器的通信策略
AllReduce AllToAll 分布式优化器
tensor_parallel/layers.py · mappings.py
Tensor Parallelism 实现详解
ColumnParallelLinear / RowParallelLinear / f-g 共轭算子、Sequence Parallelism、词表并行
ColumnParallel RowParallel SP
schedules.py · p2p_communication.py
Pipeline Parallelism 调度机制
1F1B warmup/steady/cooldown 三阶段、Interleaved VPP bubble 压缩、P2P 通信协议与奇偶交替死锁规避、Activation Checkpointing 配合
1F1B VPP P2P
schedules.py · p2p_communication.py
Interleaved Pipeline Schedule 深度源码精读
forward_backward_pipelining_with_interleaving 逐行解析：vid / schedule_table、三个共享缓冲区、P2P 四层架构、Warmup / Steady / Cooldown 全流程追踪与数据流全景图
深度精读 virtual_microbatch_id P2P overlap buffer 追踪
优化器
distrib_optimizer.py
分布式优化器（ZeRO-1）源码精读
ZeRO-1 风格优化器状态分片、混合精度主参数、Reduce-Scatter + All-Gather 通信策略与 Overlap 异步优化
ZeRO-1 Reduce-Scatter 混合精度
模型架构
gpt_model.py · transformer_layer.py
GPT 模型架构
Embedding → TransformerBlock → Output Layer 三大组件、Layer Spec 声明式定义、SwiGLU MLP、GQA Attention
GPTModel TransformerLayer SwiGLU
moe_layer.py · router.py · token_dispatcher.py · experts.py
MoE 层与 Expert Parallelism
TopKRouter 路由机制、负载均衡策略（aux_loss/sinkhorn/expert_bias）、AlltoAll Token Dispatcher、GroupedGEMM Expert 计算、Shared Expert Overlap
MoE Expert Parallelism AlltoAll DeepSeek-V3