Санитайзер потоков 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 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() [исходный код]

Включает CUDA Sanitizer.

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

torch.cuda._sanitizer.zip_arguments(schema, args, kwargs) [исходный код]

Тип возвращаемого значения:

Iterator[tuple[Argument, Any]]

torch.cuda._sanitizer.zip_by_key(a, b) [исходный код]

Тип возвращаемого значения:

Iterator[tuple[TK, TVa, TVb]]