如何用 PyTorch AoT 编译让 ZeroGPU Spaces 提速

ZeroGPU 让任何人都能在 HF Mirror Spaces 中启动强大的 Nvidia H200 硬件，而无需为闲置流量锁定 GPU。它高效、灵活，非常适合演示，但它并不总能充分利用 GPU 和 CUDA 堆栈所能提供的一切。生成图像或视频可能需要大量时间。在这种情况下，能够充分利用 H200 硬件来榨取更多性能确实很重要。

这正是 PyTorch 提前编译（AoT）的用武之地。AoT 无需即时编译模型（这与 ZeroGPU 的短生命周期进程不太兼容），而是让你一次优化、随时即时重载。

结果：更迅捷的演示和更流畅的体验，在 Flux、Wan 和 LTX 等模型上实现 1.3×–1.8× 的加速 🔥

在本文中，我们将展示如何在 ZeroGPU Spaces 中接入提前（AoT）编译。我们将探索 FP8 量化和动态形状等高级技巧，并分享你可立即尝试的可用演示。如果你等不及了，我们邀请你查看 zerogpu-aoti 组织下一些由 ZeroGPU 驱动的演示。

Pro 用户和 Team / Enterprise 组织成员可以创建 ZeroGPU Spaces，而任何人都可以免费使用它们（Pro、Team 和 Enterprise 用户可获得 8 倍 的 ZeroGPU 配额）

什么是 ZeroGPU

Spaces 是一个由 HF Mirror 提供支持的平台，让机器学习从业者可以轻松发布演示应用。

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()

这很有效，但最终会在 Space 的整个生命周期内为其保留一块 GPU——即使它没有任何用户活动。

当在这一行执行 .to('cuda') 时：

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

PyTorch 会初始化 NVIDIA 驱动，从而将该进程永久地设置在 CUDA 上。鉴于应用流量并不完全平稳，而是极度稀疏且呈尖峰状，这种做法的资源效率并不高。

ZeroGPU 采用即时（just-in-time）方式进行 GPU 初始化。它不会将主进程设置在 CUDA 上，而是自动 fork 进程，将其设置在 CUDA 上，运行 GPU 任务，并在需要释放 GPU 时最终终止该 fork。

这意味着：

当应用没有流量时，它不使用任何 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 装饰器，我们：

拦截 PyTorch API 调用以推迟 CUDA 操作
使被装饰的函数在稍后调用时在 fork 中运行
（调用内部 API 以使正确的设备对该 fork 可见，但这不在本博客文章的范围之内）

ZeroGPU 目前分配 H200 的一个 MIG 切片（3g.71gb 配置）。包括完整切片（7g.141gb 配置）在内的更多 MIG 规格将于 2025 年底推出。

PyTorch 编译

像 PyTorch 和 JAX 这样的现代机器学习框架都有编译的概念，可用于优化模型延迟或推理时间。在幕后，编译会应用一系列（通常依赖硬件的）优化步骤，例如算子融合、常量折叠等。

PyTorch（从 2.0 起）目前有两个主要的编译接口：

即时编译，使用 torch.compile
提前编译，使用 torch.export + AOTInductor

torch.compile 在标准环境中表现出色：它会在模型首次运行时进行编译，并在后续调用中复用优化后的版本。

然而，在 ZeroGPU 上，由于（几乎）每个 GPU 任务都会重新启动进程，这意味着 torch.compile 无法有效地复用编译结果，因此只能依赖其文件系统缓存来恢复已编译的模型。根据所编译模型的不同，这一过程需要几十秒到几分钟不等，这对于 Spaces 中的实际 GPU 任务来说耗时过长。

这正是提前（AoT）编译大显身手的地方。

借助 AoT，我们可以将编译后的模型导出一次、保存下来，之后在任何进程中即时重新加载，这正是 ZeroGPU 所需要的。这有助于我们降低框架开销，同时也消除了即时编译通常会产生冷启动时间的问题。

但我们如何在 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()

在下面的讨论中，我们只编译 pipe 的 transformer 组件，因为在这类生成模型中，transformer（或更一般地说，去噪器）是计算量最大的组件。

使用 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,
)

导出的 PyTorch 程序是一个计算图，它表示张量计算以及原始模型参数值。

3. 编译导出的模型

模型导出后，编译它就非常简单了。

PyTorch 中传统的 AoT 编译通常需要将模型保存到磁盘，以便之后重新加载。在我们的例子中，我们将利用 spaces 包中的一个辅助函数：spaces.aoti_compile。它是对 torch._inductor.aot_compile 的一个小型封装，负责按需保存和延迟加载模型。它的使用方式如下：

compiled_transformer = spaces.aoti_compile(exported_transformer)

这个 compiled_transformer 现在是一个已 AoT 编译的二进制文件，可直接用于推理。

4. 在流水线中使用已编译的模型

现在我们需要将编译后的 transformer 绑定到我们原来的流水线，即 pipeline。

一种朴素且几乎可行的做法是像 pipe.transformer = compiled_transformer 那样直接修补我们的 pipeline。遗憾的是，这种做法行不通，因为它会删除像 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 函数里，然后把编译好的模型取回到应用的根处。从我们最初的 demo 代码开始，就得到：

  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) # maximum duration allowed during startup
+ 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()

仅用十几行额外代码，我们就成功让 demo 快了不少（在 FLUX.1-dev 的情况下快了 1.7 倍）。

如果你想了解更多关于 AoT 编译的内容，可以阅读 PyTorch 的AOTInductor 教程

注意事项

既然我们已经展示了在 ZeroGPU 约束下所能实现的加速，接下来我们将讨论在使用这套设置时遇到的一些注意事项。

量化

AoT 可以与量化结合，带来更大的加速。对于图像和视频生成，FP8 训练后动态量化方案能提供良好的速度-质量权衡。然而，FP8 需要至少 9.0 的 CUDA 计算能力才能工作。幸运的是，对于 ZeroGPU 来说，由于它们基于 H200，我们已经可以利用 FP8 量化方案。

为了在我们的 AoT 编译工作流中启用 FP8 量化，我们可以像这样利用 torchao 提供的 API：

+ from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig

+ # Quantize the transformer just before the export step.
+ quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())

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

（你可以在这里找到更多关于 TorchAO 的细节。）

然后我们就可以继续执行上面列出的其余步骤。使用量化又带来了 1.2 倍的加速。

动态形状

图像和视频可以有各种不同的形状和尺寸。因此，在执行 AoT 编译时，考虑形状的动态性也很重要。torch.export.export 提供的原语使其可以轻松配置哪些输入应按动态形状处理，如下所示。

对于 Flux.1-Dev transformer 的情况，不同图像分辨率的变化会影响它的两个 forward 参数：

hidden_states：含噪输入 latents，transformer 应对其进行去噪。它是一个 3D 张量，表示 batch_size, flattened_latent_dim, embed_dim。当 batch size 固定时，图像分辨率的任何变化都会改变 flattened_latent_dim。
img_ids：一个 2D 数组，包含编码后的像素坐标，形状为 height * width, 3。在这种情况下，我们希望让 height * width 动态化。

我们首先定义一个希望让（latent）图像分辨率变化的范围。为了推导出这些值范围，我们检查了 pipeline 中 hidden_states 相对于不同图像分辨率的形状。具体数值取决于模型，需要人工检查并借助一些直觉。对于 Flux.1-Dev，我们最终得到：

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

然后我们定义一个参数名与其输入值中哪些维度预期为动态的映射：

transformer_dynamic_shapes = {
    "hidden_states": {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

现在，在执行导出步骤时，我们只需将 torch.export.export 提供给 transformer_dynamic_shapes：

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

看看这个 Space，它展示了如何在导出步骤中同时使用量化和动态形状。

多编译/共享权重

当动态性过于重要时，动态形状有时并不够用。

例如，对于 Wan 系列视频生成模型，如果你希望编译后的模型能够生成不同的分辨率，就会出现这种情况。在这种情况下可以这样做：为每个分辨率编译一个模型，同时保持模型参数共享，并在运行时分派正确的模型。

下面是这种方法的一个最小示例：zerogpu-aoti-multi.py。你也可以在 Wan 2.2 Space 中看到该范式的完整可用实现。

FlashAttention-3

由于 ZeroGPU 硬件和 CUDA 驱动与 Flash-Attention 3 (FA3) 完全兼容，我们可以在 ZeroGPU Spaces 中使用它来进一步加速。FA3 支持提前编译。因此，这非常适合我们的场景。

从源码编译和构建 FA3 可能需要几分钟，而且这个过程依赖硬件。作为用户，我们不想浪费宝贵的 ZeroGPU 计算时长。这正是 HF Mirror kernels 库发挥作用的地方。它提供了对特定硬件兼容的预构建内核的访问。例如，当我们尝试运行：

from kernels import get_kernel

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

它会尝试从 kernels-community/vllm-flash-attn3 仓库加载与当前设置兼容的内核。否则，它将因兼容性问题而报错。幸运的是，这在 ZeroGPU Spaces 上可以无缝运行。这意味着我们可以使用 kernels 库，在 ZeroGPU 上利用 FA3 的强大能力。

这里是 Qwen-Image 模型的一个完整可用的 FA3 attention 处理器示例。

区域编译

到目前为止，我们一直在编译整个模型。取决于模型的不同，完整模型编译可能会导致显著较长的冷启动时间。过长的冷启动时间会让开发体验变得糟糕。

我们也可以选择编译模型中的区域，显著减少冷启动时间，同时保留完整模型编译的几乎所有优势。当一个模型包含重复的计算块时，区域编译会变得很有前景。例如，标准语言模型有许多结构相同的 Transformer 块。

在我们的示例中，我们可以提前编译 Flux transformer 的重复块，并将编译后的计算图传播到其余重复块。Flux Transformer 有两种重复块：FluxTransformerBlock 和 FluxSingleTransformerBlock。

你可以查看这个 Space 获取完整示例。

💡 对于 Flux.1-Dev，切换到区域编译可将编译时间从6 分钟减少到仅30 秒，同时提供相同的加速效果。

使用来自 Hub 的编译计算图

一旦模型（甚至模型块）完成提前编译，我们就可以将编译后的图模块序列化为产物，并在之后复用。在 Spaces 上基于 ZeroGPU 驱动的演示中，这将通过跳过编译时间显著缩短演示的启动时间。

为了保持存储轻量，我们可以只保存编译后的模型图，而不在产物中包含任何模型参数。

查看这个合集，它展示了获取编译模型图、将其推送到 Hub，然后用于构建演示的完整工作流。

AoT 编译的 ZeroGPU Spaces 演示

加速对比

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

精选 AoTI Spaces

区域编译

结论

HF Mirror Spaces 中的 ZeroGPU 是一项强大的功能，它通过提供对强大计算资源的访问来赋能 AI 构建者。在这篇文章中，我们展示了用户如何从 PyTorch 的提前编译技术中受益，以加速他们利用 ZeroGPU 的应用程序。

我们展示了 Flux.1-Dev 的加速效果，但这些技术并不局限于这一个模型。因此，我们鼓励你尝试这些技术，并在这个社区讨论中向我们提供反馈。

资源

访问我们的Hub 上的 ZeroGPU-AOTI 组织，查看利用本文所讨论技术的演示合集。
浏览 spaces.aoti_* API 的源代码以了解更多关于该接口的信息
查看Hub 上的 Kernels Community 组织
从这里了解更多关于区域编译的信息
升级到 HF Mirror 上的Pro，创建你自己的 ZeroGPU Spaces（并每天获得 25 分钟的 H200 使用时长）

致谢：感谢 ChunTe Lee 为本文制作了出色的缩略图。感谢 Pedro 和 Vaibhav 对本文提供反馈。感谢 PyTorch 团队的 Angela Yi 帮助我们进行 AOT 指导。

如何用 PyTorch AoT 编译让 ZeroGPU Spaces 提速

目录

什么是 ZeroGPU

PyTorch 编译

在 ZeroGPU 上进行提前编译

1. 获取示例输入

2. 导出模型

3. 编译导出的模型

4. 在流水线中使用已编译的模型

5. 把所有内容整合到一起

注意事项

量化

动态形状

多编译/共享权重

FlashAttention-3

区域编译

使用来自 Hub 的编译计算图

AoT 编译的 ZeroGPU Spaces 演示

加速对比

精选 AoTI Spaces

区域编译

结论

资源