Подробный разбор Dynamo
Дата создания: 2 апреля 2024 г. | Последнее обновление: 6 июля 2026 г.
TorchDynamo (или просто Dynamo) — это трассировщик в torch.compile, и чаще всего именно он виноват в этих безумных трассировках стека. Однако мы не можем бездумно обвинять Dynamo в этих ошибках. Чтобы предоставить пользователю доступную гибкость, Dynamo поручена трудная задача — понимать любую программу на Python. В частности, Dynamo должен самостоятельно реализовать значительную часть языка программирования Python!
В этой статье мы подробно разберём внутреннее устройство Dynamo с самых основ. Мы обсудим его функциональность и то, как она реализована. К концу статьи вы будете лучше понимать, что пошло не так, когда вы torch.compiled программу PyTorch и компиляция завершилась ошибкой или прошла успешно, но ускорение оказалось не таким, как вы ожидали.
Краткое введение в Dynamo
Прежде чем углубляться во все детали реализации, давайте сначала разберёмся, чем занимается Dynamo.
Dynamo — это трассировщик. Это означает, что, получив функцию и её входные данные, он выполняет функцию и записывает в граф линейную последовательность инструкций (без потока управления). Например, рассмотрим следующую программу:
import torch
@torch.compile
def mse(x, y):
z = (x - y) ** 2
return z.sum()
x = torch.randn(200)
y = torch.randn(200)
mse(x, y)
Если сохранить эту программу в файл example.py и запустить
TORCH_LOGS=graph_code python example.py
мы увидим вывод трассировки Dynamo
def forward(l_x_: torch.Tensor, l_y_: torch.Tensor):
# File: example.py:5, code: z = (x - y) ** 2
sub = l_x_ - l_y_
z = sub ** 2
# File: example.py:6, code: return z.sum()
sum_1 = z.sum()
return (sum_1,)
Мы называем это графом (или трассировкой) функции для заданных входных данных. Он представлен в виде графа FX. Будем считать граф FX контейнером, хранящим список вызовов функций.
В первую очередь следует заметить, что граф представляет собой линейную последовательность операций PyTorch. [1] Dynamo записывает все операции PyTorch и сохраняет их последовательно. Например, он разделил z = (x - y) ** 2 на две составляющие операции: sub = l_x_ - l_y_ и z = sub ** 2.
Когда мы говорим, что трассировка линейна, мы имеем в виду отсутствие ветвлений и любого потока управления. Чтобы увидеть это, рассмотрим
import torch
@torch.compile
def fn(x, n):
y = x ** 2
if n >= 0:
return (n + 1) * y
else:
return y / n
x = torch.randn(200)
fn(x, 2)
которая при выполнении с TORCH_LOGS=graph_code возвращает
def forward(l_x_: torch.Tensor):
# File: example.py:5, code: y = x ** 2
y = l_x_ ** 2
# File: example.py:7, code: return (n + 1) * y
mul = 3 * y
return (mul,)
Мы видим, что Dynamo полностью удалил оператор if из трассировки и просто записал операции, выполненные с входными данными.
Таким образом, должно быть ясно, что трассировка функции зависит от входных данных. В частности, это означает, что трассировка создаётся не при написании @torch.compile, а при выполнении функции fn(x, 2) с фактическими аргументами.
Ещё одна интересная деталь: Dynamo удалил второй аргумент функции. Вместо этого он счёл его константой и записал в граф результат операции n + 1. Это ещё одна особенность Dynamo: он считает константами все значения, не являющиеся тензорами… за исключением целых чисел. Теперь посмотрим, чем отличаются целые числа.
Последнее определяющее свойство Dynamo — умение работать с динамическими формами. Символические формы позволяют Dynamo трассировать формы и, в более общем случае, целые числа, а не оставлять их константами. Это позволяет избегать повторной компиляции и развёртывать в рабочей среде универсальные модели, работающие с данными любого размера. Основные примеры использования динамических форм — размер пакета, когда модель обучают с фиксированным размером пакета, а затем выполняют вывод для произвольного размера, или переменная длина последовательности, с которой сталкиваются при обработке текста или аудио.
Это можно увидеть, если выполнить приведённый выше пример ещё несколько раз
import torch
@torch.compile
def fn(x, n):
y = x ** 2
if n >= 0:
return (n + 1) * y
else:
return y / n
x = torch.randn(200)
fn(x, 2)
fn(x, 3)
fn(x, -2)
В этом случае TORCH_LOGS=graph_code создаёт ещё два графа
# Graph for n==2 omitted
def forward(self, l_x_: torch.Tensor, l_n_: torch.SymInt):
# File: a.py:5, code: y = x ** 2
y = l_x_ ** 2
# File: a.py:7, code: return (n + 1) * y
add = l_n_ + 1
mul = add * y
return (mul,)
def forward(self, l_x_: torch.Tensor, l_n_: torch.SymInt):
# File: a.py:5, code: y = x ** 2
y = l_x_ ** 2
# File: a.py:9, code: return y / n
truediv = y / l_n_
return (truediv,)
Dynamo обнаружил, что значение одного целого числа изменилось после первого вызова, и начал трассировать его. Мы видим, что эти графы универсальны и символически трассируют переменную n с помощью объекта типа SymInt.
Если после этих вызовов вызвать fn(x, 4), Dynamo не будет выполнять повторную компиляцию, а повторно использует уже построенный граф.
Подведём итог: 1. Dynamo — это трассировщик Python. 2. Получив входные данные, он возвращает граф FX с выполненными функциями PyTorch. 3. Он также может трассировать целые числа, если обнаруживает, что их значения меняются между вызовами. 4. Все значения, не являющиеся тензорами или скалярами, он специализирует.
Разумеется, Dynamo делает гораздо больше: определяет, когда нужно повторно выполнить трассировку, переписывает байткод функции, реализует разрывы графа… Чтобы не затягивать введение, далее мы постепенно обсудим все эти возможности.
PEP 523: добавление API оценки кадров в CPython
Представим, что нам поручили реализовать Dynamo. С чего вообще начать? К счастью для нас, PEP 523 был опубликован вместе с Python 3.6. Этот PEP был создан, чтобы позволить сторонним разработчикам создавать JIT-компиляторы для Python. Посмотрим, как это работает.
Примечание о CPython: CPython внутри реализован как стековая машина. Программа на Python компилируется в байткод, который затем выполняется интерпретатором. Подробнее об этом байткоде можно узнать из модуля dis стандартной библиотеки. Введение в интерпретатор CPython также см. в документации для разработчиков. Будем считать, что читатель знаком с понятием стековой машины.
PEP 523 предоставляет API, с помощью которого пользователь может добавить собственный интерпретатор для каждой функции. Тогда CPython будет использовать этот интерпретатор вместо собственного для выполнения функции. Чтобы иметь возможность выполнить функцию, при её вызове CPython передаёт пользовательскому интерпретатору, помимо прочего: байткод функции; значения аргументов функции (то есть локальных переменных) и их имена; значения глобальных переменных и их имена; встроенные функции, такие как abs или print.
Все поля можно посмотреть здесь. [2]
Итак, CPython предоставляет пользовательскому интерпретатору всю информацию, необходимую для выполнения функции. [3]
С помощью этого API мы можем реализовать трассировщик, создав интерпретатор, который выполняет код и записывает в граф все операции PyTorch, происходящие во время выполнения. Именно это и делает Dynamo.
Dynamo использует API CPython для разбора всех этих объектов и упаковывает их в структуру Python. Выполнив это… он возвращается из C в Python. За исключением этого фрагмента кода, взаимодействующего с CPython, Dynamo полностью реализован на Python.
Должно быть ясно, что задача декоратора @torch.compile — установить необходимые механизмы, которые передадут байткод, аргументы, глобальные переменные и прочее в Dynamo при вызове функции. Напомним: @torch.compile сам по себе ничего не компилирует.
Реализация CPython на Python
Итак, мы снова в мире Python. У нас есть байткод функции и весь контекст, необходимый для её выполнения. В частности, мы оказались в функции _convert_frame_assert. Именно эту функцию возвращает декоратор torch.compile! Мы попадаем в эту функцию из _dynamo.optimize. Декоратор torch.compile — это просто удобный API для _dynamo.optimize.
Прежде чем реализовывать интерпретатор Python, определим промежуточное представление (IR). В частности, мы хотим обернуть все локальные и глобальные переменные в собственные внутренние классы. Это позволяет лучше отслеживать эти объекты и объединять объекты, с которыми Dynamo может работать одинаково.
Родительский класс внутренней структуры классов — VariableTracker, представляющий разные объекты, которые понимает Dynamo. Например, ListVariable представляет объект list и внутренне хранит список VariableTracker. Другой пример VariableTracker — это ConstantVariable. ConstantVariable оборачивает все объекты, которые Dynamo считает константами. Для объектов, требующих особого внимания, также существуют специальные подклассы, например TensorVariable. Все эти внутренние классы определены в папке torch/_dynamo/variables.
Объекты Python оборачиваются в соответствующий класс VariableTracker в VariableBuilder._wrap. Эта функция представляет собой длинную цепочку elif, которая пытается рекурсивно сопоставить входные данные Python с подходящим типом VariableTracker.
Совет по отладке. Неожиданные результаты Dynamo иногда вызваны построителем. Если его логика неверна, Dynamo может обернуть переменную в неправильный тип VariableTracker, что впоследствии приведёт к проблемам. При возникновении ошибки Dynamo полезно изучить типы VariableTracker, указанные в сообщениях об ошибках, и метод VariableTracker, вызывающий исключение. В частности, иногда объект отслеживается как UserDefinedObjectVariable (универсальный класс Dynamo), хотя должен был отслеживаться как более конкретный тип. В таких случаях часто виновата логика VariableBuilder.
Совет по отладке. При запуске программы с TORCH_LOGS=dynamo среди выводимых артефактов встречаются строки следующего вида:
TRACE LOAD_GLOBAL y [TorchInGraphFunctionVariable(<built-in method any>), TensorVariable()]
Это байткод исходной программы и состояние стека в соответствующий момент. Это очень полезно, чтобы выяснить, где объект не был преобразован в правильный VariableTracker.
Итак, у нас есть IR для трассировщика, теперь нам всего лишь нужно заново реализовать стековую машину CPython. Это делает InstructorTranslatorBase в файле symbolic_convert.py.
В InstructionTranslatorBase около 200 методов, реализующих почти все инструкции байткода Python. Например, посмотрим на реализацию BUILD_LIST
def BUILD_LIST(self, inst):
items = self.popn(inst.argval)
self.push(ListVariable(items, mutation_type=ValueMutationNew()))
Этот байткод создаётся такими конструкциями, как l = [2, 3, 4]. В данном случае, поскольку элементов три, сгенерированный байткод — BUILD_LIST 3. Это означает, что мы извлекаем с вершины стека элементы 3 и помещаем на вершину новый объект списка, образованный этими тремя элементами.
Построение выходного графа
Теперь, когда у нас есть способ символьного выполнения кода Python, мы можем извлечь операции PyTorch, выполняемые при символьном выполнении программы с заданными входными данными. В Dynamo это реализовано с помощью объекта OutputGraph. Объект OutputGraph привязан к объекту InstructionTranslator и отслеживает все данные, необходимые для создания графа FX, который будет возвращён Dynamo.
Все входные и промежуточные элементы графа FX — это fx.Node. В Dynamo объекты fx.Node обёрнуты в объекты fx.Proxy. Объекты fx.Proxy используются для построения графа FX. В частности, они записывают в граф каждую выполненную над ними операцию PyTorch. Создать новую операцию для добавления в граф можно вызовом create_proxy. Затем её можно добавить в граф с помощью функции wrap_fx_proxy.
Граф хранит операции над тензорами… и операции над символьными целыми числами. Символьные целые числа мы обсудим позже, а сначала рассмотрим, как Dynamo решает важную проблему корректности.
Обеспечение корректности Dynamo: защитные проверки
На этом этапе у нас есть способ трассировать программы, полностью игнорируя поток управления. И для этого мы заново реализовали весь CPython… Если это кажется чрезмерным, то так и есть. torch.jit.trace уже реализует это без всех этих сложностей. Так в чём же дело?
Проблема torch.jit.trace, как предупреждает документация, в том, что он работает только в том случае, если трассируемая программа не зависит от данных. Иными словами, он работает, только если сама программа линейна. Это означает, что нельзя использовать в программе конструкции if-else, циклы for и while, исключения. Более того, ни одна из используемых библиотек не должна применять поток управления! В итоге отказ от потока управления в таком динамичном языке, как Python, накладывает огромное ограничение.
JAX решает эту проблему, всегда повторно выполняя трассировку и кэшируя граф после её завершения. Dynamo, напротив, использует защитные проверки, чтобы не трассировать всю программу заново при каждом запуске.
Защитная проверка — это предположение (логическое выражение для входных данных), необходимое для специализации кадра под один набор примеров входных данных. Повторно использовать граф можно, только если эти предположения выполняются и для новых входных данных.
Например, для любого константного аргумента функции, например строки, устанавливается защитная проверка, согласно которой входные данные должны иметь тип str и совпадать с переданной строкой. Запустив
import torch
@torch.compile
def fn(a, b):
return a * len(b)
fn(torch.arange(10), "Hello")
с TORCH_LOGS=guards, мы получим (среди прочих проверок)
___check_type_id(L['b'], 94334122025024) L['b'] == 'Hello'
Это означает: «локальная переменная b должна иметь определённый тип (в данном случае str, обозначенный константой 9433...), а её значение должно быть 'Hello'». Если затем снова выполнить функцию с другим аргументом
import torch
@torch.compile
def fn(a, b):
return a * len(b)
fn(torch.arange(10), "Hello")
fn(torch.arange(10), "Hi")
можно узнать, какая проверка не прошла, запустив TORCH_LOGS=recompiles
Recompiling function fn in script.py:3
triggered by the following guard failure(s):
- L['b'] == 'Hello'
Защитные проверки накапливаются, пока входные данные функции оборачиваются построителем, а также во время выполнения программы. В следующем разделе мы покажем ещё много примеров защитных проверок, но сначала обсудим источники.
Источник отслеживает способ восстановления переменной из исходных локальных или глобальных переменных, существовавших при входе в текущий кадр. В частности, он отслеживает исходные локальные и глобальные объекты, а также содержащиеся в них объекты. В примере
def foo(x: Tensor, y: List[Tensor]):
a = x * y[0]
return a * x
источником для x и y служит LocalSource, а для y[0] — GetItemSource, который хранит внутри LocalSource. У a источника не будет, поскольку это промежуточная переменная, существующая только внутри графа FX.
Все они определены в файле torch/_dynamo/source.py. Защитную проверку, создаваемую GetItemSource, можно увидеть в следующем примере:
import torch
@torch.compile
def fn(x, l):
return x * len(l[0])
fn(torch.randn(8), ["Hi", "Hello"])
создаёт следующие защитные проверки
___check_type_id(L['l'], 94439025877664) len(L['l']) == 2 ___check_type_id(L['l'][0], 94439025840192) L['l'][0] == 'Hi' ___check_type_id(L['l'][1], 94439025840192) L['l'][1] == 'Hello'
Здесь мы видим код, сгенерированный GetItemSource ([0] и [1]), оборачивающий LocalSource (L['l']).
Теперь, имея источники и защитные проверки, мы можем реализовать систему кэширования, которая позволит избежать повторной компиляции и не потребует каждый раз заново выполнять трассировку. Далее мы подробнее обсудим эту систему кэширования.
Внимательный читатель заметил, что пока не объяснено, почему нам нужно настолько тонко контролировать интерпретатор Python, что приходится его заново реализовывать. Приведённые примеры защитных проверок зависят от входных объектов, поэтому их можно вычислить до выполнения функции. Иными словами, эту систему защитных проверок можно было бы реализовать поверх torch.jit.trace и получить те же возможности с гораздо меньшими усилиями… Но тут появляются символьные формы.
Символьные формы
Ещё один пункт, который мы обсуждали во введении: Dynamo умеет трассировать целые числа. Для этого используется символьный класс torch.SymInt, который ведёт себя как int, но записывает в выходной граф FX все выполняемые над ним операции. [4] Мы уже встречались с этим классом во введении, когда говорили о трассировке символьных целых чисел.
Теперь обсудим три свойства, определяющие трассировку символьных форм в Dynamo, и способы их реализации.
По умолчанию — статические
Dynamo считает любое целое число — будь то входное значение или размерность тензора — статическим по умолчанию. Иными словами, при первом выполнении функции целые числа не трассируются. Только если обнаруживается, что значение целого числа или размерность изменились во время выполнения, Dynamo начинает их трассировать и создаёт граф, универсальный для этой переменной.
Мы уже видели такое поведение на примере целых чисел во введении. Теперь рассмотрим пример с формами тензоров.
import torch
@torch.compile
def fn(a, b):
return a.shape[0] * a * b
fn(torch.randn(4, 3), torch.randn(4, 3))
fn(torch.randn(8, 3), torch.randn(8, 3))
Если запустить эту программу с TORCH_LOGS=graph_code, мы увидим, что эти два вызова трассируются следующим образом:
def forward(self, l_a_: torch.Tensor, l_b_: torch.Tensor):
mul = 4 * l_a_
mul_1 = mul * l_b_
return (mul_1,)
def forward(self, s0: torch.SymInt, l_a_: torch.Tensor, l_b_: torch.Tensor):
size = l_a_.size()
getitem = size[0]
mul = getitem * l_a_
mul_1 = mul * l_b_
return (mul_1,)
В первом графе форма трассируется как константа, но после её изменения Dynamo трассирует её символически с помощью SymInts. Как правило, проще всего посмотреть формы промежуточных значений, запустив программу с TORCH_LOGS=graph_sizes
TRACED GRAPH TENSOR SIZES ===== __compiled_fn_1 ===== l_a_: (s0, 3) l_a_ (concrete): (8, 3) l_b_: (s0, 3) l_b_ (concrete): (8, 3) mul: (s0, 3) mul (concrete): (8, 3) mul_1: (s0, 3) mul_1 (concrete): (8, 3)
где видно, что первая размерность обоих тензорных аргументов является динамической, поскольку она представлена переменной s0.
Узнать, как это реализовано в Dynamo, можно, запустив TORCH_LOGS=guards
# Guards first call check_tensor(L['a'], torch.float32, device=None, requires_grad=False, size=[4, 3], stride=[3, 1]) check_tensor(L['b'], torch.float32, device=None, requires_grad=False, size=[4, 3], stride=[3, 1]) # Guards second call check_tensor(L['a'], torch.float32, device=None, requires_grad=False, size=[None, 3], stride=[3, 1]) check_tensor(L['b'], torch.float32, device=None, requires_grad=False, size=[None, 3], stride=[3, 1]) L['b'].size()[0] == L['a'].size()[0] 2 <= L['a'].size()[0]
Видно, что при первом вызове защитные проверки требуют, чтобы тензоры имели определённые фиксированные размеры и шаги. При втором выполнении эти проверки не проходят, поэтому Dynamo выполняет трассировку заново. Поскольку не прошла проверка int, на этой, второй итерации Dynamo трассирует int символически и устанавливает более общие проверки для этого более универсального ядра.
Совет по повышению производительности компиляции. Если известно, что размерность будет меняться, её можно пометить как динамическую с помощью torch._dynamo.mark_dynamic перед вызовом torch.compile. Это позволит избежать первой компиляции со статической формой. Есть и другие полезные вспомогательные функции, например maybe_mark_dynamic или mark_static. Также можно настроить трассировку всех целых чисел и форм, вызвав torch.compile(dynamic=True). Это пригодится главным образом для отладки.
0 и 1 всегда специализируются
Независимо от того, помечена ли размерность как динамическая, если передать входное значение, у которого эта размерность равна 0 или 1, Dynamo будет трассировать её как статическую и создаст для неё отдельный граф. Поэтому в приведённом выше примере мы видим проверки вида 2 <= L['a'].size()[0].
У этого решения есть несколько причин. Особенно важны следующие: тензор является пустым тогда и только тогда, когда одна из его размерностей равна нулю; тензор может быть непрерывным, только если один из шагов равен единице.
Это решение не распространяется на обычные целые числа Python: если мы считаем, что целое число Python нужно компилировать динамически, оно не будет специализироваться по умолчанию; специализация зависит от того, как оно используется.
Утиная типизация форм
Dynamo применяет подход, который мы называем «утиной типизацией форм». Если два динамических целых числа имеют одинаковое значение во время трассировки, мы считаем их равными и добавляем соответствующую защитную проверку. Фактически вместо двух символов s0 и s1 в примере выше мы объединили их в s0 и добавили проверку L['b'].size()[0] == L['a'].size()[0]. Это позволяет компилятору выполнять слияния и при этом создавать достаточно универсальные ядра.
Проверки для символьных целых чисел
Теперь мы понимаем, как на высоком уровне реализованы символьные формы и какими свойствами они обладают. Но почему символьные формы заставили нас выбрать непростой путь и получить контроль над интерпретатором CPython? Рассмотрим следующий пример:
import torch
@torch.compile(dynamic=True)
def fn(a):
if a.shape[0] * 2 < 16:
return a
else:
return a + 1
fn(torch.randn(8))
В этом коде есть проверка вида 2*L['a'].size()[0] >= 16. Это нетривиальная проверка входных данных функции, но регистрируется она в середине выполнения программы. Более того, мы не можем знать, что эта проверка понадобится, пока не увидим оператор if, зависящий от аргумента SymNodeVariable. Такие условия невидимы для torch.jit.trace и требуют глубокого анализа кода Python.
Совет по отладке. Если запустить этот код с TORCH_LOGS=dynamo, можно узнать, где была добавлена эта проверка:
eval 2*s0 >= 16 [guard added] at script.py:5 in fn (_dynamo/variables/tensor.py:812 in evaluate_expr)
Установка точки останова в этом месте и просмотр трассировки стека помогают понять, откуда появилась проверка.
Завершение Dynamo: разрывы графа
Используя все рассмотренные нами инструменты, мы получили трассировщик, который может отслеживать операции PyTorch над тензорами и целыми числами, а также имеет систему кэширования, которая знает, когда можно повторно использовать ранее отслеженный граф, а когда нужно выполнить трассировку заново. И всё это при выполнении произвольного кода Python!
Здесь есть лишь одна небольшая проблема. Утверждение «выполнение произвольного кода Python» — возможно, слишком широкое обобщение. Dynamo реализует значительную часть Python, но поддерживает ли он более сложные возможности, например сопрограммы или асинхронный код? Реализует ли он всю стандартную библиотеку Python? У NumPy тоже есть API на Python. Понимает ли torch.compile и NumPy? А Django? [5]
Экосистема Python огромна, и значительная её часть написана на других, более производительных языках, таких как C++ или Rust, и лишь предоставляет привязки для Python. Dynamo не сможет трассировать объекты Python, реализованные на C++. Что может сделать трассировщик, если встречает операцию, которую он не понимает?
Обычно трассировщики машинного обучения решают эту проблему, сообщая пользователю об операции, на которой они остановились, и полностью прекращая трассировку. В случае PyTorch это создало бы серьёзные неудобства: его пользователи привыкли к гибкости, которую он им предоставляет. Например, модель doctr_det_predictor использует NumPy и библиотеку cv2, чтобы постобработать результат модели.
Вот ещё один случай, когда доступ к CPython оказывается полезен. Вместо того чтобы выдавать ошибку, Dynamo может позволить CPython выполнить проблемный код! Для этого во время трассировки Dynamo генерирует один граф со всеми операциями до проблемного кода, а другой — со всеми операциями после него. [6] Затем во время выполнения он поручает CPython выполнить первый граф, затем проблемный код и, наконец, второй граф. Этот процесс остановки трассировки и генерации нескольких графов называется разрывом графа.
Признаюсь: во введении и первых разделах я всё время говорил неправду. Dynamo генерирует не один граф, а несколько графов! Для всех практических целей повторную трассировку, начинающуюся после второго графа, можно считать трассировкой новой функции. Новый граф после разрыва графа будет иметь собственные проверки, новый набор локальных переменных и так далее.
Чтобы обсудить реализацию разрывов графа, сначала нужно ещё раз рассмотреть взаимодействие Dynamo с CPython. С помощью PEP 523 CPython позволяет пользователю использовать собственный механизм вычисления кадров. Мы не упомянули, что CPython также предоставляет собственный механизм вычисления кадров для использования другими. Dynamo использует эту возможность, чтобы позволить быстрому интерпретатору CPython выполнять скомпилированный код. Для функции без разрывов графа весь процесс трассировки и выполнения программы, которая вызывает функцию дважды с одинаковыми аргументами, выглядит так:
-
При первом вызове функции
-
Dynamo трассирует функцию в граф FX
- Компилятор (Inductor) компилирует граф FX в эффективный низкоуровневый код… но это уже история для другого раза
- Он переписывает байткод функции так, чтобы тот просто вызывал скомпилированную функцию
- Он передаёт CPython этот новый байткод и просит выполнить его здесь
-
-
При втором вызове функции
Сам по себе этот процесс выглядит чрезмерно сложным. Зачем генерировать новый байткод и просить CPython выполнить его, вместо того чтобы просто создать привязку на C++ к скомпилированной функции и выполнить её? Этот подход позволяет реализовать разрывы графа! Байткод, генерируемый при разрыве графа, имеет следующую структуру:
- Байткод, выполняющий первый граф
- Байткод, который оставляет стек в том состоянии, в котором он был бы, если бы CPython выполнил первый граф. Он также воспроизводит все изменения локальных или глобальных переменных, которые были бы видны в этот момент
- Байткод, вызвавший разрыв графа в Dynamo
- Байткод, выполняющий второй граф
Рассмотрим простой пример
import torch
@torch.compile
def fn(a):
b = a + 2
print("Hi")
return b + a
fn(torch.randn(4))
Запустив этот пример с TORCH_LOGS=bytecode, мы увидим исходный и изменённый байткод
MODIFIED BYTECODE fn script.py line 3
0 LOAD_GLOBAL 1 (__compiled_fn_0)
2 LOAD_FAST 0 (a)
4 CALL_FUNCTION 1
6 STORE_FAST 3 (graph_out_0)
8 LOAD_GLOBAL 0 (print)
10 LOAD_CONST 2 ('Hi')
12 LOAD_FAST 3 (graph_out_0)
14 LOAD_CONST 3 (0)
16 BINARY_SUBSCR
18 STORE_FAST 1 (b)
20 CALL_FUNCTION 1
22 LOAD_GLOBAL 2 (__resume_at_14_1)
24 ROT_TWO
26 LOAD_FAST 0 (a)
28 LOAD_FAST 1 (b)
30 CALL_FUNCTION 3
32 RETURN_VALUE
MODIFIED BYTECODE resume_in_fn script.py line 6
0 LOAD_GLOBAL 1 (__compiled_fn_2)
2 LOAD_FAST 2 (b)
4 LOAD_FAST 1 (a)
6 CALL_FUNCTION 2
8 UNPACK_SEQUENCE 1
10 RETURN_VALUE
Мы видим, что изменённый байткод разделён на две функции: fn, исходную функцию, и функцию с именем resume_in_fn. Эта вторая функция создана Dynamo для реализации выполнения программы начиная с места разрыва графа. Её часто называют функцией-продолжением. Эта функция-продолжение просто вызывает вторую скомпилированную функцию с нужными аргументами. Код исходной функции переписан с использованием описанной выше стратегии:
- L0-4. Вызвать скомпилированную функцию (
a + 2). - L6. Сохранить результат в локальной переменной с именем
graph_out_0.graph_out_0— это кортеж - L8-18. Оставить стек в том состоянии, в котором он был бы в месте разрыва графа
- L20. Выполнить код, вызвавший разрыв графа
- L22-32. Вызвать скомпилированную функцию-продолжение (
a + b)
Генерация кода для стека в Dynamo делегирована подклассам VariableTracker. У каждого объекта VariableTracker в Dynamo есть метод reconstruct, который генерирует необходимый байткод для создания на стеке объекта Python, который он представляет.
Совет по отладке. Разрывы графа снижают производительность, поэтому лучше избегать их. Запуск программы с TORCH_LOGS=graph_breaks — отличный способ выяснить, сколько разрывов графа произошло в нашей программе. Информация, которую он возвращает, представлена в виде объектов VariableTracker, поэтому приведённые выше советы по отладке иногда тоже помогают понять, что вызвало разрыв графа.
Заключение
Dynamo — сложная программа. Если вы решили реализовать интерпретатор CPython, будьте готовы к непростой работе. Тем не менее мы надеемся, что эта статья поможет немного развеять завесу тайны над ним.
Dynamo (в основном) реализован на Python. Мы оставили множество ссылок на рассмотренные фрагменты кода. Надеемся, что чтение этих фрагментов, поиск мест их вызова или установка точек останова и изучение стека вызовов помогут разобраться в остальной части кодовой базы.
Разумеется, лучший способ узнать, как работает программа, — расширить её. В данном случае лучше всего взглянуть на открытые задачи Dynamo на GitHub. Многие из них требуют лишь незначительных изменений в коде — нужно только найти, где именно их внести.
Сноски
Ниже приведены дополнительные сведения и ссылки на понятия, упомянутые в этом документе.
© 2026, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://docs.pytorch.org/docs/2.14/user_guide/torch_compiler/torch.compiler_dynamo_deepdive.html