DeepCompile：DeepSpeed 中面向分布式训练的编译器级优化

2026-09-05 16:21:41作者：韦蓉瑛

DeepCompile 是 DeepSpeed 将深度学习编译器能力扩展到分布式训练的子系统：它从一份标准单卡模型实现出发，在编译器 IR 层面自动注入参数切分、通信算子与 offload 逻辑，并通过剖析驱动（profile-guided）的优化 pass 动态调度计算与通信，无需修改模型代码。读完本文，你将理解 DeepCompile 的整体编译与优化流程、其内置的 ZeRO-3 / ZeRO-1 / 自适应 offload 各 pass 的实现原理与源码位置，掌握在 DeepSpeed 配置中启用 DeepCompile 的方法、关键参数与相互排斥约束，并了解官方评测中各项优化的实测收益。

1. 为什么分布式训练需要编译器级优化

深度学习编译器（如 PyTorch compiler）通过算子融合、算子调度等手段显著提升了单卡训练性能，但在分布式场景下能力有限：DeepSpeed、FSDP 等框架虽然靠先进的并行策略让大模型训练成为可能，但其优化逻辑实现于 PyTorch 框架层（Python 侧的 hook 与运行时逻辑），难以表达依赖分析、全局算子调度等编译器风格的技术。

DeepCompile 的目标正是弥补这一空白：

输入是一个标准单卡模型实现（例如 Hugging Face 模型库中的模型）；
输出是一个优化后的多卡训练图，自动包含参数切分、通信原语插入等分布式行为；
训练过程中，DeepCompile 还会基于剖析数据（算子级显存占用、执行延迟）持续调优这些并行化策略。

按官方博客的评测结论：在 GPU 资源充足时，DeepCompile 相比 ZeRO-3 基线最高可获得约 1.5× 的加速；在需要 offload 的 GPU 受限场景下最高可达约 7×。博客指出 DeepCompile 自 DeepSpeed 0.16.6 起可用，且处于活跃开发中，建议以最新版本的 DeepSpeed 或源码安装为准。

2. 设计总览：从单卡模型到分布式计算图

DeepCompile 的编译流程分两步（对应下图）：先把模型编译成中间表示（IR）形式的计算图，再依次施加一系列 优化 pass，每个 pass 负责一次特定的图变换或定向性能改进（如为切分参数插入 all-gather、插入 optimizer state 的 offload 等），整个过程保持原始计算语义不变。

其核心建立在两种能力之上：

自动并行化（Automatic parallelization）：优化 pass 可以将单卡计算图重写为分布式多卡版本，内置 ZeRO、FSDP 等策略，免去手工编写分布式逻辑的工程成本。
剖析驱动的性能调优（Profile-guided tuning）：运行时收集算子级显存、延迟等剖析数据，用于动态调度计算与通信算子，实现更充分的通信-计算 overlap 并避开显存瓶颈。

下图给出了 DeepCompile 的优化循环：编译器生成初始图后，先剖析前向/反向过程中的算子执行时间、通信开销与显存占用；据此施加 pass（插入、删除或重排算子）修改图；修改后再剖析，如此循环，直到得到一组稳定的优化，再部署到后续训练迭代。若运行期间显存等运行时特征发生变化，DeepCompile 可按预定义的 pass 时间表重新进入剖析-优化循环，让图持续自适应。

2.1 源码结构：deepspeed/compile 目录

从源码结构看，DeepCompile 位于 deepspeed/compile/ 目录，与上文的描述一一对应：

deepspeed/compile/config.py：定义 CompileConfig 配置模型（见第 5 节的参数说明）；
deepspeed/compile/init_z3.py：ZeRO-3 模式的 DeepCompile 初始化入口，负责注销原有 Python hook、注册参数到 DeepCompile 句柄、构建 pass 调度表并返回编译后端；
deepspeed/compile/passes/：各优化 pass 的实现，包括 zero3_compile.py（all-gather / reduce-scatter 注入）、prefetch.py（主动预取）、selective_gather.py（选择性 unshard）、offload_adam_states.py、offload_parameters.py、zero_1_and_2_compile.py 等；
deepspeed/compile/profilers/：graph_profile.py 与 comm_profile.py，分别为图级与通信级剖析器；
deepspeed/compile/list_schedule.py：显存感知的算子重排调度（即优化 pass 内部对图节点的重排算法）；
csrc/compile/：通信算子的 C++/CUDA 实现（z1.cpp、z2.cpp、z3.cpp），Python 侧以 torch.ops.dc.* 自定义算子形式接入 FX 图。

剖析驱动循环在代码中的落点可以看到 deepspeed/compile/passes/init.py 中的 run_opt_passes：它按序执行每个 pass，每次变换后调用 graph.lint() 与 gm.recompile()，随后用 MemoryProfilingInterpreter 重新跑一遍前向计算显存曲线并记录到 profile.fwd_mem / profile.bwd_mem——这正是博客中"修改后图被重新剖析，循环往复"描述的具体实现。

以 ZeRO-3 默认路径为例，deepspeed/compile/init_z3.py 中默认 pass 调度表为：

WARMUP = 5
schedule.append((0, [zero3_compile.add_z3_gather_release]))
schedule.append(
    (WARMUP,
     [zero3_compile.add_z3_gather_release, prefetch.schedule_prefetch, selective_gather.selective_gather]))

即第 0 步只注入 all-gather / release 算子，经过 5 步 warmup 收集到剖析数据后，再施加主动预取与选择性 unshard 两个 pass——与博客附录中"ZeRO-3 通信调度：All-gather & reduce-scatter → Proactive prefetching → Selective unsharding"的组合一致。此外，init_z3 会注销 DeepSpeed 原有 ZeRO-3 的参数 hook（如 parameter_offload._remove_module_hooks()、梯度累积 hook），并把每个参数的 ds_persist 置为 False，将参数注册进 DeepCompile 句柄 dc，使后续通信全部由编译图驱动。

2.2 显存感知的算子重排

博客中的 "memory-aware execution" 在 deepspeed/compile/list_schedule.py 中体现为具体的调度算法。create_mem_table 为图中节点建立显存表：allgather_ds_param 节点记正显存（占用 tensor_size），release_ds_param / reduce_ds_param 节点记负显存；choose_next_node 在可运行节点中优先调度显存影响最小的节点。simple_prefetch 进一步利用剖析得到的 device_time，为每个候选 all-gather 计算"解锁时间"（该 all-gather 能解锁多少后续计算），选择能解锁最多计算量的通信提前发射——这就是"主动预取"在调度器里的量化形式。fast_free_schedule 则围绕 AllgatherTask 的 free_cost / allgather_acc_mem 等指标决定 all-gather 的提前程度，使 gather 后的参数尽快释放，缩短其占用显存的窗口。

3. 优化 Pass 详解

3.1 ZeRO-3：通信算子注入、主动预取与选择性 unsharding

ZeRO-3 将参数、梯度、optimizer state 切分到各设备以降低显存。传统实现通常在运行时用 Python hook 插入 all-gather、reduce-scatter 和 buffer 释放；DeepCompile 则在编译期把这些算子直接注入计算图（deepspeed/compile/passes/zero3_compile.py 中的 add_z3_gather_release），由编译器依据图的静态结构 + 运行时剖析信息精确决定其位置。两个关键 pass：

主动预取（Proactive prefetching）：基于显存剖析把 all-gather 提前到计算流更早处发射，提高通信-计算 overlap、提升吞吐，同时避免 OOM；小的通信操作还会被融合（fused）以降低启动延迟。实现见 deepspeed/compile/passes/prefetch.py。
选择性 unsharding（Selective unsharding）：当显存允许时，让某些参数在前向/反向全程保持非切分状态，减少 all-gather 次数、避免冗余通信，在开启梯度累积的场景尤其有效。实现见 deepspeed/compile/passes/selective_gather.py。

3.2 自适应 Offloading

DeepCompile 支持自适应 offloading：只把超出显存预算的 optimizer state 部分（如 Adam 的 momentum 与 variance）搬到 pinned host 内存，并在优化器步骤前取回；传输被调度为与计算重叠，属于"选择性 + 异步"策略。与把 optimizer 计算整体放到 CPU 的传统方案不同，DeepCompile 中 optimizer 计算仍留在 GPU 上。对应配置项为 compile.offload_opt_states，实现见 deepspeed/compile/passes/offload_adam_states.py。

init_z3 中针对该模式的默认调度值得注意（deepspeed/compile/init_z3.py）：

schedule.append((0, [zero3_compile.add_z3_gather_release]))
schedule.append((1, [offload_adam_states_for_init, zero3_compile.add_z3_gather_release, move_opt_states]))

即 offload 从第 1 个 optimizer step 才开始生效（state 在第 0 步的 optimizer 更新后才存在），并用 offload_adam_states_for_init 在剖析前清空 state，使显存规划基于"floor"进行——保证只有在 offload 后才能装下的作业不会带着全部驻留状态运行任何一步。

3.3 ZeRO-1

ZeRO-1 只切分 optimizer state，参数与梯度保持完整复制，是对计算流改动最小的轻量方案。DeepCompile 通过 deepspeed/compile/passes/zero_1_and_2_compile.py 将 reduce-scatter 直接注入计算图（初始化入口见 deepspeed/compile/init_z1_and_2.py），避免 Python hook 带来的开销。

博客附录列出的当前可用 pass 与本文上述章节一一对应：

Pass	所属策略
All-gather & reduce-scatter insertion	ZeRO-3
Proactive prefetching	ZeRO-3
Selective unsharding	ZeRO-3
Reduce-scatter insertion	ZeRO-1
Adaptive offloading	通用

官方实验所用组合为：ZeRO-3 通信调度采用 "All-gather & reduce-scatter → Proactive prefetching → Selective unsharding"；ZeRO-3 offload 在其基础上追加 Adaptive offloading；ZeRO-1 采用 "Adding reduce-scatter"。

4. 性能评测

4.1 ZeRO-3

评测使用 Hugging Face 模型实现 + 参数切分，模型为 Llama-3-70B 与 Mixtral 8x7B，32 块 H100，序列长度 1024，横轴为不同梯度累积步数，指标为训练吞吐（TFLOPs/GPU）。对比基线为：(i) ZeRO3+Eager（无编译器的 eager 模式）；(ii) ZeRO3+Compile（叠加 PyTorch compiler）。DeepCompile 同时开启了主动预取与选择性 unsharding。

结果：两个模型上 DeepCompile 吞吐均稳定领先，且累积步数越大收益越明显（参数更新频率降低使选择性 unsharding 更有效），Llama-3-70B 上最高 1.28×、Mixtral 8x7B 上最高 1.54×。另一方面，ZeRO3+Compile 在部分设置下反而引入少量开销——ZeRO-3 中为支持 prefetching 等运行时特性存在大量条件分支，编译器遇到无法静态解析的分支时会把计算图切成多个片段，片段化削弱了优化机会并带来额外执行开销。DeepCompile 由于通信算子由编译器原生插入图中，不存在这类 graph break 问题。

4.2 Offloading

Llama-3 70B 用 ZeRO-3 训练通常需要 32 块 80GB GPU。DeepCompile 将 offload 实验放在 16 块 H100-80GB（一半的卡数）上：只 offload optimizer 内存（Adam 的 momentum、variance 与 fp32 master weights），optimizer 计算保留在 GPU，前向/反向剖析显存后按需转移，并在运行时依据显存模式动态调整传输时机以异步重叠。对比 ZeRO3+Eager 与 ZeRO3+Compile 两个基线，DeepCompile 最高取得约 7× 的加速；而 ZeRO3+Compile 与 ZeRO3+Eager 性能接近。

4.3 ZeRO-1

使用 Llama-3-8B，8 块 GPU，batch size 2，对比 ZeRO1+Eager 与 ZeRO1+Compile。不同序列长度下 DeepCompile 均稳定领先；batch size 1、序列长度 512 时收益最大，分别较 ZeRO1+Eager 最高 1.9×、较 ZeRO1+Compile 最高 2.5×。博客的解释是：PyTorch compiler 路线在大 batch / 长序列下可通过更高效 kernel 获益，但在 ZeRO-1 式训练中通信算子周围会产生 graph breaks，小 batch、短序列时这种开销被放大；DeepCompile 的通信算子在编译期原生插入图中，避免图碎片化，因此在小规模负载下更稳健，同时保留编译器级优化收益。

更多结果（不同 batch size、序列长度与显存占用的详细对比）见博客原文末尾引用的 arXiv 论文（此处不重复给出外部链接）。

5. 如何在 DeepSpeed 配置中启用 DeepCompile

DeepCompile 通过 DeepSpeed 配置中的 compile 段开启。仓库中的真实测试配置 tests/torch_compile/ds_config_z3_deepcompile_no_persist.json 展示了 ZeRO-3 + DeepCompile 的完整示例：

{
  "train_batch_size": 8,
  "bf16": { "enabled": true },
  "compile": {
    "deepcompile": true
  },
  "zero_optimization": {
    "stage": 3,
    "reduce_scatter": true,
    "overlap_comm": false,
    "contiguous_gradients": false,
    "stage3_param_persistence_threshold": 0,
    "stage3_model_persistence_threshold": 0
  }
}

用法说明：

在训练配置 JSON 的 compile 段中设置 "deepcompile": true；
与 zero_optimization 的 stage: 1 或 stage: 3 搭配使用，分别启用 ZeRO-1 / ZeRO-3 编译路径；overlap_comm 等 ZeRO 运行时开关在编译路径下由编译器接管调度，示例配置中将其关闭；
引擎侧的接线在 deepspeed/runtime/engine.py 中完成：is_deepcompile_enabled() 读取配置，训练时分别调用 deepcompile_backward_prologue / deepcompile_backward_epilogue，前向阶段还有 deepcompile_z3_forward_context 的 eager 回退上下文（deepspeed/compile/z3_eager_fallback.py），保证未被编译覆盖的路径仍可正常执行。

5.1 `compile` 段关键参数

以下参数与默认值取自 deepspeed/compile/config.py 的 CompileConfig 定义：

参数	默认值	说明
`deepcompile`	`false`	开启/关闭 DeepCompile 模式
`passes`	`null`	组合不同优化 pass 的列表（当前支持的取值见下方说明）
`offload_parameters`	`false`	参数 offloading；与 `offload_opt_states`、`offload_activation` 互斥
`offload_opt_states`	`false`	将 Adam 的 fp32 master weights 与动量 offload 到 pinned host 内存，optimizer 步骤前取回；为 `gradient_accumulation_steps=1` 设计
`offload_activation`	`false`	将反向需要的前向激活 offload 到 pinned host 内存；仅考虑 ≥5MB 且形状固定的张量；不在默认调度中，需通过自定义 `schedule=` 显式启用，且与 `offload_parameters` / `offload_opt_states` 互斥
`offload_activation_pin_memory`	`true`	激活 offload 的 host buffer 是否 pin；受 `ulimit -l` 限制紧张时置 `false`
`free_activation`	`false`	自由激活模式：前向保存给反向的激活在阈值内不再保留
`free_activation_threshold`	`10 * 1024 * 1024`	自由激活模式下，不小于该字节阈值的激活会被立即释放
`double_buffer`	`true`	开启/关闭双缓冲
`symmetric_memory`	`false`	开启/关闭对称内存
`debug_log`	`false`	开启/关闭图 dump 等调试日志
`sync_before_reduce` / `sync_after_reduce`	`false`	reduce 前后的同步开关
`sync_before_allgather` / `sync_after_allgather`	`false`	all-gather 前后的同步开关
`keep_int_input_tensors`	`true`	`InputStorage` 对 int 张量保留真实值而非 dummy 值
`keep_all_input_tensors`	`false`	对全部输入张量保留真实值

注意 passes 的类型定义（deepspeed/compile/config.py）为 Literal["z1", "z3", "autosp", "autotp"]：当前配置层的 passes 字段用于组合 ZeRO-1/3 与 AutoSP / AutoTP 等策略，而博客附录中列出的 all-gather 注入、主动预取、选择性 unshard 等细粒度 pass 由默认调度表（第 2.1 节）按 warmup 进度自动编排；官方配置文档 docs/_pages/config-json.md 中给出的 AutoSP / AutoTP 用法即为：

{
  "zero_optimization": {"stage": 0},
  "compile": {
    "deepcompile": true,
    "passes": ["autosp"]
  }
}

5.2 配置互斥约束（来自源码的校验逻辑）

deepspeed/compile/init_z3.py 在初始化前会做参数合法性校验，以下约束可直接从源码确认：

compile.offload_opt_states 不能与 zero_optimization.offload_optimizer 设为 cpu/nvme 同时使用——两者管理同一份 optimizer state，前者在显存允许时让其驻留、在编译图中移动，后者则全程驻留 CPU/NVMe 并在那里执行 optimizer；
compile.offload_activation 不能与 offload_parameters 或 offload_opt_states 组合——各自独立地对整块显存预算做规划，两两组合会移动远超需求的数据量；
offload_parameters 与 offload_opt_states 也不能同时启用（引擎可能在 ZeRO 配置同时 offload optimizer 和参数到 CPU 时隐式开启 offload_parameters，需留意）。

6. 展望与参考

DeepCompile 当前发布的 ZeRO-3 / ZeRO-1 / 自适应 offload pass 只是起点，博客明确列出的后续方向包括：自动化并行（序列/张量并行）、更智能的显存管理、对运行时行为的动态适配。从仓库结构看，deepspeed/compile/ 中已存在 init_sp.py、init_tp.py 及 passes/sp_compile.py、passes/tp_compile.py，与上述方向一致；相关测试可参考 tests/unit/v1/compile/（如 test_compile_zero.py、test_offload_opt_states.py、test_selective_gather.py、test_graph_profile.py）与 tests/torch_compile/（如 test_deepcompile_z3_release.py）。