torch.cuda

Генератор случайных чисел

Коллективные операции обмена данными

Потоки и события

Графы (бета-версия)

Аннотации ядер графа (прототип)

torch.cuda.graph_annotations добавляет к ядрам, захваченным в графе CUDA, пользовательские метаданные с ключами, позволяющими сопоставлять аннотации с полем graph node id событий ядер в трассировках профилировщика. Включите запись для каждого захвата с помощью аргумента enable_annotations функции torch.cuda.graph, а затем оберните нужные области захваченной рабочей нагрузки в блоки mark_kernels().

Для этих API требуется пакет cuda-bindings и драйвер CUDA с поддержкой cudaGraphNodeGetToolsId (CUDA 13.1 или новее либо совместимый пакет cuda-compat); в противном случае запись без предупреждения заменяется пустой операцией. ROCm не поддерживается. Используйте is_available(), чтобы проверить поддержку.

Полный процесс: добавьте аннотации во время захвата, профилируйте повторное воспроизведение, затем объедините аннотации с экспортированной трассировкой и просмотрите ее в Perfetto. Во время захвата:

import torch
from torch.cuda.graph_annotations import mark_kernels, get_kernel_annotations

x = torch.randn(1024, 1024, device="cuda")

# Warmup: run the workload once outside capture so lazy initialization
# (e.g. cuBLAS handles) does not end up in -- or invalidate -- the capture.
y = x @ x.t()
z = torch.relu(y) @ x
torch.cuda.synchronize()

g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g, enable_annotations=True):
    with mark_kernels("attention"):
        y = x @ x.t()
    with mark_kernels({"name": "mlp", "layer": 3}):
        z = torch.relu(y) @ x

with torch.profiler.profile() as prof:
    g.replay()
    torch.cuda.synchronize()
prof.export_chrome_trace("trace.json")

Каждое событие ядра в экспортированной трассировке содержит graph node id в поле args; записанные аннотации используют те же идентификаторы в качестве ключей, поэтому для их объединения с трассировкой достаточно обратиться к словарю:

import json

annotations = get_kernel_annotations()
with open("trace.json") as f:
    trace = json.load(f)
for event in trace["traceEvents"]:
    node_id = event.get("args", {}).get("graph node id")
    for ann in annotations.get(node_id, []):
        event["args"].update(ann)
with open("trace_annotated.json", "w") as f:
    json.dump(trace, f)

Если открыть trace_annotated.json в Perfetto (или chrome://tracing) и щелкнуть ядро из повторно воспроизводимого графа, рядом с размерами сетки и блока ядра будут показаны поля аннотаций — name: attention или name: mlp, layer: 3. Они указывают, к какой области захваченной рабочей нагрузки относится каждое ядро.

Поскольку аннотации хранятся в глобальном для процесса реестре с идентификаторами, соответствующими идентификаторам профилировщика, сериализованный с помощью pickle объект dict(get_kernel_annotations()) можно также сохранить рядом с трассировкой и объединить с ней автономно.

Этот пакет добавляет поддержку управления памятью устройства, реализованного в CUDA.

Управление памятью

class torch.cuda.use_mem_pool(pool, device=None) [исходный код]

Менеджер контекста, направляющий выделение памяти в указанный пул.

Параметры:

• pool (torch.cuda.MemPool) – объект MemPool, который нужно активировать, чтобы выделение памяти направлялось в этот пул.
• device (torch.device или int, необязательный) – выбранное устройство. Используется MemPool на текущем устройстве, определяемом с помощью current_device(), если device равно None (по умолчанию).

Примечание

Этот менеджер контекста направляет в указанный пул выделения памяти только текущего потока. Если внутри менеджера контекста запускается новый поток (например, при вызове backward), выделения в этом потоке не будут направляться в указанный пул.

Примечание

При использовании во время захвата CUDAGraph граф удерживает пул до сброса или уничтожения графа.

torch.cuda.nccl.version() [исходный код]

Возвращает версию NCCL.

Эта функция возвращает кортеж с основным, дополнительным и исправительным номерами версии NCCL. Если у версии есть суффикс, он также включается в кортеж. :returns: Информация о версии NCCL. :rtype: tuple

Расширение инструментов NVIDIA (NVTX)

Jiterator (бета-версия)

TunableOp

Некоторые операции можно реализовать с помощью нескольких библиотек или несколькими способами. Например, GEMM для CUDA или ROCm можно реализовать с помощью библиотек cublas/cublasLt или hipblas/hipblasLt соответственно. Как определить, какая реализация быстрее и какую следует выбрать? Для этого предназначен TunableOp. Некоторые операторы реализованы как настраиваемые операторы с использованием нескольких стратегий. Во время выполнения профилируются все стратегии, а самая быстрая выбирается для всех последующих операций.

Информацию об использовании см. в документации.

Stream Sanitizer (прототип)

CUDA Sanitizer — это прототип инструмента для обнаружения ошибок синхронизации между потоками в PyTorch. Информацию об использовании см. в документации.

GPUDirect Storage (прототип)

API в torch.cuda.gds предоставляют тонкие обёртки для некоторых API cuFile, позволяющих напрямую передавать данные между памятью GPU и хранилищем, избегая промежуточного буфера в CPU. Подробнее см. в документации по API cufile.

Эти API можно использовать в CUDA версии 12.6 и выше. Для их использования необходимо убедиться, что система настроена для работы с GPUDirect Storage в соответствии с документацией по GPUDirect Storage.

Пример использования см. в документации по GdsFile.

Зелёные контексты (экспериментальная функция)

torch.cuda.green_contexts предоставляет тонкие обёртки для API CUDA Green Context, позволяя более гибко распределять ресурсы SM для ядер CUDA.

Для этих API требуется пакет cuda.bindings; их можно использовать в PyTorch с CUDA версии 12.8 и выше. Для настройки Workqueue требуется CUDA 13.1 или новее.

Инструкции по установке cuda.bindings доступны здесь: https://nvidia.github.io/cuda-python/

Создайте потоки из зелёного контекста и используйте их так же, как другие пользовательские потоки CUDA:

ctx = GreenContext(...)
stream = ctx.Stream()
with torch.cuda.stream(stream):
    # torch operations here are using resources from `ctx`
    pass

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

Методы GreenContext.set_context() и GreenContext.pop_context() являются устаревшими API, оставленными для совместимости.

torch.cuda.nccl.is_available(tensors) [исходный код]

Этот пакет добавляет поддержку NVIDIA Tools Extension (NVTX), используемого при профилировании.