#51700 [2/2][Model Runner V2] FULL CUDA graph capture for microbatched steps (DBO)
原始 PR · 作者 specture724 · 合并时间 2026-09-14 13:42
MRV2 DBO 微批步骤接入 FULL CUDA 图捕获与回放
值得精读,尤其是三类设计决策:(a) `_maybe_ubatch_twin` 用与 DP 握手完全相同的阈值函数生成候选,把“可捕获”与“可分发”强行绑定,这是避免多 rank 不对称执行的关键;(b) `run()` 拆成 `begin_capturable_run()` / `finish()`,把线程创建、device 设置与 cuBLAS 预热留在图外,只让 handoff 与 join 进图,这是 DBO 能捕获成功的前提;(c) `_capture_stream` 强制捕获流与微批线程 compute stream 一致,是移植 V1 `gpu_ubatch_wrapper.py::_capture_ubatches` 时最容易漏掉的一点。阅读时请结合 #50945 一起看(本 PR 的 diff 建立在它之上),并重点关注 dp_utils.py 中 `2 * min(num_tokens_across_dp) < ubatch_desc.num_tokens` 这一守卫——它已被确认存在边界 bug,跟进补丁前不建议在生产 DP 负载高度不均的场景依赖 FULL 微批回放。此外建议跟踪 coderabbitai 提出的“永不被回放的孪生图”是否会形成显存浪费的后续处理。
参与讨论