使用 PAI-Blade 优化 Stable Diffusion 推理流程-Toy模板网

这篇具有很好参考价值的文章主要介绍了使用 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 性能对比

image size	samplesteps	Time of Pytorch(s)	Time of PAI-Blade(s)	speedup	Pytorch memory usage (GB)	PAI-Blade memory usage (GB)
1024x1024	50	13.26	4.34	3.06X	32.91	6.25
768x768	50	5.65	2.00	2.83X	14.99	5.91
512x512	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)
1024x1024	50	OOM	13.86	-	OOM	6.89
768x768	50	13.13	5.61	2.34X	12.60	6.22
512x512	50	4.53	2.11	2.15X	6.28	5.47