使用 PAI-Blade 优化 Stable Diffusion 推理流程
背景
AIGC 是人工智能计算领域里发展迅速的重要业务。Stable Diffusion 是其中最热门的开源模型,受到广泛关注。然而,随着应用场景不断扩大,Stable Diffusion 所面临的推理时延和计算成本问题也越来越突出。
简介
PAI-Blade 是 PAI 推出的通用推理优化工具,可以通过模型系统联合优化,使模型达到最优推理性能。PAI-Blade 依托于完全动态尺寸的 AI 编译器 BladeDISC 和 基于深度学习自动调度的高性能计算库 BlaDNN, 为包括图像生成模型 Stable Diffsuion, 大语言模型 LLM, 大规模稀疏推荐模型 CTR, 语音识别模型 ASR 等等在内的众多模型提供自动的高性能推理优化。
BladeDISC 是一款支持完全动态尺寸的 AI 编译器,前端支持 Pytorch 和 Tensorflow 模型。对于 Pytorch 模型能够支持 TorchScript 和 TorchDynamo 两种输入模式,后端通过 AStitch 大尺度算子融合技术和高效的 codegen 逻辑提升模型访存密集算子的执行效率。BladeDISC 现已在 github 开源,项目地址:https://github.com/alibaba/BladeDISC 。
BlaDNN 是基于深度学习自动调度的高性能计算库。BlaDNN 作为 Ansor 的升级版,不仅生成的 kernel 性能超过 Ansor,而且可以完全依赖 DNN 自动调度而不使用 Tuning 调优,使得 Dynamic Shape 业务场景的在线自动调度成为可能,基于 DNN 自动调度生成的 GPU 计算密集算子的平均性能达到极致 tuning 性能的 99.39%,通过模型系统联合优化 DNN 推理延时低至 2us, 并且只使用一个 CPU Core,从而不会对 GPU 模型本身的性能造成任何抖动。
通过采用 PAI-Blade 加速推理优化技术,对访存密集型算子进行大尺度融合及优化代码生成,对计算密集型算子进行自动调度,可以大幅度降低 Stable Diffusion 的推理延迟和显存占用,从而减少计算成本。使用 PAI-Blade 优化 Stable Diffusion 具有以下三点优势:
高性能,使用 Blade 可以降低 Text2Img、Img2Img 等推理流程的端到端延迟 2.42-3.05 倍,同时可降低省显存占用至多 5.27 倍,超过 TensorRT-8.5 等业内 SOTA 优化手段。
完全动态 shape 支持,一次优化后,可以支持任意形状、batch size 的输入。
易用性、可扩展性:仅需数行代码即可在多类 pipeline 中启用 Blade 优化,同时能支持 LoRA 等推理方案的优化。
使用示例
本文接下来以社区流行的 "runwayml/stable-diffusion-v1-5" 的 Text2Img pipeline 为例,详细介绍 PAI-Blade 在各类使用场景下的使用方法。
环境安装
下述示例完整的运行脚本及相关环境已集成到 registry.cn-beijing.aliyuncs.com/blade_demo/blade_diffusion
docker 中。在该 docker 中,直接通过 python /blade/blade_diffusion.py
即可运行推理示例。
官方模型优化
使用 PAI-Blade 优化 Stable Diffusion 模型可以分为以下几个步骤。
首先,加载预训练的模型。
第二步,使用 PAI-Blade 进行优化。注意,由于 PAI-Blade 是完全动态 shape 的优化工具,优化完成后可使用任意 shape 进行推理。
最后,使用优化好的模型替换原始模型,后续即可以原始 pipeline 同样的方式进行推理。
A100 性能对比
A10 性能对比
推理结果验证
使用 PAI-Blade 优化后,生成的图像与 Pytorch 原始输出对比,观察优化结果是否正确。左图为 Pytorch eager 模式输出,右图为 PAI-Blade 优化后的模型输出。
已验证的 pipeline 类型
StableDiffusionPipeline
StableDiffusionImg2ImgPipeline
StableDiffusionInpaintPipeline
AltDiffusionPipeline
LoRA 优化
LoRA 是指在原始模型基础上,添加额外的低秩矩阵来微调预训练的模型,并且只训练那些新添加的权重,从而大幅降低微调成本。可以通过 diffusers官方训练代码 微调得到 LoRA 权重。diffusers 加载使用 LoRA 后,模型运行方式与原始模型略有不同,带来额外计算开销。
PAI-Blade 目前已适配 huggingface/diffusers 中 LoRA 优化方式。同样的,Blade 针对同一 pipeline,只需优化一次,即可使用任意的 LoRA 权重进行推理。我们将在下一篇文章中介绍 PAI-Blade 优化 LoRA 的使用方式,敬请期待。
展望
目前,Stable Diffusion 相关技术仍在不断演化中,PAI-Blade 团队也时刻关注社区趋势,将优化适配到各种工具中去。目前团队主要集中在:
将相关优化集成到 stable-diffusion-webui 中;
优化 finetune 训练速度。
版权声明: 本文为 InfoQ 作者【阿里云大数据AI技术】的原创文章。
原文链接:【http://xie.infoq.cn/article/4d629b23ce69730721d8c584b】。文章转载请联系作者。
评论