第8章：训练框架实战(Megatron-LM、DeepSpeed)

📖 本章概述

理论知识需要落地到具体框架。本章深入 Megatron-LM 和 DeepSpeed 两大主流框架的工程实现，并讲解训练稳定性保障和 Checkpoint 策略——这些是实际大模型训练中最常踩坑的地方。

📑 章节结构

1. Megatron-LM 深度解读

代码架构总览：
- megatron/core/：核心并行原语（TP、PP、SP 的实现）
- megatron/core/tensor_parallel/：ColumnParallelLinear、RowParallelLinear
- megatron/core/pipeline_parallel/：Schedule（1F1B、Interleaved）
- megatron/core/distributed/：分布式通信组管理
TP 实现细节：
- 前向：Column Parallel 的 AllGather / Row Parallel 的 AllReduce
- 反向：自定义 autograd Function 中的通信操作
- SP 的 ReduceScatter/AllGather 插入位置
PP 实现细节：
- PipelineParallel Schedule 的状态机
- Stage 间通信：send/recv 激活值和梯度
- micro-batch 管理和 Bubble 调度
3D 并行初始化：
- initialize_model_parallel(tp_size, pp_size, dp_size) 构建通信组
- 进程 rank 到 (tp_rank, pp_rank, dp_rank) 的映射
训练 GPT 配置示例：完整的 pretrain_gpt.py 启动参数解析

2. DeepSpeed 深度解读

架构特点：以 JSON 配置驱动，用户代码改动最小化

ZeRO 配置使用：

{
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {"device": "cpu"},
    "overlap_comm": true,
    "contiguous_gradients": true
  }
}

DeepSpeed Config 关键字段：
- train_batch_size / train_micro_batch_size_per_gpu / gradient_accumulation_steps 三者关系
- fp16 / bf16 混合精度配置
- gradient_clipping 梯度裁剪
- activation_checkpointing 配置
与 HuggingFace Transformers 集成：Trainer + deepspeed 参数一键启用
DeepSpeed Chat / RLHF：PPO 训练中 Actor/Critic/Reference 多模型的 ZeRO 配置

3. 训练稳定性保障

Loss Spike 排查：
- 常见原因：数据质量问题（异常样本）、学习率过大、梯度爆炸、数值溢出
- 排查步骤：检查 grad norm 曲线 → 定位具体 batch → 检查数据
- 处理方式：跳过异常 batch、回滚到前一个 checkpoint
梯度裁剪（Gradient Clipping）：
- torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
- 分布式下需要跨 DP 组 AllReduce grad norm 后再裁剪
学习率 Warmup：
- 为什么需要：训练初期参数远离最优解，大学习率易导致发散
- 常用策略：线性 Warmup + Cosine Decay
- 典型配置：Warmup 步数 = 总步数的 1%-5%
数值稳定性：
- BF16 下 LayerNorm 的 FP32 累加
- Attention Score 的缩放（ $\frac{1}{\sqrt{d_k}}$ ）
- 避免 FP16 下的 exp 溢出

4. Checkpoint 策略

保存频率：
- 经验法则：每 1000-2000 步保存一次（取决于训练总步数和单步时间）
- 大模型保存一次可能耗时数分钟 → 需要异步保存
异步保存：
- 训练不暂停，后台线程/进程写入存储
- Megatron-LM: --async-save
- 需要额外 CPU 内存作为写入缓冲
Distributed Checkpoint（DCP）：
- 每个 rank 保存自己的分片，无需 AllGather 到单卡
- 支持 resharding：保存时 TP=4，加载时可 TP=8
断点续训：
- 恢复内容：模型参数 + 优化器状态 + 学习率调度器 + 数据加载进度 + RNG 状态
- 数据加载恢复：记录已消费的 sample 数或 DataLoader 的 state_dict
- 验证恢复正确性：续训后第一步的 loss 应与中断时一致

5. 框架选型建议

维度	Megatron-LM	DeepSpeed
并行策略	TP + PP + DP 原生深度集成	ZeRO 为主，TP/PP 需配合 Megatron
代码侵入性	需要用 Megatron 的模型定义方式	JSON 配置驱动，对用户代码改动小
性能（TP/PP）	通常更优（深度优化通信）	依赖 Megatron Core 或自有实现
易用性	学习曲线陡	低门槛（尤其配合 HuggingFace）
适用场景	大规模预训练（数百/数千卡）	中等规模微调、快速实验

6. 动手实验

实验 A：用 DeepSpeed ZeRO-2 训练小模型，对比 Stage 1/2/3 的显存和速度
实验 B：用 Megatron-LM 配置 TP=2, PP=2 的训练任务
实验 C：模拟断点续训——在第 100 步保存 checkpoint，从 checkpoint 恢复，验证 loss 一致性

🎯 本章学习目标

能描述 Megatron-LM 的代码结构和 3D 并行初始化流程
能编写 DeepSpeed 的 JSON 配置文件并解释各字段含义
能列举 3 种 Loss Spike 的常见原因和对应排查方法
能设计一套完整的 Checkpoint 策略（频率、异步保存、resharding、断点续训验证）
能根据团队规模和任务类型选择 Megatron-LM 或 DeepSpeed

搜索