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 （https://huggingface.co/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?DiffusionPipelinepipe = DiffusionPipeline.from_pretrained(...).to('cuda')def?generate(prompt):return?pipe(prompt).imagesgr.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 spacesfrom diffusers import DiffusionPipelinepipe = DiffusionPipeline.from_pretrained(...).to('cuda')+ @spaces.GPUdef generate(prompt):return pipe(prompt).imagesgr.Interface(generate, "text", "gallery").launch()

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

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

[!NOTE]
ZeroGPU 當前會分配 H200 的一個 MIG 切片（https://docs.nvidia.com/datacenter/tesla/mig-user-guide/#h200-mig-profiles?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?DiffusionPipelineMODEL_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).imagesgr.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 grimport spacesimport torchfrom diffusers import DiffusionPipelineMODEL_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.GPUdef generate(prompt):return pipe(prompt).imagesgr.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,
)

接著，我們就可以按照上面描述的步驟繼續進行。使用量化可以再帶來?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_mapdynamic_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]
可以參考?https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-Kontext-Dev-fp8-dynamic，它詳細說明了如何在導出步驟中把量化和動態形狀結合起來使用。

多重編譯 / 權重共享

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

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

這里有一個這種方法的示例：zerogpu-aoti-multi.py （https://gist.github.com/cbensimon/8dc0ffcd7ee024d91333f6df01907916）。
你也可以在 Wan 2.2 Space （https://huggingface.co/spaces/zerogpu-aoti/wan2-2-fp8da-aoti-faster/blob/main/optimization.py）中看到該范式的完整實現。

FlashAttention-3

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

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

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

from?kernels?import?get_kernelvllm_flash_attn3 = get_kernel("kernels-community/vllm-flash-attn3")

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

可以參考 Qwen-Image 模型的 FA3 注意力處理器完整示例：?https://gist.github.com/sayakpaul/ff715f979793d4d44beb68e5e08ee067#file-fa3_qwen-py。

提前編譯的 ZeroGPU Spaces 演示

加速對比

未使用 AoTI 的 FLUX.1-dev：
https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-dev-base
使用 AoTI 和 FA3 的 FLUX.1-dev （__1.75 倍__ 加速）
https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-dev-fa3-aoti

精選 AoTI Spaces

FLUX.1 Kontext：
https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-Kontext-Dev
QwenImage Edit：
https://huggingface.co/spaces/multimodalart/Qwen-Image-Edit-Fast
Wan 2.2：
https://huggingface.co/spaces/zerogpu-aoti/wan2-2-fp8da-aoti-faster
LTX Video：
https://huggingface.co/spaces/zerogpu-aoti/ltx-dev-fast

結論

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

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

資源

訪問 Hub 上的 ZeroGPU-AOTI 組織，瀏覽一系列利用文中技術的演示。https://huggingface.co/zerogpu-aoti
查看?spaces.aoti_*?API 的源代碼，了解接口細節。
https://pypi-browser.org/package/spaces/spaces-0.40.1-py3-none-any.whl/spaces/zero/torch/aoti.py
查看 Hub 上的 Kernels Community 組織。
https://huggingface.co/kernels-community
升級到 Hugging Face 的 Pro，創建你自己的 ZeroGPU Spaces（每天可獲得 25 分鐘 H200 使用時間）。
https://huggingface.co/pro

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