近年来，随着深度学习技术的迅猛发展，越来越多的模型展现出动态特性，这引发了对动态形状深度学习编译器(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
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++简介
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 技术重新生成回来，通常包含两种技术手段：
- Offloading 技术在显存即将超过阈值时将 GPU Tensor 从 GPU 内存转移到 CPU 内存，并在用到它之前转移回来
- 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
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和 %1085：expr1 = 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
Figure4: Operation Scheduling
在 AI 编译器中，通过调整算子的调度顺序，可以有效降低显存峰值。例如，在图4中，黄色曲线表示内存使用量为1，蓝色曲线表示为2。考虑以下两种执行顺序：
- A→C→E→B→D→F
- 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,996 和 expr2 = S0 × 4,096。结合计算图中推导出的关系 S0 = 12 × S1，可以较容易地判断出调度 E 的峰值内存更低。
当然，这仅是一种尽力而为的策略。在实际测试中，并非所有的内存大小都可以化简或直接比较。在大型语言模型（例如 LLama2）中，尽管少部分包含二次项的表达式无法进行比较，但大多数情况下，我们可以将表达式规约到相同的 SymbolicDim，从而实现有效的调度优化。
通过合理的算子调度，BladeDISC++ 能够在动态形状环境下有效降低显存峰值，提升深度学习模型的训练效率。尽管存在一些限制和挑战，但通过构建和简化符号形状计算图，BladeDISC++ 在大多数情况下能够实现显著的优化效果。
Just-In-Time Auto Rematerialization(JiT Remat)
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 操作。
Figure 6: Jit Remat Optimize Memory Peak
正如 Figure6 中所示，蓝色曲线表示优化前的显存使用情况，红色曲线则代表设定的内存限制。对于每一轮迭代，输入 Tensor 的 SymolicDim都会对应一个确定的值：例如S0=1024,S1=4，结合 Symbolic Shape Graph，BladeDISC++ 计算图的显存用量进行一次预估，当显存用量高于设定阈值时，系统将回溯检查之前的 Tensor，判断哪些 Tensor 需要 Remat，并依据端到端损失最小的原则，选择 offloading 还是 recompute 策略。优化后的黄色曲线尽可能地保持显存用量在设定的限制以下。然而值得注意的是这种优化算法并不能在所有情况下都确保内存使用低于限制，仅在训练集配置下，用尽量少的 GPU 资源处理更大规模的训练数据。
三、性能评估
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 issue或 TorchAcc iss或者扫描下方二维码加入钉钉群一起交流 AI Infra 相关问题：
