使用 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: https://github.com/alibaba/BladeDISC.

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 优化计算图的执行效率。

https://pytorch.org/docs/2.1/torch.compiler_deepdive.html

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

https://github.com/llvm/torch-mlir

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 镜像，具体步骤可以参考文档

https://help.aliyun.com/zh/pai/user-guide/overview-5

pai-blade-registry.cn-hangzhou.cr.aliyuncs.com/pai-blade/aicompiler:latest-stablediffusion-torch-2.0.1-cu118

1. 创建 Jupyter Notebook，启动 fine-tuning 任务

!cd /opt/StableDiffusion && bash launch_dreambooth_train.sh

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

1. 启动推理任务，并在 Jupyter Notebook 中查看生成的图片

!cd /opt/StableDiffusion && python inference.py && cp dog-bucket.png /mnt/workspace

参考文档：

• BladeDISC:

https://github.com/alibaba/BladeDISC

• TorchDynamo:

https://pytorch.org/docs/2.1/torch.compiler_deepdive.html

• Torch-MLIR Project:

https://github.com/llvm/torch-mlir

• 文档：

https://help.aliyun.com/zh/pai/user-guide/overview-5

?