深度学习框架与动态shape

Posted 2021-10-13 archimekai

从PPT中复制出来后格式乱了，PPT：  https://download.csdn.net/download/archimekai/24838917

动起来更快？
从 Nimble 和 DISC 看深度 学习 框架 与动态 shape

• 2021 年 9 月 25 日 [e-mail]

• 什么是动态 shape/ 动态模型

• 静态模型：每个算子的输入输出 shape 均在图执行前已知。

• 编译期就知道形状，并且执行期保持不变

• 张量内存预先分配

• 动态模型：包括控制流、动态数据结构（例如 tree-structured long short-term memory ）、动态形状的模型

• 张量形状、内存不能预先知道

•

• 使用动态图的主要领域

• 自然语言处理

• 动态 shape 面临的挑战

• 缺少能够表示动态结构的 IR

• 缺少能够为动态 shape 生成高效算子的生成器

• 普通的算子生成器仅支持静态 shape 的算子生成

• Shape 不确定时，难以有效优化算子性能（例如，动态 shape 时可能无法消除边界检查）

• 缺少能够处理动态结构的运行时

• 该运行时最好内存占用低，开销小

• 缺少在动态模型场景下的优化技术

• 当前的优化和内存分配技术大都假定 shape 固定

•

• Nimble 的主要贡献

• 提出和构建了一个端到端的动态模型推理系统

• 为动态模型设计了若干编译和优化技术，包括动态类型系统、内存规划、设备选择机制（ device placement ）、算子生成器、基于形状的算子选择算法

• 用于执行动态模型的平台无关的 VM

•

• Nimble 实现 方法

• Nimble 动态 shape 的表示、传播和消减

• IR 中增加一个符号来表示未知的形状： Any 。例如 Tensor[(1, 10, Any), float32]

• 类型和形状推导需要正确地传播 Any 。动态输出 shape 的算子，其相关的输出形状需要使用 Any 来描述（例如 unique ）。难以在编译时进行的形状检查，需要留到运行时进行。

• 对多余的动态 shape 需要进行消除。例如 Add 的两个输入事实上形状应该是相同的，如果一个输入形状已知，就可以确定另一个输入的形状。

• 插入 shape function 在运行时动态计算形状。（ data independent shape function (conv2d), data dependent shape function (unique) ）

• Shape function 在 HOST 上执行

• Nimble Dynamic memory planning

• 内存管理的目标：

• 减少内存碎片

• 降低内存分配对运行时性能的影响

• Nimble 的方法：

• 将 TVM 的隐式内存申请改为显式

• 在 IR 中明确体现出内存的管理，从而可以在 IR 层面跟踪和优化动态内存分配计划

• invoke_mut (op, inputs, outputs)

• alloc_storage (size, alignment, device)

• alloc_tensor (storage, offset, shape, dtype , attrs )

• kill(tensor) 释放内存

• 优化手段例如：

• Storage coalescing （多个 tensor 共享一块内存区域）

• Liveness analysis

• Graph coloring

• etc

• Nimble 设备 选择机制

• 合理分配每个节点的计算设备，从而最小化跨设备拷贝

• 使用 IR 元素 device_copy 表示在设备之间的拷贝

•

• 一些原则：

• Shape function 的输入和输出都要在 CPU 上

• Invoke_mut 的所有输入必须在相同的设备上

• 当 CPU 和 GPU 都可以时，选择速度最快的设备

• Nimble symbolic codegen

• 自动算子代码生成是提升算子性能的有效手段

• 当前的算子代码生成依赖于静态 shape 的前提。

• 动态 shape 算子代码生成的挑战：

• 如何将现有的算子性能调优手段应用到动态 shape 的场景？（ shape 不确定时，算子性能调优的时间可能会指数级增长。）

• 如何在动态 shape 的情况下达到和静态 shape 相同的性能？（ shape 不确定时，可能需要进行很多边界检查）

• 解决办法：

• 使用基于人工模板和人工定义搜索空间的算子性能调优方法来减少搜索时间。先将动态 shape 设置为一个足够大的数，例如 64 ，在这个大的 shape 上进行算子性能调优，然后保留 top k （ 100 ）个最好的结果，将这些结果应用到一批 shape 上，然后选择平均性能最好 的结果。

• 按照 tiling 的约数的余数预先准备好不同余数的算子。在余数确定的情况下，可以消除大多数的边界检查

•

•

• Nimble virtual machine

VM作用：

• 处理控制流逻辑

• 派发算子到对应的执行设备

• 支持多种平台

• 将动态模型编译为平台无关的 VM exec 和平台相关的算子代码

•

• 技术选型：

• 基于寄存器的 vm

• 无限寄存器：方便理解和原型开发，降低优化 pass 的开发难度

• CISC 风格的指令，可变长度（可变入参个数）（原因在于，输入数据的 shape 可变）

• VM 的好处

• 简单的 VM 有助于形式化验证其安全性

• 通过使用不同的 VM ，有助于隔离和调度不同的工作负载（云上场景，多个实例使用相同硬件时）

• 支持多平台

•

• Nimble evaluation

• 动态 shape 的评估要点：

• 相对于当前的 sota ，新方案在在各种硬件上的性能上如何？

• 新方案提出的优化技术效果怎样？

• Nimble 的性能提升来自以下两点

• 通过深度学习编译器获得更好的 kernel fusion 和算子实现

• 控制流带来的额外开销小

•

• Nimble micro benchmark

• Memory planning ：

• 减少 47% 的内存申请

• 降低 75% 的内存申请延迟（参考一些内存管理库提供更快的内存申请速度）

• 相较于静态分析、预分配内存的 TVM ，多使用 8% 的内存

• Symbolic codegen :

• 动态的 symbolic codegen 可以实现同 static codegen 相近的速度

• DISC 的主要贡献

• 通过在编译时生成 runtime flow 来避免 VM 的开销，进行更丰富的 HOST 和 DEVICE 联合优化

• 利用 shape propagation 和形状约束传递，在不知道完整 shape 信息的情况下进行更有效的算子融合

• 相较于 sota 性能提升 1.8 倍

•

• 问题：主要应用在推理场景。缺少针对训练场景的工业级解决方案。

•

• DISC  实现方法： AOT 风格的运行管理

•

支持动态shape的IR

• 对 HLO IR 进行扩展，扩展的重点在于静态 shape 场景中被常量折叠的部分。例如 slice 、 pad 、 broadcast 等，需要将其中的静态入参替换为运行时的动态计算。

计算shape hints以辅助编译优化

• 通过相同形状约束、相同大小约束来辅助编译优化。例如，当两个算子操作的张量大小相同时，我们可以考虑将这两个算子融合在一起。

• 利用 DHLO 中的算子语义，可以得到上述两种约束的具体情况。例如， Transpose 的输入输出张量满足相同大小约束。 Add 算子的输入和输出张量、 Add 算子的两个输入张量满足相同形状约束。

• 编译优化时，优先考虑形状泛化性好的优化手段，例如 loop fusion ， input fusion 等。对于和具体 shape 有关的优化手段，例如 loop vectorized load/store, implicit broadcast 等，则提前生成多种版本的算子，并在执行时基于输入 shape 动态选择。

•

•

动态内存管理

• 执行期动态进行内存的分配和释放操作。

• 为了降低内存分配和释放的开销，使用以下方法：

• 基于 shape hints 进行内存 liveness 分析和优化

• 使用高性能的内存分配器（例如 TensorFlow 、 PyTorch 自带的内存分配器）

• 算子融合

• 动态 shape 场景下，判断哪些算子能够融合在一起是困难的

• DISC evaluation

• 相较于 TensorFlow 、 PyTorch 、 Nimble ， DISC 均有不同程度的性能提升。

• 优化的时间主要来自：访存密集性 kernel 的内存访问时间节省。以及算子融合带来的算子数目减少。

• DISC 可以利用 形状传播和形状约束来进行更有效的融合 。

• 相较于 Nimble 的 VM ，由于能够进行 HOST 和 DEVICE 的联合优化， DISC 的 runtime flow 更有效

• DISC evaluation

• 将网络固定为静态 shape ，分别使用静态 shape 编译优化和 DISC 的动态 shape 编译优化， DISC 的性能可以达到静态 shape 编译优化的 85% 。

• 性能差距的原因是动态 shape 编译优化时 shape 信息不足，难以实施更激进的优化。