Hugging Face ZeroGPU 预编译指南

Hugging Face 为 ZeroGPU Spaces 引入了提前编译（AoT）功能，使开发者能够优化模型延迟并消除冷启动时间。通过一次编译模型并在短生命周期的 GPU 进程中即时重新加载，用户在 Flux、Wan 和 LTX 等模型上可实现 1.3Ñ–1.8Ñ 的加速。

ZeroGPU 架构与 AoT 的必要性

ZeroGPU 采用即时（just-in-time）方式进行 GPU 初始化，以最大化资源利用率。ZeroGPU 并不会为整个 Space 的生命周期保留 GPU，而是 fork 出子进程，在 Nvidia H200（当前使用 3g.71gb MIG 切片）上执行 GPU 任务，任务完成后立即终止该子进程。

虽然 torch.compile 在标准环境中有效，但在 ZeroGPU 上它依赖文件系统缓存，因为几乎每个任务都会全新启动进程。此缓存过程可能需要数十秒到数分钟，使得实时演示变得不切实际。AoT 编译通过一次导出已编译模型来解决此问题，随后可以在任何进程中即时重新加载。

在 ZeroGPU 上实现 AoT 编译

实现 AoT 编译需要使用 spaces 包和 PyTorch 工具进行五步工作流：

  1. 捕获示例输入：使用 spaces.aoti_capture 拦截传递给模型组件的参数和关键字参数（例如扩散管线中的 transformer）。
  2. 导出模型：使用 torch.export.export 将模型转换为 ExportedProgram，它是包含张量计算和模型参数的计算图。
  3. 编译导出模型：使用 spaces.aoti_compile（torch._inductor.aot_compile 的包装）生成 AoT 编译的二进制文件。
  4. 应用编译模型：使用 spaces.aoti_apply 修补模型的 forward 方法，并从内存中移除原始模型参数，以防止内存溢出（OOM）错误。
  5. 在 GPU 环境中包装：由于编译依赖硬件并需要真实 GPU，编译步骤必须在应用启动阶段用 @spaces.GPU 函数包装。

对于 FLUX.1-dev 模型，此过程实现了 1.7x 的加速。

高级优化技术

FP8 量化

AoT 编译可以结合 torchao 库的 FP8 训练后动态量化使用。由于 ZeroGPU 使用 H200 GPU（计算能力 9.0+），支持 FP8，可额外带来 1.2x 的加速。

处理动态形状

为了支持不同的图像或视频分辨率，开发者可以使用 torch.export.Dim 定义动态维度。对于 Flux.1-Dev，需要将 hidden_states 的 flattened_latent_dim 和 img_ids 的 height * width 设置为动态。这些配置通过 dynamic_shapes 映射传递给 torch.export.export。

多编译与共享权重

当动态范围过大以致单图动态形状无法覆盖（例如 Wan 视频生成系列），开发者可以为每种分辨率编译一个模型，同时共享模型参数，并在运行时调度相应的编译图。

FlashAttention-3（FA3）

ZeroGPU 与 FlashAttention-3 兼容。为避免从源码构建 FA3 的耗时过程，Hugging Face 提供了 kernels 库，用户可以加载预构建、硬件兼容的内核（例如 kernels-community/vllm-flash-attn3）。

局部编译

开发者可以仅编译模型中重复的计算块（例如 Flux 中的 FluxTransformerBlock 和 FluxSingleTransformerBlock），而不是编译整个模型。这大幅降低冷启动时间——将 Flux.1-Dev 的编译时间从 6 分钟降至 30 秒——同时保持与完整模型编译相同的加速效果。

部署与分发

编译后的计算图模块可以序列化为制品并上传至 Hugging Face Hub。仅保存编译后的模型图而不包含模型参数，可保持存储轻量。演示程序随后可以下载并加载这些预编译图，在启动时完全跳过编译阶段。

Sources