Расширение экспортеров

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

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

Патчи и исправления

Обходным решением может быть как заплатка (patch), так и исправление (fix). Разница между ними заключается в том, можно ли их откатить.

Оба находятся в реестре в exporters/utils.py, а экспортер устанавливает всё, что зарегистрировано для его бэкенда, на соответствующем этапе.

Обращайтесь к патчу, если проблема заключается в ошибке понижения (lowering) в одном конкретном бэкенде: отсутствует перевод в ONNX, есть особенность валидации ORT или разложение FX генерирует неиспользуемую операцию. Обходное решение остаётся в экспортере, а код моделирования остаётся чистым.

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

Добавить патч

Допустим, метод модели выполняет действие, которое torch.export не может проследить. Метод NllbMoeTop2Router._cast_classifier в NLLB-MoE приводит веса классификатора к другому типу данных, что не поддаётся трассировке. Замените его на пустую операцию на время экспорта.

Напишите фабрику, которая принимает исходный метод и возвращает его замену, затем зарегистрируйте фабрику по точечному пути метода:

from transformers.exporters.utils import register_patch

@register_patch("dynamo", "transformers.models.nllb_moe.modeling_nllb_moe.NllbMoeTop2Router._cast_classifier")
def _patch_classifier_cast(_original):
    # Replace the untraceable dtype cast with a no-op during export.
    return lambda self, *args, **kwargs: None

Экспортер заменяет метод перед трассировкой и восстанавливает его после, поэтому патч влияет только на экспорт. Несколько вариантов:

• Передавайте дополнительные пути, чтобы использовать один фабрика в нескольких местах вызова, например @register_patch("dynamo", path_a, path_b).
• Патчите операцию torch, указывая на неё путь, например @register_patch("onnx", "torch.where"). Фабрика получает реальную операцию в качестве аргумента, поэтому замена может вызывать её.
• Напишите фикс вместо патча, если вам нужно переписать граф после трассировки. Механизм тот же: декорированная функция в соответствующем реестре.

Этап ссылок

Исходный код каждого экспортера помечает свои этапы блоками комментариев вида # ── Stage N: … ──, поэтому структура файла и этот справочник совпадают. Обращайтесь к исходному коду, чтобы узнать точные операции и классы, обрабатываемые на каждом этапе.

DynamoExporter

Базовый экспортатор выполняет один этап патча и четыре вспомогательных функции последовательно внутри DynamoExporter.export (см. exporter_dynamo.py).

1. Патч сигнатуры forward: придаёт model.forward плоскую сигнатуру аргументов, чтобы torch.export не объединял входные данные в один кортеж **kwargs. Это внутренний механизм, а не точка расширения.
2. Патчи модели: во время трассировки заменяют нетрассируемые методы модели на эквиваленты, безопасные для экспорта. Расширяется с помощью @register_patch("dynamo", ...).
3. Регистрация Pytree: зарегистрируйте каждый Cache и ModelOutput, чтобы torch.export мог их развернуть и собрать заново (обычно это происходит автоматически). Добавьте ветку в _flatten_to_context / _unflatten_from_context для типа, до которого обход атрибутов не может добраться.
4. Динамические формы: присваивайте Dim.AUTO каждому тензору и кэшируйте лист, когда dynamic=True. Переопределите с помощью DynamoConfig.dynamic_shapes.
5. Очистка состояния: сброс атрибутов тензоров, которые модель устанавливает внутри forward, а torch.export оставляет как фиктивные тензоры. Расширьте, добавив имя атрибута в _STATEFUL_CACHE_ATTRS.

OnnxExporter

OnnxExporter добавляет пять этапов вокруг torch.onnx.export (см. exporter_onnx.py). Используйте grep для поиска полного списка патчей в файле:

grep -nE "^def (_patch_|_fix_|_aten_)" src/transformers/exporters/exporter_onnx.py

1. Патчи Torch: замена операций torch, которые экспортёр ONNX не может преобразовать напрямую. Расширяется с помощью @register_patch("onnx", ...).
2. Патчи ONNX: повторное выполнение исправлений узлов после run_decompositions, чтобы вновь введённые узлы проверки форм были исправлены до понижения. Используется тот же реестр @register_patch("onnx", ...).
3. Исправления узлов FX: переписывание узлов графа, которые экспортёр ONNX не может опустить, таких как операции alias, представления in-place и мёртвые assert. Расширение с помощью @register_fx_node_fix("onnx").
4. Переводы ONNX: предоставьте собственный lowering для aten-операции, если стандартный отсутствует или содержит ошибку (например, aten.index_put или aten._grouped_mm). Добавьте функцию _aten_* в _ONNX_TRANSLATION_TABLE.
5. Исправления ONNX IR: переписывайте готовую программу ONNX, чтобы обойти ошибки ONNX Runtime (например, принудительно задавая TopK(sorted=True)). Добавьте функцию _fix_ir_* в _IR_FIXES.

ExecutorchExporter

ExecutorchExporter добавляет пять этапов вокруг to_edge_transform_and_lower и to_executorch, начиная с подготовки бэкенда (см. exporter_executorch.py).

1. Подготовка бэкенда: переместите модель на целевое устройство и dtype, а также выберите её партиционер (prepare_for_xnnpack, prepare_for_cuda). Добавьте бэкенд, зарегистрировав prepare_for_<name> в _BACKEND_PREPARE.
2. Патчи Torch: заменяют операции torch, которые бэкенды ExecuTorch не могут обработать, например split_copy, chunk и topk(k>dim). Расширяются с помощью @register_patch("executorch", ...).
3. Патчи ExecuTorch: заменяют внутренние компоненты ExecuTorch, которые вызывают сбой на валидных графах с динамической формой. Используется тот же реестр @register_patch("executorch", ...).
4. Исправления для FX-программ: восстановление экспортированной программы в тех случаях, когда для исправления требуется контекст всей программы, например, расширение ограничений диапазона или заполнение отсутствующих метаданных плейсхолдеров. Расширяется с помощью @register_fx_program_fix("executorch").
5. Исправления узлов FX: переписывание отдельных узлов, например, отображение символьных операций Python на executorch_prim.* или переписывание pow в цепочку mul. Расширение с помощью @register_fx_node_fix("executorch").

Известные обходные решения в исходном коде

Некоторые классы моделей сталкиваются с подтверждёнными ошибками в оптимизаторе графов onnxscript (сбой при свёртке констант на SplitToSequence, потеря инициализаторов FPN). ONNX_DISABLE_OPTIMIZE отключает оптимизацию onnxscript для этих моделей. В каждой записи рядом с названием модели указан соответствующий upstream-баг. Ожидается, что список будет сокращаться по мере исправления ошибок upstream, поэтому новая запись должна ссылаться на конкретный upstream-баг, а не отключать оптимизацию произвольно.

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