新API `torch.export` で PyTorch モデルデプロイを効率化

新しい人工知能または機械学習プロジェクトに着手する際、多くの焦点は自然と巨大なタスクに集まります。例えば、膨大なデータセットのキュレーション、洗練されたモデルの設計、トレーニング用の強力なGPUクラスターの確保などです。しかし、予期せぬ障害となるのは、一見些細に見える詳細であることが多く、それがフラストレーションのたまるバグや深刻な生産遅延につながります。その典型的な例が、トレーニング済みモデルを開発環境から推論環境へ引き渡すプロセスです。このステップは簡単に見えるかもしれませんが、異なるランタイムライブラリ、ハードウェア構成、バージョン管理の現実が、それをかなりの頭痛の種に変えることがあります。開発者は、モデルの定義とそのトレーニング済みウェイトが正しくロードされ、そして何よりも、その動作が変化しないことを確認する必要があります。

従来、この重要なモデルのキャプチャとデプロイには、主に2つの方法が用いられてきました。1つ目の最もシンプルな方法は、torch.save を使用してモデルのウェイトのみを保存することです。このアプローチは最大の柔軟性を提供し、推論環境でのマシン固有の最適化を可能にします。しかし、デプロイ設定でモデルのアーキテクチャを明示的に再定義する必要があり、特にランタイムライブラリの制御が限られている制約された環境では、バージョン管理の悪夢や依存関係の不一致を引き起こす可能性があります。定義とウェイトの分離はしばしば「醜いバグ」の温床となり、厳格なバージョン管理が求められます。

長年にわたり、より包括的なソリューションはTorchScriptでした。これはモデル定義とウェイトの両方をシリアライズ可能なグラフ表現にバンドルします。TorchScriptは torch.jit.script と torch.jit.trace の2つの異なる機能を提供しました。スクリプト化はソースコードの静的解析を実行し、条件付き制御フローや動的な入力形状といった複雑な要素をキャプチャできます。対照的に、トレースはサンプル入力に対するモデルの実際の実行パスを記録するため、特定の障害には強いですが、動的な動作を処理することはできませんでした。多くの場合、両方の組み合わせが必要でしたが、それでもTorchScriptは複雑なモデルで頻繁に苦戦し、互換性を確保するために骨の折れる侵襲的なコードの書き換えを要求しました。HuggingFaceの画像からテキストへの生成モデルでの私たち自身の実験は、この限界を明確に示しました。固定入力のエンコーダーはトレースできましたが、動的な形状のデコーダーは、基盤となるライブラリコードに大幅な変更を加えない限り、スクリプト化に一貫して失敗しました。

ここに torch.export が登場します。PyTorchのモデルキャプチャのための新しい、より堅牢なソリューションです。torch.jit.trace と同様に、torch.export はモデルの実行をトレースすることで動作します。しかし、動的性や条件付き制御フローのサポートを組み込むことで、その前任者よりも大幅に改善されており、TorchScriptの多くの歴史的な制限を克服しています。出力はExport IRとして知られる中間グラフ表現であり、最小限の依存関係でスタンドアロンのPyTorchプログラムとしてロードおよび実行できます。torch.export の主要な利点は、torch.compile との互換性であり、推論環境でのさらなるオンザフライ最適化を可能にします。これはTorchScriptモデルには拡張されていない機能です。この機能は、PyTorchのグラフコンパイルソリューションのコアコンポーネントであるTorch Dynamoによって支えられています。

強力な機能を備えているにもかかわらず、torch.export はまだプロトタイプ機能であり、独自の課題を抱えています。一般的なハードルは「グラフブレイク」であり、エクスポート関数がトレース不可能なPythonコードに遭遇したときに発生します。PyTorchが熱心な実行（eager execution）にフォールバックするモデルコンパイルとは異なり、torch.export はグラフブレイクを厳しく禁止しており、開発者はそれらを回避するためにコードを書き換える必要があります。エクスポートされたグラフのデバッグもトリッキーな場合があります。標準の torch.nn.Module オブジェクトのように動作しますが、従来のデバッガーはコンパイルされた forward 関数にステップインできません。問題は、エクスポート環境からの変数が不注意にも定数としてグラフに「焼き付けられて」しまい、異なる環境でのランタイムエラーにつながることがよくあります。たとえば、私たちのエクスポートされたデコーダーは、エクスポート環境からのCPUデバイス参照がハードコードされていたため、最初にGPUで失敗し、解決のためにHuggingFaceライブラリの「モンキーパッチ」を手動で行う必要がありました。おもちゃのモデルには効果的でしたが、このような侵襲的な修正は、広範なテストなしには本番システムには推奨されません。

私たちのサンプルモデルでテストしたところ、torch.export はグラフブレイクに遭遇することなく、エンコーダーとデコーダーの両方を正常にキャプチャしました。これはTorchScriptと比較して大幅な改善です。修正されたエクスポートモデルをAmazon EC2インスタンスにデプロイしたところ、元のモデルと比較して推論時間がわずか10.7%短縮されました。興味深いことに、エクスポートモデルに torch.compile を適用すると、有望であるにもかかわらず、この特定のシナリオでは実行時間が予期せず増加しました。これは、コンパイルパラメータの慎重な調整の必要性を浮き彫りにしています。

まとめると、torch.export はPyTorchモデルデプロイにおける説得力のある飛躍的な進歩を表しています。以前TorchScriptを悩ませていたアーキテクチャのキャプチャを可能にし、複雑なモデルに対する優れたサポートを示しています。結果として得られるエクスポートモデルは、非常にポータブルで、広範なパッケージ依存関係なしにスタンドアロンで実行でき、torch.compile を介した強力なマシン固有の最適化と互換性があります。しかし、急速に進化しているプロトタイプであるため、現在、意図しない環境固有の値がグラフに焼き付けられる可能性や、まだ初期段階のデバッグツールセットなどの制限があります。これらの荒削りな点にもかかわらず、torch.export は以前のソリューションからの実質的な改善であり、AIモデル開発の重要なラストマイルを効率化する上で計り知れない可能性を秘めています。