共计 3915 个字符,预计需要花费 10 分钟才能阅读完成。
背景
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 模型能够分为以下几个步骤。
首先,加载预训练的模型。
from diffusers import StableDiffusionPipeline
device = torch.device("cuda:0")
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", revision="fp16", torch_dtype=torch.float16).to(device)
第二步,应用 PAI-Blade 进行优化。留神,因为 PAI-Blade 是齐全动静 shape 的优化工具,优化实现后可应用任意 shape 进行推理。
import torch_blade
opt_cfg = torch_blade.Config()
opt_cfg.enable_fp16 = True
with opt_cfg, torch.no_grad():
encoder = blade_optimize(pipe.text_encoder, model_inputs=encoder_inputs, allow_tracing=True)
unet = blade_optimize(pipe.unet, model_inputs=unet_inputs, allow_tracing=True)
decoder = blade_optimize(pipe.vae.decoder, model_inputs=decoder_inputs, allow_tracing=True)
最初,应用优化好的模型替换原始模型,后续即能够原始 pipeline 同样的形式进行推理。
@dataclass
class UNet2DConditionOutput:
sample: torch.FloatTensor
class TracedUNet(torch.nn.Module):
def __init__(self):
super().__init__()
self.config = pipe.unet.config
self.in_channels = pipe.unet.in_channels
self.device = pipe.unet.device
def forward(self, latent_model_input, t, encoder_hidden_states, **kwargs):
sample = unet(latent_model_input.half(), t.half(), encoder_hidden_states.half())["sample"]
return UNet2DConditionOutput(sample=sample)
class TracedEncoder(torch.nn.Module):
def __init__(self):
super().__init__()
self.config = pipe.text_encoder.config
self.device = pipe.text_encoder.device
self.dtype = torch.half
def forward(self, input_ids, **kwargs):
embeddings = encoder(input_ids.long())
return [embeddings["last_hidden_state"]]
class TracedDecoder(torch.nn.Module):
def forward(self, input):
return decoder(input.half())
pipe.text_encoder = TracedEncoder()
pipe.unet = TracedUNet()
pipe.vae.decoder = TracedDecoder()
A100 性能比照
image size | samplesteps | Time of Pytorch(s) | Time of PAI-Blade(s) | speedup | Pytorch memory usage (GB) | PAI-Blade memory usage (GB) |
---|---|---|---|---|---|---|
1024×1024 | 50 | 13.26 | 4.34 | 3.06X | 32.91 | 6.25 |
768×768 | 50 | 5.65 | 2.00 | 2.83X | 14.99 | 5.91 |
512×512 | 50 | 2.24 | 0.84 | 2.67X | 6.60 | 5.42 |
A10 性能比照
image size | samplesteps | Time of Pytorch(s) | Time of PAI-Blade(s) | speedup | Pytorch memory usage (GB) | PAI-Blade memory usage (GB) |
---|---|---|---|---|---|---|
1024×1024 | 50 | OOM | 13.86 | – | OOM | 6.89 |
768×768 | 50 | 13.13 | 5.61 | 2.34X | 12.60 | 6.22 |
512×512 | 50 | 4.53 | 2.11 | 2.15X | 6.28 | 5.47 |
推理后果验证
应用 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 训练速度。