这篇博文围绕 Megatron-LM 与 CUDA Graph 展开分享训练侧延迟优化的原理、接入路径和真实踩坑复盘。按照要求全文纯 Markdown 输出、无主标题、H2/H3 带编号、字数充足、零模板化总结以下为最终交付内容。1. 训练循环里那 3% 的 GPU 利用率去哪了先说一个我自己的经历。有段时间我负责把一批 GPT 类模型的训练性能拉上去模型规模在 7B 到 13B 之间框架用的是 Megatron-LM。当时 GPU 利用率看起来不算差但用 Nsight Systems 一抓单个 step 的 timeline发现问题非常扎眼GPU Compute 那一段并不是连续的kernel 与 kernel 之间布满了 CPU 侧的黑洞。这些黑洞不是显存访问等待而是 launch 延迟——也就是 CPU 在忙着往 GPU 上投递下一个 kernel但 GPU 已经空闲了。单个 kernel 的 launch 开销大约在 3 到 10 微秒听起来微不足道可一次 175B 级别的训练 step 里前向、反向、梯度通信、优化器更新加在一起少说也有几千个 kernel。几百步叠加起来GPU 每个 step 里真正在算的时间可能只有 70% 左右剩下全在等 CPU 排队。这就是 CUDA Graph 在 LLM 训练中被重新重视的核心原因。CUDA Graph 并不是什么新的“加速魔法”它只是把一串 kernel 的依赖关系提前录制下来然后在每次迭代时一次性回放省掉 CPU 逐个 launch 的中间过程。Megatron-LM 在训练侧接入 CUDA Graph 后我们把一个典型 7B 模型 step 的 wall time 从 2.9 秒压到了 2.2 秒上下吞吐提升接近 25%。这篇文章不打算从零科普 CUDA Graph 的 API而是想讲清楚一件事在 LLM 训练里Graph 到底该在哪里接、能包住哪些算子、以及为什么很多人接进去之后反而变慢或 loss 跑飞。如果读者之前主要接触的是推理侧优化比如 TensorRT、FasterTransformer 里用 Graph 做 token 生成加速那训练侧的思路会有所不同。训练场景下我们要处理反向传播、梯度 AllReduce、dropout 的随机性、分布式通信同步这些都会直接影响 Graph 能不能安全地录制和回放。这也是我这段时间踩坑最多的部分。2. CUDA Graph 的生成本质录制一段“算子录像”而不是修改计算逻辑理解 CUDA Graph 的关键是先放下“它是不是一种编译器优化”的念头。它不是编译优化不改变 kernel 的执行速度也不改变计算顺序的语义。更形象地说它像一台录像机capture 阶段把所有算子和依赖关系按 DAG 形式录下来replay 阶段按图一次性回放。你把一道菜的做法录成视频回放时锅还是那个锅、火还是那个火只是不再需要厨师每步都看菜谱。2.1 capture 阶段发生了什么在 PyTorch 里最简单的录制方式长这样# 先申请一个稳定的 graph pool避免 capture 前后内存指针漂移 graph_pool torch.cuda.graph_pool_handle() # 默认 capture 步 graph torch.cuda.CUDAGraph() with torch.cuda.graph(graph): loss model(input_ids) loss.backward() # 后续迭代直接回放 graph.replay()很多人第一次用会犯一个错误就是把参数初始化或者 optimizer 第一次执行也包进 capture 区域里。capture 阶段本身并不真正在 GPU 上执行计算它只是在记录 API 调用序列。如果区域里混入了依赖 CPU 数据的操作比如.item()、.cpu()、条件判断、动态形状的 tensor 分配轻则报错重则录制出一个“行为诡异”的图回放时每次跑出来的结果都一样dropout 的 mask 永远固定成一个。还有一个容易被忽略的点capture 阶段产生的内存分配行为会被固化。默认情况下进入torch.cuda.graph区域后PyTorch 的内存分配会切换到 graph 私有的内存池而不是原来的流序内存池。这会导致模型参数、中间激活这些张量的虚拟地址发生变化。如果模型权重在进入 Graph 前已经存在旧内存池里而 Graph 捕获后又放到了新池子replay 时 kernel 拿到的指针如果还是旧的就会直接踩到错误内存地址。Megatron-LM 的 cudagraph 模块里专门做了全局内存池管理和参数指针重定位就是为这个原因。2.2 replay 为什么能省时间CPU 批量投递与依赖预建单个 kernel launch 的 CPU 开销在微秒级但一次 step 可能有几千次 launch。更糟糕的是PyTorch 的 eager 模式下每个 kernel 之间的依赖要通过 event 或 stream 同步来管理这也是 CPU 开销的一部分。CUDA Graph 在做 capture 的时候把这些依赖关系全部提前计算好replay 时 GPU 驱动只需遍历一次图结构把可以并行的节点尽量并行调度把必须串行的节点通过图中已有的 event 依赖传递而不是在 CPU 上一步步排队。打个通俗比方原来 CPU 是那个要站在流水线旁边每拧一颗螺丝都要看一眼图纸的工人CUDA Graph 相当于把整条流水线的图纸提前刻在机器里工人只需要按一次开始键。所以它的收益上限取决于 CPU launch 开销在整体 step time 中的占比。推理场景里这个占比很高因为生成一个 token 往往只需要几十个 kernel而训练侧 kernel 数量多、单 kernel 计算量大收益会被稀释但绝对值依然可观。2.3 训练侧相反的方向不是所有算子都值得包进 Graph这里要和推理场景做一个明确区分。推理用 Graph 通常是把整条 forward 路径固化输入输出 shape 固定逻辑简单。训练侧的前向还相对好说反向加进去之后问题就来了反向里包含梯度计算梯度计算依赖中间激活值激活值在每次 iteration 里可能因为 batch 内不同样本的 padding 而产生 shape 变化如果这个模型还开了梯度 checkpointing反向时的重计算逻辑会导致一部分算子在每个 step 里重复出现Graph 的录制区间也相应变得不稳定。所以我在做 Megatron-LM 接入时没有天真地试图把整个 step 全包进去。比较务实的做法是区分“稳定可捕获区域”和“必须留在动态路径上的区域”。稳定可捕获的包括前向的核心 transformer 部分、反向的梯度计算、以及不涉及跨 rank 通信的局部规约必须留在动态路径上的包括pipeline 并行阶段的跨 stage 同步、NCCL 的 AllGather/ReduceScatter、优化器 step尤其是涉及动态内存重分配的时候。这个区分是后续所有配置的基础也是理解 Megatron-LM 里--cudagraph-modules这类参数含义的前提。3. Megatron-LM 里 Graph 与分布式训练的共处之道Megatron-LM 官方代码仓库很早就加了 CUDA Graph 的支持入口集中在cudagraph.py和几个训练脚本的启动参数里。但真正要把它用起来不能只靠开一个开关需要理解分布式并行和 Graph 之间的几条硬约束。3.1 启动参数与模块级 capture 的意图以 GPT 预训练脚本的常见参数为例社区里常见的用法是这样组合的--use-cudagraph \ --cudagraph-steps-per-module 1 \ --cudagraph-modules bert.language_model.encoder,lm_head \ --cudagraph-warmup-steps 3--cudagraph-steps-per-module这个参数初看容易误解。它并不是说一个模块要连续跑几步才录一次 Graph而是指为同一个模块维护的独立 Graph 实例数量。设成 1 意味着每个模块只保存一个 Graph内存开销最小设成更大的值适合 step 之间形状会略有波动的场景多几个候选 Graph 可以轮流切换。不过我实测下来如果模型本身没有动态形状问题设成 1 就够了多设反而白白增加显存占用。--cudagraph-modules控制哪些子模块会被包进 Graph。这里有个经验不是模块越多收益越大。transformer 层里的 attention 和 MLP 是最值得包的因为它们 kernel 密集且 shape 稳定而嵌入层和输出层很多时候涉及 vocab 维度的 gather或者和分布式通信耦合包进去反而增加内存池压力。我自己在 7B 模型上把lm_head排除在 Graph 外后性能没有明显下降因为那一层本来就不是 kernel launch 密集的地方。3.2 张量并行与数据并行Graph 的舒适区先说结论纯张量并行加数据并行是最适合 CUDA Graph 的分布式组合。原因在于这两种并行方式下transformer 的每个层在每一次 iteration 里执行的算子序列是稳定重复的。张量并行里的 AllGather 和 ReduceScatter 是在 transformer 层内部发生的每次调用的通信量、参与的 rank 数都固定可以完整包进 Graph数据并行里的梯度 AllReduce 发生在全部层都计算完后位置固定、方向固定也适合作为 Graph 的收尾节点。Megatron 在处理这个场景时走的是“模块级局部捕获”路线不要求从 input_ids 到 optimizer step 整条链路都在一个 Graph 里而是把前向的每个 transformer 层模块、反向的梯度模块分别捕获再在 Graph 之间用普通的 stream 依赖串起来。这样每一块内部都是稳定的 kernel 序列块与块之间的连接如果用上了原生的 tensor 传递而没有任何 CPU 同步整体的 launch 开销依然能被大幅压缩。我们 7B 模型的测试就是在这种组合下做的收益最干净容易复现。3.3 pipeline 并行下 Graph 的“循环依赖”难题pipeline 并行PP是让 CUDA Graph 掉头发的地方。原因不在 kernels而在 pipeline schedule 本身经典 1F1B 调度下每个 micro-batch 的前向、反向、跨 stage 通信相互交错尤其是在 virtual pipeline stage 开启后一个 stage 上要执行的 micro-batch 序列存在跨 iteration 的循环依赖。一个 stage 的 forward 可能要等上一个 stage 的 forward 结果而那个 stage 同样也在等待。Graph 回放是一个 DAG不支持这种带环路的等待关系。强行包住可能会导致某些 rank 在回放时拿到的是上一轮 iteration 的中间数据loss 曲线直接飞掉。所以碰到 PP 场景我的建议是要么干脆不配 Graph要么只对 PP stage 内部的“无跨 stage 通信”的算子块捕获。换句话说把 Graph 当作用来加速局部 kernel launch 的工具而不是试图统一整个 pipeline step。有人可能会问那 PP 的同步等待开销还在收益会不会大打折扣会但如果你是在 PP8 以上跑超大规模模型本来单 stage 的计算量就很大launch 开销占比被稀释Graph 的绝对收益有限。我见过一些项目组在 PP2 或 PP4 时依然能拿到 8% 到 15% 的 step time 下降但这种收益不稳定需要根据实际 profile 结果判断不要想当然。3.4 优化器 step 和梯度裁剪留在 Graph 外面更安全梯度裁剪和 Adam step 有一个共同特点它们依赖全局梯度范数或者跨 rank 的梯度统计。全局梯度范数需要先做一次 AllReduce拿到范数后 CPU 端才能决定是否缩放梯度、是否走更新分支。这种依赖 CPU 结果的控制流是 CUDA Graph 天然排斥的——capture 期间 CPU 必须“无脑”地录制 API 调用一旦中途要读一个 GPU 标量来决定走向整个 graph 就废了。我见过一种折中做法在 capture 之前算好范数然后把缩放后的梯度作为常量输入传进去让 Adam 更新也被包进 Graph。理论可行但工程上很脆尤其是混合精度场景下 master weights 与 fp32 gradient 的指针重定位、动态 loss scaling 的调整都会引入隐藏分支。更稳妥的做法是让 Graph 止步于梯度计算完成之后后面 AllReduce、clip、optimizer step 全留在动态路径上。这样 CPU launch 开销依然有一大块被消除了但保留了灵活性。这也是 Megatron-LM 团队在默认实现里更推荐的路径。4. 实测数据与形状约束一个规模翻倍就失效的方案很多博客讲 CUDA Graph 只讲机制不讲边界。这里我把我实际跑出来的一组对照数据放上来顺便也说说哪些场景下这套方案会迅速失效。4.1 一组可复现的基准测试测试环境8 卡 A100 80GMegatron-LM 的 GPT 模型参数量 7B张量并行 TP8数据并行 DP8不做 pipeline 并行。序列长度 2048global batch size 512即每个 rank 每个 step 用 64 个样本。混合精度 bf16。配置step time平均吞吐tokens/秒/卡GPU 利用率基线无 CUDA Graph2.9s约 36100约 78%开启 CUDA Graph模块级捕获2.2s约 47700约 93%开启 CUDA Graph但把全部模块包入2.3s约 45600约 91%第二行和第三行的对比说明一个现象全包裹并不总是最优。把嵌入层、输出层这些和通信耦合紧密的模块也硬塞进 Graph内存池压力变大偶尔会触发 cudaFree 或者额外的同步反而抵消了一部分收益。所以模块选择性捕获不是保守而是理性。GPU 利用率能从 78% 提到 93%主要就是因为 step 内部 kernel 间隙被 Graph 回放的批量调度填上了。如果你用nsys profile抓开启 Graph 后的 timeline会发现 kernel 序列几乎连成一片不再有一个个小缺口。4.2 形状固定规则训练里的“静态化约束”CUDA Graph 有一个硬规则录制的图和回放的图输入 tensor 的 shape、dtype、layout 必须完全一致。推理侧通常好满足因为 prompt 长度虽然变但很多服务会把请求 padding 到固定长度或者用 varlen API 让内核本身支持动态长度。训练侧不一样很多人用 padding 把输入序列对齐后以为万事大吉结果 rollout 阶段如果做 RL 训练或者 packing 数据出现时每个样本长度一变整个 Graph 匹配失败代码会退回到 eager 路径性能一夜回到解放前。所以实践中必须给训练数据链路加一道“静态形状护栏”固定 seq len固定 micro-batch 大小关闭任何形式的动态 padding。如果业务确实需要不同长度混合训练一个折中是准备多个 Graph每个对应一组长度 bucket比如 512、1024、2048 各录一张图样本按长度路由。这个方案我们在一个多模态混合训练任务里用过效果还可以代价是显存里多出两三张图的容量。另外要注意 FlashAttention 的兼容性。varlen 版 FlashAttention 在输入长度变化时内部 kernel 的 grid 配置会变这和 Graph 的完全固定语义天然冲突。建议在 capture 区域里使用固定 seq len 的 FlashAttention 路径或者干脆在捕获期间切换到标准的 attention 实现。4.3 为什么会“规模翻倍就失效”我们把同样的方案从 7B 挪到 13BTP8 不变结果 step time 只下降了 12% 左右而不是 25%。原因不难理解模型越大单层内计算量占比越高kernel launch 开销占比下降Graph 收益自然缩水。这是所有 launch 优化技术的共同规律带宽和计算量堆起来之后CPU 投递不再是瓶颈。反过来如果你跑的是小模型、高吞吐、短序列比如 1B 模型配极短的 prompt 做大量小 step 的训练Graph 的收益会非常夸张有时候能省掉 40% 的 wall time。这也可以作为选择模型的判断依据先用 profile 工具算出单 step 中 CPU launch 开销占比如果低于 10%Graph 大概率帮不上忙如果高于 25%那它应该是第一优先级优化手段。5. 比文档多一点我踩过的四个 Graph 坑这一节是全文最想留给大家的部分。文档里的 API 用法是死的真正导致项目返工的往往是那些文档里没写、但运行时会准时出现的问题。以下是我在 Megatron-LM 里给训练 step 套 Graph 时踩过的四个坑顺序按“踩坑时间线”排。5.1 没做 warmup 导致的“假收敛”第一次接入时我图省事直接 skip 前几步就跑 capture结果训练 loss 掉了两天都降不下来。后来定位到是 warmup 步数不足。为什么训练必须 warmup因为很多算子第一次调用时会有 lazy initializationcuBLAS 会分配 workspace、cuDNN 会找最优算法并缓存、NCCL 会建立通信 plan。这些初始化动作如果发生在 capture 期间会被当成正常的 kernel 录进 Graph于是每一次 replay 都会重复执行一遍“初始化”这些初始化 kernel 里的随机状态和指针分配行为完全无法预期轻则 loss 抖动重则显存不断膨胀。解决办法在开启 Graph capture 前先跑 2 到 5 个普通 step 作为 warmup让所有算子的缓存和 workspace 就位。Megatron 里--cudagraph-warmup-steps 3干的就是这件事。我自己后来会把 warmup 步数加到 5不是因为 3 不够而是为了确保 lazy 初始化的各种异步分支都执行过一遍。5.2 RNG 状态固化导致的 dropout mask 问题这是所有训练侧 Graph 用户都会撞到的经典坑dropout。PyTorch 的 dropout 在 eager 模式下每次迭代都会用当前的 RNG 状态生成一个新的 mask但 capture 时 RNG 状态只取一次replay 时会反复使用同一个 offset 生成的 mask。换句话说模型在 Graph 区内所有 dropout 层的 mask 从此固定正则化效果直接失效过拟合风险上升。我有一次就是没注意这个训练 loss 一直在降但验证集指标一路走坏。处理这个问题的正确姿势是把 dropout mask 的生成放到 capture 之外。比如在进入 Graph 前用当前 RNG 状态生成一组 mask然后把 mask 作为额外输入传给 Graph 内部的算子如果模型本身不允许这样改造就退而求其次把 dropout 层排除在 Graph 捕获模块之外让它们继续走 eager 路径的随机逻辑。这个取舍会牺牲少量收益但保住训练的统计学正确性是必要的。5.3 capture 区域里的 CPU 同步点性能反而变慢的元凶还有一种情况比 loss 跑飞更难排查Graph 开启后 step time 不降反升。我遇到过一回原因是模型里某个算子当时是 embedding 的 sparse 梯度处理内部有.sum().item()这样的 CPU 同步点。在 eager 模式下这种同步点虽然慢但不影响 kernel 执行流一旦被包进 Graphcapture 会强制把同步点前所有算子落盘等待然后 replay 时又重复等待一次形成 CPU-GPU 乒乓效应。排查方法很简单在 capture 区域代码里搜索.item()、.cpu()、float()、int()、.tolist()这类一切会产生 device-to-host 同步的调用逐个移除。Megatron 的 loss 计算通常发生在 Graph 外问题集中在自定义 loss 函数、自定义 metric 和部分 debug 打印代码上。建议把 capture 看成“录音棚”任何需要现场互动的行为都不允许发生。5.4 全局内存池配置和显存碎片capture 后的内存分配行为切换到 graph 私有池后如果某个模块的中间激活在每次 step 里大小都一样那就很稳但一旦某个 step 因为微批次切分产生了额外的临时 buffergraph 池内部就会出现碎片。碎片多了之后下次 capture 新 Graph 时可能申请不到连续显存直接 OOM。Megatron 里会试图通过全局 pool handle 来复用显存但实践中我发现一个更直接的策略控制 Graph 实例数量。--cudagraph-steps-per-module值越大显存占用越高在动态 shape 优化里它有用但静态 shape 场景下设成 1 最省心。另外如果发现开启 Graph 后跨 step 显存持续上涨优先怀疑是不是有模块被反复 capture 了而不是模型代码泄漏。6. 同一套思路的另一面把训练侧经验迁移到推理与 CV 训练最后再说点延伸。CUDA Graph 在 LLM 推理侧已经很成熟了很多部署框架默认就把解码循环包进 Graph。开发者在训练侧掌握的原理完全可以平移过去只是方向不同推理侧更关注动态 token 长度下的缓存管理训练侧更关注分布式通信与 RNG 语义。把两边的经验交叉起来能少走很多弯路。6.1 推理侧Graph 解决的是“记忆体”和“调度”双重问题推理时generate 循环每个 step 都调用一次模型kernel 数量少但 launch 密集Graph 把逻辑上完全相同的一串 kernel 固化下来后延迟能降一个量级。和训练不同的是推理侧常常要多张 Graph 覆盖不同 input length bucket这对显存比较敏感可以将训练侧的内存池复用思路带进来用一个全局 Cache 管理多个 Graph 实例而不是每个请求新建一个。6.2 CV 训练YOLO 这类任务用 Graph 要格外小心我也有同事拿这套思路去试 YOLO 系列模型训练。理论上可行但 CV 训练里常见的 mosaic、mixup 数据增强会导致每个 batch 的 target 数量和样本 shape 在不同 step 间变化这就直接撞上了 Graph 的静态形状约束。如果目标是给 CV 训练加速更值得投入的反而是 torch.compile 这类编译器方案因为它会根据每一步的实际输入重新生成优化后的计算图不需要手工维护形状 bucket。相较之下手工 CUDA Graph 更适合形状稳定、算子重复度高的 NLP 训练场景。6.3 和 torch.compile 的边界手工 Graph 不等于编译器很多人问有了 torch.compile 何必再手写 CUDA Graph这里要分清楚torch.compile 在modereduce-overhead下本身就会生成 CUDA Graph但它是编译器自动产出的且优化目标是“通用计算图”。Megatron 这种重度自定义分布式并行代码里很多通信和重计算逻辑是编译器无法完美建模的手工控制的模块级 Graph 反而更精准。如果项目里有大量自定义 CUDA kernel 或者奇怪的分布式 schedule我建议保留手工 Graph 的入口如果模型本身是标准 transformer、并行度不高直接上 torch.compile 会更省人力。我个人在实操中的体会是CUDA Graph 不是银弹它是一个需要“持续维护”的性能资产。每次改动网络结构、调整并行策略、增加新的通信算子都要重新审视哪些模块适合被包进 Graph、哪些应该移出。如果你能把这种审视变成常规性能检查的一部分它对 LLM 训练的收益是稳定且可观的。最后再分享一个小技巧在跑长训练任务前先开一版带 Graph 和不带 Graph 的各跑 50 step比较 loss 曲线是否一致一旦有细微偏离就不要急着让它上生产。训练正确性永远是第一位的性能优化只是锦上添花。