性能优化：fused kernel / flash-attn / recompute / fp8
10.1 优化的三个维度
Megatron 的性能优化可以归到三笔账：
$$ \text{MFU} = \frac{\text{有效 FLOPs}}{\text{理论 FLOPs}} = f(\underbrace{\text{算子效率}}_{\text{fused/flash}}, \underbrace{\text{通信效率}}_{\text{overlap, bucket}}, \underbrace{\text{显存效率}}_{\text{recompute, fp8}}) $$
- 算子效率：fused kernel、FlashAttention、torch.compile；
- 通信效率：overlap、distributed optimizer、bucket size；
- 显存效率：activation recompute、fp8、selective recompute、SP/CP（第 9 章）。
10.2 算子优化：fused kernels
Megatron 在 megatron/core/fusions/ 维护一组 fused kernel：
| kernel | 融合了什么 |
|---|---|
fused_bias_gelu / fused_bias_geglu | bias + GeLU/SwiGLU 一个 launch |
fused_bias_dropout | bias + dropout + residual |
fused_layer_norm | LayerNorm 全流程一个 kernel |
fused_rms_norm | RMSNorm fused 版 |
fused_softmax | scaling + mask + softmax（attention 内） |
fused_cross_entropy | vocab parallel CE 内部 fused |
fused_swiglu | SwiGLU fused |
大部分会通过 --transformer-impl transformer_engine 自动启用（TransformerEngine 提供）。不开 fused kernel 会损失 20-40% 速度。
10.3 FlashAttention
FlashAttention 把 attention 的 $O(S^2)$ 内存降到 $O(S)$，速度也快 2-3 倍。Megatron 通过 --attention-backend 选择：
--attention-backend | 说明 |
|---|---|
auto（默认） | 自动选最快的 |
flash | FlashAttention 2/3（推荐） |
fused | TransformerEngine fused attention |
unfused | PyTorch 原生 SDPA |
local | Megatron 自家实现（不推荐） |
FlashAttention-3 在 H100 上 fp16 / fp8 进一步加速，是2024 年后训长上下文的标配。
10.4 activation recompute（激活重算）
原理：前向时不存中间激活，反向时重算一遍。用时间换显存。
| flag | 说明 |
|---|---|
--recompute-activations | 启用 selective recompute（默认级别） |
--recompute-granularity selective | 推荐：只对 attention 内最贵的几个 op 重算 |
--recompute-granularity full | 整层重算（费时 30%+，但显存最省） |
--recompute-method uniform | 每 N 层 ckpt 一次 |
--recompute-method block | 每个 stage 内的前 K 层 ckpt |
--recompute-num-layers | full 模式时每 N 层 ckpt |
--distribute-saved-activations | 把保存的激活也按 TP 切（叠加 SP 用） |
selective vs full
| 方式 | 额外时间 | 额外显存收益 |
|---|---|---|
| selective recompute | +5-10% | 降 attention 矩阵 |
| full recompute (uniform 1) | +30-35% | 几乎全部激活省下 |
经验法则：先 selective，不够再 full。
10.5 fp8 训练
H100 引入 fp8 张量核，吞吐是 fp16 的 2 倍。Megatron 通过 TransformerEngine 支持：
--fp8 \
--fp8-format hybrid # E4M3 forward + E5M2 backward
--fp8-margin 0 # 自动 scaling 的边界
--fp8-interval 1 # 多少步重新 calibrate实战注意：
- fp8 训练目前只在 H100/H200/Blackwell 上能用；
- 必须配 TransformerEngine（
--transformer-impl transformer_engine）； - 需要 fp16/bf16 master weight 兜底；
- 训练前几百 step 可能数值波动比 bf16 大，loss spike 概率略高；
- 实际加速大约 1.4-1.6 倍（不到理论 2 倍，因为大量算子不在 GEMM 上）。
10.6 distributed optimizer：ZeRO-1 等价
--use-distributed-optimizer \
--overlap-grad-reduce \
--overlap-param-gather原理：把 Adam 的 $m, v, $ fp32 master 切到 $\mathrm{DP}$ 个 rank 上（类似 ZeRO-1）。每张卡只更新自己负责的 1/DP 优化器状态，optimizer.step 后再 all-gather 参数。
显存收益：optimizer state 占用从 $12\Psi$ 降到 $12\Psi/\mathrm{DP}$，对 70B 模型 8 卡训练能省接近 80 GB。
10.7 通信 overlap
Megatron 支持多种通信 / 计算 overlap：
| flag | 说明 |
|---|---|
--overlap-grad-reduce | 反向中梯度 reduce 与下一层 backward 重叠 |
--overlap-param-gather | 前向中参数 all-gather 与上一层 forward 重叠 |
--tp-comm-overlap | TP 内 all-reduce 与 GEMM 重叠（需 TE） |
--tp-comm-bulk-wgrad | bulk weight gradient 通信 |
--tp-comm-overlap-ag | TP all-gather 重叠 |
--tp-comm-overlap-rs | TP reduce-scatter 重叠 |
--gradient-accumulation-fusion | 累积梯度时 fuse |
所有 overlap 默认都该开，少数情况下因为 stream 资源紧会反而变慢，需要单独 benchmark。
10.8 torch.compile
--use-torch-compile把模型 forward 用 PyTorch 2.x 的 compile() 包一次。注意：
- 第一次 step 编译 1-3 分钟；
- 变长 sequence 会反复重编译，必须配合 packing；
- fp8 + compile 同开偶有 bug，多看 issue。
10.9 NCCL 与硬件层调优
| 环境变量 | 用途 |
|---|---|
NCCL_IB_DISABLE=0 | 启用 InfiniBand |
NCCL_SOCKET_IFNAME=ibX | 指定网卡 |
NCCL_IB_HCA=mlx5 | 指定 HCA |
NCCL_NVLS_ENABLE=1 | NVLink SHARP（H100） |
NCCL_ASYNC_ERROR_HANDLING=1 | 异步错误处理（debug 友好） |
CUDA_DEVICE_MAX_CONNECTIONS=1 | Megatron 必设（保证 NCCL 通信顺序） |
10.10 性能体检清单
| 序 | 检查项 |
|---|---|
| 1 | 用了 --transformer-impl transformer_engine？ |
| 2 | --attention-backend flash？ |
| 3 | --bf16 而非 --fp16（A100+）？ |
| 4 | --use-distributed-optimizer？ |
| 5 | --overlap-grad-reduce --overlap-param-gather？ |
| 6 | TP > 1 时配 --sequence-parallel？ |
| 7 | --recompute-activations --recompute-granularity selective？ |
| 8 | 设了 CUDA_DEVICE_MAX_CONNECTIONS=1？ |
| 9 | NCCL 走 NVLink + InfiniBand？ |
| 10 | num_microbatches ≥ 4 × PP？ |
| 11 | 样本长度齐整 / 开 packing？ |
| 12 | H100 用 --fp8？ |
10.11 性能参考
NVIDIA 公开过 175B 训练的 MFU 数字（H100，2048 卡）：
| 配置 | MFU | tokens/sec/GPU |
|---|---|---|
| bf16 + 全套优化 | ~58% | ~5000 |
| fp8 + 全套优化 | ~52% | ~8000 |
如果你的训练 MFU 低于 40%，大概率是上面 12 条没勾全。用 tools/report_theoretical_memory.py 算理论显存，对比实际占用看差异。
10.12 FP8 训练（H100/H200/B100）
FP8 在 Megatron 里走 TransformerEngine。两个 recipe：
| recipe | 含义 | 典型场景 |
|---|---|---|
hybrid | fwd 用 E4M3、bwd grad 用 E5M2 | 默认；通用 |
e4m3 | 全程 E4M3 | 推理 / 数值范围小的模型 |
CLI 全套：--fp8-hybrid --fp8-amax-history-len 1024 --fp8-amax-compute-algo max --fp8-margin 0。注意：FP8 需要 TE ≥ 1.x 且模型必须用 TE 版的 layer norm / linear（即 spec 选 TE 那一套）。
10.13 Kitchen Quantization（2026 新）
arguments.py:3335 _add_kitchen_quantization_arguments 注册的训练期量化路径。区别于推理后量化：
- 训练时部分 layer / 部分 tensor 直接以低精度（INT4/INT8/MXFP4）参与计算；
- QAT（Quantization Aware Training）的工业落地；
- 典型 CLI：
--kitchen-quantize --kitchen-quant-method mxfp4 --kitchen-quant-target weights； - 与 FP8 不互斥（weight INT4 + activation FP8 是常见组合）。
10.14 Multi-Token Prediction（MTP）
DeepSeek-V3 论文带火的训练目标。Megatron 实现见 core/transformer/multi_token_prediction.py，TransformerConfig.mtp_num_layers / mtp_loss_scaling_factor。要点：
- 主 head 之外再加 N 个浅层 head，每个预测后 1..N 个 token；
- 训练 loss = main_loss + λ·sum(mtp_loss_i)；
- 提速来源：相同步数下信号密度提升 → 收敛快；
- 推理时可丢掉 MTP head，也可保留做 speculative decoding。
10.15 这章你需要带走什么
- 性能优化 = "算子 / 通信 / 显存" 三笔账；
- fused kernel + FlashAttention + TransformerEngine 是默认套餐；
- selective recompute 是性价比最高的显存招；
- distributed optimizer 是 Megatron 自带的 ZeRO-1；
- NCCL 环境变量是多机性能的"暗坑"；
- H100 + fp8 能再多 30-50% 吞吐。