Черновой экспорт

Какие ошибки он обнаруживает?

Черновой экспорт помогает обнаруживать и отлаживать следующие ошибки:

• Проверки условий для ошибок, зависящих от данных
• Ошибки нарушения ограничений
• Отсутствующие фиктивные ядра
• Неправильно написанные фиктивные ядра

Как это работает?

При обычном экспорте мы преобразуем входные примеры в FakeTensors и используем их для записи операций и трассировки программы в граф. Формы входных тензоров, которые могут изменяться (задаются с помощью dynamic_shapes), или значения внутри тензоров (обычно полученные в результате вызова .item()) будут представлены символической формой (SymInt), а не конкретным целым числом. Однако во время трассировки могут возникнуть проблемы: например, мы можем столкнуться с проверками условий, которые не удаётся вычислить, например, если нужно проверить, больше ли некоторый элемент тензора нуля (u0 >= 0). Поскольку трассировщик ничего не знает о значении u0, он выдаст ошибку, зависящую от данных. Если в модели используется пользовательский оператор, для которого не определено фиктивное ядро, возникнет ошибка fake_tensor.UnsupportedOperatorException, поскольку экспорт не знает, как применить его к FakeTensors. Если фиктивное ядро пользовательского оператора реализовано неправильно, экспорт без предупреждения создаст неверный граф, не соответствующий поведению при eager-выполнении.

Для устранения перечисленных выше ошибок черновой экспорт использует трассировку реальными тензорами, чтобы определить, как действовать во время трассировки. При трассировке модели с фиктивными тензорами черновой экспорт также выполняет оператор над сохранёнными реальными тензорами, полученными из входных примеров, переданных в экспорт, для каждой операции с фиктивным тензором. Это позволяет устранить перечисленные выше ошибки: встретив проверку условия, которую не удаётся вычислить, например u0 >= 0, мы используем сохранённые значения реального тензора для вычисления этого условия. В граф будут добавлены проверки времени выполнения, чтобы гарантировать, что граф проверяет то же условие, которое мы предполагали во время трассировки. Если мы столкнёмся с пользовательским оператором без фиктивного ядра, то выполним обычное ядро оператора с сохранёнными реальными тензорами и вернём фиктивный тензор с тем же рангом, но с неподкреплёнными формами. Поскольку у нас есть результат в виде реального тензора для каждой операции, мы сравним его с результатом фиктивного тензора, полученным от фиктивного ядра. Если фиктивное ядро реализовано неправильно, мы обнаружим это поведение и сгенерируем более корректное фиктивное ядро.

Как использовать черновой экспорт?

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

class M(torch.nn.Module):
    def forward(self, x, y, z):
        res = torch.ops.mylib.foo2(x, y)

        a = res.item()
        a = -a
        a = a // 3
        a = a + 5

        z = torch.cat([z, z])

        torch._check_is_size(a)
        torch._check(a < z.shape[0])

        return z[:a]

inp = (torch.tensor(3), torch.tensor(4), torch.ones(3, 3))

ep = torch.export.export(M(), inp)

Это приводит к ошибке «отсутствует фиктивное ядро» для mylib.foo2, а затем к ошибке GuardOnDataDependentExpression из-за среза z с использованием a — неподкреплённого symint.

Чтобы вызвать draft-export, можно заменить строку torch.export следующим кодом:

ep = torch.export.draft_export(M(), inp)

ep — это допустимый ExportedProgram, который теперь можно передавать в следующие инструменты!

Отладка с помощью чернового экспорта

В выводе терминала при запуске чернового экспорта должно появиться следующее сообщение:

#########################################################################################
WARNING: 2 issue(s) found during export, and it was not able to soundly produce a graph.
To view the report of failures in an html page, please run the command:
    `tlparse /tmp/export_angelayi/dedicated_log_torch_trace_axpofwe2.log --export`
Or, you can view the errors in python by inspecting `print(ep._report)`.
########################################################################################

Черновой экспорт автоматически сохраняет журналы для tlparse. Просмотреть ошибки трассировки можно с помощью print(ep._report) или передать журналы в tlparse, чтобы создать HTML-отчёт.

Запуск команды tlparse в терминале создаст HTML-отчёт tlparse. Ниже приведён пример отчёта tlparse:

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

• Трассировку стека в месте возникновения ошибки
• Список локальных переменных и их форм
• Информацию о том, как было создано это условие

Возвращаемая экспортированная программа

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

Например, если в графе есть ветвление в зависимости от того, больше ли значение 5, и входные примеры для чернового экспорта были больше 5, то возвращаемый ExportedProgram будет специализирован для этой ветви и будет проверять, что значение больше 5. Это означает, что программа выполнится успешно, если передать ей другое значение больше 5, но завершится с ошибкой, если передать значение меньше 5. Это надёжнее, чем torch.jit.trace, который без предупреждения специализируется на этой ветви. Чтобы torch.export корректно поддерживал обе ветви, нужно переписать код с помощью torch.cond, который захватит обе ветви.

Благодаря проверкам времени выполнения в графе возвращённую экспортированную программу также можно повторно трассировать с помощью torch.export или torch.compile, добавив небольшое изменение в случае отсутствия фиктивного ядра у пользовательского оператора.

Генерация фиктивных ядер

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

Чтобы решить эту проблему, после чернового экспорта мы создадим профиль оператора для каждого вызова пользовательского оператора и сохраним его в отчёте, прикреплённом к экспортированной программе: ep._report.op_profiles. Затем пользователи смогут использовать менеджер контекста torch._library.fake_profile.unsafe_generate_fake_kernels, чтобы сгенерировать и зарегистрировать фиктивную реализацию на основе этих профилей операторов. Благодаря этому последующая повторная трассировка с фиктивными тензорами будет работать.

Рабочий процесс может выглядеть примерно так:

class M(torch.nn.Module):
    def forward(self, a, b):
        res = torch.ops.mylib.foo(a, b)  # no fake impl
        return res

ep = draft_export(M(), (torch.ones(3, 4), torch.ones(3, 4)))

with torch._library.fake_profile.unsafe_generate_fake_kernels(ep._report.op_profiles):
    decomp = ep.run_decompositions()

new_inp = (
    torch.ones(2, 3, 4),
    torch.ones(2, 3, 4),
)

# Save the profile to a yaml and check it into a codebase
save_op_profiles(ep._report.op_profiles, "op_profile.yaml")
# Load the yaml
loaded_op_profile = load_op_profiles("op_profile.yaml")

Профиль оператора представляет собой словарь, сопоставляющий имена операторов с наборами профилей, которые описывают входные и выходные данные оператора. Профиль можно написать вручную, сохранить в YAML-файл и добавить в кодовую базу. Вот пример профиля для mylib.foo.default:

"mylib.foo.default": {
    OpProfile(
        args_profile=(
            TensorMetadata(
                rank=2,
                dtype=torch.float32,
                device=torch.device("cpu"),
                layout=torch.strided,
            ),
            TensorMetadata(
                rank=2,
                dtype=torch.float32,
                device=torch.device("cpu"),
                layout=torch.strided,
            ),
        ),
        out_profile=TensorMetadata(
            rank=2,
            dtype=torch.float32,
            device=torch.device("cpu"),
            layout=torch.strided,
        ),
    )
}

Профиль mylib.foo.default содержит только один вариант: для двух входных тензоров ранга 2, с типом данных torch.float32 и устройством cpu будет возвращён один тензор ранга 2, с типом данных torch.float32 и устройством cpu. Менеджер контекста затем сгенерирует фиктивное ядро, которое для двух входных тензоров ранга 2 (и остальных метаданных тензоров) будет возвращать один тензор ранга 2 (и остальные метаданные тензора).

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

Что делать дальше?

Теперь, когда с помощью чернового экспорта мы успешно создали ExportedProgram, можно использовать следующие компиляторы, например AOTInductor, чтобы оптимизировать его производительность и создать исполняемый артефакт. Эту оптимизированную версию затем можно использовать для развёртывания. Параллельно можно использовать отчёт, созданный черновым экспортом, чтобы выявить и исправить ошибки torch.export, благодаря чему исходную модель можно будет напрямую трассировать с помощью torch.export.