使用 BentoML 部署 DeepFloyd IF

TL;DR

Hugging Face 已详细介绍了使用 BentoML（一个开源模型服务框架）部署 DeepFloyd IF 的工作流程。此集成使开发者能够通过在不同 GPU 资源上独立扩展推理运行器来管理 DeepFloyd IF 的复杂多阶段架构，以优化生产性能。

DeepFloyd IF 架构和功能

DeepFloyd IF 是一个开源的文本到图像模型，直接在像素空间运行，这使其区别于如 Stable Diffusion 这样的潜在扩散模型。它采用由冻结文本编码器和三个级联像素扩散模块组成的模块化结构：

  • 阶段 1：生成 64x64 像素的基础图像。
  • 阶段 2 和 3：逐步将图像上采样至最终分辨率 1024x1024 像素。

为了增强对复杂提示的语言理解，DeepFloyd IF 将 T5-XXL-1.1 大型语言模型 (LLM) 作为其文本编码器。

使用 BentoML 进行生产部署

BentoML 提供了一个统一框架，用于将模型从 Hugging Face Hub 过渡到生产就绪的 AI 应用。部署过程遵循五步生命周期：

  1. 定义模型：使用在 PyTorch 或 TensorFlow 等库中训练的模型。
  2. 保存模型：将模型存储在 BentoML 本地模型存储中，以便管理和服务访问。
  3. 创建 BentoML 服务：使用 service.py 文件包装模型并通过 Runners 定义服务逻辑，Runners 负责大规模模型推理并公开用于输入/输出处理的 API。
  4. 构建 Bento：将服务和模型打包成一个“Bento”——一个可部署的工件，通过 YAML 配置文件包含所有必要的代码和依赖项。
  5. 部署 Bento：将 Bento 容器化为 Docker 镜像以进行 Kubernetes 部署，或使用 Yatai 在 Kubernetes 上实现自动扩展。

技术实现和资源管理

部署 DeepFloyd IF 需要大量计算资源。推荐的硬件包括至少 2x16GB VRAM GPU 或 1x40GB VRAM GPU。对于生产级服务，不建议使用单个 Tesla T4。

GPU 分配策略

BentoML 使得可以独立扩展 DeepFloyd IF 管道每个阶段的 Runners，使开发者能够根据每个阶段的具体需求分配资源：

  • 高 VRAM (40GB+)：所有模型可以在单个 GPU 上运行。
  • 双 Tesla T4（每个 15GB）：阶段 1 可以分配给第一个 GPU，而阶段 2 和 3 分配给第二个 GPU。
  • 混合 GPU 配置：阶段 1 可以分配给主 GPU（例如 T4），而阶段 2 和 3 分布在两个额外的 VRAM 较小的 GPU 上。

部署工作流程

为了实现此设置，开发者使用 import_models.py 将所需阶段下载到 BentoML 模型存储，并使用 service.py 定义接受 JSON 提示和负面提示以返回生成图像的 API。最终工件使用 bentoml build 打包，可通过 bentoml serve 在本地提供服务，或使用 bentoml containerize 部署到云端。

Sources