# エクスポート機能の拡張

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

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

## パッチと修正

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

|              | Patch                                                                          | 修正                                                                                   |
| ------------ | ------------------------------------------------------------------------------ | ------------------------------------------------------------------------------------- |
| その機能 | エクスポートの期間中、属性(`torch`のオペレーション、ExecuTorchの内部処理、またはモデルのメソッド)を別のものに置き換えます | トレースされたグラフまたはプログラムを書き換えます                                                  |
| 取り消しました     | はい、その後で元の状態に復元されます                                        | いいえ、次のステージが実行される前に、そのアーティファクトを修復します。                                |
| 登録 | `@register_patch(backend, *paths)`                                            | `@register_fx_node_fix(backend)` または `@register_fx_program_fix(backend)`               |

両者は[exporters/utils.py](https://github.com/huggingface/transformers/blob/main/src/transformers/exporters/utils.py)内のレジストリに格納されており、exporter はそのバックエンド用に登録されたすべての処理を適切な段階でインストールします。

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

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

## パッチを追加

モデルのメソッドが`torch.export`でトレースできない処理を実行する場合を想定します。NLLB-MoEの`NllbMoeTop2Router._cast_classifier`は、分類器の重みを別のdtypeにキャストしますが、これはトレースできません。エクスポートの間だけ、これを何もしない処理（no-op）に置き換えてください。

元のメソッドを受け取り、その代替メソッドを返すファクトリを作成し、メソッドのドット区切りパスに対してファクトリを登録します:

```python
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](https://github.com/huggingface/transformers/blob/main/src/transformers/exporters/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](https://github.com/huggingface/transformers/blob/main/src/transformers/exporters/exporter_onnx.py) を参照）。パッチの完全なリストについては、ファイル内で grep してください：

```bash
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](https://github.com/huggingface/transformers/blob/main/src/transformers/exporters/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 sym opsのマッピングや、`pow`を`mul`のチェーンとして書き換えるなど、個々のノードを書き換えます。`@register_fx_node_fix("executorch")`で拡張します。

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

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

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