Shortcuts

CUDA 流清理器

注意

这是一个原型功能,意味着它处于早期阶段,用于反馈和测试,其组件可能会发生变化。

概述

本模块介绍CUDA Sanitizer,这是一个用于检测在不同流上运行的内核之间的同步错误的工具。

它存储有关张量访问的信息,以确定它们是否同步。当在Python程序中启用并且检测到可能的数据竞争时,将打印详细的警告并退出程序。

可以通过导入此模块并调用 enable_cuda_sanitizer() 或通过导出 TORCH_CUDA_SANITIZER 环境变量来启用。

用法

这是一个PyTorch中简单同步错误的示例:

import torch

a = torch.rand(4, 2, device="cuda")

with torch.cuda.stream(torch.cuda.Stream()):
    torch.mul(a, 5, out=a)

张量 a 在默认流上初始化,并且没有任何同步方法,在新流上进行了修改。这两个内核将在同一个张量上并发运行,这可能导致第二个内核在第一个内核能够写入数据之前读取未初始化的数据,或者第一个内核可能会覆盖第二个内核的部分结果。当在命令行上运行此脚本时:

TORCH_CUDA_SANITIZER=1 python example_error.py

以下输出由CSAN打印:

============================
CSAN 检测到在数据指针为 139719969079296 的张量上可能存在数据竞争
由流 94646435460352 在核函数期间访问:
aten::mul.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)
写入参数 self, out 和输出
堆栈跟踪:
  File "example_error.py", line 6, in 
    torch.mul(a, 5, out=a)
  ...
  File "pytorch/torch/cuda/_sanitizer.py", line 364, in _handle_kernel_launch
    stack_trace = traceback.StackSummary.extract(

之前的访问由流 0 在核函数期间进行:
aten::rand(int[] size, *, int? dtype=None, Device? device=None) -> Tensor
写入输出
堆栈跟踪:
  File "example_error.py", line 3, in 
    a = torch.rand(10000, device="cuda")
  ...
  File "pytorch/torch/cuda/_sanitizer.py", line 364, in _handle_kernel_launch
    stack_trace = traceback.StackSummary.extract(

张量分配时的堆栈跟踪:
  File "example_error.py", line 3, in 
    a = torch.rand(10000, device="cuda")
  ...
  File "pytorch/torch/cuda/_sanitizer.py", line 420, in _handle_memory_allocation
    traceback.StackSummary.extract(

这提供了对错误来源的深入了解:

  • 张量被错误地从以下流中访问:0(默认流)和94646435460352(新流)

  • 张量是通过调用 a = torch.rand(10000, device="cuda") 分配的

  • The faulty accesses were caused by operators
    • a = torch.rand(10000, device="cuda") 在流 0 上

    • torch.mul(a, 5, out=a) 在流 94646435460352

  • 错误消息还会显示被调用操作符的模式,并附带一个说明,显示操作符的哪些参数对应于受影响的张量。

    • 在示例中,可以看到张量 a 对应于参数 selfout 以及调用的操作符 torch.muloutput 值。

另请参阅

支持的 torch 运算符及其模式的列表可以查看 这里

可以通过强制新流等待默认流来修复该错误:

with torch.cuda.stream(torch.cuda.Stream()):
    torch.cuda.current_stream().wait_stream(torch.cuda.default_stream())
    torch.mul(a, 5, out=a)

当脚本再次运行时,没有报告任何错误。

API参考

torch.cuda._sanitizer.enable_cuda_sanitizer()[源代码]

启用CUDA Sanitizer。

清理器将开始分析由torch函数调用的低级CUDA调用中的同步错误。所有发现的数据竞争将被打印到标准错误输出,并附带可疑原因的堆栈跟踪。为了获得最佳结果,清理器应在程序的开始时启用。