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, на предмет ошибок синхронизации. Все обнаруженные гонки данных будут выведены в стандартный поток ошибок вместе со стеками вызовов предполагаемых причин. Для наилучших результатов саннитайзер следует включить в самом начале программы.