exporters 확장하기

exporters 확장하기 (Extending the exporters)

torch.export는 모델을 그래프로 추적(trace)하고, 이후 단계들이 그 그래프를 변환하며, 마지막 단계가 타깃 런타임을 위해 그래프를 내리거나 방출(emit)해요. 대부분의 모델은 손대지 않고 그대로 통과해요. 백엔드가 처리할 수 없는 PyTorch 패턴이 있다면, exporter가 모델을 수정하는 대신 그 단계에서 작은 우회책(workaround)을 적용해요.

첫 번째 동작(함수) 하나를 작성하고 데코레이터로 등록하면 우회책을 추가할 수 있어요. 각 우회책은 그것을 깔끔하게 표현할 수 있는 가장 낮은 단계에 속해야 해요.

출처: 문서

본문

Patch와 Fix

우회책은 patch 또는 fix 중 하나예요. 둘의 차이는 되돌릴 수 있는지(revert) 여부예요.

Patch Fix
하는 일 내보내기(export) 기간 동안 속성(torch op, ExecuTorch 내부, 또는 모델 메서드)을 교체해요 추적된 그래프나 프로그램을 다시 작성해요
되돌리기 네, 이후 원본이 복원돼요 아니오, 다음 단계가 실행되기 전에 결과물을 수리해요
등록 방법 @register_patch(backend, *paths) @register_fx_node_fix(backend) 또는 @register_fx_program_fix(backend)

둘 다 exporters/utils.py의 레지스트리에 살고 있어요. exporter는 해당 백엔드에 등록된 모든 것을 올바른 단계에서 설치해요.

문제가 단일 백엔드의 하강(lowering) 버그일 때, 즉 ONNX 번역이 빠져 있거나, ORT 검증의 특이한 점, 또는 죽은 op를 방출하는 FX 분해(decomposition)일 때는 patch를 사용해요. 우회책은 exporter 안에 남고, 모델링 코드는 깨끗하게 유지돼요.

데이터 의존 루프, Cache 밖의 상태 저장 캐시, 손으로 작성한 split-loop attention처럼 여러 백엔드에서 export를 막는 패턴이라면 모델을 대신 고쳐야 해요. 모델에서 한 번 고치면 모든 exporter가 도움을 받을 수 있어요.

patch 추가하기

모델 메서드가 torch.export가 추적할 수 없는 일을 한다고 가정해 봐요. NLLB-MoE의 NllbMoeTop2Router._cast_classifier는 분류기 가중치를 다른 dtype으로 캐스팅하는데, 이는 추적할 수 없어요. export 기간 동안 no-op로 교체해요.

원본 메서드를 받아 교체물을 반환하는 팩토리를 작성하고, 메서드의 점으로 구분된 경로(dotted path)에 대해 팩토리를 등록해요.

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는 추적 전에 메서드를 교체하고, 추적 후에 복원해요. 그래서 patch는 export에만 영향을 줘요. 몇 가지 변형이 있어요:

  • 하나의 팩토리를 여러 호출 지점에서 공유하려면 추가 경로를 전달해요. 예: @register_patch("dynamo", path_a, path_b).
  • 경로를 그 op를 가리키게 해서 torch op를 patch할 수 있어요. 예: @register_patch("onnx", "torch.where"). 팩토리는 실재 op를 인자로 받으므로, 교체물이 이를 통해 호출할 수 있어요.
  • 추적 후 그래프를 다시 작성해야 한다면 patch 대신 fix를 작성해요. 메커니즘은 동일하고, 일치하는 레지스트리의 데코레이션된 함수를 쓰면 돼요.

단계 참조 (Stage reference)

각 exporter의 소스는 자신의 단계를 # ── Stage N: … ── 주석 블록으로 표시하므로, 파일과 이 참조가 맞춰져요. 각 단계가 처리하는 정확한 op와 클래스는 거기서 확인해요.

DynamoExporter

기본 exporter는 DynamoExporter.export 안에서 patch 단계 하나와 도우미 4개를 순서대로 실행해요 (exporter_dynamo.py 참고).

  1. Forward-signature patch: model.forward에 평평한 인자 시그니처를 줘서 torch.export가 입력을 하나의 **kwargs 튜플로 묶지 않게 해요. 이것은 내부적이며 확장 지점이 아니에요.
  2. Model patches: 추적 중에 추적 불가능한 모델 메서드를 export에 안전한 등가물로 교체해요. @register_patch("dynamo", ...)로 확장해요.
  3. Pytree 등록: 각 Cache와 ModelOutput을 등록해서 torch.export가 이를 펼치고(flatten) 다시 만들 수 있게 해요 (보통 자동으로 일어나요). 속성 탐색이 도달할 수 없는 타입은 _flatten_to_context / _unflatten_from_context에 브랜치를 추가해요.
  4. Dynamic shapes: dynamic=True일 때 모든 텐서와 캐시 리프에 Dim.AUTO를 할당해요. DynamoConfig.dynamic_shapes로 덮어쓸 수 있어요.
  5. State cleanup: 모델이 forward 안에서 설정하고 torch.export가 fake tensor로 남겨두는 텐서 속성을 재설정해요. 속성 이름을 _STATEFUL_CACHE_ATTRS에 추가해서 확장해요.

OnnxExporter

OnnxExporter는 torch.onnx.export 주위에 다섯 단계를 추가해요 (exporter_onnx.py 참고). 전체 patch 목록은 파일에서 grep으로 찾아봐요:

grep -nE "^def (_patch_|_fix_|_aten_)" src/transformers/exporters/exporter_onnx.py
  1. Torch patches: ONNX exporter가 그대로 번역할 수 없는 torch op를 교체해요. @register_patch("onnx", ...)로 확장해요.
  2. ONNX patches: run_decompositions 후에 노드 fix를 다시 실행해서 새로 생긴 shape-guard 노드가 하강 전에 수리되게 해요. 동일한 @register_patch("onnx", ...) 레지스트리를 사용해요.
  3. FX node fixes: ONNX exporter가 내릴 수 없는 그래프 노드(alias op, in-place view, 죽은 assert 등)를 다시 작성해요. @register_fx_node_fix("onnx")로 확장해요.
  4. ONNX translations: 기본 번역이 없거나 버그가 있는 aten op에 대해 커스텀 하강을 제공해요 (예: aten.index_put 또는 aten._grouped_mm). _ONNX_TRANSLATION_TABLE에 _aten_* 함수를 추가해요.
  5. ONNX IR fixes: ONNX Runtime 버그를 우회하려고 완성된 ONNX 프로그램을 다시 작성해요 (예: TopK(sorted=True) 강제). _IR_FIXES에 _fix_ir_* 함수를 추가해요.

ExecutorchExporter

ExecutorchExporter는 백엔드 준비로 시작해서 to_edge_transform_and_lower와 to_executorch 주위에 다섯 단계를 추가해요 (exporter_executorch.py 참고).

  1. Backend preparation: 모델을 타깃 디바이스와 dtype으로 옮기고 파티셔너를 선택해요 (prepare_for_xnnpack, prepare_for_cuda). _BACKEND_PREPARE에 prepare_for_<name>을 등록해서 백엔드를 추가해요.
  2. Torch patches: ExecuTorch 백엔드가 받아들일 수 없는 torch op(split_copy, chunk, topk(k>dim) 등)를 교체해요. @register_patch("executorch", ...)로 확장해요.
  3. ExecuTorch patches: 유효한 dynamic-shape 그래프에서 크래시나는 ExecuTorch 내부를 교체해요. 동일한 @register_patch("executorch", ...) 레지스트리를 사용해요.
  4. FX program fixes: 전체 프로그램 컨텍스트가 필요한 수정(범위 제약 넓히기, 빠진 placeholder 메타데이터 채우기 등)에서 내보낸 프로그램을 수리해요. @register_fx_program_fix("executorch")로 확장해요.
  5. FX node fixes: 개별 노드를 다시 작성해요. 예를 들어 Python sym op를 executorch_prim.*로 매핑하거나 pow를 mul 체인으로 다시 쓰는 것. @register_fx_node_fix("executorch")로 확장해요.

알려진 업스트림 우회책

몇몇 모델 클래스는 onnxscript 그래프 옵티마이저의 확인된 버그(SplitToSequence에서 constant folding 크래시, FPN initializer 누락 등)에 부딪혀요. ONNX_DISABLE_OPTIMIZE는 그 모델들에 대해 onnxscript 최적화를 비활성화해요. 각 항목은 모델 이름 옆에 업스트림 이슈를 기록해요. 업스트림 버그가 해결되면 목록이 줄어들 것으로 기대되므로, 새 항목은 임의로 최적화를 비활성화하는 게 아니라 특정 업스트림 버그를 참조해야 해요.

EXPORT_SKIPS는 모델이 근본적으로 그대로 내보낼 수 없을 때(벡터화할 수 없는 데이터 종속 제어 흐름, 또는 forward 인자로 간주되는 모듈) 소수의 모델 클래스를 export 스윕(sweep)에서 완전히 제외해요. 각 항목은 필요한 모델 측 변경을 지목하는 이유를 담아요. 이 목록도 늘어나기보다 줄어들 것으로 기대돼요.

더 알아보기 (Learn more)

  • exporters의 시작 방법과 기본 사용법은 Exporters 문서를 확인해 보세요.