Chapter 11
OOM、NCCL hang、loss spike 排查指南
Megatron 训练大模型时，大部分时间是在"训练前 + 训练中卡住 + 训练崩溃后"中度过的。本章按这三个时段组织。
11.1 训练前：理论显存自查
Megatron 自带 tools/report_theoretical_memory.py，能根据你的配置算理论显存：
python tools/report_theoretical_memory.py \
--num-layers 80 --hidden-size 8192 --num-attention-heads 64 \
--num-query-groups 8 --kv-channels 128 --ffn-hidden-size 28672 \
--seq-length 8192 --max-position-embeddings 8192 \
--micro-batch-size 1 --global-batch-size 1024 \
--tensor-model-parallel-size 8 \
--pipeline-model-parallel-size 4 \
--bf16 \
--use-distributed-optimizer输出包含每张卡的参数 / 梯度 / 优化器 / 激活预估。跑训练前先跑这个，比"启动后等 OOM"快得多。
11.2 OOM：按概率从高到低排查
| 现象 | 常见原因 | 修复 |
|---|---|---|
| 第 1 步就 OOM | 基础配置就装不下 | 开 selective recompute；升 TP / PP；开 distributed optimizer；降 micro batch |
| 第 50-200 步 OOM | 遇到长样本 | 开 packing；降 --seq-length；过滤异常长样本 |
| save_checkpoint 时 OOM | ZeRO-3 / dist optim 时 gather 显存 | 用 --ckpt-format torch_dist；或 --async-save |
| 反向 OOM 但前向 OK | 梯度没 distribute | --use-distributed-optimizer + --overlap-grad-reduce |
| 有的 stage OOM 别的 stage 没事（PP） | layer layout 不均 | --standalone-embedding-stage / 自定义 layout |
| fp8 + flash-attn 突然 OOM | TE 内部缓存 | 查 TE 版本；试 NVTE_FUSED_ATTN=0 |
| 长上下文 OOM | 没开 CP | --context-parallel-size 4 |
| VPP 模式 OOM | interleaved 每卡持多份激活 | 降 virtual_pipeline_size 或减 num_microbatches |
OOM 工具箱
# 1. 看每卡显存使用峰值
nvidia-smi -l 1
# 2. 用 PyTorch 显存 profiler
torch.cuda.memory._record_memory_history(max_entries=10000)
# ...训练几步...
torch.cuda.memory._dump_snapshot("memdump.pkl")
# 在 https://pytorch.org/memory_viz 打开
# 3. Megatron 自带 wall clock + memory profile
--profile \
--profile-step-start 5 \
--profile-step-end 611.3 NCCL hang：调度类问题
"NCCL hang" 是 Megatron 分布式训练的标志性故障 —— 多张卡卡住，没有报错，看起来在干活但其实在干等。原因几乎都是跨 rank 调度不对称。
常见原因（按发生率）
| 原因 | 调试 | 修复 |
|---|---|---|
| 某 rank Python 异常但没 propagate | NCCL_DEBUG=INFO | TORCH_DISTRIBUTED_DEBUG=DETAIL |
| 跨 rank 模型 forward 路径不一致 | print 每 rank shape | 检查 if rank==0 这种条件分支 |
| 数据 sampler 在某 rank 提前耗尽 | 查 dataloader length | 确认 sampler 走 DistributedSampler |
| NCCL 拓扑发现失败 | nvidia-smi topo -m | 检查 NVLink 是否完整 |
| 跨节点 ib 抖动 | 跑 nccl-tests 测带宽 | 查交换机；隔离故障节点 |
| NCCL 超时太短 | 默认 30 分钟 | NCCL_TIMEOUT=3600 让它 hang 出错而不是直接退出 |
没设 CUDA_DEVICE_MAX_CONNECTIONS=1 | — | 必须设 |
调 hang 的杀手锏
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=ALL
export TORCH_DISTRIBUTED_DEBUG=DETAIL
export TORCH_NCCL_BLOCKING_WAIT=1
export TORCH_NCCL_TRACE_BUFFER_SIZE=10000
# 启动训练；hang 后 ctrl-C 会打印每个 rank 卡在哪个 NCCL call11.4 loss spike：训练飞了
大模型预训练经常碰到：训了几千步、loss 缓慢下降，某一步突然飙到 10+ 然后开始震荡发散。原因谱系：
| 现象 | 多半原因 | 修复路线 |
|---|---|---|
| fp16 时频繁飞 | 数值范围窄 | 换 bf16 |
| bf16 也偶尔飞 | 梯度爆炸 / 坏样本 | 开 --clip-grad 1.0；--check-for-nan-in-loss-and-grad；查坏样本 |
| 每次 ramp up batch 时飞 | batch 突变 → 优化器没准备好 | 用 --rampup-batch-size 平滑过渡 |
| 训了 1000 步飞、回滚也飞 | 数据集脏 / 异常长序列 | 过滤异常样本；查 token 分布 |
| fp8 训练飞 | fp8 scaling 还没收敛 | 前 100 步用 bf16，之后切 fp8 |
| 启用 SP 后飞 | RNG context 错 | 查 dropout RNG 实现 |
| 大 LR + 大 batch 飞 | — | 把 LR 减半；warmup_steps 升 1.5× |
恢复策略
- 从最近一个 healthy checkpoint 恢复；
- 跳过引发 spike 的 batch（用
--skip-train-iteration-range）； - 临时降 LR 训几百步，再恢复正常 LR；
- 开
--check-for-nan-in-loss-and-grad提前发现。
11.5 收敛奇怪
| 现象 | 原因 | 修复 |
|---|---|---|
| loss 完全不降 | LR 太小 / tokenizer 不匹配 / template 错 | 查 LR；验证数据；最小 example dry-run |
| loss 慢慢降，验证集不降 | 过拟合 / 数据重复 | 查数据去重；缩小 epoch |
| training loss 比 val loss 高 | train 没开 eval mode（少见） | 检查 dropout |
| perplexity 异常高 | position embedding 类型错 | 核对 --position-embedding-type |
| resume 后 loss 跳变 | scheduler / optimizer 状态没恢复 | 不要加 --no-load-optim |
11.6 通信 / 性能"哪一步慢"
用 Megatron 自带 wall clock：
--log-throughput \
--log-progress \
--log-timers-to-tensorboard训几十步后 tensorboard 上能看到分段时间：
| 计时器 | 含义 |
|---|---|
forward-compute | 前向计算 |
backward-compute | 反向计算 |
forward-send-backward-recv | PP 通信 |
backward-params-all-reduce | DP 梯度同步 |
optimizer | optimizer.step |
哪一项占比异常就专项排查。
11.7 启动期常见报错
| 报错 | 修复 |
|---|---|
RuntimeError: NCCL error 5: invalid usage | 没设 CUDA_DEVICE_MAX_CONNECTIONS=1 |
AssertionError: num_attention_heads must be divisible by TP | 调 TP 或 num_heads |
cannot find megatron.core._C | fused kernel 没编；cd megatron/core/datasets && make |
flash_attn import error | pip install flash-attn --no-build-isolation |
TransformerEngine not found | 用 NGC 容器；或 pip install transformer-engine[pytorch] |
data path .bin not found | 第 6 章 6.13 |
11.8 多机训练专属问题
| 现象 | 修复 |
|---|---|
| 跨节点 PP 慢 | 查 IB；NCCL_IB_HCA / NCCL_SOCKET_IFNAME 正确 |
| 某节点 GPU 不参与 | torchrun 的 --node-rank 写错 |
| 不同节点起步时间差几分钟 | 共享文件系统慢 / 模型从 HF 拉太慢 |
| resume 跨节点失败 | checkpoint 没在共享文件系统 |
| NCCL 跨节点带宽极低 | iptables / VxLAN 阻挡了 GDR；走 NCCL_NET_GDR_LEVEL |
11.9 提 issue 的最小信息清单
- Megatron commit hash；
- PyTorch / CUDA / NCCL / TransformerEngine 版本；
- 启动命令完整复制；
- NCCL_DEBUG=INFO + TORCH_DISTRIBUTED_DEBUG=DETAIL 下的完整 log；
- 能否在更小配置（如 TP=1 PP=1）复现？
- 能否用 docker.io NGC 镜像复现？
11.10 心法
三步走
- 能用 NGC 容器就用——绝大多数"奇怪错"在标准镜像里不出现；
- 先跑
tools/report_theoretical_memory.py——OOM 不是猜的； - 打开三层 verbose log（NCCL、Distributed、Megatron timers）——比看 stderr 蒙猜效率高 100 倍。