关于人工智能:使用-PAIBlade-优化-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 具备以下三点劣势：

1. 高性能，应用 Blade 能够升高 Text2Img、Img2Img 等推理流程的端到端提早 2.42-3.05 倍，同时可升高省显存占用至少 5.27 倍，超过 TensorRT-8.5 等业内 SOTA 优化伎俩。
2. 齐全动静 shape 反对，一次优化后，能够反对任意形态、batch size 的输出。
3. 易用性、可扩展性：仅需数行代码即可在多类 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 性能比照

A10 性能比照

推理后果验证

应用 PAI-Blade 优化后，生成的图像与 Pytorch 原始输入比照，察看优化后果是否正确。左图为 Pytorch eager 模式输入，右图为 PAI-Blade 优化后的模型输入。

已验证的 pipeline 类型

1. StableDiffusionPipeline
2. StableDiffusionImg2ImgPipeline
3. StableDiffusionInpaintPipeline
4. AltDiffusionPipeline

LoRA 优化

LoRA 是指在原始模型根底上，增加额定的低秩矩阵来微调预训练的模型，并且只训练那些新增加的权重，从而大幅升高微调老本。能够通过 diffusers 官网训练代码 微调失去 LoRA 权重。diffusers 加载应用 LoRA 后，模型运行形式与原始模型略有不同，带来额定计算开销。

PAI-Blade 目前已适配 huggingface/diffusers 中 LoRA 优化形式。同样的，Blade 针对同一 pipeline，只需优化一次，即可应用任意的 LoRA 权重进行推理。咱们将在下一篇文章中介绍 PAI-Blade 优化 LoRA 的应用形式，敬请期待。

瞻望

目前，Stable Diffusion 相干技术仍在一直演变中，PAI-Blade 团队也时刻关注社区趋势，将优化适配到各种工具中去。目前团队次要集中在：

1. 将相干优化集成到 stable-diffusion-webui 中；
2. 优化 finetune 训练速度。