ZeroGPU Spaces 加速實踐：PyTorch 提前編譯全解析

ZeroGPU 讓任何人都能在 Hugging Face Spaces 中使用強大的 Nvidia H200 硬件，而不需要因為空閑流量而長期占用 GPU。它高效、靈活，非常適合演示，不過需要注意的是，ZeroGPU 并不能在所有場景下完全發揮 GPU 與 CUDA 棧的全部潛能，比如生成圖像或視頻可能需要相當多的時間。在這種情況下，充分利用 H200 硬件，使其發揮極致性能就顯得尤為重要。

這就是 PyTorch 提前編譯（AoT）的用武之地。與其在運行時動態編譯模型（這和 ZeroGPU 短生命周期的進程配合得并不好），提前編譯允許你一次優化、隨時快速加載。

結果：演示 Demo 更流暢、體驗更順滑，在 Flux、Wan 和 LTX 等模型上有 1.3×–1.8× 的提速 ??

在這篇文章中，我們將展示如何在 ZeroGPU Spaces 中接入提前編譯（AoT）。我們會探索一些高級技巧，如 FP8 量化和動態形狀，并分享你可以立即嘗試的可運行演示。如果你想盡快嘗試，可以先去 zerogpu-aoti 中體驗一些基于 ZeroGPU 的 Demo 演示。

[!TIP]
Pro 用戶和 Team / Enterprise 組織成員可以創建 ZeroGPU Spaces，而任何人都可以免費使用（Pro、Team 和 Enterprise 用戶將獲得 8 倍 的 ZeroGPU 配額）

什么是 ZeroGPU

Spaces 是一個由 Hugging Face 提供的平臺，讓機器學習從業者可以輕松發布演示應用。

典型的 Spaces 演示應用看起來像這樣：

import gradio as gr
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(...).to('cuda')

def generate(prompt):
    return pipe(prompt).images

gr.Interface(generate, "text", "gallery").launch()

這樣做雖可行，卻導致 GPU 在 Space 的整個運行期間被獨占，即使是在沒有用戶訪問的情況下。

當執行這一行中的 .to('cuda') 時：

pipe = DiffusionPipeline.from_pretrained(...).to('cuda')

PyTorch 在初始化時會加載 NVIDIA 驅動，使進程始終駐留在 CUDA 上。由于應用流量并非持續穩定，而是高度稀疏且呈現突發性，這種方式的資源利用效率并不高。

ZeroGPU 采用了一種即時初始化 GPU 的方式。它不會在主進程中直接配置 CUDA，而是自動 fork 一個子進程，在其中配置 CUDA、運行 GPU 任務，并在需要釋放 GPU 時終止這個子進程。

這意味著：

當應用沒有流量時，它不會占用任何 GPU
當應用真正執行任務時，它會使用一個 GPU
當需要并發執行任務時，它可以使用多個 GPU

借助 Python 的 spaces 包，實現這種行為只需要如下代碼改動：

  import gradio as gr
+ import spaces
  from diffusers import DiffusionPipeline

  pipe = DiffusionPipeline.from_pretrained(...).to('cuda')

+ @spaces.GPU
  def generate(prompt):
      return pipe(prompt).images

  gr.Interface(generate, "text", "gallery").launch()

通過引入 spaces 并添加 @spaces.GPU 裝飾器 (decorator)，我們可以做到：

攔截 PyTorch API 調用，以延遲 CUDA 操作
讓被裝飾的函數在 fork 出來的子進程中運行
（調用內部 API，使正確的設備對子進程可見 —— 這不在本文范圍內）

[!NOTE]
ZeroGPU 當前會分配 H200 的一個 MIG 切片（3g.71gb 配置）。更多的 MIG 配置（包括完整切片 7g.141gb）預計將在 2025 年底推出。

PyTorch 編譯

在現代機器學習框架（如 PyTorch 和 JAX）中，“編譯”已經成為一個重要概念，它能夠有效優化模型的延遲和推理性能。其背后通常會執行一系列與硬件相關的優化步驟，例如算子融合、常量折疊等，以提升整體運行效率。

從 PyTorch 2.0 開始，目前有兩種主要的編譯接口：

即時編譯（Just-in-time）：torch.compile
提前編譯（Ahead-of-time）：torch.export + AOTInductor

torch.compile 在標準環境中表現很好：它會在模型第一次運行時進行編譯，并在后續調用中復用優化后的版本。

然而，在 ZeroGPU 上，由于幾乎每次執行 GPU 任務時進程都是新啟動的，這意味著 torch.compile 無法高效復用編譯結果，因此只能依賴文件系統緩存來恢復編譯模型。根據模型的不同，這個過程可能需要幾十秒到幾分鐘，對于 Spaces 中的實際 GPU 任務來說，這顯然太慢了。這正是 提前編譯（AoT） 大顯身手的地方。

通過提前編譯，我們可以在一開始導出已編譯的模型，將其保存，然后在任意進程中即時加載。這不僅能減少框架的額外開銷，還能消除即時編譯通常帶來的冷啟動延遲。

但是，我們該如何在 ZeroGPU 上實現提前編譯呢？讓我們繼續深入探討。

ZeroGPU 上的提前編譯

讓我們回到 ZeroGPU 的基礎示例，來逐步解析啟用 AoT 編譯所需要的內容。在本次演示中，我們將使用 black-forest-labs/FLUX.1-dev 模型：

import gradio as gr
import spaces
import torch
from diffusers import DiffusionPipeline

MODEL_ID = 'black-forest-labs/FLUX.1-dev'

pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
pipe.to('cuda')

@spaces.GPU
def generate(prompt):
    return pipe(prompt).images

gr.Interface(generate, "text", "gallery").launch()

[!NOTE]
在下面的討論中，我們只編譯 pipe 的 transformer 組件。
因為在這類生成模型中，transformer（或者更廣義上說，denoiser）是計算量最重的部分。

使用 PyTorch 對模型進行提前編譯通常包含以下幾個步驟：

1. 獲取示例輸入

請記住，我們要對模型進行提前編譯。因此，我們需要為模型準備示例輸入。這些輸入應當與實際運行過程中所期望的輸入類型保持一致。

為了捕獲這些輸入，我們將使用 spaces 包中的 spaces.aoti_capture 輔助函數：

with spaces.aoti_capture(pipe.transformer) as call:
    pipe("arbitrary example prompt")

當 aoti_capture 作為上下文管理器使用時，它會攔截對任意可調用對象的調用（在這里是 pipe.transformer），阻止其實際執行，捕獲本應傳遞給它的輸入參數，并將這些值存儲在 call.args 和 call.kwargs 中。

2. 導出模型

既然我們已經得到了 transformer 組件的示例參數（args 和 kwargs），我們就可以使用 torch.export.export 工具將其導出為一個 PyTorch ExportedProgram：

exported_transformer = torch.export.export(
    pipe.transformer,
    args=call.args,
    kwargs=call.kwargs,
)

3. 編譯導出的模型

一旦模型被導出，編譯它就非常直接了。

在 PyTorch 中，傳統的提前編譯通常需要將模型保存到磁盤，以便后續重新加載。在我們的場景中，可以利用 spaces 包中的一個輔助函數：spaces.aoti_compile。
它是對 torch._inductor.aot_compile 的一個輕量封裝，能夠根據需要管理模型的保存和延遲加載。其使用方式如下：

compiled_transformer = spaces.aoti_compile(exported_transformer)

這個 compiled_transformer 現在是一個已經完成提前編譯的二進制，可以直接用于推理。

4. 在流水線中使用已編譯模型

現在我們需要將已編譯好的 transformer 綁定到原始流水線中，也就是 pipeline。接下來，我們需要將編譯后的 transformer 綁定到原始的 pipeline 中。一個看似簡單的做法是直接修改：pipe.transformer = compiled_transformer。但這樣會導致問題，因為這種方式會丟失一些關鍵屬性，比如 dtype、config 等。如果只替換 forward 方法也不理想，因為原始模型參數依然會常駐內存，往往會在運行時引發 OOM（內存溢出）錯誤。

因此spaces 包為此提供了一個工具 —— spaces.aoti_apply：

spaces.aoti_apply(compiled_transformer, pipe.transformer)

這樣以來，它會自動將 pipe.transformer.forward 替換為我們編譯后的模型，同時清理舊的模型參數以釋放內存。

5. 整合所有步驟

要完成前面三個步驟（攔截輸入示例、導出模型，以及用 PyTorch inductor 編譯），我們需要一塊真實的 GPU。在 @spaces.GPU 函數之外得到的 CUDA 仿真環境是不夠的，因為編譯過程高度依賴硬件，例如需要依靠微基準測試來調優生成的代碼。這就是為什么我們需要把所有步驟都封裝在一個 @spaces.GPU 函數中，然后再將編譯好的模型傳回應用的根作用域。從原始的演示代碼開始，我們可以得到如下實現：

  import gradio as gr
  import spaces
  import torch
  from diffusers import DiffusionPipeline
  
  MODEL_ID = 'black-forest-labs/FLUX.1-dev'
  
  pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
  pipe.to('cuda')
  
+ @spaces.GPU(duration=1500) # 啟動期間允許的最大執行時長
+ def compile_transformer():
+     with spaces.aoti_capture(pipe.transformer) as call:
+         pipe("arbitrary example prompt")
+ 
+     exported = torch.export.export(
+         pipe.transformer,
+         args=call.args,
+         kwargs=call.kwargs,
+     )
+     return spaces.aoti_compile(exported)
+ 
+ compiled_transformer = compile_transformer()
+ spaces.aoti_apply(compiled_transformer, pipe.transformer)
  
  @spaces.GPU
  def generate(prompt):
      return pipe(prompt).images
  
  gr.Interface(generate, "text", "gallery").launch()

只需增加十幾行代碼，我們就成功讓演示運行得更快（在 FLUX.1-dev 的情況下提升了 1.7 倍）。

如果你想進一步了解提前編譯，可以閱讀 PyTorch 的 AOTInductor 教程。

注意事項

現在我們已經展示了在 ZeroGPU 條件下可以實現的加速效果，接下來將討論在這一設置中需要注意的一些問題。

量化（Quantization）

提前編譯可以與量化結合，從而實現更大的加速效果。對于圖像和視頻生成任務，FP8 的訓練后動態量化方案提供了良好的速度與質量平衡。不過需要注意，FP8 至少需要 9.0 的 CUDA 計算能力才能使用。
幸運的是，ZeroGPU 基于 H200，因此我們已經能夠利用 FP8 量化方案。要在提前編譯工作流中啟用 FP8 量化，我們可以使用 torchao 提供的 API，如下所示：

+ from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig

+ # 在導出步驟之前對 transformer 進行量化
+ quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())

exported_transformer = torch.export.export(
    pipe.transformer,
    args=call.args,
    kwargs=call.kwargs,
)

（你可以在這里找到更多關于 TorchAO 的詳細信息。）

接著，我們就可以按照上面描述的步驟繼續進行。使用量化可以再帶來 1.2 倍 的加速。

動態形狀（Dynamic shapes）

圖像和視頻可能具有不同的形狀和尺寸。因此，在執行提前編譯時，考慮形狀的動態性也非常重要。torch.export.export 提供的原語讓我們能夠很容易地配置哪些輸入需要被視為動態形狀，如下所示。

以 Flux.1-Dev 的 transformer 為例，不同圖像分辨率的變化會影響其 forward 方法中的兩個參數：

hidden_states：帶噪聲的輸入潛變量，transformer 需要對其去噪。它是一個三維張量，表示 batch_size, flattened_latent_dim, embed_dim。當 batch size 固定時，隨著圖像分辨率變化，flattened_latent_dim 也會變化。
img_ids：一個二維數組，包含編碼后的像素坐標，形狀為 height * width, 3。在這種情況下，我們希望讓 height * width 是動態的。

我們首先需要定義一個范圍，用來表示（潛變量）圖像分辨率可以變化的區間。為了推導這些數值范圍，我們檢查了 pipeline 中 hidden_states 的形狀在不同圖像分辨率下的變化。這些具體數值依賴于模型本身，需要人工檢查并結合一定直覺。對于 Flux.1-Dev，我們最終得到：

transformer_hidden_dim = torch.export.Dim('hidden', min=4096, max=8212)

接下來，我們定義一個映射，指定參數名稱，以及在其輸入值中哪些維度需要被視為動態：

transformer_dynamic_shapes = {
    "hidden_dim": {1: transformer_hidden_dim}, 
    "img_ids": {0: transformer_hidden_dim},
}

然后，我們需要讓動態形狀對象的結構與示例輸入保持一致。對于不需要動態形狀的輸入，必須將其設置為 None。這可以借助 PyTorch 提供的 tree_map 工具輕松完成：

from torch.utils._pytree import tree_map

dynamic_shapes = tree_map(lambda v: None, call.kwargs)
dynamic_shapes |= transformer_dynamic_shapes

現在，在執行導出步驟時，我們只需將 transformer_dynamic_shapes 傳遞給 torch.export.export：

exported_transformer = torch.export.export(
    pipe.transformer,
    args=call.args,
    kwargs=call.kwargs,
    dynamic_shapes=dynamic_shapes,
)

[!NOTE]
可以參考這個 Space，它詳細說明了如何在導出步驟中把量化和動態形狀結合起來使用。

多重編譯 / 權重共享

當模型的動態性非常重要時，僅依靠動態形狀有時是不夠的。

例如，在 Wan 系列視頻生成模型中，如果你希望編譯后的模型能夠生成不同分辨率的內容，就會遇到這種情況。在這種情況下，可以采用的方法是：為每種分辨率編譯一個模型，同時保持模型參數共享，并在運行時調度對應的模型。

這里有一個這種方法的示例：zerogpu-aoti-multi.py。
你也可以在 Wan 2.2 Space 中看到該范式的完整實現。

FlashAttention-3

由于 ZeroGPU 的硬件和 CUDA 驅動與 Flash-Attention 3（FA3）完全兼容，我們可以在 ZeroGPU Spaces 中使用它來進一步提升速度。FA3 可以與提前編譯（AoT）配合使用，因此非常適合我們的場景。

從源碼編譯和構建 FA3 可能需要幾分鐘時間，并且這個過程依賴于具體硬件。作為用戶，我們當然不希望浪費寶貴的 ZeroGPU 計算時間。這時 Hugging Face 的 kernels 庫就派上用場了，因為它提供了針對特定硬件的預編譯內核。

例如，當我們嘗試運行以下代碼時：

from kernels import get_kernel

vllm_flash_attn3 = get_kernel("kernels-community/vllm-flash-attn3")

它會嘗試從 kernels-community/vllm-flash-attn3 倉庫加載一個內核，該內核與當前環境兼容。
否則，如果存在不兼容問題，就會報錯。幸運的是，在 ZeroGPU Spaces 上這一過程可以無縫運行。這意味著我們可以在 ZeroGPU 上借助 kernels 庫充分利用 FA3 的性能。

這里有一個 Qwen-Image 模型的 FA3 注意力處理器完整示例。

提前編譯的 ZeroGPU Spaces 演示

加速對比

未使用 AoTI 的 FLUX.1-dev
使用 AoTI 和 FA3 的 FLUX.1-dev （1.75 倍 加速）

結論

Hugging Face Spaces 中的 ZeroGPU 是一項強大的功能，它為 AI 構建者提供了高性能算力。在這篇文章中，我們展示了用戶如何借助 PyTorch 的提前編譯（AoT）技術，加速他們基于 ZeroGPU 的應用。

我們用 Flux.1-Dev 展示了加速效果，但這些技術并不僅限于這一模型。因此，我們鼓勵你嘗試這些方法，并在社區討論中向我們提供反饋。

資源

訪問 Hub 上的 ZeroGPU-AOTI 組織，瀏覽一系列利用文中技術的演示。
查看 spaces.aoti_* API 的源代碼，了解接口細節。
查看 Hub 上的 Kernels Community 組織。
升級到 Hugging Face 的 Pro，創建你自己的 ZeroGPU Spaces（每天可獲得 25 分鐘 H200 使用時間）。

致謝：感謝 ChunTe Lee 為本文制作了精彩的縮略圖。感謝 Pedro 和 Vaibhav 對文章提供的反饋。

英文原文: https://huggingface.co/blog/zerogpu-aoti
原文作者: Charles Bensimon, Sayak Paul, Linoy Tsaban, Apolinário Passos

譯者: AdinaY

posted @ 2025-09-09 10:16 HuggingFace 閱讀(163) 評論(0) 收藏舉報

刷新頁面返回頂部

Hugging Face 博客

The AI community building the future.