Hugging Face ZeroGPU 前置編譯指南

Hugging Face 為 ZeroGPU Spaces 引入了前置編譯（AoT），讓開發者能優化模型延遲並消除冷啟動時間。透過一次編譯模型並在短暫的 GPU 程序中即時重新載入，使用者可在 Flux、Wan、LTX 等模型上獲得 1.3Ñ–1.8Ñ 的加速。

ZeroGPU 架構與前置編譯的需求

ZeroGPU 採用即時（just-in-time）方式初始化 GPU，以最大化資源效率。ZeroGPU 不會為整個 Space 的生命週期保留 GPU，而是分叉進程，在 Nvidia H200（目前使用 3g.71gb MIG 切片）上執行 GPU 任務，任務完成後即終止分叉的進程。

雖然 torch.compile 在標準環境中有效，但在 ZeroGPU 上它依賴檔案系統快取，因為幾乎每個任務都會重新啟動進程。此快取過程可能需要數十秒至數分鐘，使得即時示範變得不切實際。前置編譯透過一次匯出已編譯的模型，之後即可在任何進程中即時重新載入，解決了此問題。

在 ZeroGPU 上實作前置編譯

實作前置編譯涉及使用 spaces 套件與 PyTorch 工具的五步工作流程：

  1. 捕獲範例輸入：使用 spaces.aoti_capture 攔截傳遞給模型元件的參數與關鍵字參數（例如擴散管線中的 transformer）。
  2. 匯出模型：使用 torch.export.export 將模型轉換為 ExportedProgram，這是一個包含張量運算與模型參數的計算圖。
  3. 編譯匯出模型：使用 spaces.aoti_compile（torch._inductor.aot_compile 的包裝器）產生前置編譯的二進位檔。
  4. 套用編譯模型：使用 spaces.aoti_apply 修補模型的 forward 方法，並從記憶體中移除原始模型參數，以防止記憶體不足（OOM）錯誤。
  5. 在 GPU 環境中包裹：由於編譯依賴硬體且需要實體 GPU，編譯步驟必須在應用程式啟動階段以 @spaces.GPU 函式包裹。

對於 FLUX.1-dev 模型，此流程帶來了 1.7x 的加速。

進階最佳化技術

FP8 量化

前置編譯可結合 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 函式庫，讓使用者載入預先建好的、相容硬體的 kernel（例如 kernels-community/vllm-flash-attn3）。

區域編譯

開發者可以僅編譯重複的計算區塊（例如 Flux 中的 FluxTransformerBlock 與 FluxSingleTransformerBlock），而非整個模型。這大幅縮短冷啟動時間——將 Flux.1-Dev 的編譯時間從 6 分鐘縮減至 30 秒——同時保有與完整模型編譯相同的加速效果。

部署與發佈

編譯好的圖形模組可序列化為工件並上傳至 Hugging Face Hub。僅儲存編譯模型圖而不包含模型參數，可保持儲存空間輕量。示範程式可下載並載入這些預編譯圖，於啟動時完全跳過編譯階段。

SUMMARY: Hugging Face 為 ZeroGPU Spaces 引入前置編譯（AoT），透過消除即時編譯開銷，使 Flux、Wan、LTX 等生成模型的速度提升 1.3 倍至 1.8 倍。

TITLE: Hugging Face ZeroGPU 前置編譯指南

Sources