에クスポート拡張

torch.export는 모델을 그래프로 추적하고, 이후 단계에서 해당 그래프를 변환하며, 최종 단계에서는 대상 런타임에 맞게 낮추거나 출력합니다. 대부분의 모델은 그대로 통과합니다. 백엔드가 처리할 수 없는 PyTorch 패턴이 있는 경우, 익스포터는 모델을 수정하는 대신 해당 단계에서 작은 우회 조치를 적용합니다.

워크어라운드를 추가하려면 함수 하나를 작성하고 데코레이터로 등록합니다. 각 워크어라운드는 이를 깔끔하게 표현할 수 있는 가장 낮은 단계에 배치해야 합니다.

패치 및 수정

우회 방법은 패치 또는 수정입니다. 두 방식은 되돌릴 수 있는지 여부에 따라 구별됩니다.

두 개 모두 exporters/utils.py의 레지스트리에 있으며, exporter는 해당 백엔드에 등록된 모든 항목을 적절한 단계에서 설치합니다.

문제가 단일 백엔드의 lowering 버그, 즉 ONNX 변환 누락, ORT 검증의 특이한 동작, 또는 dead op를 출력하는 FX decomposition인 경우 패치를 사용하세요. 이 경우 워크어라운드는 exporter에 유지되고 모델링 코드는 깔끔하게 유지됩니다.

Cache 외부의 상태ful 캐시, 데이터 의존 루프, 수동으로 작성된 split-loop attention과 같은 여러 백엔드로의 패턴 블록 내보내기가 발생하는 경우, 모델을 수정해야 합니다. 모델을 한 번만 수정해도 모든 exporter에 도움이 됩니다.

패치 추가

모델 메서드가 torch.export로 추적할 수 없는 작업을 수행한다고 가정해 보겠습니다. NLLB-MoE의 NllbMoeTop2Router._cast_classifier는 분류기 가중치를 다른 dtype으로 캐스팅하는데, 이는 추적할 수 없습니다. export가 진행되는 동안 이를 no-op로 대체하세요.

…

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

exporter는 트레이싱 전에 메서드를 교체하고 이후에 원래대로 복원하므로, 패치는 export에만 영향을 미칩니다. 몇 가지 변형:

• 여러 호출 지점에서 하나의 팩토리를 공유하려면 추가 경로를 전달합니다. 예를 들어 @register_patch("dynamo", path_a, path_b).
• torch op의 경로를 지정하여 패치합니다. 예를 들어 @register_patch("onnx", "torch.where")와 같이 사용합니다. 팩토리는 실제 op를 인수로 받으므로, 대체 로직에서 해당 op를 호출할 수 있습니다.
• 트레이싱 후 그래프를 다시 작성해야 하는 경우 패치가 아닌 수정을 작성하세요. 메커니즘은 동일하며, 매칭 레지스트리에 있는 데코레이터가 적용된 함수입니다.

Stage 참조

각 exporter의 소스에는 # ── Stage N: … ── 주석 블록으로 단계가 표시되어 있어, 해당 파일과 이 참고 문서가 일치합니다. 각 단계가 처리하는 정확한 연산과 클래스는 해당 소스에서 확인하세요.

DynamoExporter

기본 exporter는 DynamoExporter.export 안에서 패치 단계 1개와 헬퍼 4개를 순서대로 실행합니다(자세한 내용은 exporter_dynamo.py를 참조하세요).

1. Forward-signature 패치: model.forward에 평탄한 인자 시그니처를 부여하여 torch.export가 입력을 하나의 **kwargs 튜플로 묶지 않도록 합니다. 이는 내부 전용이며 확장 지점이 아닙니다.
2. 모델 패치: 트레이싱 중에 추적할 수 없는 모델 메서드를 내보내기 안전한 대안으로 교체합니다. @register_patch("dynamo", ...)로 확장할 수 있습니다.
3. Pytree 등록: torch.export가 Cache와 ModelOutput을 평탄화하고 다시 구성할 수 있도록 각각 등록합니다(보통 자동으로 수행됩니다). 속성 탐색으로 도달할 수 없는 타입에 대해서는 _flatten_to_context / _unflatten_from_context에 분기를 추가합니다.
4. 동적 형상: dynamic=True인 경우 모든 텐서와 캐시 리프에 Dim.AUTO를 할당합니다. DynamoConfig.dynamic_shapes로 오버라이드할 수 있습니다.
5. 상태 정리: torch.export가 가짜 텐서로 남기는, 모델이 forward 내부에서 설정하는 텐서 속성을 초기화합니다. _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 패치: ONNX 익스포터가 그대로 변환할 수 없는 torch 연산을 교체합니다. @register_patch("onnx", ...)로 확장할 수 있습니다.
2. ONNX 패치: run_decompositions 실행 후 노드 수정을 다시 실행하여, lowering 전에 새로 도입된 shape-guard 노드가 복구되도록 합니다. 동일한 @register_patch("onnx", ...) 레지스트리를 사용합니다.
3. FX 노드 수정: ONNX exporter가 lower할 수 없는 alias ops, in-place views, dead asserts와 같은 그래프 노드를 재작성합니다. @register_fx_node_fix("onnx")로 확장할 수 있습니다.
4. ONNX 번역: 기본 lowering이 없거나 버그가 있는 aten op에 대해 사용자 정의 lowering을 제공합니다(예: aten.index_put 또는 aten._grouped_mm). _ONNX_TRANSLATION_TABLE에 _aten_* 함수를 추가합니다.
5. ONNX IR 수정: 완성된 ONNX 프로그램을 다시 작성하여 ONNX Runtime의 버그를 우회합니다(예: TopK(sorted=True) 강제 적용). _IR_FIXES에 _fix_ir_* 함수를 추가합니다.

ExecutorchExporter

ExecutorchExporter는 to_edge_transform_and_lower와 to_executorch 주위에 5개의 단계를 추가하며, 백엔드 준비 단계로 시작합니다(자세한 내용은 exporter_executorch.py를 참조).

1. 백엔드 준비: 모델을 대상 디바이스와 dtype으로 이동하고 파티셔너를 선택합니다(prepare_for_xnnpack, prepare_for_cuda). _BACKEND_PREPARE에 prepare_for_<name>를 등록하여 백엔드를 추가합니다.
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 노드 수정: executorch_prim.*으로 Python 심볼 연산을 매핑하거나 pow를 mul 체인으로 재작성하는 등 개별 노드를 재작성합니다. @register_fx_node_fix("executorch")로 확장합니다.

알려진 상위 프로젝트 우회 방법

일부 모델 클래스는 onnxscript 그래프 최적화기에서 확인된 버그(예: SplitToSequence에서 상수 접기 시 크래시, FPN 초기화 값 누락)에 영향을 받습니다. ONNX_DISABLE_OPTIMIZE는 해당 모델에 대해 onnxscript 최적화를 비활성화합니다. 각 항목은 모델 이름 옆에 해당 업스트림 이슈를 기록합니다. 업스트림 버그가 수정됨에 따라 목록은 축소될 것으로 예상되므로, 새로운 항목은 임의로 최적화를 비활성화하는 것이 아니라 특정 업스트림 버그를 참조해야 합니다.

EXPORT_SKIPS, 모델이 근본적으로 그대로 내보낼 수 없는 경우(벡터화할 수 없는 데이터 의존 제어 흐름, 또는 forward 인수로 취급되는 모듈)에 일부 모델 클래스를 export sweep에서 완전히 제외하는 옵션입니다. 각 항목에는 모델 측에서 필요한 변경 사항을 명시하는 이유가 포함됩니다. 이 목록은 커지기보다 줄어드는 것이 기대됩니다.