CUDA Потоковый Сантайзер

Обзор

Этот модуль представляет Потоковый Сантайзер CUDA, инструмент для обнаружения ошибок синхронизации между ядрами, запущенными в разных потоках. Он сохраняет информацию о доступах к тензорам, чтобы определить, синхронизированы ли они или нет. При включении в программе 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 detected a possible data race on tensor with data pointer 139719969079296
Access by stream 94646435460352 during kernel:
aten::mul.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)
writing to argument(s) self, out, and to the output
With stack trace:
  File "example_error.py", line 6, in <module>
    torch.mul(a, 5, out=a)
  ...
  File "pytorch/torch/cuda/_sanitizer.py", line 364, in _handle_kernel_launch
    stack_trace = traceback.StackSummary.extract(

Previous access by stream 0 during kernel:
aten::rand(int[] size, *, int? dtype=None, Device? device=None) -> Tensor
writing to the output
With stack trace:
  File "example_error.py", line 3, in <module>
    a = torch.rand(10000, device="cuda")
  ...
  File "pytorch/torch/cuda/_sanitizer.py", line 364, in _handle_kernel_launch
    stack_trace = traceback.StackSummary.extract(

Tensor was allocated with stack trace:
  File "example_error.py", line 3, in <module>
    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")

• Ошибочные доступы были вызваны операторами

  • a = torch.rand(10000, device="cuda") в потоке 0
  • torch.mul(a, 5, out=a) в потоке 94646435460352

• Сообщение об ошибке также отображает схемы вызываемых операторов, а также примечание, показывающее, какие аргументы операторов соответствуют затронутому тензору.

  • В примере видно, что тензор a соответствует аргументам self, out и значению output вызываемого оператора torch.mul.

Ошибка может быть исправлена, заставив новый поток дождаться стандартного потока:

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() [source]

Включает Потоковый Сантайзер CUDA.

Сантайзер начнёт анализировать низкоуровневые вызовы CUDA, вызванные функциями torch, на предмет ошибок синхронизации. Все обнаруженные гонки данных будут выводиться в стандартный поток ошибок вместе с трассировками стека предполагаемых причин. Для наилучших результатов сантайзер следует включать в самом начале программы.