性能调优：bucket size / overlap / fused optimizer

11.1 调优前先量化"瓶颈在哪"

打开 ds_config.json：

{
  "wall_clock_breakdown": true,
  "steps_per_print": 50,
  "flops_profiler": {
    "enabled": true,
    "profile_step": 5,
    "module_depth": -1,
    "top_modules": 3,
    "detailed": false,
    "output_file": null
  }
}

跑 100 步后能看到一份 breakdown：

step= 50  fwd= 240ms  bwd= 380ms  step= 95ms
       allgather=  60ms (24% of fwd)
       reduce_scatter= 110ms (28% of bwd)
       optim_step= 95ms

诊断思路：

11.2 通信 bucket：reduce / allgather 的"批"

DS 通信不是一个张量一个张量地发，而是凑齐一个 bucket（≈ 几百 MB）后一次发出。bucket 太小 → call 多，启动开销大；bucket 太大 → 显存被占；最佳值通常 $2 \times 10^8 \sim 5 \times 10^8$。

11.3 overlap_comm：必须开

不开 overlap，整个 step 串行执行：

$$T_\text{step} = T_\text{compute} + T_\text{comm}$$

开 overlap：

$$T_\text{step} \approx \max(T_\text{compute}, T_\text{comm}) + \text{少量开销}$$

差距通常 30%+。任何场景都建议开：

"zero_optimization": {
  "stage": 3,
  "overlap_comm": true,
  "contiguous_gradients": true
}

11.4 fused optimizer

DS 内置的 fused 实现比 PyTorch 自带快 1.5–3 倍：

用法：直接在 ds_config 写 "optimizer": {"type": "AdamW"}，DS 自动选 FusedAdam（如果 op 已编）；offload 时自动选 CPUAdam。不需要你手写 import。

11.5 activation checkpointing

大模型激活内存约：

$$M_\text{act} \approx B \cdot L \cdot H \cdot (\text{constant})$$

$B$ batch、$L$ sequence length、$H$ hidden size。开 checkpointing 后激活内存降到 $\approx \sqrt{L}$ 倍数级。代价：反向时要重做一次 forward，训练时间 +25%–35%。

三种打开方式：

同时开多个不会冲突，DS 内部去重。

11.6 torch.compile 加持

PyTorch 2.x 的 compile() 能给训练带 10–30% 提速。DS 0.14+ 支持：

{
  "torch_compile": {
    "enabled": true,
    "backend": "inductor"
  }
}

注意：

• 第一次 step 会卡几分钟编译；
• 跟 ZeRO-3 + activation checkpoint 同开偶有 bug，多看 issue；
• 动态 shape（变长 sequence）会触发重编译，建议配合 packing 把长度齐整。

11.7 autotuning：让 DS 帮你调

DS 自带 autotuner：

{
  "autotuning": {
    "enabled": true,
    "fast": true,
    "results_dir": "autotune_results"
  }
}

启动命令加 --autotuning run：

deepspeed --autotuning run --num_gpus 8 train.py --deepspeed --deepspeed_config ds_config.json

它会扫一组合理的 micro batch / accum / zero stage 组合，跑短训练测吞吐，给一份最佳推荐。第一次接触新硬件时极有用。

11.8 一份"性能体检"清单

11.9 实战吞吐参考（A100 80G）

具体数字会因模型架构、batch size、context length 浮动 ±30%；用作"我的训练是不是远低于预期"的参考标尺。

11.10 Coordinator 八个 timer 的实战用法

ch05.11.3 列了 Stage 3 PartitionedParameterCoordinator 的八个 timer。wall_clock_breakdown: true 时它们会逐步打印——下表给出"这个 timer 大代表哪类问题"：

一个经验顺序：先看 *_WAIT 大不大；其次看 ALL_GATHER 大不大；最后看 SUBMIT。WAIT 大几乎都是带宽 / 拓扑问题，调字段救不回来。

11.11 ZeRO++（HPZ 分层 + 量化通信）

ZeRO++ 的核心 idea：节点内通信便宜（NVLink ~ 600 GB/s）、节点间通信贵（IB ~ 50 GB/s）。所以——

• 把 ZeRO-3 的 all-gather 拆成两层：节点内 all-gather + 节点间 broadcast，节点间数据量减少；
• 节点间传输前量化：权重 int8、梯度 int4，带宽再省 4×。

对应字段（ds_config）：

陷阱：

• 单机不要开：HPZ 分区数 = WORLD_SIZE 时跟原 ZeRO-3 等价，但多绕一层 wrapper 反而慢；
• 量化梯度会引入数值误差，对 fp16 训练影响明显（bf16 一般 OK）；
• 跨架构 checkpoint 不可移植 —— 量化路径写下的 checkpoint 只能在 ZeRO++ 下加载。

11.12 DeepCompile：把 ZeRO 和 torch.compile 真正打通

torch.compile 单独用不难，但跟 ZeRO-3 的"参数动态 all-gather / release" 配合不是显然的——参数地址会变，graph 重编频率高。DeepCompile（deepspeed/compile/）的工作是把 ZeRO 的 hook 行为转成编译器能理解的 IR pass。

关键源码：compile/backend.py:44 的 GraphOrder 跟踪 frame ID 与图执行顺序；line 79 register_compile_pass(name, fn) 注册 custom pass；line 95 launch_compile_passes(global_steps) 按 step 触发 pass。pass 调度逻辑在 compile/list_schedule.py。

启用前置：PyTorch 2.x 带 torch._dynamo、装好 functorch、AOT autograd 可用。第一步 step 会卡几分钟编译。

11.13 Autotuner 搜索空间

deepspeed/autotuning/autotuner.py:42 的 Autotuner 在跑一组短训练实验找最佳配置。搜索维度（你不写也会自动展开）：

三种 tuner（autotuning/tuner/）：

• gridsearch：穷举（小空间 OK）；
• random：随机抽样（大空间 + 时间紧）；
• modelbased：用代理模型（Gaussian Process / RF）逐步收敛；

配 "max_evals": 30 控制最多跑多少组合。结果写到 results_dir，里面 exps_dir/exps_*.json 是每次实验的 config，results.csv 是吞吐对比。第一次接触新硬件时极有用——比手调省两天。