关于人工智能:使用-PAIBlade-优化-Stable-Diffusion-推理流程

28次阅读

共计 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 具备以下三点劣势:

  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 性能比照

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 类型

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

正文完
 0