Spec-Zone.ru › PyTorch 2.14

CUDAGraph

class torch.cuda.CUDAGraph(keep_graph=False) [исходный код]

Обёртка для графа CUDA.

Параметры:

keep_graph (bool, необязательно) – Если keep_graph=False, cudaGraphExec_t будет создан на GPU в конце capture_end, а нижележащий cudaGraph_t будет уничтожен. Пользователи, которым нужно запросить или иным образом изменить нижележащий cudaGraph_t до создания экземпляра, могут задать keep_graph=True и получить к нему доступ через raw_cuda_graph после capture_end. Обратите внимание, что в этом случае cudaGraphExec_t не будет создан в конце capture_end. Вместо этого он будет создан явным вызовом instantiate или автоматически при первом вызове replay, если instantiate ещё не был вызван. Рекомендуется вручную вызвать instantiate до replay, чтобы избежать увеличения задержки при первом вызове replay. После первого вызова instantiate разрешается изменять исходный cudaGraph_t, но пользователь должен снова вручную вызвать instantiate, чтобы убедиться, что созданный экземпляр графа учитывает эти изменения. PyTorch не отслеживает эти изменения.

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

Self

Предупреждение

Этот API находится на стадии бета-тестирования и может измениться в будущих выпусках.

capture_begin(pool=None, capture_error_mode='global', check_input_liveness=False) [исходный код]

Начать захват операций CUDA в текущем потоке.

Обычно не следует самостоятельно вызывать capture_begin. Используйте graph или make_graphed_callables(), которые вызывают capture_begin внутри себя.

Параметры:
  • pool (необязательно) – Токен (возвращаемый graph_pool_handle() или other_Graph_instance.pool()) или MemPool, указывающий, что этот граф может совместно использовать память с указанным пулом. См. раздел Управление памятью графа.
  • capture_error_mode (str, необязательно) – задаёт cudaStreamCaptureMode для потока захвата графа. Может принимать значения “global”, “thread_local” или “relaxed”. Во время захвата графа CUDA некоторые действия, например cudaMalloc, могут быть небезопасны. Режим “global” вызывает ошибку при таких действиях в других потоках, “thread_local” — только при действиях в текущем потоке, а “relaxed” не вызывает ошибку при таких действиях. НЕ меняйте эту настройку, если вы не знакомы с cudaStreamCaptureMode
  • check_input_liveness (bool, необязательно) –

    Если True, отслеживает внешние входные тензоры во время захвата графа и вызывает ошибку, если какой-либо из них освобождён до повторного воспроизведения. Это помогает выявлять ошибки «использование после освобождения», когда входные тензоры удаляются сборщиком мусора между захватом и повторным воспроизведением. По умолчанию: False.

    Примечание

    Пользовательские ядра CUDA, добавленные вне PyTorch (например, с помощью cuLaunchKernel или DLPack), этим механизмом не отслеживаются.

capture_end() [исходный код]

Завершить захват графа CUDA в текущем потоке.

После capture_end для этого экземпляра можно вызвать replay.

Обычно не следует самостоятельно вызывать capture_end. Используйте graph или make_graphed_callables(), которые вызывают capture_end внутри себя.

capture_end_post() [исходный код]

Завершить захват, начатый с помощью capture_end_pre(): уничтожить шаблон при вызове keep_graph=False (граф уже должен быть создан; это выполняют capture_end() и менеджер контекста).

capture_end_pre() [исходный код]

Завершить захват, но не выполнять окончательную обработку: оставить захваченный cudaGraph_t активным (в обоих режимах keep_graph), чтобы его можно было проверить до того, как capture_end_post() создаст его экземпляр и/или уничтожит его.

debug_dump(debug_path, *, verbose=True) [исходный код]

Сохранить захваченный граф в debug_path в формате Graphviz DOT.

Шаблон графа должен быть активен: keep_graph=True (или enable_debug_mode()) либо вызван из обработчика завершения захвата. Требуется пакет cuda.bindings.

Параметры:
  • debug_path (обязательный) – Путь для сохранения графа.
  • verbose (bool) – Если True (значение по умолчанию), использовать максимально подробный вывод DOT.
enable_debug_mode() [исходный код]

Сохранить захваченный граф (эквивалентно keep_graph=True), чтобы его можно было проверить, например с помощью debug_dump(). Оставлен для обратной совместимости.

get_graph_data() [исходный код]

Возвращает словарь с описанием топологии графа и метаданных узлов.

Значение keep_graph должно быть True. Перед вызовом этого метода необходимо создать экземпляр графа (с помощью instantiate()). Требуется пакет cuda.bindings.

Возвращается словарь со следующей структурой:

{
    "exec_graph_id": int,
    "nodes": [
        {
            "index": int,
            "node_type": str,
            "tools_id": int,
            "graph_id": int,
            "node_id": int,
            "kernel_name": str or None,
            "event_ptr": int,
            "host_fn_addr": int,
            "host_fn_name": str or None,
            "dependencies": [int, ...],
            "dependents": [int, ...],
        },
        ...,
    ],
}

event_ptr — это дескриптор cudaEvent_t (в виде int), на который узел записи или ожидания события записывает событие или ожидает его. Эти узлы не создают временных записей CUPTI, поэтому сопоставление ожидания с записью, которая его сигнализирует, — единственный способ понять закодированную в нём межпоточную синхронизацию. Для узлов других типов значение равно 0.

host_fn_addr / host_fn_name заполняются для узлов хоста (обратный вызов CPU, выполняемый как узел графа): адрес обратного вызова и его деманглированное имя символа, определённое по возможности (None, если для него не найден экспортируемый символ). Для узлов других типов значения равны 0 / None.

Значение graph_id каждого узла сопоставляется с идентификатором графа exec, чтобы значения tools_id совпадали со значениями, указанными профилировщиками на основе CUPTI. dependencies и dependents — это списки индексов узлов в списке nodes.

Эта структура полезна для анализа трассировки профилировщика и определения того, является ли конкретная зависимость, обнаруженная в профиле, истинной зависимостью (закодированной в графе) или ложной зависимостью, возникшей из-за сопоставления независимых потоков с одним и тем же аппаратным каналом.

Дочерние графы и условные узлы отображаются как самостоятельные узлы, однако этот обход не проходит внутрь них (это отдельные объекты cudaGraph_t), поэтому содержащаяся в них работа отсутствует в nodes и выдаётся предупреждение. При этом идентификаторы, которые указываются, остаются действительными: граф exec сохраняет идентификаторы узлов верхнего уровня.

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

dict

instantiate() [исходный код]

Создать экземпляр графа CUDA. Вызывается методом capture_end, если keep_graph=False, или методом replay, если keep_graph=True и instantiate ещё не были вызваны явно. Не уничтожает cudaGraph_t, возвращаемый методом raw_cuda_graph.

pool() [исходный код]

Возвращает непрозрачный токен, представляющий идентификатор пула памяти этого графа.

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

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

_POOL_HANDLE

pools() [исходный код]

Возвращает непрозрачные токены для всех пулов памяти, сохраняемых этим графом.

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

list[_POOL_HANDLE]

raw_cuda_graph() [исходный код]

Возвращает нижележащий cudaGraph_t. Шаблон должен быть активен: для этого требуется keep_graph=True (он сохраняется после capture_end) либо доступ из обработчика завершения захвата (до уничтожения шаблона для keep_graph=False).

Сведения об API для управления этим объектом см. в разделах Управление графами и привязки управления графами cuda-python

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

int

raw_cuda_graph_exec() [исходный код]

Возвращает нижележащий cudaGraphExec_t. Если keep_graph имеет значение True, должен быть вызван instantiate; если keep_graph имеет значение False, должен быть вызван capture_end. Если вызвать instantiate() после raw_cuda_graph_exec(), ранее возвращённый cudaGraphExec_t будет уничтожен. Вы несёте ответственность за то, чтобы не использовать этот объект после уничтожения.

Сведения об API для управления этим объектом см. в разделах Выполнение графов и привязки выполнения графов cuda-python

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

int

register_capture_end_hook(hook) [исходный код]

Зарегистрировать hook(graph) для выполнения после завершения захвата, но до окончательной обработки графа. Захваченный cudaGraph_t активен (доступен через raw_cuda_graph()) в обоих режимах keep_graph. Обработчики вызываются в порядке регистрации. Возвращает дескриптор, метод remove() которого отменяет регистрацию обработчика.

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

RemovableHandle

register_capture_start_hook(hook) [исходный код]

Зарегистрировать hook(graph) для выполнения при начале захвата этим графом, сразу после начала захвата в текущем потоке. Обработчики вызываются в порядке регистрации. Возвращает дескриптор, метод remove() которого отменяет регистрацию обработчика.

Предупреждение

Обработчик выполняется во время захвата: любые выполняемые им операции CUDA попадают в граф, а при режиме захвата по умолчанию "global" небезопасный вызов вызывает ошибку. См. torch.cuda.graphs.register_graph_capture_start_hook().

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

RemovableHandle

register_destroy_callback(cb, *, synchronize_before_release=False) [исходный код]

Зарегистрировать cb() для выполнения при уничтожении (окончательной обработке) этого графа или при явном вызове reset(), непосредственно перед освобождением его ресурсов CUDA. Обратные вызовы выполняются один раз за цикл захвата, в порядке регистрации; исключения подавляются, чтобы сбой одного обратного вызова не прерывал выполнение остальных. cb НЕ должен ссылаться на этот граф: финализатор, запускающий его, хранится в глобальном реестре, поэтому обратный вызов, доступный через граф, не даст графу освободиться до завершения интерпретатора (он никогда не будет собран сборщиком мусора, а значит, обратный вызов не будет выполнен). Возвращает дескриптор, метод remove() которого отменяет регистрацию обратного вызова.

Завершение работы не синхронизирует CUDA, а cudaGraphExecDestroy освобождает граф, выполняющийся в данный момент, только асинхронно. Поэтому обратный вызов, освобождающий память устройства, которую граф читает или записывает, приведёт к использованию после освобождения, если повторное воспроизведение ещё выполняется. Передайте synchronize_before_release=True, чтобы синхронизировать каждый поток, в котором воспроизводился этот граф, до вызова обратного вызова. В противном случае обратные вызовы не должны освобождать объекты, на которые ссылается граф.

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

RemovableHandle

register_post_instantiate_hook(hook) [исходный код]

Зарегистрировать hook(graph) для выполнения после каждого создания экземпляра (включая повторное создание, в результате которого создаётся новый граф exec). Созданный экземпляр графа доступен через raw_cuda_graph_exec(). Обработчики вызываются в порядке регистрации. Возвращает дескриптор, метод remove() которого отменяет регистрацию обработчика.

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

RemovableHandle

register_replay_end_hook(hook) [исходный код]

Зарегистрировать hook(graph) для выполнения в конце каждого вызова replay(), сразу после запуска графа. Запуск выполняется асинхронно, поэтому обработчик запускается, когда повторное воспроизведение поставлено в очередь, а не после завершения работы GPU. Обработчики вызываются в порядке регистрации. Возвращает дескриптор, метод remove() которого отменяет регистрацию обработчика. См. примечание о критическом пути в описании register_replay_start_hook().

Обработчики завершения вызываются, даже если запуск завершается ошибкой, поэтому каждому обработчику начала соответствует обработчик завершения; затем ошибка запуска передаётся дальше. (Если обработчик начала вызывает исключение, повторное воспроизведение прерывается до запуска и обработчик завершения не вызывается.)

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

RemovableHandle

register_replay_start_hook(hook) [исходный код]

Зарегистрировать hook(graph) для выполнения в начале каждого вызова replay(), непосредственно перед запуском графа (после создания экземпляра по запросу, поэтому raw_cuda_graph_exec() действителен). Обработчики вызываются в порядке регистрации. Возвращает дескриптор, метод remove() которого отменяет регистрацию обработчика.

Примечание

Повторное воспроизведение — это критический путь, и зарегистрированный обработчик запускается при каждом повторном воспроизведении, поэтому он должен быть лёгким. Если обработчик не зарегистрирован, выполняется только одна проверка пустоты словаря.

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

RemovableHandle

replay() [исходный код]

Повторно воспроизвести операции CUDA, захваченные этим графом.

reset() [исходный код]

Удалить граф, хранящийся в этом экземпляре.

retain_object(obj, *, synchronize_before_release=False) [исходный код]

Сохранять obj активным в течение текущего цикла захвата графа и освобождать его при уничтожении (окончательной обработке) графа или при явном вызове reset(). Обратный вызов не выполняется; обычный подсчёт ссылок освобождает obj после снятия сохранённой ссылки. Возвращает дескриптор, вызов remove() которого досрочно снимает сохранённую ссылку. Как и в случае с register_destroy_callback(), obj НЕ должен ссылаться на этот граф, иначе граф останется активным до завершения интерпретатора, а obj никогда не будет освобождён.

synchronize_before_release имеет то же значение и те же ограничения, что и в register_destroy_callback(): задайте его, если освобождение obj освободит память устройства, которую граф читает или записывает (например, если obj — последняя ссылка на используемый графом тензор), а повторное воспроизведение ещё может выполняться.

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

RemovableHandle

© 2026, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://docs.pytorch.org/docs/2.14/generated/torch.cuda.CUDAGraph.html

Spec-Zone.ru

Настройки Оффлайн Что нового Помощь О нас
Spec-Zone .ru
спецификации, руководства, описания, API