前些天在给某个 PyTorch 深度学习模型应用 full iteration CUDA graph 时，碰到了一个随机出现的 bug。CUDA graph 是一个威力很大的性能调优手段，特别是在大模型训练的场景中，但用过的人都知道，能把 CUDA graph 用起来是相对比较困难的，特别是把整个 iteration 捕获为一张 CUDA graph。这个 bug 前前后后花了 3 天时间才解决，分享以下自己的 debug 经过。

  1. 症状是在捕获 CUDA graph 时会随机抛出异常，错误消息是 CUDA error: operation failed due to a previous error during capture，单 GPU 上出现概率比较低，但在 512 个 GH100 上出现概率大概是 1/3。相对比较难 debug，一是单 GPU 上难以复现，而集群上又不是交互式的，无法使用 cuda-gdb 这类调试工具，二是 CUDA graph 的错误消息比较隐晦。

  2. 因为 CUDA 的异步行为，最初报错的位置并不可信，使用 CUDA_LAUNCH_BLOCKING=1 调试，发现错误随机出现在开始捕获计算图之后的第一个或第二个 CUDA kernel 后。

  3. 怀疑存在某种形式的资源竞争，第一时间想到了 compute-sanitizer，但 compute-sanitizer 主要针对单个 kernel，用于整个 iteration 的场景下速度慢到令人发指。

  4. 查阅 CUDA Graph 文档，graph capture 时的错误 cudaErrorStreamCaptureInvalidated 表明有某种不可被捕获的操作，比如 cudaStreamSynchronize () 和 cudaMalloc()。

    我们非常确信我们的模型在这方面已经做到极致优化，在 capture graph 时不存在同步和内存分配。但保险起见，还是用 nsys profile --stats=true 检查了在出错前后的代码中涉及到的 CUDA API 调用，包含了 4 个 cudaLaunchKernel、8 个 cudaEventQuery、4 个 cudaEventDestroy、1 个 cuCtxSynchronize。这个 cuCtxSynchronize 来自 NSight System 本身，其余似乎并没有什么异常的 CUDA API 调用。

  5. 在 Slack 上搜索历史上 CUDA graph 团队帮忙解决的 cudaErrorStreamCaptureInvalidated 问题，注意到出现该错误时，原 CUDA API 也会报错，但此前并没注意到 PyTorch 在出错时有提供更多的细节。通过查看 PyTorch 的源码发现，想要让 AT_CUDA_CHECK 在出错时报告 C++ stack，需要设置环境变量 TORCH_SHOW_CPP_STACKTRACES=1。得到 C++ stack 后，检查了相关 PyTorch 代码，并无异常。

  6. 把错误消息发给了美国那边 PyTorch 和 CUDA Graph 团队的同事，次日从 PyTorch 同事那里得到意见，有可能问题发生在 graph capture 的中间再次进入 graph capture。为了验证是否成立，在捕获的代码前后通过 torch.cuda.is_current_stream_capturing() 检查是否已经处在捕获状态，但并未处在捕获状态。

  7. PyTorch 同事反馈如果其他进程上有 CUDA context，可能也会有潜在问题。修改 multiprocessing 中 BaseProcess 的代码，检查是否有有进程初始化了 CUDA context，但是并没有。

  8. CUDA Graph 团队同事反馈如果有其他线程在执行同步，就会打断 graph capture。然而我们的模型已经移除了所有的同步，再次确认后发现确实没有任何同步。

  9. 虽然 PyTorch 和 CUDA Graph 团队的反馈没有直接帮我们解决问题，但受此启发我开始怀疑可能有未注意到线程在主线程 graph capture 的同时在调用 CUDA API。

    用 NSight System profile 了 20 个 iteration，每个 GPU 上只有一个进程有 CUDA API 调用。但在 nsys 界面上发现除了 PyTorch 的正向传播和反向传播两个线程，还有其他几个线程出现在界面上，虽然几乎都是空白。仔细检查每个 iteration 中其他线程的行为，发现有 3 个线程在每个 iteration 都有一两个 cudaEventQuery()，这里要仔细去看，否则难以发现其他线程上的CUDA API调用，除此之外没有其他类型的 CUDA API 调用。

  10. 怀疑正是其他线程上的 cudaEventQuery() 导致主线程上 graph capture 失败，祭出自己在公司写的拦截 CUDA API 的工具，添加拦截 cudaEventQuery 和 cuStreamQuery 的功能，抓到完整 Python/C++ 的调用栈。

    查看调用栈，发现每个 iteration 有上千个 cudaEventQuery()，它们主要来自 3 个地方:

    • ncclCommWatchdog: 用于判断异步的 NCCL kernel 有没有完成;
    • CUDACachingAllocator: 申请 CUDA tensor 时;
    • CUDAHostAllocator: 申请 pinned memory 时;

    排除 ncclCommWatchdog，PyTorch DDP 针对 ProcessGroupNCCL 有个 ncclCommWatchdog，它执行 workCleanupLoop，定期查询 work 的完成状态并清理，其中涉及到 cudaEventQuery, 默认每 1000 毫秒执行一次，但每个 ProcessGroup 都有自己的 stream。

    最后问题定位到 CUDAHostAllocator，CUDACachingAllocator 和 CUDAHostAllocator 在 allocate memory 时会先检查 event list，所以 dataloader 的 worker thread 在执行 pin_memory 时会触发 cudaEventQuery。如果此时恰好正在 stream capturing 就会报错，而如果此时 dataloader 已经完成 pin_memory 则不会报错。

    因此，一个变通方案是在 graph capture 之前 sleep 1000 毫秒，确保其他线程上的 CUDA API 调用完成，这样就解决了 stream capturing 时的随机错误。

  11. 这个问题的本质是其他 thread 在 query 一个 record 在 default stream 上的 event，导致 stream capturing 出现随机失败。用最简短的代码重现问题，跟 PyTorch 和 CUDA graph 团队一起研究如何从根本上解决类似问题。

整整 debug 了 3 天，得到的教训是: 在 stream capturing 时，最好保证其他线程/进程都没有任何 CUDA API 调用，比如在 dataloader 启动 worker 之前进行 stream capturing。