Hugging Face 教程：用 PyTorch 提前编译加速 ZeroGPU Spaces 生成模型

ZeroGPU 让任何人都能在 Hugging Face 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
• PyTorch 编译
• ZeroGPU 上的提前编译
• Gotchas
  • 量化
  • 动态形状
  • 多重编译 / 共享权重
  • FlashAttention-3
  • 区域编译
  • 使用来自 Hub 的已编译图
• AoT 编译的 ZeroGPU Spaces 演示
• 结论
• 资源

什么是 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()

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

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

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

PyTorch 会初始化 NVIDIA 驱动，从而将进程永久设置在 CUDA 上。考虑到应用流量并不完全平稳，而是极其稀疏且呈尖峰状，这并不十分高效。

ZeroGPU 采用即时方式初始化 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 编译

现代 ML 框架如 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. 在 pipeline 中使用编译后的模型

现在我们需要将编译后的 transformer 绑定到我们原始的 pipeline，即 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) # 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()

仅用十几行额外代码，我们就成功让演示变得相当快（在 FLUX.1-dev 的情况下快了 1.7 倍）。

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

注意事项

现在我们已经展示了在 ZeroGPUs 约束下可以实现的加速，我们将讨论在使用此设置时出现的一些注意事项。

量化

AoT 可以与量化结合，以实现更大的加速。 对于图像和视频生成，FP8 训练后动态量化方案提供了良好的速度-质量权衡。 然而，FP8 需要至少 9.0 的 CUDA 计算能力才能工作。 幸运的是，对于 ZeroGPUs，由于它们基于 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：带噪声的输入潜在表示，transformer 应对其进行去噪。它是一个 3D 张量，表示 batch_size, flattened_latent_dim, embed_dim。当批量大小固定时，对于图像分辨率的任何更改，变化的是 flattened_latent_dim。
• img_ids：编码像素坐标的 2D 数组，形状为 height * width, 3。在这种情况下，我们希望让 height * width 动态化。

我们首先定义一个希望让（潜在）图像分辨率变化的范围。 为了推导这些值范围，我们检查了流水线中 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

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

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 计算时长。这正是 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 上可以无缝运行。这意味着我们可以使用 kernels 库在 ZeroGPU 上利用 FA3 的强大能力。

这里是 Qwen-Image 模型的一个完全可用的 FA3 注意力处理器示例。

区域编译

到目前为止，我们一直在编译整个模型。根据模型的不同，全模型编译可能会导致显著较长的冷启动时间。较长的冷启动时间会让开发体验变得不愉快。

我们也可以选择编译模型内的区域，从而显著减少冷启动时间，同时保留全模型编译的几乎所有好处。当模型具有重复的计算块时，区域编译就变得很有前景。例如，标准语言模型有许多结构相同的 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

• FLUX.1 Kontext
• QwenImage Edit
• Wan 2.2

区域编译

• 区域编译配方
• 使用 AOT 的区域编译
• Diffusers 中的原生集成
• 更多性能数据

结论

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

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

资源

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

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