Устранение неполадок PyTorch 2.0

Автор: Михаил Лазос

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

В дополнение к информационному и отладочному журналам вы можете использовать torch._logging для более тонкой настройки ведения журнала.

Диагностика ошибок во время выполнения

В целом, стек TorchDynamo состоит из захвата графа из кода Python (TorchDynamo) и компилятора бэкэнда. Например, компилятор бэкэнда может состоять из отслеживания обратного графа (AOTAutograd) и понижения графа (TorchInductor)*. Ошибки могут возникать в любом компоненте стека и будут предоставлять полные трассировки стека.

Для определения компонента, в котором произошла ошибка, вы можете использовать ведение журнала уровня информации torch._logging.set_logs(dynamo = logging.INFO) или TORCH_LOGS="dynamo" и искать Step #: ... выводы. Журналы создаются в начале и в конце каждого шага, поэтому шаг, которому должна соответствовать ошибка, — это последний залогированный шаг, конец которого еще не залогирован. Шаги соответствуют следующим частям стека:

Если информационного ведения журнала недостаточно, вы можете использовать доступные варианты бэкэнда. Эти варианты включают:

• "eager": выполняется только захват прямого графа TorchDynamo, а затем захваченный граф выполняется с помощью PyTorch. Это дает указание о том, генерирует ли TorchDynamo ошибку.
• "aot_eager": выполняет захват прямого графа TorchDynamo, а затем AOTAutograd отслеживает обратный граф без дополнительных шагов компилятора бэкэнда. Затем для выполнения прямых и обратных графов используется PyTorch eager. Это полезно для сужения проблемы до AOTAutograd.

Общая процедура сужения проблемы выглядит следующим образом:

1. Запустите свою программу с бэкэндом "eager" . Если ошибка больше не возникает, проблема в компиляторе бэкэнда, используемом (если используется TorchInductor, перейдите к шагу 2. Если нет, см. эту секцию эту секцию). Если ошибка все еще возникает с бэкэндом "eager", это ошибка при выполнении torchdynamo.
2. Этот шаг необходим только в том случае, если TorchInductor используется в качестве компилятора бэкэнда. Запустите модель с бэкэндом "aot_eager" . Если этот бэкэнд генерирует ошибку, ошибка возникает во время отслеживания AOTAutograd. Если ошибка больше не возникает с этим бэкэндом, то ошибка находится в TorchInductor*.

Каждый из этих случаев анализируется в следующих разделах.

Ошибки Torchdynamo

Если сгенерированная ошибка возникает с бэкэндом "eager", то TorchDynamo, скорее всего, является источником ошибки. Вот пример кода, который сгенерирует ошибку.

import torch

import torch._dynamo as dynamo


def test_assertion_error():
    y = torch.ones(200, 200)
    z = {y: 5}
    return z

compiled_test_assertion_error = torch.compile(test_assertion_error, backend="eager")

compiled_test_assertion_error()

Код выше генерирует следующую ошибку:

torch._dynamo.convert_frame: [ERROR] WON'T CONVERT test_assertion_error /scratch/mlazos/torchdynamo/../test/errors.py line 26
due to:
Traceback (most recent call last):
  File "/scratch/mlazos/torchdynamo/torchdynamo/symbolic_convert.py", line 837, in BUILD_MAP
    assert isinstance(k, ConstantVariable) or (
AssertionError

from user code:
   File "/scratch/mlazos/torchdynamo/../test/errors.py", line 34, in test_assertion_error
    z = {y: 5}

Set torch._dynamo.config.verbose=True for more information
==========

Как предлагает сообщение, вы можете установить torch._dynamo.config.verbose=True для получения полной трассировки стека как ошибки в TorchDynamo, так и кода пользователя. В дополнение к этому флагу вы также можете установить уровень log_level TorchDynamo через torch._dynamo.config.log_level. Эти уровни включают:

• logging.DEBUG: Вывод каждой встречаемой инструкции в дополнение ко всем уровням журнала, перечисленным ниже.
• logging.INFO: Вывод каждой компилируемой функции (исходный и измененный байткод) и захваченного графика в дополнение ко всем уровням журнала, перечисленным ниже.
• logging.WARNING (по умолчанию): Вывод разрывов графика в дополнение ко всем уровням журнала, перечисленным ниже.
• logging.ERROR: Вывод только ошибок.

Если модель очень большая, журналы могут стать чрезмерно объемными. Если ошибка возникает глубоко внутри кода Python модели, может быть полезно выполнить только тот фрейм, в котором произошла ошибка, чтобы облегчить отладку. Есть два доступных инструмента для этого:

• Установление переменной окружения TORCHDYNAMO_DEBUG_FUNCTION с желаемым именем функции запустит torchdynamo только для функций с этим именем.
• Включение инструмента записи/воспроизведения (установка torch._dynamo.config.replay_record_enabled = True) который сохраняет запись выполнения при возникновении ошибки. Затем эту запись можно воспроизвести, чтобы выполнить только фрейм, в котором произошла ошибка.

Диагностика ошибок TorchInductor

Если ошибка не возникает с бэкэндом "eager", значит, источником ошибки является компилятор бэкэнда (пример ошибки). Существует несколько вариантов компиляторов бэкэнда для TorchDynamo, при этом TorchInductor удовлетворяет потребностям большинства пользователей. Этот раздел сосредоточен на TorchInductor в качестве мотивационного примера, но некоторые инструменты также могут быть использованы с другими компиляторами бэкэнда.

Ниже представлена часть стека, на котором мы фокусируемся:

При выборе TorchInductor в качестве бэкэнда используется AOTAutograd для генерации обратного графа из прямого графа, захваченного torchdynamo. Важно отметить, что ошибки могут возникать во время этого отслеживания, а также во время понижения TorchInductor прямого и обратного графов в код GPU или C++. Модель часто может состоять из сотен или тысяч узлов FX, поэтому сужение до точных узлов, вызвавших проблему, может быть очень сложным. К счастью, доступны инструменты, которые автоматически минимизируют эти входные графы до узлов, вызывающих проблему. Первый шаг — определить, возникает ли ошибка во время отслеживания обратного графа с помощью AOTAutograd или во время понижения TorchInductor. Как упоминалось выше в шаге 2, бэкэнд "aot_eager" может использоваться для запуска только AOTAutograd в изоляции без понижения. Если ошибка все еще возникает с этим бэкэндом, это указывает на то, что ошибка возникает во время отслеживания AOTAutograd.

Вот пример:

import torch

import torch._dynamo as dynamo

model = torch.nn.Sequential(*[torch.nn.Linear(200, 200) for _ in range(5)])

def test_backend_error():

    y = torch.ones(200, 200)
    x = torch.ones(200, 200)
    z = x + y
    a = torch.ops.aten._foobar(z)  # dummy function which errors
    return model(a)


compiled_test_backend_error = torch.compile(test_backend_error, backend="inductor")
compiled_test_backend_error()

При запуске этого вы получите эту ошибку с более длинной трассировкой стека ниже:

Traceback (most recent call last):
  File "/scratch/mlazos/torchdynamo/torchinductor/graph.py", line 246, in call_function
    return lowerings[target](*args, **kwargs)
  File "/scratch/mlazos/torchdynamo/torchinductor/lowering.py", line 185, in wrapped
    return decomp_fn(*args, **kwargs)
  File "/scratch/mlazos/torchdynamo/torchinductor/lowering.py", line 810, in _foobar
    assert False
AssertionError
...

ошибка с полной трассировкой стека

Если затем изменить torch.compile(backend="inductor") на torch.compile(backend="aot_eager"), он будет выполняться без ошибок, потому что проблема заключается в процессе понижения TorchInductor, а не в AOTAutograd.

Минимизация ошибок TorchInductor

Отсюда запустим минимизатор, чтобы получить минимальное воспроизведение. Установка переменной окружения TORCHDYNAMO_REPRO_AFTER=“aot” (или установка torch._dynamo.config.repro_after="aot" напрямую) сгенерирует программу Python, которая уменьшит граф, полученный AOTAutograd, до наименьшего подграфа, воспроизводящего ошибку. (См. ниже пример, где мы минимизируем граф, полученный TorchDynamo) Запуск программы с этой переменной окружения должен показать почти идентичный вывод с дополнительной строкой, указывающей, куда minifier_launcher.py был записан. Каталог вывода настраивается путем установки torch._dynamo.config.base_dir на допустимое имя каталога. Последний шаг — запустить минимизатор и проверить, что он выполняется успешно. Успешный запуск выглядит как это. Если минимизатор выполняется успешно, он генерирует исполняемый код Python, воспроизводящий точную ошибку. В нашем примере это следующий код:

import torch
from torch import tensor, device
import torch.fx as fx
from torch._dynamo.testing import rand_strided
from math import inf
from torch.fx.experimental.proxy_tensor import make_fx

# torch version: 1.13.0a0+gitfddfc44
# torch cuda version: 11.6
# torch git version: fddfc4488afb207971c54ad4bf58130fdc8a4dc5


# CUDA Info:
# nvcc: NVIDIA (R) Cuda compiler driver
# Copyright (c) 2005-2022 NVIDIA Corporation
# Built on Thu_Feb_10_18:23:41_PST_2022
# Cuda compilation tools, release 11.6, V11.6.112
# Build cuda_11.6.r11.6/compiler.30978841_0

# GPU Hardware Info:
# NVIDIA A100-SXM4-40GB : 8

from torch.nn import *

class Repro(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, add):
        _foobar = torch.ops.aten._foobar.default(add);  add = None
        return (_foobar,)

args = [((200, 200), (200, 1), torch.float32, 'cpu')]
args = [rand_strided(shape, stride, dtype, device) for shape, stride, dtype, device in args]
mod = make_fx(Repro())(*args)
from torch._inductor.compile_fx import compile_fx_inner

compiled = compile_fx_inner(mod, args)
compiled(*args)

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

Минимизация ошибок компилятора бэкенда

При использовании компиляторов бэкенда, отличных от TorchInductor, процесс поиска подграфа, вызывающего ошибку, практически идентичен процедуре в ошибках TorchInductor с одним важным нюансом. А именно, что минимизатор теперь будет запущен на графе, прослеженном TorchDynamo, а не на выходном графе AOTAutograd. Давайте рассмотрим пример.

import torch

import torch._dynamo as dynamo

model = torch.nn.Sequential(*[torch.nn.Linear(200, 200) for _ in range(5)])
# toy compiler which fails if graph contains relu
def toy_compiler(gm: torch.fx.GraphModule, _):
    for node in gm.graph.nodes:
        if node.target == torch.relu:
            assert False

    return gm


def test_backend_error():
    y = torch.ones(200, 200)
    x = torch.ones(200, 200)
    z = x + y
    a = torch.relu(z)
    return model(a)


compiled_test_backend_error = torch.compile(test_backend_error, backend=toy_compiler)
compiled_test_backend_error()

Для запуска кода после прохождения трассировки TorchDynamo по прямому графу можно использовать переменную среды TORCHDYNAMO_REPRO_AFTER. Запуск этой программы с TORCHDYNAMO_REPRO_AFTER=“dynamo” (или torch._dynamo.config.repro_after="dynamo") должен привести к этому выводу и следующему коду в {torch._dynamo.config.base_dir}/repro.py.

import torch
import torch._dynamo as dynamo
from torch import tensor, device
import torch.fx as fx
from torch._dynamo.testing import rand_strided
from math import inf
from torch._dynamo.debug_utils import run_fwd_maybe_bwd

from torch.nn import *

class Repro(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, add):
        relu = torch.relu(add);  add = None
        return (relu,)


mod = Repro().cuda()
opt_mod = torch.compile(mod, backend="None")


args = [((200, 200), (200, 1), torch.float32, 'cpu', False)]
args = [rand_strided(sh, st, dt, dev).requires_grad_(rg) for (sh, st, dt, dev, rg) in args]


with torch.cuda.amp.autocast(enabled=False):
    ref = run_fwd_maybe_bwd(mod, args)
    res = run_fwd_maybe_bwd(opt_mod, args)

Минимизатор успешно уменьшил граф до операции, вызвавшей ошибку в toy_compiler. Другое отличие от процедуры в ошибках TorchInductor заключается в том, что минимизатор автоматически запускается после обнаружения ошибки компилятора бэкенда. После успешного выполнения минимизатор записывает repro.py в torch._dynamo.config.base_dir.

Профилирование производительности

Доступ к профайлеру TorchDynamo

TorchDynamo имеет встроенную функцию stats для сбора и отображения времени, затраченного на каждой фазе компиляции. К этим данным можно получить доступ, вызвав torch._dynamo.utils.compile_times() после выполнения Torch._Dynamo. По умолчанию это возвращает строковое представление времени компиляции, затраченного на каждую функцию TorchDynamo по имени.

Отладка TorchInductor с помощью TORCH_COMPILE_DEBUG

TorchInductor имеет встроенные функции stats и trace для отображения времени, затраченного на каждой фазе компиляции, кода вывода, визуализации выходного графа и вывода IR. Это инструмент отладки, призванный упростить понимание и устранение неполадок в TorchInductor.

Давайте выполним пример со следующей тестовой программой (repro.py):

import torch

@torch.compile()
def test_model(x):
    model = torch.nn.Sequential(
        torch.nn.Linear(10, 10),
        torch.nn.LayerNorm(10),
        torch.nn.ReLU(),
    )
    return model(x)


y = test_model(torch.ones(10, 10))

Указание переменной среды TORCH_COMPILE_DEBUG=1 приведет к созданию каталога отладочных трассировок. По умолчанию этот каталог будет расположен в текущей директории и называться torch_compile_debug (это можно переопределить в поле конфигурации torchdynamo debug_dir_root и также в env var TORCH_COMPILE_DEBUG_DIR). Внутри этого каталога каждый запуск будет иметь отдельную папку, названную с временной меткой и идентификатором процесса запуска:

$ env TORCH_COMPILE_DEBUG=1 python repro.py
$ cd torch_compile_debug
$ ls
run_2023_03_01_08_20_52_143510-pid_180167

В папке запуска будет каталог torchdynamo, содержащий отладочные журналы, и каталог torchinductor, содержащий подкаталог для каждого скомпилированного ядра с артефактами отладки inductor.

$ cd
run_2023_03_01_08_20_52_143510-pid_180167
$ ls
torchinductor  torchdynamo

Переходя дальше в каталог torchinductor, файлы \*.log являются логами фазы AOT Autograd компиляции, model__0_forward_1.0 содержит артефакты отладки inductor.

$ cd torchinductor
$ ls
aot_model___0_debug.log  model__0_forward_1.0
$ cd model__0_forward_1.0
$ ls
debug.log  fx_graph_readable.py  fx_graph_runnable.py  fx_graph_transformed.py  ir_post_fusion.txt  ir_pre_fusion.txt  output_code.py

Вот краткое содержание содержимого:

• fx_graph_readable.py и fx_graph_runnable.py — читаемые и исполняемые версии fx_graph, полученные inductor.
• fx_graph_transformed.py — граф fx после выполнения всех fx-пасов inductor.
• ir\*.txt — inductor ir до и после слияния.
• output_code.py — скомпилированное ядро triton для подграфа.

Вот пример содержимого каталога отладки для тестовой программы:

import torch

@torch.compile()
def test_model(x):
    model = torch.nn.Sequential(
        torch.nn.Linear(10, 10),
        torch.nn.LayerNorm(10),
        torch.nn.ReLU(),
    )
    return model(x)


y = test_model(torch.ones(10, 10))

Каждый файл в этой отладочной трассировке можно включить и выключить через torch._inductor.config.trace.*. Профилирование и диаграммы по умолчанию отключены, так как их генерация дорогостояща.

Один узел в этом новом формате отладки выглядит следующим образом:

buf1: SchedulerNode(ComputedBuffer)
buf1.writes =
    {   MemoryDep(name='buf1', index=0, size=()),
        MemoryDep(name='buf1', index=0, size=(s0,))}
buf1.unmet_dependencies = {MemoryDep(name='buf0', index=c0, size=(s0,))}
buf1.met_dependencies = {MemoryDep(name='primals_2', index=c0, size=(s0,))}
buf1.group.device = cuda:0
buf1.group.iteration = (1, s0)
buf1.sizes = ([], [s0])
class buf1_loop_body:
    var_ranges = {z0: s0}
    index0 = z0
    index1 = 0
    def body(self, ops):
        get_index = self.get_index('index0')
        load = ops.load('buf0', get_index, False)
        get_index_1 = self.get_index('index0')
        load_1 = ops.load('primals_2', get_index_1, False)
        add = ops.add(load, load_1)
        get_index_2 = self.get_index('index1')
        reduction = ops.reduction('buf1', torch.float32, torch.float32, 'sum', get_index_2, add)
        return reduction

См. пример выходного каталога отладки для получения дополнительных примеров.

Разрывы графа

Учитывая программу такого вида:

def some_fun(x):
    ...

compiled_fun = torch.compile(some_fun, ...)
...

TorchDynamo попытается скомпилировать все операции torch/tensor в пределах some_fun в один граф FX, но может не захватить все в один граф.

Некоторые причины разрыва графика непреодолимы для TorchDynamo и не могут быть легко исправлены. - Вызов C-расширения, отличного от torch, невидим для torchdynamo и может выполнять произвольные действия без возможности TorchDynamo ввести необходимые защиты для обеспечения безопасности повторного использования скомпилированной программы. Разрывы графика могут препятствовать производительности, если полученные фрагменты невелики. Для максимальной производительности важно иметь как можно меньше разрывов графика.

Выявление причины разрыва графика

Чтобы определить все разрывы графика в программе и соответствующие причины разрывов, можно использовать torch._dynamo.explain. Этот инструмент запускает TorchDynamo для заданной функции и агрегирует обнаруженные разрывы графика. Вот пример использования:

import torch
import torch._dynamo as dynamo
def toy_example(a, b):
    x = a / (torch.abs(a) + 1)
    print("woo")
    if b.sum() < 0:
        b = b * -1
    return x * b
explanation, out_guards, graphs, ops_per_graph, break_reasons, explanation_verbose = (
    dynamo.explain(toy_example, torch.randn(10), torch.randn(10))
)
print(explanation_verbose)
"""
Dynamo produced 3 graphs, with 2 graph breaks and 6 ops.
 Break reasons:
1. call_function BuiltinVariable(print) [ConstantVariable(str)] {}
   File "t2.py", line 16, in toy_example
    print("woo")

2. generic_jump
   File "t2.py", line 17, in toy_example
    if b.sum() < 0:
 """

Выходные данные включают:

• out_guards — список списков, где каждый подсписок содержит проверки, которые должны быть пройдены для обеспечения валидности прослеженных графов.
• graphs — список модулей графов, которые были успешно прослежены.
• ops_per_graph — список списков, где каждый подсписок содержит операции, выполняемые в графе.

Чтобы выдать ошибку при первом обнаружении разрыва графика, используйте режим nopython. Этот режим отключает Python-обработку TorchDynamo и работает только в случае, если вся программа может быть преобразована в один граф. Пример использования:

def toy_example(a, b):
   ...

compiled_toy = torch.compile(toy_example, fullgraph=True, backend=<compiler>)

Чрезмерная перекомпиляция

Когда TorchDynamo компилирует функцию (или часть ее), он делает определенные предположения о локальных и глобальных переменных, чтобы позволить оптимизации компилятора, и выражает эти предположения как проверки, которые проверяют определенные значения во время выполнения. Если какая-либо из этих проверок терпит неудачу, Dynamo перекомпилирует эту функцию (или часть) до torch._dynamo.config.cache_size_limit раз. Если ваша программа достигает предела кэша, вам сначала нужно определить, какая проверка терпит неудачу и какая часть вашей программы ее вызывает.

Профилировщик компиляции автоматизирует процесс установки предела кэша TorchDynamo на 1 и запуска вашей программы под наблюдением «компилятора», который записывает причины любых сбоев проверок. Вы должны быть уверены, что запускаете программу не менее долго (много итераций), чем вы запускали, когда столкнулись с проблемой, и профилировщик будет накапливать статистику в течение этого времени.

Если ваша программа демонстрирует ограниченное количество динамики, вы можете настроить предел кэша TorchDynamo, чтобы позволить компиляцию и кэширование каждой вариации, но если предел кэша слишком высок, вы можете обнаружить, что затраты на перекомпиляцию перевешивают преимущества оптимизации.

torch._dynamo.config.cache_size_limit = <your desired cache limit>

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

from torch._dynamo.utils import CompileProfiler

def my_model():
    ...

with CompileProfiler() as prof:
    profiler_model = torch.compile(my_model, backend=prof)
    profiler_model()
    print(prof.report())

Отладка точности

Проблемы с точностью также могут быть минимизированы, если вы зададите переменную среды TORCHDYNAMO_REPRO_LEVEL=4. Она работает по принципу бинарного поиска git и полный репро может выглядеть примерно так TORCHDYNAMO_REPRO_AFTER="aot" TORCHDYNAMO_REPRO_LEVEL=4. Причина, по которой нам это нужно, заключается в том, что компиляторы ниже по потоку будут генерировать код, будь то код Triton или C++ бэкенд, числовые значения из этих компиляторов ниже по потоку могут отличаться незначительно, но оказывать драматическое влияние на стабильность вашей подготовки. Поэтому отладчик точности очень полезен для обнаружения ошибок в нашем кодогенерировании или с компилятором бэкенда.

Если вы хотите гарантировать, что генерация случайных чисел будет одинаковой как в torch, так и в triton, вы можете включить torch._inductor.config.fallback_random = True.