CUDA graph in the collector: from full-capture to update-state-only, and what we learned

[TATP-233]

Collector CUDA Graph 复盘：从 Full Capture 到 Update-State-Only，以及为什么停下来

概要

这篇 discussion 回顾 collector CUDA graph 这条线的工作（#2307 → #2311 → #2338 → #2350/#2351）。它是一份工程复盘，而不是把 graph 模式推向生产的决策。

结论可以概括为四点：

• Full collector graph 是错误的边界。 它试图把推理相关的控制流和 reset 编排一起捕获，而这些部分本质上由 host 拥有。
• Update-state-only graph 是正确的边界，但它仍然不是一个通用优化。 它捕获的是 Manager kernel 序列，但前提是移除或绕过了真实任务需要的大部分动态语义。
• 测到的性能是上界，不是生产语义下的收益。 graph 组需要关闭 corruption、关闭 term noise、绕过 RNG、启用 deferred validation，并绕过 allocator proof。
• 剩余成本是结构性的，不是一次性的。 graph 安全性会变成每一个 reward / observation / command term 作者的永久负担，而换来的是 2%–19% 且依赖 workload 的收益。

因此当前结论是：保留其中协议层面正确的部分（backend preflight、owner boundary classification、receipt 修复），继续推荐 MPS 作为共享执行方式，但不把 update-state graph 作为默认能力或通用能力推广。

起点：为什么要做

早期 Phase A profiling 在标准的 MJWarp / FlashSAC / G1 motion workload 上发现，collector 侧可处理的 overhead 大约为：

• 独立 context：约 19.5%
• MPS 下：约 32.7%

主要集中在：

• Manager update 和 episode reset 的 Python dispatch
• 热路径上的 finite / shape validation reduction
• 动态 reset row compaction
• reset 路径上的临时分配

MJWarp 本身已经有固定 shape 的 physics step/reset graph，但这些 graph 不覆盖 Manager-based observation、reward、termination、command、packed reset 和 transition publication。于是自然出现一个问题：Manager update body 能否也用同样的方式捕获。

第一次尝试：Full Collector Graph

第一种设计尝试把 collector tick 里尽可能多的部分都放进 graph：

• action / control 读取
• backend physics
• Manager update
• reset 处理
• transition publication

这条路失败的原因是结构性的，而不是实现细节没写好。

1. Collector tick 不是一个封闭的 device function

K=1 off-policy tick 有一个无法回避的依赖：

observation[t] -> learner action[t] -> env.step(action[t]) -> observation[t+1]

Policy inference 和 action publication 都在 env 外。任何把它们包含进 capture 的边界，都必须在 graph 内解决 inference IPC 的顺序关系。这会立刻把 graph 和 uni_rl 的 transport 细节耦合在一起。

2. Reset 是 mask-driven 且由 host 编排

Reset 需要 fixed-width mask、owner transaction、scene 写入和 event application。mask 本身是合法的 device state，但围绕它的编排是 host lifecycle。把它放进 graph，要么把 mask 语义冻结在 capture 时刻，要么引入 CUDA graph 无法表达的数据依赖分支。

3. Validation 和 diagnostics 天然希望对 host 可见

Finite / shape check 被写成 reduction 后，很容易变成 host bool read。在 eager 路径里，这只是性能问题；在 capture 里，它要么变成非法同步，要么变成被冻结的分支。graph-safe 的做法——persistent device verdict、deferred publication、只在失败边界做 host diagnostics——方向是正确的，但意味着必须先改变所有 manager 的 validation contract，然后才能评估 graph 本身。

4. 边界随着真实任务不断变大

每换一个 workload，就会暴露出一条新的动态路径：reset row compaction、event mode、command resample、recorder term、curriculum transition。每一条路径都必须被安全 capture、移到 graph 外，或者被 preflight 拒绝。这个设计最终收敛到的状态是：“capture 这个特定任务恰好不需要的那部分”。

第二次尝试：Update-State-Only Graph

第二种设计把边界收窄到真正属于固定 shape tensor program 的部分：

advance episode steps
compute termination / reset flags
compute rewards
compute commands
compute observations
map observations
publish into persistent carriers

Physics stepping、reset 编排、inference 和 transition publication 都留在 graph 外。这是正确的边界，也是最终成功 capture 的版本。

这个阶段产出了应该保留的东西

即便最终不推广 graph 模式，下面这些工作仍然有价值：

• 为 update boundary 建立了独立的 public backend preflight，与 backend physics graph 是否开启无关。physics graph READY 并不能证明 Manager update body 可 capture。

• 建立了 owner preflight，对不支持的动态语义 fail closed，而不是尝试 capture 后在运行时崩掉。

• 建立了 persistent carrier contract：update 输入和输出的 address、shape、dtype、device 都固定。

• 建立了 stream / completion ordering contract：

  backend step receipt
  -> update stream wait
  -> update graph replay
  -> independent update completion receipt
  

• 引入了 deferred finite verdict：validation 结果保存在 device carrier 中，在明确的 completion boundary 再发布。

• 对 curriculum term 做了 reset-owned classification：global / step-owned curriculum 变化留在 graph 外；只有 reset-owned boundary 可以和 graph body 共存。

• 在 warmup、capture、replay 周围加入了 allocator invariant check。

• run manifest 中加入了 capture identity、state、replay counter、disable reason、RNG state / draw version、carrier layout 等证据。

这份清单是整个尝试中最重要的沉淀。即使最终性能答案是零，这些工作本身也是有意义的。

但为了跑通，仍然需要禁用大量内容

成功 capture 的那次运行并不是生产语义。为了让 motion tracking 达到 graph-ready，测试路径需要：

• 关闭 observation corruption
• 关闭 term noise
• 拒绝 interval event
• 使用 RNG shortcut，包括一个只在 lower == upper 时成立的 command resample 路径
• 通过环境变量 bridge 强制启用 deferred validation
• 绕过 allocator proof

motion tracking 本身的 resampling_time_range: [1e9, 1e9] 也接近“几乎从不 resample”。因此这次 graph run 实际上是一个接近静态的 kernel 序列。

它证明的是：CUDA graph replay 可以加速一个静态化的 Manager body。
它没有证明的是：按生产配置运行时，motion tracking 任务可以被 graph 模式加速。

这个区别比任何单个性能数字都重要。

性能对比实际说明了什么

2×2×2 训练路径对比覆盖了：

• FlashSAC / G1 motion tracking
• SAC / G1 walk flat
• independent context 与 MPS
• disabled / eager / graph 三种模式

测试使用真实训练参数，只缩短 iterations，并且临时禁用了阻塞 graph 的 item。

Workload	Sharing	Graph vs eager
FlashSAC motion	independent	+16.1%
FlashSAC motion	MPS	+19.2%
SAC walk	independent	+8.8%
SAC walk	MPS	+2.3%

MPS 本身是更大、更普适的收益：

• FlashSAC eager：MPS 约 +17%
• SAC eager：MPS 约 +32%

怎么解读

• Graph 收益和被捕获的 update body 在整个 iteration 中所占的时间比例强相关。FlashSAC 的 learner update 更少，collector body 占比更高；SAC walk 的 updates_per_step: 8 使 learner work 占主导，Amdahl 定律大幅压低了 graph 的上限。
• MPS 已经消除了部分 serialization / context overhead，留给 graph 恢复的空间更小。SAC + MPS 下 graph 基本接近中性。
• 当前 graph 数字是 上界。恢复 noise、RNG 和 interval event 后，两个组的 kernel workload 都会增加，相对收益会下降。SAC + MPS 下 2.3% 的优势很可能会完全消失。
• MPS 不需要对 term 作者提出任何约束，也不需要维护 eager / graph 等价性。作为默认建议，它优于 graph 模式。

真正的阻塞：开发成本，而不是 capture 报错

最终决定性的问题不是 capture 失败，而是为了让 capture 通过，未来需要持续付出什么。

如果 update-state graph 变成一个受支持的通用路径，那么每一个 reward、observation、command term 都需要回答：

• 这个实现会不会分配内存？
• 会不会在 Python 里根据 device tensor 分支？
• 所有输出是否写入 persistent carrier？
• 所有 reduction 是否都被推迟而不是立刻变成 host-visible？
• 在固定 seed 下，这个实现是否 eager / graph 等价？
• 有没有对应的等价性测试？
• 新增这个 term 会不会使已有 capture 失效？

这意味着一个二十行的 reward term，会变成二十行实现，加上一份永久的 graph-compatibility 义务。它还会持续诱导出一些数学上并不正确、只是为了满足 capture 的 shortcut。lower == upper 的 resample 分支正是这一类改动，后来也被正确回退了。

这个成本不是一次性的。它会随着每一个新 term、每一次 manager 改动、每一个新任务族重复出现。而换来的收益是恢复生产语义前的 2%–19%，在 MPS SAC 场景下接近零。

这不是一个好的默认基础设施 trade。

保留什么

以下内容与 graph 是否推广无关，应当保留：

1. Replay IPC 路径中的 receipt operation 处理。 consumer 必须按 receipt 的实际 operation wait，而不是假设它是 step。这是一个真实的 baseline bug，eager 路径同样会命中。
2. Replay sync worker 的 first-error propagation。 没有它，原始错误会被后续的 non-contiguous ingress 次生错误遮蔽。
3. Backend-first public preflight。 capability negotiation 应该位于 SimBackend 上；env 代码不应探测 backend 私有状态。
4. Owner preflight 与 boundary classification。 尤其是 reset-owned curriculum 分类和对不支持动态 term 的 fail-closed 拒绝。
5. Persistent carrier 与 deferred verdict 基础设施。 即使在 eager 执行中，它们也能减少 host sync 和临时分配。
6. MPS 作为推荐的共享执行方式。 对支持的 CUDA workload，它是更大、更普适、且没有约束的收益。

不继续携带什么

• 为了避免 allocation 而引入的 out= graph-only 机制
• lower == upper 的 command resample shortcut
• 临时的 corruption / noise / RNG 禁用
• allocator proof bypass
• 从 training 通过环境变量进入 env graph mode 的 bridge
• 任何“graph 模式是 Manager-Based 任务通用优化”的暗示

经验教训

1. 边界应该从 ownership 推导，而不是从 profiling 百分比推导

profiling 能告诉我们 update body 很热，但真正可 capture 的边界，是输入输出有单一 owner 且 shape 稳定的那部分。决定 capture 可行性的是 ownership，不是 kernel 数量。

2. 一个需要禁用语义才能完成的 capture，还没有被真正评估

每一个被禁用的 feature 都应该视为一个未测试的交互，而不是简化条件。当这个列表包含 noise、RNG、event 和 validation 时，benchmark 测量的是合成 workload。

3. Fail-closed preflight 值得坚持

preflight 把模糊的 “graph 坏了” 变成了明确原因：不支持的 event mode、非 reset-owned curriculum、active recorder、carrier layout 不稳定。即使最终决定不使用 graph，这种清晰度仍然有价值。

4. 错误传播是实验的一部分

大量诊断时间花在了一个由 daemon thread 吞掉第一错误后引发的次生错误上。根因只是 IPC receipt 的 operation mismatch。如果第一错误能直接暴露，整个调查会大幅缩短。

5. 评估维护模型，而不只是吞吐

性能数字本身不是验收标准。真正的问题是：未来的每一个 term 作者现在必须保证什么？这个永久约束是否值得？在这里，答案是否定的。

6. 优先选择没有作者契约的优化

MPS 在不改变 term 写法的情况下带来了更大、更普适的收益。当两个优化重叠时，没有语义契约的那个应该优先成为默认。

7. workload 依赖应该被预期，而不是当作意外

CUDA graph replay 作为机制是通用的，但端到端收益遵循 Amdahl 定律。被捕获 body 在 iteration 中所占比例、learner update 比例、host orchestration 的多少，共同决定收益上限。单一任务的 benchmark 不够。

当前建议

• update-state graph 保持 experimental、opt-in，不作为默认，也不作为通用建议。
• 当前数字只视为上界证据。
• collector / learner sharing 优先使用 MPS。
• 不要求 Manager term 作者写 graph-safe 代码，除非未来某个任务族在完整生产语义下证明了足够大且持续的收益。
• 如果重新评估 graph 模式，至少要求：
  1. 恢复 allocator proof；
  2. 恢复 noise、RNG、event 语义；
  3. 增加 fixed-seed eager / graph 等价性测试；
  4. 建立 certified term subset 和显式 opt-in contract，保证 subset 之外的 term 不受影响；
  5. 生产配置入口走 owner YAML，而不是环境变量；
  6. 在完整任务语义恢复后重新测量性能。