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 中体验一些基于 ZeroGPU 的 Demo 演示。 [!TIP] Pro 用户和 Team / Enterprise 组织成员可以创建 ZeroGPU Spaces，而任何人都可以免费使用（Pro、Team 和 Enterprise 用户将获得 8 倍 的 ZeroGPU 配额） 目录 什么是 ZeroGPU PyTorch 编译 ZeroGPU 上的提前编译 注意事项 量化 动态形状 多重编译 / 权重共享 FlashAttention-3 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() 这样做虽可行，却导致 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 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)，我们可以做到： 拦截 PyTorch API 调用，以延迟 CUDA 操作 让被装饰的函数在 fork 出来的子进程中运行 （调用内部 API，使正确的设备对子进程可见 —— 这不在本文范围内） [!NOTE] ZeroGPU 当前会分配 H200 的一个 MIG 切片（3g.71gb 配置）。更多的 MIG 配置（包括完整切片 7g.141gb）预计将在 2025 年底推出。