1. 背景交代为什么偏偏是 CUDA Graph1.1 一个训练卡顿的现场复盘第一次意识到 CUDA Graph 在 LLM 训练里的价值是我在把 Megatron-LM 从单机几卡玩具规模向多节点千亿参数规模迁移的时候。当时模型 forward 和 backward 都能正常跑GPU 利用率从 nvidia-smi 上看还挺漂亮稳定在 80% 上下可实际吞吐只有预期的一半左右。最明显的一个现象是训练曲线是“一顿一顿”的日志里的 interval time 忽高忽低GPU 时间线一打开密密麻麻的全是小空隙。这种空隙最常见的原因是 CPU 调度跟不上 GPU 的执行速度。PyTorch 默认是 eager 模式每个算子都要经历“CPU 端构造 kernel 描述 → 写入 CUDA stream → GPU 端解析执行 → 返回”这一整条链路。模型规模一大算子数量动辄上千上万CPU 开销被无限放大。当 GPU 比 CPU 快的时候GPU 只能停下来等 CPU 发下一条指令。你说它是 CPU bound 吧cpu 占用好像也不高你说它是 GPU bound 吧SM 里面又明显有大量 idle。这种“两头都不满”的状态在长序列、大 batch 的 LLM 训练里特别常见。当时我试过很多粗暴的办法调大 batch size、把 PyTorch 的编译选项换掉、甚至手动 reduce 一下 Python 层面的开销。这些手段有效但都不够彻底。因为瓶颈不在某个算子里而在算子之间的“调度空档”里。后来换了思路干脆跳过 CPU 的逐算子调度把整套执行关系一次性提交给 GPU也就是 CUDA Graph。1.2 CUDA Graph 到底解决了训练里的什么问题CUDA Graph 做的事情从使用者的视角看很朴素把一串依赖明确的 CUDA kernel 执行关系捕获成一张图然后一次提交、整体回放。回放的时候 CPU 只发出一个很小的 launch 指令剩下的 kernel 调度顺序都在 GPU 侧预先编排好了。类比一下就是以前是 CPU 当导演每拍一个镜头喊一次“action”现在是把剧本里的所有分镜先录好导演只要按一次播放键整个流程自己跑。这里面最关键的一点是LLM 训练里的执行拓扑其实是高度稳定的。一个标准的 transformer layerforward 路径上就是 embedding、QKV 投影、attention、MLP、dropout、layer norm 这一套backward 路径则是对应的梯度计算。虽然参数值每个 step 都在变但算子的类型、数量、依赖关系几乎不变。这种“结构固定参数变化”的场景正好是 CUDA Graph 最舒服的适用对象。但要注意CUDA Graph 不是无代价的。它要求执行图里没有动态分支、没有 CPU 依赖、没有跨 stream 的不可控同步本质上要求你把训练循环里最“规律”的那段切下来固化成静态图。这也是为什么 CUDA Graph 在推理场景早就用烂了但在训练场景里落地还需要仔细设计因为训练里有随机性、有优化器状态更新、有分布式通信不能只把前向一包就完事。2. Megatron-LM 训练里的 CUDA Graph 落点分析2.1 Megatron-LM 的核心关注点不是单卡小模型而是大规模分布式先说清楚 Megatron-LM 是什么它是一套面向大规模语言模型预训练的分布式训练框架核心卖点是张量并行、流水线并行、序列并行以及配套的分布式优化器和 checkpoint 机制。它不像 HuggingFace 那样开箱即用跑个小模型而是假设你要在几十上百张 GPU 上训练一个几十亿到几千亿参数的模型。在这种框架下做 CUDA Graph 改造第一反应不能是“我只要把某个 layer 包一层 Graph 就行”而是得先想清楚图的捕获范围到底该划在哪里。Megatron-LM 的训练循环里每一步大概由这几部分组成数据加载和搬运、前向计算、反向计算、梯度 AllReduce 或 Pipeline 通信、优化器更新、日志采样。这里面前向和反向的计算部分是纯 GPU operator 的集合节点拓扑比较稳定适合进图。数据搬运动作涉及 CPU 和 GPU 之间的拷贝优化器更新涉及大量参数和动量状态的操作这些虽然也能在一定程度上静态化但参数规模和超参数组合经常变强行进图反而会引入新的维护成本。我个人的经验是默认把“单个 micro batch 的前向 反向”作为 CUDA Graph 的捕获单元优化器更新留在图外梯度同步根据并行模式单独处理。这个边界最干净也能解决训练吞吐里最头痛的“小算子调度”问题。2.2 训练阶段哪些环节值得捕获哪些不值得不是所有计算都适合进图。我踩过一段时间后发现区分标准其实很简单看这段计算是否具备三个特征——拓扑固定、shape 固定、没有 CPU 同步。拓扑固定指的是算子之间的依赖关系不会因为输入内容不同而改变。普通 dense attention 和 MLP 都没问题但如果你的模型里有动态路由、条件判断、稀疏模块那就不适合整体捕获。shape 固定指的是每个 tensor 的维度必须在 capture 前后一致。LLM 训练里最理想的组合是“固定 seq_len 固定 micro batch size”。如果跑的是 packed sequence 或者变长输入那么捕获时用的 mask、位置编码、batch 排布都得跟着变这会让图捕获变得非常痛苦。我一般采取的策略是训练数据统一 padding 到某个最大长度牺牲一点计算量换取整体的静态性。没有 CPU 同步指的是这段计算内部不能有其他线程或进程在等待某个 CPU 结果。比如某些 kernel 需要根据 tensor 数值来决定后续执行逻辑或者某个自定义 op 里混了 C 端的 lock这类都不适合进图。反向传播其实也天然满足这三个特征因为 autograd 引擎展开后的反向执行顺序在模型结构不变的前提下是确定的。所以我在 Megatron-LM 里常用的做法是把 forward 和 backward 一起捕获而不是只捕 forward。这样能顺手把 autograd 引擎的调度开销也消掉一部分。2.3 静态图与动态形状的冲突讲到 shape 固定这是 LLM 训练里最容易碰墙的地方。很多团队习惯用 dynamic padding、sample packing 来提高样本利用率或者用 gradient accumulation 时把不同 micro batch 的 shape 做成不完全一样。这种灵活性在 eager 模式下毫无问题但一旦启用 CUDA Graph任何 shape 变化都意味着必须重新捕获。一个比较实用的解法是分层设计外层保留动态 shape 和动态数据处理内层把“固定 batch size、固定 seq_len”的 transformer 计算包进图里。具体到代码上可以给每个不同的 bucket shape 保存一张图运行时按实际输入 shape 选择对应的图来 replay。比如训练数据集中在 [2048, 4096, 8192] 三个序列长度附近那就分别捕获三张图动态切换。这个方案看起来简单但要注意 GPU 显存开销。图捕获会固定住它捕获期间分配的所有显存多张图就意味着多份显存预留。我实际测试下来每张额外的图大概会多占几百 MB 到几个 GB 不等具体要看模型规模和激活显存。所以在设计 bucket 数量时不能太贪最好是两到三个典型 shape别想着为每一种样本长度都配一张图。3. 在 Megatron-LM 里接入 CUDA Graph 的实操过程3.1 环境与基准我的测试环境是一台 8 卡 A100 80G 节点PyTorch 2.1CUDA 12.1Megatron-LM 版本基于 core 分支数据并行 梯度累积暂不开启流水线并行。模型规模大概 13B隐藏层 512032 层序列长度 4096micro batch size 4。启动参数大概是--tensor-model-parallel-size 1 \ --pipeline-model-parallel-size 1 \ --num-layers 32 \ --hidden-size 5120 \ --num-attention-heads 40 \ --seq-length 4096 \ --micro-batch-size 4 \ --global-batch-size 128 \ --bf16先跑通基准记录 baseline 的吞吐和 GPU kernel 时间线。我的 baseline 大约在 3800 tokens/s 左右nvidia-smi 显示利用率波动明显许多 kernel 的 launch 间隔在 5 到 20 微秒。这个数字看起来很吓人但在大模型里很常见因为大量小算子比如 norm、elementwise、reshape 的开销都是 launch bound 的。3.2 捕获前需要做的准备工作不要直接上来就开 capture否则会撞见一堆内存分配和 stream 问题。我总结下来有三个前置步骤必须做足第一预热显存分配器。CUDA Graph 捕获时不允许从默认缓存分配器里新分配显存因为图捕获期间“申请显存”这个动作本身也会成为图里的节点但显存地址后续可能发生变化这会产生不可预期的结果。PyTorch 的torch.cuda.graph上下文管理器会自动切换到一个 graph pool但前提是捕获前所有需要的显存块已经分配好了。我通常会先跑两到三个真实 micro batch 的 forwardbackward让分配器把常用 tensor 的显存块都预分配一遍。第二固定输入缓冲区。把每个 micro batch 的 input_ids、position_ids、attention_mask 放在预先分配好的 pinned memory 或固定显存 buffer 里。每次训练迭代时往这个 buffer 里拷贝新的数据然后触发 graph replay。第三处理好随机性。图捕获时 dropout mask 会被固定下来如果你不复现随机状态那么每次 replay 都会用同样的 mask。我的建议是若你的模型对 dropout 很敏感要么每 N 步重新捕获一次图要么在 graph 外部单独生成一组随机数塞进去。对大多数预训练任务来说dropout mask 固定几个 step 影响很小但如果你做的是需要精细随机性的小模型实验这个点就必须额外注意。3.3 核心代码改造点下面给出一个简化但可用的伪代码框架展示如何把 Megatron-LM 的 forwardbackward 包进 CUDA Graphimport torch # 假定 megatron 已经构建好了 model、optimizer、dataloader # model 里包含 embedding 所有 transformer layer 输出层 # 这里只展示图捕获的核心思路 def warmup_and_capture(model, static_input, static_label, steps3): # 预热 for _ in range(steps): out model(static_input) loss loss_fn(out, static_label) loss.backward() # 清空梯度准备捕获 model.zero_grad(set_to_noneTrue) # 捕获 forward backward g torch.cuda.CUDAGraph() with torch.cuda.graph(g): out model(static_input) loss loss_fn(out, static_label) loss.backward() return g, loss # 每次训练迭代 def train_step(g, static_input, real_input, real_label): # 拷贝新数据到静态输入 buffer static_input.copy_(real_input) static_label.copy_(real_label) # 回放图 g.replay() # 这里的 loss 值可以读取但需要先 synchronize 再取 torch.cuda.current_stream().synchronize() return current_loss这里有个容易被忽略的细节loss.backward()在 Megatron-LM 里默认会把梯度写到 parameter.grad 上。如果你开启了 gradient accumulation就要把每个 micro batch 的梯度累加做到 graph 内部或者在 graph 外部用一个显式的 buffer 去累加。我推荐后者因为更容易控制显存和数值精度。另外torch.cuda.graph(g)这个上下文管理器会自动处理 stream 切换它会从当前 stream 开始捕获内部再做一次 stream capture。捕获完成后replay 是在当前 stream 上执行的。如果你在分布式训练里注意不要让 NCCL 的 collective 出现在 capture 区域内这个后面会专门讲。3.4 与分布式并行组件的配合这才是 Megatron-LM 接入 CUDA Graph 最麻烦的地方。Megatron-LM 默认会有 tensor parallel、pipeline parallel 和 data parallel这些并行模式里都涉及跨卡通信操作。先说最简单的情况——纯数据并行tensor parallel size 和 pipeline parallel size 都是 1。这种情况最友好前向和反向计算完全在单卡内完成图里没有跨卡通信。反向结束后会有一个梯度 AllReduce这个操作放在图外执行。流程就变成graph replay 计算出梯度 → 等所有卡都算完 → 梯度 AllReduce → optimizer.step()。这样一来图的捕获区域非常干净replay 开销也很低。如果你的配置里有 tensor parallel情况就复杂了。tensor parallel 的 AllReduce 发生在 transformer 层的内部比如 attention 输出的归约、MLP 的归约这些操作如果进入 capture 区域需要 NCCL 支持 stream capture。较新版本的 NCCL 对 CUDA Graph 有支持但实测中还是容易踩到“NCCL 操作无法在 stream capture 中启动”这类错误。我的建议是如果你的主目标是稳妥提升吞吐可以先不考虑 TPgraph 的激进组合如果一定要用就升级到最新版本 PyTorch 和 NCCL并且把通信操作放在主 stream 之外的独立 stream 上预先定义再把图捕获范围尽量缩小到单卡计算部分。流水线并行同理。pipeline 的 micro batch 调度天然带有动态性每一 stage 的启动顺序会随 micro batch 数量变化强行整段捕获收益不明显反而容易引入故障。更合理的做法是只对“单 stage、单 micro batch”的局部计算做图捕获而不是把整个 pipeline schedule 都套进去。4. 我踩过的坑和排查思路4.1 内存拷贝导致的“假加速”第一次跑通图回放后我满心欢喜去对比吞吐结果发现只提升了不到 3%而且显存涨了 8 个多 G。后来查下来才发现问题出在“输入数据拷贝”这一步。我用的是input_tensor.cuda(non_blockingTrue)每次动态搬运这一搬就把 CPU 与 GPU 之间的传输延迟又加回来了而且频繁的小拷贝会打断 graph replay 的连续性。解决办法很简单把输入统一放到预先分配的固定显存 buffer 里训练迭代时用static_input.copy_(real_input)覆盖再附带一个必要的synchronize。这样虽然也要花时间但开销比重新走一遍数据传输和 kernel 调度小很多。另外要注意数据从 CPU 搬运到 Host Pin Memory 的动作最好提前做和 GPU 计算重叠。4.2 CUDA error: operation not permitted during stream capture这是新手最容易碰见的报错。意思是说在捕获期间有某个操作尝试从 CPU 端启动新的 CUDA 任务但 stream capture 状态下不允许这么做。常见的触发点包括tensor.item()、print(tensor)、动态维度计算、某些 Python 层面的 assert以及数据加载器里隐式的.cpu()操作。排查方法没有捷径只能把 capture 区域里的代码一行行过凡是可能把 tensor 转成 Python 标量的操作全部挪到图外。我还会在捕获前用torch.cuda.set_sync_debug_mode(1)开启同步调试它会在捕获失败时给出更详细的调用栈。4.3 图捕获完以后梯度不对另一个让我头疼的问题是图捕获成功后replay 出来的梯度数值和 eager 模式对不上。查了半天发现是 autograd 的 gradient buffer 没有提前分配好图捕获时把“参数梯度地址”也固化进去了。如果后续 optimizer.step() 之后梯度 buffer 被原地更新那还没问题但一旦梯度 buffer 被释放或重新分配graph replay 就会往一个已经失效的地址上写。解决方案是在捕获前确保所有参数的.grad都已经拥有稳定的存储并且用model.zero_grad(set_to_noneTrue)不行得用model.zero_grad()这样参数梯度 buffer 会被固定住。或者干脆手动给每个参数分配一个grad_buffer并把 backward 的梯度通过torch.autograd.backward(..., grad_tensors...)导向这个固定 buffer。这个细节很多人不说但我实际调试了很久。4.4 常见问题速查表问题现象可能原因处理建议捕获时报 “operation not permitted during stream capture”捕获区域内混入 CPU 同步操作把.item()、.cpu()、动态分支挪出图吞吐提升不明显输入拷贝或数据传输占了主要开销用固定显存 buffer避免非阻塞拷贝打断 replay显存增加过多图捕获固定了整个捕获期的显存分配预热分配器、合理限制 bucket 数量梯度数值与 eager 不一致梯度 buffer 在捕获后地址变化捕获前固定梯度 buffer并保持参数 grad 稳定多个图切换卡顿bucket 数量过多或显存碎片只保留 2-3 个典型 shape切换前同步NCCL 相关报错图捕获区域内含跨卡通信升级 PyTorch/NCCL或把通信排除在图外4.5 NCCL 和 CUDA Graph 的兼容性实战我不建议在一开始就把 NCCL collective 放进 capture 区域。虽然从原理上讲只要 NCCL 内核能正确注册进图它也可以随图回放但 NCCL 内部有很多基于环境变量和连接池的状态在多机多卡环境下稍有变化就可能爆出各种奇怪的 host-side hang 或 “invalid usage of NCCL group”。我自己的稳定组合是DP 梯度 AllReduce 放图外TP/PP 通信不放图内。如果你的场景必须使用 TP可以在 Megatron-LM 里先把--tensor-model-parallel-size调成 1只验证 DP 场景的收益。等确认图捕获机制没问题后再去逐一尝试打开 TP 并评估稳定性。好在很多用户场景里纯 DP 加超大单卡模型已经能吃到 CUDA Graph 的大部分红利。5. 效果与体会5.1 我实测下来的一组数据在 13B 模型、单节点 8 卡纯 DP 配置下我把 forwardbackward 整体做成 CUDA Graph 后吞吐从约 3800 tokens/s 提升到了约 4400 tokens/s提升在 15% 左右。GPU 利用率从波动不定的 80% 出头稳定到了接近 96%。这个提升幅度和模型规模关系很大模型越小算子越碎CUDA Graph 的帮助也越明显模型大到一定程度后单个算子本身耗时很长调度开销占比下降收益可能会回落到 5%-10%。显存方面因为固定了输入 buffer 和图的完整分配峰值显存增加了约 4 到 6 个 G。对于 80G 的 A100 来说还能接受但如果你用的是 40G 甚至更小显存的卡就要提前算好余量。开启梯度累积后显存压力会更大因为多个 micro batch 的激活值空间都被图固定了不能像 eager 模式那样灵活释放。5.2 什么时候别用 CUDA Graph我也要泼一点冷水。不是所有 LLM 训练场景都应该上 CUDA Graph。如果你在小模型或瓶颈根本不在调度上收益会非常有限。比如模型层数少、单算子耗时长、或训练时大量时间花在数据读取上那么即使图捕获成功整体吞吐也几乎不变但你会多承担捕获调试和显存固定的成本。如果你的训练流程里有频繁改结构的需求比如动态调整层数、插入自定义分支、经常改 attention mask 的类型那就别急着用图。图捕获本身是一次性开销不小的操作经常重建图会反倒拖慢训练。交给编译器类技术如torch.compile可能更省事因为它在底层会主动做类似 kernel fusion同时保留了一部分动态性。如果你的形状经常变比如不同批次序列长度差异很大那也要非常谨慎。除非你愿意做 padding 和 bucket 方案否则 CUDA Graph 的“静态 shape”限制会把灵活的训练流程逼成一场噩梦。5.3 一点个人经验我最后想说的是CUDA Graph 不是一个能让你“一键起飞”的银弹它更像是一块需要嵌进训练流水线里的定制零件。理解它的核心价值不是为了炫技而是为了知道 GPU 在等什么、CPU 在忙什么。在 Megatron-LM 这种大型分布式训练框架里每个看似微小的调度优化放大到成千上万个 step 和多卡并行之后都会变成实实在在的电力成本和时间成本。我现在的习惯做法是每个大模型实验开始前先用 profiler 采一段训练时间线如果发现 kernel launch 间隙占整体时间超过 15%我就认真考虑 CUDA Graph如果间隙很小我就不折腾了。这套判断标准陪我避开了很多次“折腾半天没提升”的尴尬局面。你如果准备在 Megatron-LM 里做类似改造建议也先把自己的 bottleneck 量化清楚再决定要不要上这张“静态图”。