使用 PAI-Blade 加速 StableDiffusion Fine-Tuning

2023-12-15 18:45

本文主要是介绍使用 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


 

这篇关于使用 PAI-Blade 加速 StableDiffusion Fine-Tuning的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



http://www.chinasem.cn/article/497469

相关文章

C++使用栈实现括号匹配的代码详解

《C++使用栈实现括号匹配的代码详解》在编程中,括号匹配是一个常见问题,尤其是在处理数学表达式、编译器解析等任务时,栈是一种非常适合处理此类问题的数据结构,能够精确地管理括号的匹配问题,本文将通过C+... 目录引言问题描述代码讲解代码解析栈的状态表示测试总结引言在编程中,括号匹配是一个常见问题,尤其是在

Java中String字符串使用避坑指南

《Java中String字符串使用避坑指南》Java中的String字符串是我们日常编程中用得最多的类之一,看似简单的String使用,却隐藏着不少“坑”,如果不注意,可能会导致性能问题、意外的错误容... 目录8个避坑点如下:1. 字符串的不可变性:每次修改都创建新对象2. 使用 == 比较字符串,陷阱满

Python使用国内镜像加速pip安装的方法讲解

《Python使用国内镜像加速pip安装的方法讲解》在Python开发中,pip是一个非常重要的工具,用于安装和管理Python的第三方库,然而,在国内使用pip安装依赖时,往往会因为网络问题而导致速... 目录一、pip 工具简介1. 什么是 pip?2. 什么是 -i 参数?二、国内镜像源的选择三、如何

使用C++实现链表元素的反转

《使用C++实现链表元素的反转》反转链表是链表操作中一个经典的问题,也是面试中常见的考题,本文将从思路到实现一步步地讲解如何实现链表的反转,帮助初学者理解这一操作,我们将使用C++代码演示具体实现,同... 目录问题定义思路分析代码实现带头节点的链表代码讲解其他实现方式时间和空间复杂度分析总结问题定义给定

Linux使用nload监控网络流量的方法

《Linux使用nload监控网络流量的方法》Linux中的nload命令是一个用于实时监控网络流量的工具,它提供了传入和传出流量的可视化表示,帮助用户一目了然地了解网络活动,本文给大家介绍了Linu... 目录简介安装示例用法基础用法指定网络接口限制显示特定流量类型指定刷新率设置流量速率的显示单位监控多个

JavaScript中的reduce方法执行过程、使用场景及进阶用法

《JavaScript中的reduce方法执行过程、使用场景及进阶用法》:本文主要介绍JavaScript中的reduce方法执行过程、使用场景及进阶用法的相关资料,reduce是JavaScri... 目录1. 什么是reduce2. reduce语法2.1 语法2.2 参数说明3. reduce执行过程

如何使用Java实现请求deepseek

《如何使用Java实现请求deepseek》这篇文章主要为大家详细介绍了如何使用Java实现请求deepseek功能,文中的示例代码讲解详细,感兴趣的小伙伴可以跟随小编一起学习一下... 目录1.deepseek的api创建2.Java实现请求deepseek2.1 pom文件2.2 json转化文件2.2

python使用fastapi实现多语言国际化的操作指南

《python使用fastapi实现多语言国际化的操作指南》本文介绍了使用Python和FastAPI实现多语言国际化的操作指南,包括多语言架构技术栈、翻译管理、前端本地化、语言切换机制以及常见陷阱和... 目录多语言国际化实现指南项目多语言架构技术栈目录结构翻译工作流1. 翻译数据存储2. 翻译生成脚本

C++ Primer 多维数组的使用

《C++Primer多维数组的使用》本文主要介绍了多维数组在C++语言中的定义、初始化、下标引用以及使用范围for语句处理多维数组的方法,具有一定的参考价值,感兴趣的可以了解一下... 目录多维数组多维数组的初始化多维数组的下标引用使用范围for语句处理多维数组指针和多维数组多维数组严格来说,C++语言没

在 Spring Boot 中使用 @Autowired和 @Bean注解的示例详解

《在SpringBoot中使用@Autowired和@Bean注解的示例详解》本文通过一个示例演示了如何在SpringBoot中使用@Autowired和@Bean注解进行依赖注入和Bean... 目录在 Spring Boot 中使用 @Autowired 和 @Bean 注解示例背景1. 定义 Stud