Hugging Face ZeroGPU 事前コンパイルガイド

Hugging Face は ZeroGPU Spaces 向けに事前コンパイル（AoT）を導入し、開発者がモデルのレイテンシを最適化し、コールドスタート時間を排除できるようにしました。モデルを一度コンパイルして短命な GPU プロセス間で即座に再ロードすることで、Flux、Wan、LTX などのモデルで 1.3Ñ–1.8Ñ の速度向上が得られます。

ZeroGPU アーキテクチャと AoT が必要な理由

ZeroGPU は GPU 初期化にジャストインタイム方式を採用し、リソース効率を最大化します。Space の全ライフタイムで GPU を予約する代わりに、ZeroGPU はプロセスをフォークし、Nvidia H200（現在は 3g.71gb MIG スライスを使用）上で GPU タスクを実行し、タスク完了後にフォークを終了します。

標準環境で有効な torch.compile は ZeroGPU ではファイルシステムキャッシュに依存します。なぜなら、ほぼすべてのタスクでプロセスが新たに起動されるためです。このキャッシュ処理には数十秒から数分かかり、リアルタイムデモには実用的でありません。AoT コンパイルは、コンパイル済みモデルを一度エクスポートすれば、任意のプロセスで即座に再ロードできるようにすることでこの問題を解決します。

ZeroGPU で AoT コンパイルを実装する手順

AoT コンパイルの実装は、spaces パッケージと PyTorch ユーティリティを用いた 5 ステップのワークフローです。

サンプル入力の取得: spaces.aoti_capture を使用して、モデルコンポーネント（例: ディフュージョンパイプラインのトランスフォーマー）に渡される引数とキーワード引数を捕捉します。
モデルのエクスポート: torch.export.export を使用して、モデルを ExportedProgram に変換します。これはテンソル計算とモデルパラメータを含む計算グラフです。
エクスポートされたモデルのコンパイル: spaces.aoti_compile（torch._inductor.aot_compile のラッパー）を使って AoT コンパイル済みバイナリを生成します。
コンパイル済みモデルの適用: spaces.aoti_apply を用いてモデルの forward メソッドをパッチし、元のモデルパラメータをメモリから除去して OOM エラーを防ぎます。
GPU コンテキストでラップ: コンパイルはハードウェア依存で実際の GPU が必要なため、コンパイルステップはアプリ起動時に @spaces.GPU 関数でラップする必要があります。

FLUX.1‑dev モデルの場合、このプロセスにより 1.7x の速度向上が得られました。

高度な最適化手法

FP8 量子化

AoT コンパイルは torchao ライブラリによる FP8 ポストトレーニング動的量子化と組み合わせることができます。ZeroGPU は H200 GPU（コンピュート能力 9.0 以上）を使用しているため FP8 をサポートし、追加で 1.2x の速度向上が得られます。

動的形状の取り扱い

画像や動画の解像度が可変になるケースに対応するため、開発者は torch.export.Dim を使って動的次元を定義できます。Flux.1‑Dev では hidden_states の flattened_latent_dim と img_ids の height * width を動的にします。これらの設定は dynamic_shapes マップとして torch.export.export に渡されます。

マルチコンパイルと重み共有

単一グラフの動的形状で対応しきれないほどの変動がある場合（例: Wan の動画生成ファミリー）、解像度ごとにモデルをコンパイルし、モデルパラメータは共有、実行時に適切なコンパイル済みグラフをディスパッチする手法が取れます。

FlashAttention-3（FA3）

ZeroGPU は FlashAttention-3 に対応しています。FA3 をソースからビルドする手間を省くため、Hugging Face は kernels ライブラリを提供しており、kernels-community/vllm-flash-attn3 などの事前ビルド済みハードウェア対応カーネルをロードできます。

局所コンパイル

モデル全体をコンパイルする代わりに、計算の繰り返しブロック（例: Flux の FluxTransformerBlock と FluxSingleTransformerBlock）だけをコンパイルできます。これによりコールドスタート時間が大幅に短縮され、Flux.1‑Dev のコンパイル時間は 6 minutes to 30 seconds から 30 seconds に削減され、フルモデルコンパイルと同等の速度向上が得られます。

デプロイと配布

コンパイル済みグラフモジュールはアーティファクトとしてシリアライズし、Hugging Face Hub にアップロードできます。モデルパラメータを除いたコンパイル済みグラフだけを保存すれば、ストレージ負荷は軽くなります。デモはこれらの事前コンパイル済みグラフをダウンロードしてロードするだけで、起動時のコンパイルフェーズを完全にスキップできます。