ZeroGPU Spaces 加速實踐：PyTorch 提(ti)前編(bian)譯全解析
ZeroGPU 讓任何人都能在 Hugging Face Spaces 中使用強大的 Nvidia H200 硬件，而不(bu)需(xu)要(yao)因(yin)為空閑流(liu)量而長期占用 GPU。 它高效(xiao)、靈活，非常適合演(yan)示(shi)，不(bu)過需(xu)要(yao)注意的是，ZeroGPU 并不(bu)能(neng)在(zai)所有場景下完全發揮 GPU 與 CUDA 棧的全部潛(qian)能(neng)，比如生成圖像或(huo)視頻可能(neng)需(xu)要(yao)相當多的時間。在(zai)這種情況下，充分利用 H200 硬件，使其發揮極致(zhi)性能(neng)就顯得尤為重要(yao)。
這就是 PyTorch 提(ti)(ti)前(qian)編譯（AoT）的(de)用(yong)武(wu)之地(di)。與其在運行(xing)時動態編譯模型（這和(he) ZeroGPU 短生命(ming)周期的(de)進程配合得并不好），提(ti)(ti)前(qian)編譯允許你一次(ci)優化、隨時快速加載。
結果：演示 Demo 更流暢、體驗更順滑，在 Flux、Wan 和 LTX 等模型上有 1.3×–1.8× 的提速 ??
在這篇(pian)文章中(zhong)，我們將(jiang)展示如(ru)何(he)在 ZeroGPU Spaces 中(zhong)接入提(ti)前(qian)編(bian)譯(yi)（AoT）。我們會探索一(yi)些高級技(ji)巧，如(ru) FP8 量(liang)化和(he)動態形狀，并分享你可(ke)以(yi)(yi)立即嘗試(shi)的可(ke)運行演示。如(ru)果你想盡(jin)快嘗試(shi)，可(ke)以(yi)(yi)先去 中(zhong)體驗一(yi)些基于 ZeroGPU 的 Demo 演示。
[!TIP]
用戶和 組織成員可以創建 ZeroGPU Spaces，而任何人都可以免費使用（Pro、Team 和 Enterprise 用戶將獲得 8 倍 的 ZeroGPU 配額）
目錄
什么是 ZeroGPU
是一個由 Hugging Face 提供的(de)平臺，讓機器(qi)學(xue)習(xi)從業者可以輕松發布演示應(ying)用。
典型的(de) Spaces 演示應(ying)用看起來像(xiang)這(zhe)樣(yang)：
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()
這樣做(zuo)雖可行，卻導致 GPU 在(zai) Space 的整個(ge)運(yun)行期(qi)間被獨占，即使(shi)是在(zai)沒有用戶訪問的情況下。
當執行這一行中的 .to('cuda') 時：
pipe = DiffusionPipeline.from_pretrained(...).to('cuda')
PyTorch 在(zai)初始化時(shi)會加載 NVIDIA 驅動(dong)，使進程(cheng)始終(zhong)駐(zhu)留在(zai) CUDA 上。由于應用流量并非持續(xu)穩定，而(er)是高(gao)度稀疏且呈現突發(fa)性(xing)，這(zhe)種(zhong)方式(shi)的資源利用效(xiao)率并不高(gao)。
ZeroGPU 采用了一(yi)種即時(shi)初始化 GPU 的方(fang)式。它不會(hui)在(zai)(zai)主進程中(zhong)直接配置 CUDA，而(er)是自動 fork 一(yi)個子(zi)進程，在(zai)(zai)其中(zhong)配置 CUDA、運(yun)行(xing) GPU 任務，并在(zai)(zai)需(xu)要釋放 GPU 時(shi)終止(zhi)這個子(zi)進程。
這意味著：
- 當應用沒有流量時，它不會占用任何 GPU
- 當應用真正執行任務時，它會使用一個 GPU
- 當需要并發執行任務時，它可以使用多個 GPU
借助 Python 的 spaces 包，實現(xian)這種行為只需要如下(xia)代碼改(gai)動：
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)，我(wo)們可以做(zuo)到：
- 攔截 PyTorch API 調用，以延遲 CUDA 操作
- 讓被裝飾的函數在 fork 出來的子進程中運行
- （調用內部 API，使正確的設備對子進程可見 —— 這不在本文范圍內）
[!NOTE]
ZeroGPU 當前會分配 H200 的一個 切片（3g.71gb配置）。更多的 MIG 配置（包括完整切片7g.141gb）預計將(jiang)在(zai) 2025 年(nian)底(di)推出。
PyTorch 編譯
在現代機器學習(xi)框(kuang)架（如 PyTorch 和 JAX）中，“編譯”已經成為一個(ge)重要概念，它能夠有效(xiao)優化(hua)模型的(de)延遲(chi)和推理性能。其背后通(tong)常會執行(xing)一系列與(yu)硬件(jian)相關的(de)優化(hua)步驟，例如算子融合、常量折疊等，以提升整體運行(xing)效(xiao)率(lv)。
從 PyTorch 2.0 開始，目(mu)前(qian)有兩(liang)種(zhong)主要的編譯接口(kou)：
- 即時編譯（Just-in-time）：
torch.compile - 提前編譯（Ahead-of-time）：
torch.export+AOTInductor
在(zai)標準環境(jing)中(zhong)表現很好：它會在(zai)模型第一次運行時進行編譯，并在(zai)后續調(diao)用中(zhong)復用優化(hua)后的版本。
然而，在 ZeroGPU 上，由于幾乎每次執行 GPU 任務時進程都是新啟動的，這意味著 torch.compile 無法高效復用編譯結果，因此只能依賴 來恢復編譯模型。 根據模型的不同，這個過程可能需要幾十秒到幾分鐘，對于 Spaces 中的實際 GPU 任務來說，這顯然太慢了。 這正是 提前編譯（AoT） 大顯身手的地方。
通(tong)過提前編譯，我們(men)可(ke)以(yi)在(zai)一(yi)開始導出已編譯的(de)(de)模型，將其保存(cun)，然后在(zai)任意進程中即(ji)時加載。這不(bu)僅能(neng)(neng)減少(shao)框架的(de)(de)額外開銷，還能(neng)(neng)消除即(ji)時編譯通(tong)常帶來的(de)(de)冷啟動延遲。
但(dan)是，我(wo)們(men)該如何在 ZeroGPU 上實現提前(qian)編譯呢？讓我(wo)們(men)繼續深入探討。
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組件。
因為在(zai)這類生成模型中(zhong)，transformer（或者更廣義上說(shuo)，denoiser）是計算量最重的部(bu)分。
使(shi)用(yong) PyTorch 對模型進行提前(qian)編譯通常包(bao)含(han)以(yi)下幾個步驟：
1. 獲取示例輸入
請記住，我們要對模型進行 提前 編(bian)譯。因此(ci)，我們需(xu)要(yao)為模型準備示例(li)輸入(ru)。這些輸入(ru)應當與實際運行過程中所期望的輸入(ru)類(lei)型保持(chi)一致(zhi)。
為了捕獲這些輸入，我們將使用 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. 導出模型
既然(ran)我(wo)們已(yi)經得(de)到了 transformer 組(zu)件的示(shi)例參數（args 和 kwargs），我(wo)們就(jiu)可以使用(yong) 工具將其導出為一個(ge) PyTorch ：
exported_transformer = torch.export.export(
pipe.transformer,
args=call.args,
kwargs=call.kwargs,
)
3. 編譯導出的模型
一旦模型被導(dao)出，編譯它就非常直接了。
在 PyTorch 中，傳統的提前編譯通常需要將模型保存到磁盤，以便后續重新加載。 在我們的場景中，可以利用 spaces 包中的一個輔助函數：spaces.aoti_compile。
它是對 torch._inductor.aot_compile 的一個輕量封裝，能夠(gou)根據需要管理模型的保存和延遲加載。其使用方(fang)式如下：
compiled_transformer = spaces.aoti_compile(exported_transformer)
這個 compiled_transformer 現在是一個已經(jing)完(wan)成提前編譯的二進制，可以直(zhi)接用于(yu)推理。
4. 在流水線中使用已編譯模型
現在我們需要將已編譯好的 transformer 綁定到原始流水線中，也就是 pipeline。 接下來，我們需要將編譯后的 transformer 綁定到原始的 pipeline 中。 一個看似簡單的做法是直接修改：pipe.transformer = compiled_transformer。但這樣會導致問題，因為這種方式會丟失一些關鍵屬性，比如 dtype、config 等。 如果只替換 forward 方法也(ye)不理想，因為(wei)原始模(mo)型參數(shu)依然會常駐內存，往(wang)往(wang)會在(zai)運行時引發(fa) OOM（內存溢出）錯誤。
因此spaces 包為此提供了一個工具 —— spaces.aoti_apply：
spaces.aoti_apply(compiled_transformer, pipe.transformer)
這樣以來，它會自動將 pipe.transformer.forward 替換為我們編譯后的(de)模(mo)型(xing)，同(tong)時清理舊的(de)模(mo)型(xing)參數以(yi)釋放內存。
5. 整合所有步驟
要完成前面三個步驟（攔截輸入示例、導出模型，以及用 PyTorch inductor 編譯），我們需要一塊真實的 GPU。 在 @spaces.GPU 函數之外得到的 CUDA 仿真環境是不夠的，因為編譯過程高度依賴硬件，例如需要依靠微基準測試來調優生成的代碼。這就是為什么我們需要把所有步驟都封裝在一個 @spaces.GPU 函數中，然后再(zai)將編譯好的模型傳回應用的根作用域。 從原始(shi)的演(yan)示代碼開(kai)始(shi)，我們可以得(de)到如下實現(xian)：
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 倍）。
如果(guo)你想進一步了解提前編譯，可以閱讀 PyTorch 的 。
注意事項
現(xian)在(zai)(zai)我們已經展示了(le)在(zai)(zai) ZeroGPU 條件下(xia)可以實(shi)現(xian)的加速效果，接下(xia)來將(jiang)討論在(zai)(zai)這一設置中需要注意(yi)的一些(xie)問(wen)題。
量化（Quantization）
提前編譯可以與量化結合，從而實現更大的加速效果。對于圖像和視頻生成任務，FP8 的訓練后動態量化方案提供了良好的速度與質量平衡。不過需要注意，FP8 至少需要 9.0 的 CUDA 計算能力才能使用。
幸(xing)運的(de)是，ZeroGPU 基于 H200，因(yin)此我們(men)已經能夠(gou)利(li)用(yong) FP8 量化方案。 要在提前編譯(yi)工作流中啟用(yong) FP8 量化，我們(men)可以使用(yong) 提供(gong)的(de) API，如(ru)下(xia)所示：
+ from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig
+ # 在導出步驟之前對 transformer 進行量化
+ quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
exported_transformer = torch.export.export(
pipe.transformer,
args=call.args,
kwargs=call.kwargs,
)
（你可以在 找(zhao)到更多(duo)關于 TorchAO 的(de)詳細信息。）
接著，我們就可以按照上面描述的步驟繼續進行。使用量化可以再帶來 1.2 倍 的加速。
動態形狀（Dynamic shapes）
圖像和視頻可能具有不同的形狀和尺寸。因此，在執行提前編譯時，考慮形狀的動態性也非常重要。torch.export.export 提(ti)供的(de)原語讓我們能夠很容(rong)易地配置哪(na)些輸入需要被視為動態形狀，如下(xia)所(suo)示。
以 Flux.1-Dev 的 transformer 為例，不同圖像分辨率的變化會影響其 forward 方(fang)法中的兩個參數：
-
hidden_states：帶噪聲的輸入潛變量，transformer 需要對其去噪。它是一個三維張量，表示batch_size, flattened_latent_dim, embed_dim。當 batch size 固定時，隨著圖像分辨率變化，flattened_latent_dim也會變化。 -
img_ids：一個二維數組，包含編碼后的像素坐標，形狀為height * width, 3。在這種情況下，我們希望讓height * width是動態的。
我(wo)們首先(xian)需(xu)要定義一個(ge)范圍，用來表示(shi)（潛變量(liang)）圖(tu)像分辨(bian)率可以變化的(de)區(qu)間。為了推(tui)導這些數值范圍，我(wo)們檢查(cha)了 pipeline 中 的(de)形狀在(zai)不同(tong)圖(tu)像分辨(bian)率下(xia)的(de)變化。這些具體(ti)數值依(yi)賴于(yu)模型本(ben)身，需(xu)要人工(gong)檢查(cha)并結合(he)一定直(zhi)覺。 對(dui)于(yu) Flux.1-Dev，我(wo)們最終得(de)到：
transformer_hidden_dim = torch.export.Dim('hidden', min=4096, max=8212)
接下來，我們(men)定義一個映射，指(zhi)定參數名(ming)稱，以及(ji)在其輸入值中哪些維(wei)度需(xu)要被視為動(dong)態(tai)：
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]
可以參考 ，它詳細說明了如何在導(dao)出步驟中把量化和動態形狀結合(he)起來使用。
多重編譯 / 權重共享
當模型(xing)的(de)動態性(xing)非(fei)常重要時，僅依靠動態形狀有時是不夠的(de)。
例如，在 Wan 系列視頻生成模型中，如果你希望編譯后的模型能夠生成不同分辨率的內容，就會遇到這種情況。在這種情況下，可以采用的方法是：為每種分辨率編譯一個模型，同時保持模型參數共享，并在運行時調度對應的模型。
這里有一個這種方法的示例：。
你也可(ke)以在 中看到該范式的完(wan)整(zheng)實(shi)現。
FlashAttention-3
由于 ZeroGPU 的硬件(jian)和(he) CUDA 驅動與 Flash-Attention 3（FA3）完全兼(jian)容，我(wo)們可(ke)以在 ZeroGPU Spaces 中(zhong)使用(yong)它來進一步提升速度。FA3 可(ke)以與提前編譯(yi)（AoT）配合(he)使用(yong)，因(yin)此非常適(shi)合(he)我(wo)們的場景(jing)。
從源碼編譯和構建 FA3 可能需要幾分(fen)鐘時(shi)間，并且這個過(guo)程依賴于具體硬(ying)件(jian)。作為用戶，我們當然不希(xi)望浪費寶貴的 ZeroGPU 計算(suan)時(shi)間。這時(shi) Hugging Face 的 就派上用場了，因為它提供了針對(dui)特定硬(ying)件(jian)的預編譯內核。
例(li)如，當(dang)我們嘗試運行以(yi)下代碼時(shi)：
from kernels import get_kernel
vllm_flash_attn3 = get_kernel("kernels-community/vllm-flash-attn3")
它會嘗試從 倉庫加載一個內核，該內核與當前環境兼容。
否則，如果存在不兼容問題，就會報錯。幸運的是，在 ZeroGPU Spaces 上這一過程可以無縫運行。這意味著我們可以在 ZeroGPU 上借助 kernels 庫充分利用 FA3 的性(xing)能(neng)。
這里有一個 。
提前編譯的 ZeroGPU Spaces 演示
加速對比
- （1.75 倍 加速）
精選 AoTI Spaces
結論
Hugging Face Spaces 中的 ZeroGPU 是一(yi)項(xiang)強(qiang)大的功能，它為 AI 構(gou)建者提供(gong)了高(gao)性(xing)能算力。在(zai)這篇文章中，我們(men)展(zhan)示了用戶如何借助 PyTorch 的提前編譯（AoT）技術，加速他們(men)基(ji)于 ZeroGPU 的應用。
我(wo)們(men)(men)用 Flux.1-Dev 展示了加速效(xiao)果(guo)，但(dan)這些技術并不僅限于這一模型。因此，我(wo)們(men)(men)鼓勵(li)你嘗試這些方法(fa)，并在 中向我(wo)們(men)(men)提供反饋(kui)。
資源
- 訪問 ，瀏覽一系列利用文中技術的演示。
- 查看
spaces.aoti_*API 的，了解接口細節。 - 查看 。
- 升級到 Hugging Face 的 ，創建你自己的 ZeroGPU Spaces（每天可獲得 25 分鐘 H200 使用時間）。
致謝：感謝 ChunTe Lee 為本文制作了精彩的縮略圖。感謝 Pedro 和 Vaibhav 對文章提供的反饋。
英文原文:
原(yuan)文作者: Charles Bensimon, Sayak Paul, Linoy Tsaban, Apolinário Passos譯者: AdinaY