Деревья CUDAGraph

Обзор CUDAGraph

Для более подробного ознакомления с CUDAGraphs, прочитайте ускорение PyTorch с помощью CUDAGraphs.

CUDA Graphs, появившиеся в CUDA 10, позволяют определять и инкапсулировать серию CUDA ядер как единый блок, т.е. график операций, а не как последовательность отдельных операций. Это предоставляет механизм запуска нескольких операций на GPU с помощью одной операции на CPU, а значит, уменьшает издержки запуска.

CUDA Graphs могут обеспечить значительное ускорение, особенно для моделей с высокой нагрузкой на CPU или небольшими вычислениями. Существует ряд ограничений, связанных с требованием, чтобы одни и те же ядра запускались с одними и теми же аргументами и зависимостями, а также адресами памяти.

• Невозможна обработка потоков управления
• Ядра, которые вызывают синхронизацию хост-устройство (например, .item()), приводят к ошибкам
• Все входные аргументы ядер фиксируются на значениях, зафиксированных при записи
• Адреса CUDA памяти фиксированы, однако значения памяти по этим адресам могут меняться
• Нет существенных операций CPU или эффектов на стороне CPU

Интеграция PyTorch с CUDAGraph

PyTorch предоставляет обёртку удобного использования вокруг CUDAGraphs, которая обрабатывает несколько сложных взаимодействий с кэширующим аллокатором PyTorch.

CachingAllocator использует отдельный пул памяти для всех новых выделений. Во время записи CUDAGraph память учитывается, выделяется и освобождается точно так же, как и во время выполнения eager. При воспроизведении вызываются только ядра, и нет изменений в аллокаторе. После начальной записи аллокатор не знает, какая память активно используется в пользовательских программах.

Использование отдельного пула памяти между eager-выделениями и выделениями cudagraph может увеличить объём памяти вашей программы, если выделено значительное количество памяти для обоих типов.

Создание вызываемых функций с графиками

Создание вызываемых функций с графиками — это абстракция PyTorch для совместного использования одного пула памяти над рядом вызываемых функций. Вызываемые функции с графиками используют тот факт, что при записи CUDA Graph память точно учитывается кэширующим аллокатором, чтобы безопасно разделять память между отдельными записями CUDA Graph. В каждом вызове выходы сохраняются как активная память, предотвращая перезапись активной памяти одной вызываемой функцией другой. Вызываемые функции с графиками могут вызываться только в одном порядке; адреса памяти из первого запуска записываются во второй и так далее.

Интеграция TorchDynamo с предыдущими CUDA Graphs

Запуск с cudagraph_trees=False не повторно использует память между отдельными захватами графика, что может привести к значительным проблемам с памятью. Даже для модели без разрывов графика это создаёт проблемы. Вперёд и назад — это отдельные захват графика, поэтому пулы памяти для вперёд и назад не совмещаются. В частности, память для активаций, сохранённых во вперёд, не может быть освобождена в назад.

Интеграция деревьев CUDAGraph

Как и вызываемые функции с графиками, деревья CUDA Graph используют один пул памяти для всех захватов графика. Однако вместо требования единственной последовательности вызовов, деревья CUDA Graph создают отдельные деревья захватов CUDA Graph. Давайте рассмотрим иллюстративный пример:

@torch.compile
def foo(x):
    # GRAPH 1
    y = x * x * x
    # graph break triggered here
    if y.sum() > 0:
        # GRAPH 2
        z = y ** y
    else:
        # GRAPH 3
        z = (y.abs() ** y.abs())
    torch._dynamo.graph_break()
    # GRAPH 4
    return z * torch.rand_like(z)

    # the first run warms up each graph, which does things like CuBlas or Triton benchmarking
    foo(torch.arange(0, 10), device="cuda")
    # The second run does a CUDA Graph recording, and replays it
    foo(torch.arange(0, 10), device="cuda")
    # Finally we hit the optimized, CUDA Graph replay path
    foo(torch.arange(0, 10), device="cuda")

В этом примере есть два отдельных пути, которые мы проходим через функцию: 1 -> 2 -> 4 или 1 -> 3 -> 4.

Мы совмещаем всю память в одном пуле памяти между отдельными записями, создавая ленту записей CUDA Graph, в данном случае 1 -> 2 -> 4. Мы добавляем инварианты, чтобы гарантировать, что память всегда находится в том же месте, что и при записи, и нет активных тензоров в пользовательских программах, которые могли бы быть перезаписаны.

• Применимы те же ограничения от CUDA Graphs: одни и те же ядра должны вызываться с одними и теми же аргументами (статические размеры, адреса и т.д.)
• Такая же структура памяти должна наблюдаться между записью и воспроизведением: если тензор, являющийся выходом одного графика, погибает после другого графика во время записи, то он также должен погибнуть во время воспроизведения.
• Активная память в пуле CUDA требует зависимости между двумя записями
• Эти записи могут быть вызваны только в одном порядке 1 -> 2 -> 4

Вся память совмещается в одном пуле памяти, поэтому дополнительной нагрузки на память по сравнению с eager нет. Что происходит, если мы пройдём новый путь и запустим Graph 3?

График 1 воспроизводится, а затем мы достигаем графика 3, который мы ещё не записали. При воспроизведении графика частный пул памяти не обновляется, поэтому y не отражается в аллокаторе. Без предосторожности мы перезапишем его. Чтобы поддержать повторное использование того же пула памяти после воспроизведения других графиков, мы сохраняем пул памяти в состоянии по окончании графика 1. Теперь, когда наши активные тензоры отражены в кэширующем аллокаторе, мы можем безопасно запустить новый график.

Сначала мы бы достигли оптимизированного пути CUDAGraph.replay(), который мы уже записали в графике 1. Затем мы бы достигли графика 3. Как и прежде, нам нужно будет один раз разогреть график перед записью. При разогреве адреса памяти не фиксируются, поэтому график 4 также вернётся к вызову inductor, не использующему cudagraph.

Во второй раз, когда мы достигнем графика 3, он разогрет и готов к записи. Мы запишем график 2, а затем снова запишем график 4, поскольку адреса входной памяти изменились. Это создаёт дерево записей CUDA Graph. Дерево CUDA Graph!

  1
 / \\
2   3
 \\   \\
  4   4

Ограничения

Так как CUDA Graph фиксирует адреса памяти, CUDA Graphs не имеют хорошего способа обработки активных тензоров из предыдущего вызова.

Предположим, мы проводим бенчмаркинг запуска инференса с помощью следующего кода:

import torch

def my_model(x):
    y = torch.matmul(x, x)
    return y

x = torch.randn(10, 10)
y1 = my_model(x)
y2 = my_model(x)

В реализации с отдельными CUDA Graph выход из первого вызова будет перезаписан вторым вызовом. Аналогично, в деревьях CUDA Graph, по умолчанию, активный выход от первого запуска потребовал бы зависимости между первым и вторым запуском, и мы никогда бы не достигли оптимизированного вызова cudagraph replay. Деревья CUDA Graph проигнорируют выходы от предыдущего запуска torch.compile и не будут создавать зависимость памяти. При обучении мы не будем игнорировать выходы от предыдущего запуска torch.compile, если у нас есть ожидающие обратные вычисления, которые ещё не были вызваны. TODO — добавить API для ручного увеличения поколения, выводить ошибку при доступе к предыдущему хранилищу

Сравнение