BladeDISC++:Dynamic Shape AI 编译器下的显存优化技术

本文涉及的产品
交互式建模 PAI-DSW,每月250计算时 3个月
模型在线服务 PAI-EAS,A10/V100等 500元 1个月
模型训练 PAI-DLC,100CU*H 3个月
简介: 本文介绍了阿里云 PAI 团队近期发布的 BladeDISC++项目,探讨在动态场景下如何优化深度学习训练任务的显存峰值,主要内容包括以下三个部分:Dynamic Shape 场景下显存优化的背景与挑战;BladeDISC++的创新解决方案;Llama2 模型的实验数据分析

近年来,随着深度学习技术的迅猛发展,越来越多的模型展现出动态特性,这引发了对动态形状深度学习编译器(Dynamic Shape AI Compiler)的广泛关注。本文将介绍阿里云 PAI 团队近期发布的 BladeDISC++项目,探讨在动态场景下如何优化深度学习训练任务的显存峰值,主要内容包括以下三个部分:

  • Dynamic Shape 场景下显存优化的背景与挑战
  • BladeDISC++的创新解决方案
  • Llama2 模型的实验数据分析

本文内容来自 NeurIPS WorkShop 2024 论文:BladeDISC++: Memory Optimizations Based On Symbolic Shape



一、背景与挑战

动态形状深度学习编译器的挑战

随着模型架构的不断演进,其动态性日益增强。例如,传统的计算机视觉(CV)模型中,图像尺寸和批量大小(batch size)在训练过程中会不断变化;大型语言模型的序列长度和批量大小也呈动态调整状态;多模态模型中的图像、视频长度及序列长度同样变化不定。此外,一些更为复杂的混合专家(MoE)模型还涉及与数据相关的动态形状,这些都体现了模型的动态特性。

相应地,Dynamic Shape AI 编译器的发展也在加速。目前一些最先进的编译器,如TVM,OpenXLA主要侧重于静态形状下的优化,包括高效代码生成和图层优化等,类似 Torch Inductor 已经在 Dynamic Shape 场景做了一些工作,但在内存优化方面仍存在不足。

其主要挑战有:

1. 张量(Tensor)形状的不确定性

Tensor 形状的动态变化给代码生成带来了麻烦,由于无法在编译期提前买预知 Tensor 的大小和布局,编译器必须采取一些措施使得不同形状下都可以生成高效的代码。

2. 显存分配的动态性

在静态形状下,可以通过整体图计算提前确定显存分配,预先分配适当大小的显存块,并在使用缓存分配器时尽量减少显存碎片。然而,在动态形状下,无法提前预分配显存块的大小,这将导致显存碎片的增加。

3. 优化算法的复杂度增加

常用的图优化算法,比如算子调度,算子融合等需要根据具体的 Tensor 形状来预估内存用量从而对调度、计算策略进行调整。例如,判断两个相似算子是否能融合通常依赖相同的形状信息或相同的代码生成策略,若缺乏形状信息,难以判断融合的可能性,导致性能下降。针对这些挑战,PAI 团队提出了 BladeDISC++方案,旨在动态形状的工作负载下,尽量将显存峰值控制在某一个阈值之下,使深度学习训练任务可以用更少的显存资源处理更多的训练数据。总结:随着深度学习模型动态性的增强,动态形状 AI 编译器的重要性日益凸显。然而,这也带来了形状不确定性、优化算法复杂度增加以及显存分配动态性等多重挑战。BladeDISC++通过专注于优化显存峰值,力求在动态形状环境中提供高效的性能表现。

BladeDISC with TorchAcc

image.png

Figure1: BladeDISC with TorchAcc

BladeDISC[1] 这一名称来源于Dynamic Shape Compiler,它是一个基于MLIR构建的端到端编译器。由于 BladeDISC 主要作为编译器的后端,前端则支持多种Tracing系统。TorchAcc[2]作为 PyTorch 加速器对大规模深度学习模型训练推理提供了完整的优化方案。在本次实验中,BladeDISC++ 作为 TorchAcc 编译器的后端,对 LLM 训练效率优化进行了完整的验证。

整体架构来看,TorchAcc 的前端支持更多类型的模型,通过 GraphCapture 技术(包括 PyTorch Lazy Tensor Core 和 Dynamo 等 PyTorch Tracing 系统)捕捉 PyTorch 的计算图,并将其统一转换为 StableHLO 的中间表示层。BladeDISC 接收一张 StableHLO 计算图通过一系列优化算法最终生成不同硬件上运行的可执行文件,BladeDISC 的优化主要涵盖以下几个方面:

  • 显存优化:基于对 Buffer 生存周期的分析,实现自动的 offloading,recomputation 等显存相关的优化。
  • 计算优化:例如算子的融合、计算图化简 以及 Custom Operator 的支持等优化计算效率。
  • 通信优化:例如 Multiple CUDA Stream 的异步通信等。

BladeDISC 还支持多种主流芯片架构,目前已涵盖Nvidia,AMD以及 Intel CPU 等主流硬件平台。


二、BladeDISC++的创新方案

正如前文中提到的面向 Dynamic Shape 计算图下的诸多挑战,PAI 团队创新性地提出了 BladeDISC++ 解决方案。

BladeDISC++简介

image.png

Figure2: BladeDISC++ Overview

在 Dynamic Shape 场景下,最核心的挑战在于缺乏具体的形状信息,这导致无法静态地完成计算图的优化算法。BladeDISC++ 采取了结合编译时和运行时进行联合优化的技术方案。正如 Figure2 中展示的一样,BladeDISC接收到的是 TorchAcc GraphCapture 记录下来的动态形状计算图,这个计算图包含两个特点:一是输入张量的形状是未知的;二是计算图的拓扑结构是固定的。随后 BladeDISC++ 开始对 Dynamic Shape 计算图进行优化,主要包含以下核心阶段:

  • Operation Fusion

将访存算子或GEMM算子合并在一起,可以减少 Kernel 发射带来的额外开销并增加共享内存的利用率,提升系统吞吐。

  • Operation Scheduling

调整算子的执行顺序,通过优化 Tensor 生存周期降低显存峰值。

  • Automatic Rematerialization(Auto Remat)

Auto Remat 技术将临时不用的 Tensor 从 GPU HBM 中释放掉,在使用它时通过 Regenerate 技术重新生成回来,通常包含两种技术手段:

  1. Offloading 技术在显存即将超过阈值时将 GPU Tensor 从 GPU 内存转移到 CPU 内存,并在用到它之前转移回来
  2. Recomputation 技术将暂时不用的 GPU Tensor 释放掉,在用到它之前重新计算回来,是一种以时间换空间的策略。

当然这些阶段是当前 AI 编译器中较为典型的优化步骤。然而,由于缺乏具体的张量形状信息,使得现有的方法无法生效。BladeDISC++ 首先构建一个基于 Symbolic Shape 的计算图,用以表示未知的形状信息以及其关联关系。基于 Symbolic Shape Graph,BladeDISC++ 可以尽量判断张量之间所占内存的大小关系,从而调整算子执行顺序(Operation Scheduling)来降低显存峰值。对于 Auto Remat,由于缺少了具体的形状信息,编译器无法在编译期静态的决定哪些张量需要被 Offloading 或 Recomputation ,BladeDISC++ 采取了编译期和运行时联合优化的方法,在编译期选取可能被暂时释放的 Tensor,在运行期,根据实际的显存用量,动态的决定 Regenreate 的策略:offloading 或者 recomputation 。

综上所述,BladeDISC++ 通过结合编译期和运行时的联合优化方法,成功应对了动态形状下的诸多挑战,特别是在优化显存峰值方面,实现了高效的深度学习训练任务优化。

Symbolic Shape Graph

image.png

Figure3: Symbolic Shape Graph

Figure3 中的 MLIR Intermediate Representation (IR) 展示了 Dynamic Shape Graph 及其对应的 Symbolic Shape。Graph 输入 %arg0 的形状是未知的,使用 “?” 表示,为了便于全局分析,BladeDISC++ 引入了 SymbolicDimOp 通过 Tensor Attribution,将 Symbolic Shape 与未知 Shape 进行绑定,例如: %arg0: tensor

为了评估 Tensor 之间占用显存的大小关系 BladeDISC++ 基于 Symbolic Shape 构建了 Symbolic Expression (SymbolicExpr) 来描述 Tensor 的元素个数,例如:%1084%1085expr1 = 11008 * @S1; expr2 = 1024 * @S0。为了比较 SymbolicExpr 之间的大小关系,BladeDISC++ 通过对算子的语义以及输入输出之间的约束关系进行分析,进而化简表达式。例如对于 ReshapeOp,其算子语义是将一个输入 Tensor 变换形状并输出一个新的 Tensor。同时 ReshapeOp 隐含了一个约束,即输入和输出 Tensor 的元素数量是一致的。基于这些信息,可以推导出形状之间的关系,例如 Figure3 中 @S0 = 12 * @S1, BladeDISC++ 可以将 expr2 化简为 expr2=132096 * @S0 ,从而推断出 expr1 小于 expr2。

Operation Scheduling

image.png

Figure4: Operation Scheduling

在 AI 编译器中,通过调整算子的调度顺序,可以有效降低显存峰值。例如,在图4中,黄色曲线表示内存使用量为1,蓝色曲线表示为2。考虑以下两种执行顺序:

  1. A→C→E→B→D→F
  2. A→C→B→D→E→F

当执行到 A→C 后,系统有两个选择:调度 B 或 E。如果选择 C→E→B 的顺序,此时活跃张量(Living Tensor)的总大小为5;而若选择 C→B→D,则活跃张量的总大小为6。由此可见,不同的调度序列会显著影响显存的峰值。

然而,在 Dynamic Shape 环境下,由于无法预先确定具体张量的大小,难以准确预估不同调度序列对活跃张量总大小的影响。为了解决这一问题,BladeDISC++ 利用前面介绍的符号表达式(SymbolicExpr)对 Tensor 占用的显存大小进行比较。例如,在上述例子中,选择调度 B 或 E 分别会导致不同的内存需求:expr1=S1 × 10,996expr2 = S0 × 4,096。结合计算图中推导出的关系 S0 = 12 × S1,可以较容易地判断出调度 E 的峰值内存更低。

当然,这仅是一种尽力而为的策略。在实际测试中,并非所有的内存大小都可以化简或直接比较。在大型语言模型(例如 LLama2)中,尽管少部分包含二次项的表达式无法进行比较,但大多数情况下,我们可以将表达式规约到相同的 SymbolicDim,从而实现有效的调度优化。

通过合理的算子调度,BladeDISC++ 能够在动态形状环境下有效降低显存峰值,提升深度学习模型的训练效率。尽管存在一些限制和挑战,但通过构建和简化符号形状计算图,BladeDISC++ 在大多数情况下能够实现显著的优化效果。

Just-In-Time Auto Rematerialization(JiT Remat)

image.png

Figure5: Jit Remat

静态形状计算图上的 Auto Remat 技术通过预估显存峰值,在计算图上找到哪些 Living Tensor 会导致显存超过某一阈值,进而搜索出哪些 Tensor 需要被释放以及 Regenerate 策略:比如 offloading 或 recomputation ,需要注意的是无论 offloading 还是 recomputation 会引入额外显存带宽占用或增加计算量,这会对端到端性能造成一定损失,所以只有在判断显存峰值可能超过阈值时才会执行 Rematerialization 操作。但是在 Dynamic Shape 环境下,由于缺少具体的 Shape 信息,无法对显存占用进行预估,BladeDISC++ 基于 Symbolic Shape 计算图采取了编译期-运行期联合优化的思路解决了这个问题。

  • 在编译期,BladeDISC++ 搜索出所有可能被释放的 Living Tensor,以及对应的 Regenerate 子图。通过在计算图上插入 EvictOp 和 RegenerateOp,表示可能对这个 Tensor 进行 Remat 操作的语义:offloading,recompute 或者什么都不做。
  • 在运行时,BladeDISC++ 根据实际 Tensor Shape 结合 Symbolic Shape Graph 对显存峰值进行预估,从而判断需要对哪些 Living Tensor 进行 Remat 操作。

image.png

Figure 6: Jit Remat Optimize Memory Peak

正如 Figure6 中所示,蓝色曲线表示优化前的显存使用情况,红色曲线则代表设定的内存限制。对于每一轮迭代,输入 Tensor 的 SymolicDim都会对应一个确定的值:例如S0=1024,S1=4,结合 Symbolic Shape Graph,BladeDISC++ 计算图的显存用量进行一次预估,当显存用量高于设定阈值时,系统将回溯检查之前的 Tensor,判断哪些 Tensor 需要 Remat,并依据端到端损失最小的原则,选择 offloading 还是 recompute 策略。优化后的黄色曲线尽可能地保持显存用量在设定的限制以下。然而值得注意的是这种优化算法并不能在所有情况下都确保内存使用低于限制,仅在训练集配置下,用尽量少的 GPU 资源处理更大规模的训练数据。


三、性能评估

image.png

Figure7: Evaluation on Llama2

为了方便在单卡上进行验证,我们在 Llama2-7B 上进行了一些裁剪(hidden_layers=4)使其变为1B 模型。如 Figure7 所示,我们做了如下实验:

  • Dynamic Shape BladeDISC 在 Dynamic Shape 环境没有显存优化
  • Static Shape BladeDISC 在 Static Shape 环境显存优化的实验
  • BladeDISC++ 在 Dynamic Shape 环境开启了显存优化(本篇工作)

注:为了避免 Static Shape 下不同的输入 Tensor Shape 都编译一张计算图,所以实验中对数据采取了分桶 padding 的策略避免编译次数过多带来的overhead[3].

实验数据表明,在 BatchSize=14 时,BladeDISC++ 吞吐基本和 Dynamic Shape 持平,显存峰值和 Static Shape 情况持平。当 BatchSize 增大到 18 时,Dynamic Shape 由于没有显存优化策略会导致 Out of Memory (OOM), Static Shape 开启了显存优化能够将训练任务跑起来,但由于 padding 带来的冗余计算,导致吞吐只有 5103 tokens/sec,而 BladeDISC++ 显存峰值和 Static Shape 下基本一致,但吞吐提升了 11%。


四、总结

Dynamic Shape 环境下核心的挑战是 Tensor Shape 的不确定性,BladeDISC++ 构建了 Symbolic Shape Graph 来描述 Dynamic Shape 之间的关系,在此基础上提出了基于 Symbolic Shape 的 Operation Scheduling 以及编译-运行时联合优化的 Jit Auto Rematerialization 方法。在 Llama2 模型上的实验数据表明,BladeDISC++ 能够有效降低训练时的显存占用,优化效果可以与 Static Shape 环境相媲美,同时由于避免了冗余计算,吞吐优于 Static Shape。以上就是 PAI 团队在近期在 Dynamic Shape 环境下的显存优化方向的一些探索,本文分享了我们在优化动态形状计算图显存方面的行业经验。我们提出了基于符号形状的算子调度和重材化方法,并开发了 BladeDISC++。评估结果表明,BladeDISC++ 能够有效降低动态形状训练的显存使用,并且在显存优化效果上可与静态形状训练相媲美。据我们所知,这是该领域的开创性工作,我们希望它能够支持编译器社区管理动态形状工作负载,促进动态形状编译器的更广泛应用。

加入我们

PAI 核心引擎团队团队社招、实习火热招聘中,包括三个方向:

  • AI 编译器
  • 基于 MLIR 的自研编译器 BladeDISC,极致优化 GPU 计算效率。
  • OpenAI Triton,训练、推理场景中的算子调优以及运行时优化。
  • 大模型训练优化
  • 包括 LLM,VLM、 MoE 以及视频图像生成等主流模型大规模训练任务的优化,包括算子优化、分布式并行策略优化、框架优化等。
  • 大模型推理优化
  • LLM、VLM 以及MoE 等主流模型的推理系统的端到端优化,包括且不限于算子调优、量化压缩、调度优化以及框架优化等。

如果您对以上任何一个方向感兴趣,欢迎投递简历到 PAI 核心引擎团队招聘贴 或直接发邮件到 yancey.yx@alibaba-inc.com 。也欢迎通过 Github 上 BladeDISC issueTorchAcc iss或者扫描下方二维码加入钉钉群一起交流 AI Infra 相关问题:

image.png image.png

引用

相关实践学习
使用PAI-EAS一键部署ChatGLM及LangChain应用
本场景中主要介绍如何使用模型在线服务(PAI-EAS)部署ChatGLM的AI-Web应用以及启动WebUI进行模型推理,并通过LangChain集成自己的业务数据。
机器学习概览及常见算法
机器学习(Machine Learning, ML)是人工智能的核心,专门研究计算机怎样模拟或实现人类的学习行为,以获取新的知识或技能,重新组织已有的知识结构使之不断改善自身的性能,它是使计算机具有智能的根本途径,其应用遍及人工智能的各个领域。 本课程将带你入门机器学习,掌握机器学习的概念和常用的算法。
相关文章
|
3天前
|
人工智能 前端开发 Serverless
阿里云《AI 剧本生成与动画创作》解决方案技术评测
随着人工智能技术的发展,越来越多的工具和服务被应用于内容创作领域。阿里云推出的《AI 剧本生成与动画创作》解决方案,利用函数计算 FC 构建 Web 服务,结合百炼模型服务和 ComfyUI 工具,实现了从故事剧本撰写、插图设计、声音合成和字幕添加到视频合成的一站式自动化流程。本文将对该方案进行全面的技术评测,包括实现原理及架构介绍、部署文档指引、具体耗时分析以及实际使用体验。
48 16
|
13天前
|
机器学习/深度学习 人工智能 API
Aligner:自动修正AI的生成结果,北大推出残差修正模型对齐技术
介绍北大团队提出的 Aligner 模型对齐技术,通过学习对齐答案与未对齐答案之间的修正残差,提升大语言模型的性能。
74 28
|
1天前
|
存储 人工智能 Serverless
阿里云《AI 剧本生成与动画创作》技术解决方案测评
本问是对《AI 剧本生成与动画创作》的用心体验。结论不是特别理想,在实际使用中仍存在一些问题。
50 22
|
27天前
|
人工智能 达摩院 计算机视觉
SHMT:体验 AI 虚拟化妆!阿里巴巴达摩院推出自监督化妆转移技术
SHMT 是阿里达摩院与武汉理工等机构联合研发的自监督化妆转移技术,支持高效妆容迁移与动态对齐,适用于图像处理、虚拟试妆等多个领域。
67 9
SHMT:体验 AI 虚拟化妆!阿里巴巴达摩院推出自监督化妆转移技术
|
11天前
|
机器学习/深度学习 人工智能 自然语言处理
DeepSeek逆天,核心是 知识蒸馏(Knowledge Distillation, KD),一项 AI 领域的关键技术
尼恩架构团队推出《LLM大模型学习圣经》系列,涵盖从Python开发环境搭建到精通Transformer、LangChain、RAG架构等核心技术,帮助读者掌握大模型应用开发。该系列由资深架构师尼恩指导,曾助力多位学员获得一线互联网企业的高薪offer,如网易的年薪80W大模型架构师职位。配套视频将于2025年5月前发布,助你成为多栖超级架构师。此外,尼恩还提供了NIO、Docker、K8S等多个技术领域的学习圣经PDF,欢迎领取完整版资源。
|
5天前
|
人工智能 负载均衡 搜索推荐
谷歌发布双思维AI Agent:像人类一样思考,重大技术突破!
谷歌近日推出基于“快慢思维”理论的双思维AI Agent系统,模仿人类大脑的两种思维模式:快速直观的Talker(系统1)和深思熟虑的Reasoner(系统2)。Talker负责日常对话与快速响应,Reasoner则处理复杂推理任务。该系统模块化设计,灵活高效,已在睡眠教练等场景中展现应用潜力,但仍面临工作负载平衡与推理准确性等挑战。论文详情见:https://arxiv.org/abs/2410.08328v1
31 1
|
24天前
|
存储 人工智能 安全
AI时代的网络安全:传统技术的落寞与新机遇
在AI时代,网络安全正经历深刻变革。传统技术如多因素身份认证、防火墙和基于密码的系统逐渐失效,难以应对新型攻击。然而,AI带来了新机遇:智能化威胁检测、优化安全流程、生物特征加密及漏洞管理等。AI赋能的安全解决方案大幅提升防护能力,但也面临数据隐私和技能短缺等挑战。企业需制定清晰AI政策,强化人机协作,推动行业持续发展。
57 16
|
1月前
|
人工智能 缓存 Ubuntu
AI+树莓派=阿里P8技术专家。模拟面试、学技术真的太香了 | 手把手教学
本课程由阿里P8技术专家分享,介绍如何使用树莓派和阿里云服务构建AI面试助手。通过模拟面试场景,讲解了Java中`==`与`equals`的区别,并演示了从硬件搭建、语音识别、AI Agent配置到代码实现的完整流程。项目利用树莓派作为核心,结合阿里云的实时语音识别、AI Agent和文字转语音服务,实现了一个能够回答面试问题的智能玩偶。课程展示了AI应用的简易构建过程,适合初学者学习和实践。
97 22
|
30天前
|
人工智能 Java 程序员
通义灵码AI编码助手和AI程序员背后的技术
通义灵码AI编码助手和AI程序员背后的技术,由通义实验室科学家黎槟华分享。内容涵盖三部分:1. 编码助手技术,包括构建优秀AI编码助手及代码生成补全;2. 相关的AI程序员技术,探讨AI程序员的优势、发展情况、评估方法及核心难点;3. 代码智能方向的展望,分析AI在软件开发中的角色转变,从辅助编程到成为开发主力,未来将由AI执行细节任务,开发者负责决策和审核,大幅提升开发效率。
146 12
|
1月前
|
人工智能 搜索推荐
AI视频技术的发展是否会影响原创内容的价值
AI视频技术的发展显著降低了视频制作的门槛与成本,自动完成剪辑、特效添加等繁琐工作,大大缩短创作时间。它提供个性化创意建议,帮助创作者突破传统思维,拓展创意边界。此外,AI技术使更多非专业人士也能参与视频创作,注入新活力与多样性,丰富了原创内容。总体而言,AI视频技术不仅提升了创作效率,还促进了视频内容的创新与多样化。