使用 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...

相关推荐
肥猪猪爸11 分钟前
使用卡尔曼滤波器估计pybullet中的机器人位置
数据结构·人工智能·python·算法·机器人·卡尔曼滤波·pybullet
LZXCyrus39 分钟前
【杂记】vLLM如何指定GPU单卡/多卡离线推理
人工智能·经验分享·python·深度学习·语言模型·llm·vllm
我感觉。1 小时前
【机器学习chp4】特征工程
人工智能·机器学习·主成分分析·特征工程
YRr YRr1 小时前
深度学习神经网络中的优化器的使用
人工智能·深度学习·神经网络
DieYoung_Alive1 小时前
一篇文章了解机器学习(下)
人工智能·机器学习
夏沫的梦1 小时前
生成式AI对产业的影响与冲击
人工智能·aigc
goomind1 小时前
YOLOv8实战木材缺陷识别
人工智能·yolo·目标检测·缺陷检测·pyqt5·木材缺陷识别
只怕自己不够好1 小时前
《OpenCV 图像基础操作全解析:从读取到像素处理与 ROI 应用》
人工智能·opencv·计算机视觉
幻风_huanfeng1 小时前
人工智能之数学基础:线性代数在人工智能中的地位
人工智能·深度学习·神经网络·线性代数·机器学习·自然语言处理
嵌入式大圣1 小时前
嵌入式系统与OpenCV
人工智能·opencv·计算机视觉