背景
AIGC是人工智能計(jì)算領(lǐng)域里發(fā)展迅速的重要業(yè)務(wù)。Stable Diffusion 是其中最熱門的開源模型，受到廣泛關(guān)注。然而，隨著應(yīng)用場(chǎng)景不斷擴(kuò)大，Stable Diffusion所面臨的推理時(shí)延和計(jì)算成本問題也越來越突出。
簡(jiǎn)介
PAI-Blade是 PAI 推出的通用推理優(yōu)化工具，可以通過模型系統(tǒng)聯(lián)合優(yōu)化，使模型達(dá)到最優(yōu)推理性能。PAI-Blade依托于完全動(dòng)態(tài)尺寸的AI編譯器BladeDISC?和?基于深度學(xué)習(xí)自動(dòng)調(diào)度的高性能計(jì)算庫BlaDNN， 為包括圖像生成模型Stable Diffsuion, 大語言模型LLM, 大規(guī)模稀疏推薦模型CTR, 語音識(shí)別模型ASR等等在內(nèi)的眾多模型提供自動(dòng)的高性能推理優(yōu)化。
BladeDISC?是一款支持完全動(dòng)態(tài)尺寸的AI編譯器，前端支持Pytorch和Tensorflow模型。對(duì)于Pytorch模型能夠支持 TorchScript 和 TorchDynamo 兩種輸入模式，后端通過 AStitch 大尺度算子融合技術(shù)和高效的 codegen 邏輯提升模型訪存密集算子的執(zhí)行效率。BladeDISC現(xiàn)已在github開源，項(xiàng)目地址：https://github.com/alibaba/BladeDISC?。
BlaDNN?是基于深度學(xué)習(xí)自動(dòng)調(diào)度的高性能計(jì)算庫。BlaDNN 作為Ansor的升級(jí)版，不僅生成的kernel性能超過Ansor，而且可以完全依賴DNN自動(dòng)調(diào)度而不使用Tuning調(diào)優(yōu)，使得Dynamic Shape業(yè)務(wù)場(chǎng)景的在線自動(dòng)調(diào)度成為可能，基于DNN自動(dòng)調(diào)度生成的GPU計(jì)算密集算子的平均性能達(dá)到極致tuning性能的99.39%，通過模型系統(tǒng)聯(lián)合優(yōu)化DNN推理延時(shí)低至2us, 并且只使用一個(gè)CPU Core，從而不會(huì)對(duì)GPU模型本身的性能造成任何抖動(dòng)。
通過采用 PAI-Blade 加速推理優(yōu)化技術(shù)，對(duì)訪存密集型算子進(jìn)行大尺度融合及優(yōu)化代碼生成，對(duì)計(jì)算密集型算子進(jìn)行自動(dòng)調(diào)度，可以大幅度降低Stable Diffusion的推理延遲和顯存占用，從而減少計(jì)算成本。使用 PAI-Blade 優(yōu)化Stable Diffusion 具有以下三點(diǎn)優(yōu)勢(shì)：
- 高性能，使用Blade可以降低 Text2Img、Img2Img 等推理流程的端到端延遲 2.42-3.05 倍，同時(shí)可降低省顯存占用至多 5.27 倍，超過TensorRT-8.5等業(yè)內(nèi)SOTA優(yōu)化手段。
- 完全動(dòng)態(tài)shape支持，一次優(yōu)化后，可以支持任意形狀、batch size的輸入。
- 易用性、可擴(kuò)展性：僅需數(shù)行代碼即可在多類pipeline中啟用 Blade優(yōu)化，同時(shí)能支持LoRA等推理方案的優(yōu)化。
使用示例
本文接下來以社區(qū)流行的 "runwayml/stable-diffusion-v1-5" 的 Text2Img pipeline 為例，詳細(xì)介紹 PAI-Blade 在各類使用場(chǎng)景下的使用方法。
環(huán)境安裝
下述示例完整的運(yùn)行腳本及相關(guān)環(huán)境已集成到?registry.cn-beijing.aliyuncs.com/blade_demo/blade_diffusion?docker 中。在該docker中，直接通過?python /blade/blade_diffusion.py?即可運(yùn)行推理示例。
官方模型優(yōu)化
使用 PAI-Blade 優(yōu)化 Stable Diffusion 模型可以分為以下幾個(gè)步驟。
首先，加載預(yù)訓(xùn)練的模型。
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 進(jìn)行優(yōu)化。注意，由于 PAI-Blade 是完全動(dòng)態(tài)shape的優(yōu)化工具，優(yōu)化完成后可使用任意shape進(jìn)行推理。
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)最后，使用優(yōu)化好的模型替換原始模型，后續(xù)即可以原始 pipeline 同樣的方式進(jìn)行推理。
@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 性能對(duì)比
| 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 性能對(duì)比
| 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 |
推理結(jié)果驗(yàn)證
使用PAI-Blade優(yōu)化后，生成的圖像與Pytorch原始輸出對(duì)比，觀察優(yōu)化結(jié)果是否正確。左圖為Pytorch eager模式輸出，右圖為PAI-Blade優(yōu)化后的模型輸出。
已驗(yàn)證的pipeline類型
- StableDiffusionPipeline
- StableDiffusionImg2ImgPipeline
- StableDiffusionInpaintPipeline
- AltDiffusionPipeline
LoRA優(yōu)化
LoRA 是指在原始模型基礎(chǔ)上，添加額外的低秩矩陣來微調(diào)預(yù)訓(xùn)練的模型，并且只訓(xùn)練那些新添加的權(quán)重，從而大幅降低微調(diào)成本?？梢酝ㄟ^?diffusers官方訓(xùn)練代碼?微調(diào)得到 LoRA 權(quán)重。diffusers 加載使用 LoRA 后，模型運(yùn)行方式與原始模型略有不同，帶來額外計(jì)算開銷。
PAI-Blade 目前已適配 huggingface/diffusers 中 LoRA 優(yōu)化方式。同樣的，Blade 針對(duì)同一pipeline，只需優(yōu)化一次，即可使用任意的 LoRA 權(quán)重進(jìn)行推理。我們將在下一篇文章中介紹PAI-Blade 優(yōu)化 LoRA 的使用方式，敬請(qǐng)期待。
展望
目前，Stable Diffusion相關(guān)技術(shù)仍在不斷演化中，PAI-Blade 團(tuán)隊(duì)也時(shí)刻關(guān)注社區(qū)趨勢(shì)，將優(yōu)化適配到各種工具中去。目前團(tuán)隊(duì)主要集中在：
- 將相關(guān)優(yōu)化集成到 stable-diffusion-webui 中；
- 優(yōu)化 finetune 訓(xùn)練速度。
原文鏈接文章來源：http://www.zghlxwxcb.cn/news/detail-489631.html
本文為阿里云原創(chuàng)內(nèi)容，未經(jīng)允許不得轉(zhuǎn)載文章來源地址http://www.zghlxwxcb.cn/news/detail-489631.html
到了這里，關(guān)于使用 PAI-Blade 優(yōu)化 Stable Diffusion 推理流程的文章就介紹完了。如果您還想了解更多內(nèi)容，請(qǐng)?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！
