使用 PAI-Blade 加速 StableDiffusion Fine-Tuning

01

背景

Stable Diffusion 模型自从发布以来在互联网上发展迅猛,它可以根据用户输入的文本描述信息生成相关图片,用户也可以提供自己喜爱的风格的照片,来对模型进行微调。例如当我们输入 "A photo of sks dog in a bucket" ,StableDiffusion 模型会生成类似下面的图片:

02

PAI-Blade 加速 PyTorch 训练

PAI-Blade 使用编译优化技术提高 PyTorch 程序的执行效率,其代码已经开源在

Github: github.com/alibaba/Bla....

PAI-Blade API

使用 PAI-Blade 对 PyTorch 程序进行加速非常简单,只需要在原有程序上插入两行代码即可:

1. import PAI-Blade Python package import torch_blade# 2. compile and accelerate 'model' performancemodel = torch.compile(backend='aot_disc')(model)for batch, label in data_loader(): output = model(**batch) loss = compute_loss(output, label) loss.backward() optimizer.step()

torch.compile(backend='aot_disc')(model) 使用 BladeDISC 作为 TorchDynamo 的编译器后端,加速 PyTorch 模型的的前向和反向计算,其中 model也可以是一段 PyTorch 实现的 Python 函数。

PAI-Blade 编译 Pipeline

TorchDynamo 将 PyTorch 程序记录到一个或多个 FX Graph 上,PAI-Blade 通过一系列 Pass 优化计算图的执行效率。

pytorch.org/docs/2.1/to...

MHLO Conversion PAI-Blade 引入了 Torch-MLIR Project 将 PyTorch IR 转换为 MLIR 世界中的 MHLO Dialect,以便进一步使用 BladeDISC 编译器进行性能优化,同时 PAI-Blade 开发团队也将 MHLO 转换相关代码贡献给了社区。

github.com/llvm/torch-...

BlaDNN Library 提供了高性能计算密集型算子库,PAI-Blade 会根据计算图上的一些典型 Pattern,自动的将一部分子图替换为等价的,有极致性能的 BlaDNN 算子。

Memory Intensive Kernel Fusion

算子融合是图层面编译优化最重要的收益来源,一个典型的 workload 上,可能会包含 element-wise 算子,动态 shape 的 broadcast/reshape/reduce 算子以及计算密集型算子,例如 GEMM 等。在 PyTorch 中,每一个算子都是一个独立的 kernel,而过多的 kernel 会导致 Tensor 在 Cache 中频繁的交换,导致显存带宽的浪费,而频繁的发射 kernel 也会造成一定的额外的开销。

对于如上图的一个典型的访存类算子 workload ,类似 XLA 做法会将 schedule 相同的算子合并在一起,从而将 7 个 kernel 合并为 3 个 kernel。BladeDISC 会采用更为激进的 fusion 策略,从而进一步提高 workload 性能:

  • 每个 fusion block 表示为独立的 ww 结构,使用 shared-memory 进行粘连,从而将 kernel 数量由 3 减少到 1
  • 使用 AStitch 技术,将不同的 loop 结构黏贴在一起,通过 index 推导生成一个 loop 结构,同时引入了 index_cache, value_cache 消除冗余的 index 计算。

在上面 workload 中,BladeDISC 的 fusion 策略可以将 kernel 数量从 7 减少到 1,并且在 kernel 内部使用 index 推导和 cache 来减少冗余的计算,从而逼近硬件的理论峰值。

Inplace Mutation 优化

在 PyTorch Eager 模式下,通过 inplace 算子 (aten.add_) 可以实现对输入的 tensor (w) 进行更新,而不需要一个额外的输出 Tensor。但是在 MLIR 世界里,IR 必须是符合 SSA 形式的,所以没有办法直接表示 inplace 语义,通常的做法是增加一个 D2D memcpy 算子来将输出的 buffer (w') 覆盖输入 buffer (w)。但这样做会造成额外的一次显存拷贝。

BladeDISC 的做法是找到需要 inplace 更新的两个buffer,在 MHLO IR 上进行标记,将 w和 w' 标记为相同的 buffer,在生成 gpu.store指时,将输出直接写回 wbuffer,从而节省一次显存拷贝所造成的额外开销。

03

Benchmark

PAI-Blade 在 A10 和 A100 上最大可获得 41.6 % 和 28.4% 的性能收益(batchsize=1)。

04

在 DSW 上使用 PAI-Blade

  1. 在 PAI 平台中创建 DSW 实例,并使用如下自定义 Docker 镜像,具体步骤可以参考文档

help.aliyun.com/zh/pai/user...

bash 复制代码
pai-blade-registry.cn-hangzhou.cr.aliyuncs.com/pai-blade/aicompiler:latest-stablediffusion-torch-2.0.1-cu118
  1. 创建 Jupyter Notebook,启动 fine-tuning 任务
bash 复制代码
!cd /opt/StableDiffusion && bash launch_dreambooth_train.sh

在看到如下日志时,表示微调任务执行完成:

  1. 启动推理任务,并在 Jupyter Notebook 中查看生成的图片
bash 复制代码
!cd /opt/StableDiffusion && python inference.py && cp dog-bucket.png /mnt/workspace

参考文档:

  • BladeDISC:

github.com/alibaba/Bla...

  • TorchDynamo:

pytorch.org/docs/2.1/to...

  • Torch-MLIR Project:

github.com/llvm/torch-...

  • 文档:

help.aliyun.com/zh/pai/user...

相关推荐
昨日之日200632 分钟前
Moonshine - 新型开源ASR(语音识别)模型,体积小,速度快,比OpenAI Whisper快五倍 本地一键整合包下载
人工智能·whisper·语音识别
浮生如梦_34 分钟前
Halcon基于laws纹理特征的SVM分类
图像处理·人工智能·算法·支持向量机·计算机视觉·分类·视觉检测
深度学习lover35 分钟前
<项目代码>YOLOv8 苹果腐烂识别<目标检测>
人工智能·python·yolo·目标检测·计算机视觉·苹果腐烂识别
热爱跑步的恒川2 小时前
【论文复现】基于图卷积网络的轻量化推荐模型
网络·人工智能·开源·aigc·ai编程
阡之尘埃4 小时前
Python数据分析案例61——信贷风控评分卡模型(A卡)(scorecardpy 全面解析)
人工智能·python·机器学习·数据分析·智能风控·信贷风控
孙同学要努力6 小时前
全连接神经网络案例——手写数字识别
人工智能·深度学习·神经网络
Eric.Lee20216 小时前
yolo v5 开源项目
人工智能·yolo·目标检测·计算机视觉
其实吧37 小时前
基于Matlab的图像融合研究设计
人工智能·计算机视觉·matlab
丕羽7 小时前
【Pytorch】基本语法
人工智能·pytorch·python
ctrey_7 小时前
2024-11-1 学习人工智能的Day20 openCV(2)
人工智能·opencv·学习