torch.cuda.graph_annotations.mark_kernels

torch.cuda.graph_annotations.mark_kernels(annotation, *, backward=True) [source]

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

Должен использоваться внутри активного захвата torch.cuda.graph с enable_annotations=True. Каждый узел ядра, memcpy и memset, добавленный захватом в этой области действия, помечается с помощью annotation. Вне захвата, при отключенных аннотациях или если is_available() имеет значение False, менеджер контекста ничего не делает.

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

По умолчанию аннотируется также работа обратного прохода: узлы autograd, созданные операциями прямого прохода внутри области, получают хуки (через torch.autograd.graph.node_creation_hook), которые ограничивают выполнение обратного прохода. Поэтому, когда обратный проход сам захватывается — в том же захвате, что и прямой проход, или в более позднем, — его ядра помечаются той же аннотацией, а также ключом "autograd_phase": "backward", указывающим, что это работа обратного прохода (поэтому "autograd_phase" зарезервирован: аннотация обратного прохода перезаписывает его). Если обратный проход выполняется вне захвата, хуки ничего не делают. Владение распространяется и на градиенты более высокого порядка: узлы, созданные во время выполнения узла с хуком (create_graph=True, повторное вычисление контрольной точки), наследуют его аннотации, поэтому последующий захват градиента от градиента также получает соответствующую атрибуцию. Узлы AccumulateGrad никогда не аннотируются: узел листа создается один раз и кэшируется, поэтому его принадлежность к области действия зависела бы от случайного первого использования. Передайте backward=False, чтобы аннотировать только работу прямого прохода, например, если обертка реализует собственную атрибуцию обратного прохода. Наличие этого ключевого слова также служит проверкой поддержки обратного прохода на нативном уровне: "backward" in inspect.signature(mark_kernels).parameters.

Реализация: при входе фиксирует границу захвата текущего потока и существующие непосредственные зависимые от нее узлы; при выходе из области действия обходит только зависимые узлы, добавленные после входа (если область содержит первую захваченную работу, вместо них используются вновь созданные корневые узлы графа).

Параметры:

• annotation (str или dict) – Метаданные для добавления к каждому захваченному узлу. Строковое значение s записывается как {"name": s}. Значения словаря должны поддерживать сериализацию с помощью pickle. Ключ "name" задает имя области в инструментах трассировки; "stream" зарезервирован для назначения полос потоков.
• backward (bool) – Нужно ли также аннотировать ядра обратного прохода узлов autograd, созданных внутри области действия. По умолчанию: True.

Примечание

Аннотируемые узлы должны быть достижимы от границы захвата потока, текущего при входе в область действия. Работу в другом потоке, который уже выполняет захват, сначала необходимо синхронизировать с текущим потоком.

Примечание

Вложенные графы и условные узлы содержат тела в отдельном cudaGraph_t, в которое этот обход не переходит, поэтому их работа остается без аннотаций и выдается предупреждение. Переход внутрь возможен (cudaGraphNodeGetParams предоставляет доступ к графам тел), но сам по себе он не решил бы проблему: узлы тела нумеруются в пространстве идентификаторов этого графа и получают новую нумерацию при встраивании в граф exec, а способа узнать эти новые номера нет. Поэтому remap_to_exec_graph() не смог бы сопоставить аннотации с данными профилировщика. По той же причине область действия внутри тела условного узла (torch.cond / torch.while_loop) вообще ничего не записывает.

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

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

Пример:

>>> g = torch.cuda.CUDAGraph()
>>> x = torch.randn(8, device="cuda")
>>> with torch.cuda.graph(g, enable_annotations=True):
...     with torch.cuda.graph_annotations.mark_kernels("phase_A"):
...         y = x + 1