使用 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 模型可以分为以下几个步骤。

首先，加载预训练的模型。

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 训练速度。

原文链接

本文为阿里云原创内容，未经允许不得转载