掘金 后端 ( ) • 2024-04-26 21:40
释放Stable Diffusion 无限可能

最近在整理大语言模型的系列内容，Stable Diffusion 是我下一篇博客的主题。关注 Stable Diffusion，是因为它是目前最受欢迎和影响力最大的多模态生成模型之一。Stable Diffusion 于 2022 年 8 月发布，主要用于根据文本的描述产生详细图像，当然它也可以应用于其他任务，比如内补绘制、外补绘制，以及在提示词指导下，对现有图形进行风格化或转变。Stable Diffusion 模型版本正在快速迭代，开源生态也在逐步扩展，对行业生产力带来了巨大的变革。如今出现了很多的开源软件，通过调用 Stable Diffusion 来支持各种功能，并提供简洁的用户界面以方便设计师和爱好者使用。然而 Stable Diffusion 的大规模部署不是一件简单的事情，需要考虑多种因素：

亚马逊云科技开发者社区为开发者们提供全球的开发技术资源。这里有技术文档、开发案例、技术专栏、培训视频、活动与竞赛等。帮助中国开发者对接世界最前沿技术，观点，和项目，并将中国优秀开发者或技术推荐给全球云社区。如果你还没有关注/收藏，看到这里请一定不要匆匆划过，点这里让它成为你的技术宝库！

  • 个性化和可扩展性：Stable Diffusion 生态广泛，仅广泛使用的基础模型就有 1.5 ，XL 和 2 三个版本。同时大量的插件和附加模型（如 LoRA ，ControlNet 等）可被附加到基础模型上。还可针对特定工作场景（如人像生成）进行精细化调优。这些模型和插件也都在不断迭代。在大规模部署时，针对不同的工作场景可以使用不同的模型进行模型推理，这对整个系统可扩展性要求很高。

  • 推理速度：Stable Diffusion 的基础模型是在亚马逊云计算服务上使用 256 个 NVIDIA A100 GPU 训练，共花费 15 万个 GPU 小时，成本为 60 万美元。而这只是基础模型，对于调用基础模型并加载个性化数据进行推理的应用场景来说，需要使用加速计算芯片（如 NVIDIA GPU 或亚马逊云科技的 Inferentia 芯片）来提升单任务推理速度，降低用户等待时间，提升用户体验。

  • 弹性伸缩：在多种业务场景中，使用者的请求有较大的不确定性。从平衡成本的角度出发，需要考虑在请求较多时快速增加推理实例数量以应对请求，而在请求较少时降低实例数量以降低成本。 1.png

上图是一个常见的大语言模型在容器集群上的部署方式，这种部署方式存在以下问题：

  1. 所有请求都是同步的。由于模型推理相对比较耗时，每个请求耗时可达几十秒，甚至几分钟。这不仅要求客户与后端之间的网络绝对稳定，在流量突增且没有限流手段时，甚至会导致系统雪崩。
  2. 常见的自动扩容策略是基于 CPU 或 GPU 利用率的指标跟踪，无法直观反应系统负载，且触发时间长，无法应对突增请求。但如果为了避免冷启动而保留大量的空闲容量，则资源在低谷期大量闲置，空置成本高昂。
  3. 在弹性伸缩拉起新实例后，还需要加载 Stable  Diffusion 运行时和模型才能对外提供服务。Stable Diffusion 运行时的容器镜像普遍在 10 GB 以上，新实例下载镜像和解压耗费的时间过长，导致冷启动时间过长，大大影响使用者的体验。
  4. Stable Diffusion 模型通常使用存放在块存储或文件存储中，每次加载模型时候拉取性能受限，成本也较高。
解决方案

为了解决传统部署方式带来的业务痛点，亚马逊云科技中国区的技术团队基于使用者的需求和对多种业务场景的深刻理解提出了一套解决方案。该方案充分利用云上资源的弹性，无服务器产品的简便，和 Kubernetes 生态强大的任务调度能力，展示了如何构建端到端低成本、快速扩展的异步图像生成架构。

方案的具体工作流程如下图所示：

2.png

  1. 使用者将请求（模型，Prompt 等）发送至业务应用，业务应用将请求发送至 Amazon API Gateway 提供的 API 端点。
  2. 请求通过 Amazon Lambda 进行校验，并投送至 Amazon SNS 主题。
  3. Amazon SNS 根据请求中的运行时名称，基于请求过滤机制，将请求投送至对应运行时的 SQS 队列。
  4. 在 EKS 集群中，KEDA 会根据队列内消息数量扩充运行时的副本数。
  5. Stable Diffusion 运行时启动时会通过 Mountpoint for Amazon S3 CSI Driver，直接从 S3 存储桶中加载模型。
  6. Queue Agent 会从 Amazon SQS 队列里接收任务，并发送给 Stable Diffusion 运行时生成图像。
  7. 生成的图片由 Queue Agent 存储至 Amazon S3 存储桶中，并将完成通知投送至 Amazon SNS 主题，SNS 可将响应投送至 SQS 或其他目标中。
  8. 该解决方案提供完整的可观测性和管理组件，包含基于 CloudWatch 和 ADOT 的数值监控和日志，以及基于 Amazon X-Ray 的全链路跟踪。
  9. 该解决方案通过基于 Amazon CDK 的基础设施即代码部署方式进行部署和配置，通过 IAM 和 API Key 提供安全和访问控制。
方案的生产实践
API 接入和事件分发

该解决方案使用基于 Amazon API Gateway 的 API 端点对外提供服务，并基于 Amazon SNS，Amazon SQS 进行任务分发。用户将请求（模型，Prompt 等）发送至 Amazon API Gateway 提供的 API 端点，请求通过 Amazon Lambda 进行校验，并投送至 Amazon SNS 主题。Amazon SNS 根据请求中的运行时名称，将请求投送至对应运行时的 SQS 队列。

3.png

使用 Queue Agent 组件进行异步推理

部署时，每个运行时有独立的 Amazon SQS 队列以接收请求。而考虑到不同的的 Stable Diffusion 运行时的 API 接口和行为不同，对异步处理和队列的支持也不同，故我们构建了名为 Queue Agent 的组件。Queue Agent 会从 Amazon SQS 队列里接收任务，并转换成对应运行时的 API 请求格式，发送给 Stable Diffusion 运行时生成图像。Queue Agent 还兼顾了从 Internet 或 Amazon S3 上获取图像，和任务跟踪，错误处理的功能，将更多任务从 Stable Diffusion 运行时解耦，从而支持更多的 Stable Diffusion 运行时。

4.png

利用 KEDA 自动扩缩容器副本

当 Amazon SQS 队列中积压过多消息时，需要增加容器副本数以快速处理这些积压任务。我们使用开源的 Kubernetes Event-Driven Autoscaling（简称 KEDA）进行弹性伸缩。KEDA 会轮询 SQS 队列的队列内消息数量，并根据消息数量和预设的阈值来决定运行时的副本数。由于 KEDA 持续监控队列，在没有负载的时候，可以不启动任何运行时，而在有任务时再启动运行时实例，有效节省成本。除了 SQS 队列长度外，KEDA 还支持根据许多其他指标（被称为 Scaler）进行扩缩容。例如部分业务场景有明确的业务高峰期，则可以使用 Cron 在预估的业务高峰期提前进行扩容，而在业务低谷期不预置资源。通过多个 Scaler 的组合，可以支持更多的工作场景。

5.png

利用 Karpenter 自动扩缩实例

在运行时副本数扩展后，需要节点弹性伸缩组件创建实例承载新的运行时副本。传统的节点弹性伸缩组件是 Kubernetes 社区的 Cluster Autoscaler。而亚马逊云科技开源的弹性伸缩组件 Karpenter 可以智能选择实例类型和购买方式。根据同时启动的 Pod 数量，选择最合适的实例类型。例如在同时启动的 Pod 数量很多时，可以选择包含 2 张或 4 张 GPU 的实例，减少实例启动时间对冷启动时间的影响。Karpenter 使用 EC2 Fleet，实例启动速度快于基于 Autoscaling Group 的节点组。在推理任务的成本划分中，GPU 实例的成本占了总成本的 95% 以上。Karpenter 支持创建 Spot 实例，且在创建时会智能选择回收几率最低的实例类型，降低 Spot 实例回收对任务造成的影响。

6.png

基于 Bottlerocket 实现容器镜像缓存

在存储方面，方案针对 Stable Diffusion 运行时和模型使用了不同的存储。新创建的实例采用了 Bottlerocket 操作系统。Bottlerocket 是一款基于 Linux 的开源操作系统，由亚马逊云科技专门构建用于运行容器。Bottlerocket 仅包括运行容器所需的基本软件，启动速度极快。而在 Stable Diffusion 运行时镜像方面，由于 Docker 拉取镜像是单线程，且拉取完成后还需要解压缩，导致镜像拉取耗时较长。我们将 Stable Difussion 运行时镜像预加载到磁盘，并抓取 EBS 快照。在每次创建新实例时，EBS 快照会恢复到新实例的数据盘，新的实例在加入 EKS 集群后可以立即启动容器，节省拉取镜像和解压的时间。

7.png

利用对象存储支撑大量模型动态读取

常见的模型存储会使用块存储或者文件存储，但对于模型读取这种连续读的工作负载，挂载对象存储存在优势。Mountpoint for Amazon S3 可以将 S3 存储桶挂载到文件系统。在读取模型时支持多线程读取，充分利用网卡资源。在 Amazon S3 上的性能测试显示，使用 Mountpoint for Amazon S3 加载 Stable Diffusion 1.5 的基础模型仅需 14.5 秒，进一步缩减冷启动耗时。

8.png

推送处理结果

生成的图片由 Queue Agent 存储至 Amazon S3 存储桶中，并将完成通知投送至 Amazon SNS 主题。不论任务成功或失败，运行时的返回也会存储到 S3 存储桶，方便进行调试。而只有任务真正处理完成后，Queue Agent 才会从队列中删除对应的任务。当出现意外状况（如 SD 运行时崩溃）时，任务会在设定的一段时间（被称为 Visibility Timeout）后，重新出现在队列中，其他运行时可以重试这个任务。Amazon SQS 这个“两阶段确认”的特性有效的提升了整体方案的可靠性。

9.png

X-Ray 端到端全链路追踪

而在可观测性方面，除了传统的基于 CloudWatch 数值监控和日志之外，通过分布式跟踪，我们能了解每一步的耗时和性能瓶颈，方便运维人员快速排查问题。

10.png

方案的性能与成本

这套方案已经在多个 Amazon Web Services 使用者的生产环境上部署。而我们在 Amazon Web Services 美国西部（俄勒冈）区域的内部实测，端到端的扩容速度（即请求发送到处理完成）最低可达 60 秒，完全可以有效应对需要快速扩容的场景。推理成本由于使用了 Spot 实例最高可节省 70%，有效降低了方案的整体成本，适应大规模部署。

开始使用

目前这套方案已经开源，支持部署到亚马逊云科技海外和中国区域。您可以从 https://github.com/aws-samples/stable-diffusion-on-eks/?trk=cndc-detail  获取源代码并开始在您的环境中部署。

资源与参考文献：

文章来源：https://dev.amazoncloud.cn/column/article/662b99cc4e3b7d0f2fb5257f?sc_medium=regulartraffic&sc_campaign=crossplatform&sc_channel=juejin