エクスポート機能の拡張

torch.export はモデルをグラフにトレースし、その後のステージでそのグラフを変換し、最終ステージでターゲットランタイム向けにローワーまたはエミットします。ほとんどのモデルはそのまま通過します。バックエンドが処理できない PyTorch のパターンがある場合、エクスポートアプライはモデルを編集するのではなく、そのステージで小さなワークアラウンドを適用します。

デコレーターで1つの関数を登録してワークアラウンドを追加します。各ワークアラウンドは、それを明確に表現できる最も低いステージに配置します。

パッチと修正

回避策には、パッチと修正の2種類があります。両者の違いは、元に戻せるかどうかです。

両者はexporters/utils.py内のレジストリに存在し、exporter はそのバックエンド用に登録されたすべてのものを適切な段階でインストールします。

問題が単一のバックエンドのローリングのバグ（ONNX変換の欠落、ORTの検証のクイック、デッドオペレーションを出力するFX分解など）である場合は、パッチを使用してください。ワークアラウンドはエクスポート側に留まり、モデリングコードはクリーンな状態を保てます。

パターンブロックがデータ依存ループ、Cache 外部の状態付きキャッシュ、手書きの分割ループアテンションなど、複数のバックエンドにまたがってエクスポートされる場合は、モデル側を修正してください。モデルを一度修正しておくと、すべてのエクスポートツールで恩恵を受けることができます。

パッチを追加

モデルのメソッドがtorch.exportでトレースできない処理を実行する場合を想定します。NLLB-MoEのNllbMoeTop2Router._cast_classifierは、分類器の重みを別のdtypeにキャストしますが、これはトレースできません。エクスポートの間だけ、これを何もしない処理（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

エクスポートツールは、トレースの前にメソッドを差し替え、トレース後に元に戻すため、このパッチはエクスポートにのみ影響します。いくつかの変形例を以下に示します。

• 呼び出し箇所間で1つのファクトリを共有するために、追加のパスを渡します。例えば、@register_patch("dynamo", path_a, path_b)のように指定します。
• torch のオペレーションをパッチするには、そのパスを指定します。例えば @register_patch("onnx", "torch.where") のようにします。 ファクトリは実際のオペレーションを引数として受け取るため、置換側から元のオペレーションを呼び出すことができます。
• トレース後にグラフを書き直す必要がある場合は、パッチではなく修正（fix）を作成してください。仕組みは同じで、対応するレジストリにデコレータ付きの関数を登録します。

ステージの参照

各エクスポート元のソースでは、# ── Stage N: … ── というコメントブロックで各ステージがラベル付けされているため、ファイルとこのリファレンスは対応しています。各ステージが処理する正確なオペレーションとクラスについては、そちらを参照してください。

DynamoExporter

ベースのエクスポート処理は、DynamoExporter.export 内で、1つのパッチステージと4つのヘルパーを順に実行します（exporter_dynamo.py を参照）。

1. Forwardシグネチャのパッチ: 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 の前後に5つのステージを追加します（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 の実行後にノード修正を再実行し、新たに導入された shape-guard ノードが lowering 前に修復されるようにします。同じ @register_patch("onnx", ...) レジストリを使用します。
3. FXノードの修正: ONNXエクスポート機能が低レベル化できないグラフノード(エイリアス演算子、インプレースビュー、デッドアサートなど)を書き換えます。@register_fx_node_fix("onnx")で拡張できます。
4. ONNX 変換: デフォルトが欠落しているかバグがある aten オペレーションに対して、カスタムなローリングを指定します(例: 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ノードの修正: powをmulのチェーンとして書き換えるなど、Pythonのsym opsをexecutorch_prim.*にマッピングする個々のノードを書き換えます。@register_fx_node_fix("executorch")で拡張します。

既知の上流ワークアラウンド

一部のモデルクラスでは、onnxscriptのグラフオプティマイザーで確認済みのバグ（SplitToSequenceでの定数フォールディングのクラッシュ、FPNのイニシャライザーの欠落）が発生します。ONNX_DISABLE_OPTIMIZEは、これらのモデルに対してonnxscriptの最適化を無効にします。各エントリには、モデル名の横にアップストリームの問題を記録しています。アップストリームのバグが修正されるにつれてリストは縮小すると予想されるため、新しいエントリは恣意的に最適化を無効にするのではなく、特定のアップストリームのバグを参照する必要があります。

EXPORT_SKIPS、 モデルがそのままでは根本的にエクスポート不可能な場合（ベクトル化できないデータ依存の制御フロー、または forward 引数として扱われるモジュールなど）、一部のモデルクラスをエクスポートスウィープから完全に除外するためのオプションです。各エントリには、必要なモデル側の修正を特定する理由が記載されています。このリストは増加するのではなく、縮小することが期待されています。