ZeroGPU Spaces 加速实践：PyTorch 提前编译全解析
<p>ZeroGPU 让任何人都能在 Hugging Face Spaces 中使用强大的 <strong>Nvidia H200</strong> 硬件，而不需要因为空闲流量而长期占用 GPU。它高效、灵活，非常适合演示，不过需要注意的是，ZeroGPU 并不能在所有场景下完全发挥 GPU 与 CUDA 栈的全部潜能，比如生成图像或视频可能需要相当多的时间。在这种情况下，充分利用 H200 硬件，使其发挥极致性能就显得尤为重要。</p><p>这就是 PyTorch 提前编译（AoT）的用武之地。与其在运行时动态编译模型（这和 ZeroGPU 短生命周期的进程配合得并不好），提前编译允许你一次优化、随时快速加载。</p>
<p><strong>结果</strong>：演示 Demo 更流畅、体验更顺滑，在 Flux、Wan 和 LTX 等模型上有 <strong>1.3×–1.8×</strong> 的提速 🔥</p>
<p>在这篇文章中，我们将展示如何在 ZeroGPU Spaces 中接入提前编译（AoT）。我们会探索一些高级技巧，如 FP8 量化和动态形状，并分享你可以立即尝试的可运行演示。如果你想尽快尝试，可以先去 zerogpu-aoti 中体验一些基于 ZeroGPU 的 Demo 演示。</p>
<blockquote>
<p>[!TIP]<br>
Pro 用户和 Team / Enterprise 组织成员可以创建 ZeroGPU Spaces，而任何人都可以免费使用（Pro、Team 和 Enterprise 用户将获得 <strong>8 倍</strong> 的 ZeroGPU 配额）</p>
</blockquote>
<h2 id="目录">目录</h2>
<ul>
<li>什么是 ZeroGPU</li>
<li>PyTorch 编译</li>
<li>ZeroGPU 上的提前编译</li>
<li>注意事项
<ul>
<li>量化</li>
<li>动态形状</li>
<li>多重编译 / 权重共享</li>
<li>FlashAttention-3</li>
</ul>
</li>
<li>AoT 编译的 ZeroGPU Spaces 演示</li>
<li>结论</li>
<li>资源</li>
</ul>
<h2 id="什么是-zerogpu">什么是 ZeroGPU</h2>
<p>Spaces 是一个由 Hugging Face 提供的平台，让机器学习从业者可以轻松发布演示应用。</p>
<p>典型的 Spaces 演示应用看起来像这样：</p>
<pre><code class="language-python">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()
</code></pre>
<p>这样做虽可行，却导致 GPU 在 Space 的整个运行期间被独占，即使是在没有用户访问的情况下。</p>
<p>当执行这一行中的 <code>.to('cuda')</code> 时：</p>
<pre><code class="language-python">pipe = DiffusionPipeline.from_pretrained(...).to('cuda')
</code></pre>
<p>PyTorch 在初始化时会加载 NVIDIA 驱动，使进程始终驻留在 CUDA 上。由于应用流量并非持续稳定，而是高度稀疏且呈现突发性，这种方式的资源利用效率并不高。</p>
<p>ZeroGPU 采用了一种即时初始化 GPU 的方式。它不会在主进程中直接配置 CUDA，而是自动 fork 一个子进程，在其中配置 CUDA、运行 GPU 任务，并在需要释放 GPU 时终止这个子进程。</p>
<p>这意味着：</p>
<ul>
<li>当应用没有流量时，它不会占用任何 GPU</li>
<li>当应用真正执行任务时，它会使用一个 GPU</li>
<li>当需要并发执行任务时，它可以使用多个 GPU</li>
</ul>
<p>借助 Python 的 <code>spaces</code> 包，实现这种行为只需要如下代码改动：</p>
<pre><code class="language-diff">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()
</code></pre>
<p>通过引入 <code>spaces</code> 并添加 <code>@spaces.GPU</code> 装饰器 (decorator)，我们可以做到：</p>
<ul>
<li>拦截 PyTorch API 调用，以延迟 CUDA 操作</li>
<li>让被装饰的函数在 fork 出来的子进程中运行</li>
<li>（调用内部 API，使正确的设备对子进程可见 —— 这不在本文范围内）</li>
</ul>
<blockquote>
<p>[!NOTE]<br>
ZeroGPU 当前会分配 H200 的一个 MIG 切片（<code>3g.71gb</code> 配置）。更多的 MIG 配置（包括完整切片 <code>7g.141gb</code>）预计将在 2025 年底推出。</p>
</blockquote>
<h2 id="pytorch-编译">PyTorch 编译</h2>
<p>在现代机器学习框架（如 PyTorch 和 JAX）中，“编译”已经成为一个重要概念，它能够有效优化模型的延迟和推理性能。其背后通常会执行一系列与硬件相关的优化步骤，例如算子融合、常量折叠等，以提升整体运行效率。</p>
<p>从 PyTorch 2.0 开始，目前有两种主要的编译接口：</p>
<ul>
<li>即时编译（Just-in-time）：<code>torch.compile</code></li>
<li>提前编译（Ahead-of-time）：<code>torch.export</code> + <code>AOTInductor</code></li>
</ul>
<p><code>torch.compile</code> 在标准环境中表现很好：它会在模型第一次运行时进行编译，并在后续调用中复用优化后的版本。</p>
<p>然而，在 ZeroGPU 上，由于几乎每次执行 GPU 任务时进程都是新启动的，这意味着 <code>torch.compile</code> 无法高效复用编译结果，因此只能依赖 文件系统缓存 来恢复编译模型。根据模型的不同，这个过程可能需要几十秒到几分钟，对于 Spaces 中的实际 GPU 任务来说，这显然太慢了。这正是 <strong>提前编译（AoT）</strong> 大显身手的地方。</p>
<p>通过提前编译，我们可以在一开始导出已编译的模型，将其保存，然后在任意进程中即时加载。这不仅能减少框架的额外开销，还能消除即时编译通常带来的冷启动延迟。</p>
<p>但是，我们该如何在 ZeroGPU 上实现提前编译呢？让我们继续深入探讨。</p>
<h2 id="zerogpu-上的提前编译">ZeroGPU 上的提前编译</h2>
<p>让我们回到 ZeroGPU 的基础示例，来逐步解析启用 AoT 编译所需要的内容。在本次演示中，我们将使用 <code>black-forest-labs/FLUX.1-dev</code> 模型：</p>
<pre><code class="language-python">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()
</code></pre>
<blockquote>
<p>[!NOTE]<br>
在下面的讨论中，我们只编译 <code>pipe</code> 的 <code>transformer</code> 组件。<br>
因为在这类生成模型中，transformer（或者更广义上说，denoiser）是计算量最重的部分。</p>
</blockquote>
<p>使用 PyTorch 对模型进行提前编译通常包含以下几个步骤：</p>
<h3 id="1-获取示例输入">1. 获取示例输入</h3>
<p>请记住，我们要对模型进行 <em>提前</em> 编译。因此，我们需要为模型准备示例输入。这些输入应当与实际运行过程中所期望的输入类型保持一致。</p>
<p>为了捕获这些输入，我们将使用 <code>spaces</code> 包中的 <code>spaces.aoti_capture</code> 辅助函数：</p>
<pre><code class="language-python">with spaces.aoti_capture(pipe.transformer) as call:
pipe("arbitrary example prompt")
</code></pre>
<p>当 <code>aoti_capture</code> 作为上下文管理器使用时，它会拦截对任意可调用对象的调用（在这里是 <code>pipe.transformer</code>），阻止其实际执行，捕获本应传递给它的输入参数，并将这些值存储在 <code>call.args</code> 和 <code>call.kwargs</code> 中。</p>
<h3 id="2-导出模型">2. 导出模型</h3>
<p>既然我们已经得到了 transformer 组件的示例参数（args 和 kwargs），我们就可以使用 <code>torch.export.export</code> 工具将其导出为一个 PyTorch <code>ExportedProgram</code>：</p>
<pre><code class="language-python">exported_transformer = torch.export.export(
pipe.transformer,
args=call.args,
kwargs=call.kwargs,
)
</code></pre>
<h3 id="3-编译导出的模型">3. 编译导出的模型</h3>
<p>一旦模型被导出，编译它就非常直接了。</p>
<p>在 PyTorch 中，传统的提前编译通常需要将模型保存到磁盘，以便后续重新加载。 在我们的场景中，可以利用 <code>spaces</code> 包中的一个辅助函数：<code>spaces.aoti_compile</code>。<br>
它是对 <code>torch._inductor.aot_compile</code> 的一个轻量封装，能够根据需要管理模型的保存和延迟加载。其使用方式如下：</p>
<pre><code class="language-python">compiled_transformer = spaces.aoti_compile(exported_transformer)
</code></pre>
<p>这个 <code>compiled_transformer</code> 现在是一个已经完成提前编译的二进制，可以直接用于推理。</p>
<h3 id="4-在流水线中使用已编译模型">4. 在流水线中使用已编译模型</h3>
<p>现在我们需要将已编译好的 transformer 绑定到原始流水线中，也就是 <code>pipeline</code>。接下来，我们需要将编译后的 transformer 绑定到原始的 pipeline 中。 一个看似简单的做法是直接修改：<code>pipe.transformer = compiled_transformer</code>。但这样会导致问题，因为这种方式会丢失一些关键属性，比如 <code>dtype</code>、<code>config</code> 等。如果只替换 <code>forward</code> 方法也不理想，因为原始模型参数依然会常驻内存，往往会在运行时引发 OOM（内存溢出）错误。</p>
<p>因此<code>spaces</code> 包为此提供了一个工具 —— <code>spaces.aoti_apply</code>：</p>
<pre><code class="language-python">spaces.aoti_apply(compiled_transformer, pipe.transformer)
</code></pre>
<p>这样以来，它会自动将 <code>pipe.transformer.forward</code> 替换为我们编译后的模型，同时清理旧的模型参数以释放内存。</p>
<h3 id="5-整合所有步骤">5. 整合所有步骤</h3>
<p>要完成前面三个步骤（拦截输入示例、导出模型，以及用 PyTorch inductor 编译），我们需要一块真实的 GPU。 在 <code>@spaces.GPU</code> 函数之外得到的 CUDA 仿真环境是不够的，因为编译过程高度依赖硬件，例如需要依靠微基准测试来调优生成的代码。这就是为什么我们需要把所有步骤都封装在一个 <code>@spaces.GPU</code> 函数中，然后再将编译好的模型传回应用的根作用域。 从原始的演示代码开始，我们可以得到如下实现：</p>
<pre><code class="language-diff">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()
</code></pre>
<p>只需增加十几行代码，我们就成功让演示运行得更快（在 FLUX.1-dev 的情况下提升了 <strong>1.7 倍</strong>）。</p>
<p>如果你想进一步了解提前编译，可以阅读 PyTorch 的 AOTInductor 教程。</p>
<h2 id="注意事项">注意事项</h2>
<p>现在我们已经展示了在 ZeroGPU 条件下可以实现的加速效果，接下来将讨论在这一设置中需要注意的一些问题。</p>
<h3 id="量化quantization">量化（Quantization）</h3>
<p>提前编译可以与量化结合，从而实现更大的加速效果。对于图像和视频生成任务，FP8 的训练后动态量化方案提供了良好的速度与质量平衡。不过需要注意，FP8 至少需要 9.0 的 CUDA 计算能力才能使用。<br>
幸运的是，ZeroGPU 基于 H200，因此我们已经能够利用 FP8 量化方案。要在提前编译工作流中启用 FP8 量化，我们可以使用 <code>torchao</code> 提供的 API，如下所示：</p>
<pre><code class="language-diff">+ from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig
+ # 在导出步骤之前对 transformer 进行量化
+ quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
exported_transformer = torch.export.export(
pipe.transformer,
args=call.args,
kwargs=call.kwargs,
)
</code></pre>
<p>（你可以在 这里 找到更多关于 TorchAO 的详细信息。）</p>
<p>接着，我们就可以按照上面描述的步骤继续进行。使用量化可以再带来 <strong>1.2 倍</strong> 的加速。</p>
<h3 id="动态形状dynamic-shapes">动态形状（Dynamic shapes）</h3>
<p>图像和视频可能具有不同的形状和尺寸。因此，在执行提前编译时，考虑形状的动态性也非常重要。<code>torch.export.export</code> 提供的原语让我们能够很容易地配置哪些输入需要被视为动态形状，如下所示。</p>
<p>以 Flux.1-Dev 的 transformer 为例，不同图像分辨率的变化会影响其 <code>forward</code> 方法中的两个参数：</p>
<ul>
<li>
<p><code>hidden_states</code>：带噪声的输入潜变量，transformer 需要对其去噪。它是一个三维张量，表示 <code>batch_size, flattened_latent_dim, embed_dim</code>。当 batch size 固定时，随着图像分辨率变化，<code>flattened_latent_dim</code> 也会变化。</p>
</li>
<li>
<p><code>img_ids</code>：一个二维数组，包含编码后的像素坐标，形状为 <code>height * width, 3</code>。在这种情况下，我们希望让 <code>height * width</code> 是动态的。</p>
</li>
</ul>
<p>我们首先需要定义一个范围，用来表示（潜变量）图像分辨率可以变化的区间。为了推导这些数值范围，我们检查了 pipeline 中 <code>hidden_states</code> 的形状在不同图像分辨率下的变化。这些具体数值依赖于模型本身，需要人工检查并结合一定直觉。 对于 Flux.1-Dev，我们最终得到：</p>
<pre><code class="language-python">transformer_hidden_dim = torch.export.Dim('hidden', min=4096, max=8212)
</code></pre>
<p>接下来，我们定义一个映射，指定参数名称，以及在其输入值中哪些维度需要被视为动态：</p>
<pre><code class="language-python">transformer_dynamic_shapes = {
"hidden_dim": {1: transformer_hidden_dim},
"img_ids": {0: transformer_hidden_dim},
}
</code></pre>
<p>然后，我们需要让动态形状对象的结构与示例输入保持一致。对于不需要动态形状的输入，必须将其设置为 <code>None</code>。这可以借助 PyTorch 提供的 <code>tree_map</code> 工具轻松完成：</p>
<pre><code class="language-python">from torch.utils._pytree import tree_map
dynamic_shapes = tree_map(lambda v: None, call.kwargs)
dynamic_shapes |= transformer_dynamic_shapes
</code></pre>
<p>现在，在执行导出步骤时，我们只需将 <code>transformer_dynamic_shapes</code> 传递给 <code>torch.export.export</code>：</p>
<pre><code class="language-python">exported_transformer = torch.export.export(
pipe.transformer,
args=call.args,
kwargs=call.kwargs,
dynamic_shapes=dynamic_shapes,
)
</code></pre>
<blockquote>
<p>[!NOTE]<br>
可以参考 这个 Space，它详细说明了如何在导出步骤中把量化和动态形状结合起来使用。</p>
</blockquote>
<h3 id="多重编译--权重共享">多重编译 / 权重共享</h3>
<p>当模型的动态性非常重要时，仅依靠动态形状有时是不够的。</p>
<p>例如，在 Wan 系列视频生成模型中，如果你希望编译后的模型能够生成不同分辨率的内容，就会遇到这种情况。在这种情况下，可以采用的方法是：<strong>为每种分辨率编译一个模型，同时保持模型参数共享，并在运行时调度对应的模型</strong>。</p>
<p>这里有一个这种方法的示例：zerogpu-aoti-multi.py。<br>
你也可以在 Wan 2.2 Space 中看到该范式的完整实现。</p>
<h3 id="flashattention-3">FlashAttention-3</h3>
<p>由于 ZeroGPU 的硬件和 CUDA 驱动与 Flash-Attention 3（FA3）完全兼容，我们可以在 ZeroGPU Spaces 中使用它来进一步提升速度。FA3 可以与提前编译（AoT）配合使用，因此非常适合我们的场景。</p>
<p>从源码编译和构建 FA3 可能需要几分钟时间，并且这个过程依赖于具体硬件。作为用户，我们当然不希望浪费宝贵的 ZeroGPU 计算时间。这时 Hugging Face 的 <code>kernels</code> 库 就派上用场了，因为它提供了针对特定硬件的预编译内核。</p>
<p>例如，当我们尝试运行以下代码时：</p>
<pre><code class="language-python">from kernels import get_kernel
vllm_flash_attn3 = get_kernel("kernels-community/vllm-flash-attn3")
</code></pre>
<p>它会尝试从 <code>kernels-community/vllm-flash-attn3</code> 仓库加载一个内核，该内核与当前环境兼容。<br>
否则，如果存在不兼容问题，就会报错。幸运的是，在 ZeroGPU Spaces 上这一过程可以无缝运行。这意味着我们可以在 ZeroGPU 上借助 <code>kernels</code> 库充分利用 FA3 的性能。</p>
<p>这里有一个 Qwen-Image 模型的 FA3 注意力处理器完整示例。</p>
<h2 id="提前编译的-zerogpu-spaces-演示">提前编译的 ZeroGPU Spaces 演示</h2>
<h3 id="加速对比">加速对比</h3>
<ul>
<li>未使用 AoTI 的 FLUX.1-dev</li>
<li>使用 AoTI 和 FA3 的 FLUX.1-dev （<strong>1.75 倍</strong> 加速）</li>
</ul>
<h3 id="精选-aoti-spaces">精选 AoTI Spaces</h3>
<ul>
<li>FLUX.1 Kontext</li>
<li>QwenImage Edit</li>
<li>Wan 2.2</li>
<li>LTX Video</li>
</ul>
<h2 id="结论">结论</h2>
<p>Hugging Face Spaces 中的 ZeroGPU 是一项强大的功能，它为 AI 构建者提供了高性能算力。在这篇文章中，我们展示了用户如何借助 PyTorch 的提前编译（AoT）技术，加速他们基于 ZeroGPU 的应用。</p>
<p>我们用 Flux.1-Dev 展示了加速效果，但这些技术并不仅限于这一模型。因此，我们鼓励你尝试这些方法，并在 社区讨论 中向我们提供反馈。</p>
<h2 id="资源">资源</h2>
<ul>
<li>访问 Hub 上的 ZeroGPU-AOTI 组织，浏览一系列利用文中技术的演示。</li>
<li>查看 <code>spaces.aoti_*</code> API 的源代码，了解接口细节。</li>
<li>查看 Hub 上的 Kernels Community 组织。</li>
<li>升级到 Hugging Face 的 Pro，创建你自己的 ZeroGPU Spaces（每天可获得 25 分钟 H200 使用时间）。</li>
</ul>
<p><em>致谢：感谢 ChunTe Lee 为本文制作了精彩的缩略图。感谢 Pedro 和 Vaibhav 对文章提供的反馈。</em></p>
<blockquote>
<p>英文原文: https://huggingface.co/blog/zerogpu-aoti<br>
原文作者: Charles Bensimon, Sayak Paul, Linoy Tsaban, Apolinário Passos</p>
<p>译者: AdinaY</p>
</blockquote><br><br>
来源：https://www.cnblogs.com/huggingface/p/19081312
頁:
[1]