随着 AI 应用的快速发展和广泛应用，AI 应用开发和工程部署及模型优化也日益引起了产业界关注。如何更高效地完成 AI 算法从研究领域到生产部署的应用开发流程是一个常见的问题。本次 GTC 2022 S41395 中，阿里云计算平台 PAI 团队分享了通过模型系统优化工具 PAI-Blade 更好地应用 TensorRT，BladeDISC 等相关技术方案，优化云上应用场景。
背景及动机
AI 应用的主要特点是拥有大量、高度密集和并行的张量数值计算。在深度学习 AI 的热潮之下，国内外也涌现了不少支持 AI 数值计算的加速器及新硬件。但是就目前而言， NVIDIA GPU 在云端依然是最主流的 AI 加速硬件；无论从性能、生态、还是普及程度，NVIDIA GPU 都是佼佼者。TensorRT 是 NVIDIA GPU 平台上提供的一款免费模型推理加速工具，它在许多深度学习模型上都取得了领先的加速效果。然而，在许多情况下顺滑地通过TensorRT 完成从研究到生产的部署却成为了用户最头疼的问题。我们把这些困难总结为如下三个挑战:
- 从研究到部署的切换
- ONNX 的缺点
- TensorRT 的局限性
挑战 1: 从研究到部署切换
- 转换链路和流程太长，许多转换和优化过程需要人工参与和介入
- 在广泛的应用场景中，从训练到生产迁移过程中的每个步骤所依赖的开发技能不同
- 算法建模、模型训练、模型转换、模型优化部署
- Python，C/C++, CUDA 等
- 不同训练框架，不同的硬件加速器
- 平台和框架迁移的成本
挑战 2: ONNX 的缺点
ONNX 是一个开源的社区项目，它的目标是成为各个机器学习框架之间的可交换格式。因为 ONNX 格式具有严格的标准，并且在框架之间保持中立，在 AI 应用的发展过程中得到了许多权威机构的支持。目前 ONNX 的社区仍然比较繁荣和活跃。在 NVIDIA 平台上面 ONNX-TensorRT 得到了官方支持和广泛使用。
尽管如此，ONNX 社区不同工具之间的兼容性存在明显的问题。主要原因可能包括:
- 深度学习算法和框架仍然在持续发展，ONNX 标准难以覆盖新生 Operator；在 OpSet 覆盖率上面目前无法做到 100%
- 标准的制定依赖于产业的共识和努力，然而 ONNX 社区上下游之间对 ONNX 标准的支持水平参差不齐；ONNX 标准语义和框架及后端之间的语义存在不对齐的情况
- 主流框架的导出 ONNX 方式主要依赖于 tracing 机制；难以覆盖 control flow 等一些复杂的情况
挑战 3: TensorRT 的局限性
- 对前端框架算子操作(Operator)和数据类型的仅提供有限的支持
- ONNX-TensorRT 支持了 129 个 ONNX 算子，具体信息可以参考 support-matrix
- ONNX 接近 160 个算子，TensorFlow 拥有超过 750 个算子， PyTorch 拥有超过 2000 个算子
- 仅提供对动态 shape 的有限支持
- 动态 batching、文本/语音序列长度、图像/视频分辨率的动态性在推理场景中很常见
- TensorRT 需要通过设置 OptimizationProfile 来确定动态 shape 的范围；这在一些场景中并不现实，而且容易引起错误和鲁棒性问题
- 已构建的 Engine 无法跨 GPU compute capabilities 使用，同时优化无法迁移到非 NVIDIA GPU 的设备
这些挑战在现实中成为了困扰开发用户完成 AI 算法研究到 AI 生产部署的问题，并且用户经常为此付出了更多的时间和成本。基于用户的调研和反馈，PAI-Blade 旨在提供更加易用，更加鲁棒，更加自动的模型优化工具。
PAI-Blade 的设计原则
通过统一的流程在不同框架上自动完成多样化的优化策略
- Performance: 优异的执行性能仍然是我们考虑的最主要问题；没有性能的优化工具，不拥有价值。
- Generality: 广泛的通用性能够适配更多场景，同时提供更多的开发和部署自由。
- 支持 TensorFlow/PyTorch/ONNX
- 支持多样的加速硬件，包括 GPU/CPU/DCU 等等
- 支持不同的加速方案，BladeDISC/TensorRT/TVM/MNN 等等
- Usability and Robustness: 易用性和鲁棒性能够降低用户的学习成本在生产系统中集成的成本。
- 统一的简洁的 API
- 最大程度的自动化
- 轻量的集成负担(仅需要几行代码)
PAI-Blade 和以往其他的工具主要的不同点在于它不仅仅考虑性能，而是更加强调性能、通用性、易用性与鲁棒性的协同设计。
PAI-Blade 的设计考虑
可插拔的设计
PAI-Blade 将优化功能实现为原始框架的扩展形式。在这种设计考虑之下，优化部署开发的用户不需要迁移框架和编程语言，只需要几行（甚至完全没有）代码改动就可以完成自动化的优化部署。
支持不同后端接入的可扩展架构
PAI-Blade 的性能来有赖于高性能的后端加速器的支持。为此，我们设计了可扩展、可复用的统一优化流程。以 PyTorch 为例，下面图例分别展示了接入 TensorRT、量化和 BladeDISC 的方法。
接入 ONNX-TensorRT 的流程
接入量化的流程
接入 BladeDISC 的流程
在这种统一的优化流程中，所有信息的传递，层级与层级的交互更加通畅了，也为我们组合多种优化工具打开了空间。
多层 fallback 支持
到目前为止，深度学习模型优化流程的鲁棒性和易用性并没有得到足够的重视。因为一个算子不支持，而放弃整个模型的优化的事情常有发生。许多 AI 应用开发者不得不因此而止步。
为了避免这种情况发生，降低用户的使用门槛，PAI-Blade 引入了多层的 fallback 机制:
- 操作(Operation) 不支持时 fallback
- 子图转换失败时 fallback
- 运行时异常 fallback
如上图，在优化执行过程中，PAI-Blade 会主动在 Graph IR(Intermediate Representation) 上根据每个操作的输入情况，判定当前操作是否会被后端加速器支持，如果不支持则 fallback 到原端框架的 Runtime 执行。
在此基础之上，PAI-Blade 通过并查集算法合并能够被支持的操作成为若干子图；然后将子图转换到后端加速，如果子图转换失败则会进一步 fallback。最后，转换成功的子图会运行在加速器上。如果在运行时发现异常，或者输入不匹配的情况(数据类型、shape 范围不支持)，则会 fallback 到框架的 Runtime 上运行。
自动化优化流程
由于 PAI-Blade 协同设计了从训练框架前端到硬件加速器后端的转换流程，优化过程中不需要用户介入:
- 不支持的算子自动 fallback，不需要开发 plugin
- 自动圈图转换和子图优化
- 自动选择最佳的优化策略
借助 fallback 保证了模型转换的成功率 ，并且让过程更加自动化，进一步降低了使用门槛。同时优化所需要的信息在同一个系统流程中自动获取和传递，相比导出到其他框架部署，用户不需要干预和介入优化中间流程。
BladeDISC 动态 shape 优化
在推理应用场景中，动态 shape 的使用场景很多样。期待用户通过限制自己应用场景中的动态 shape 并不实际。比如，组 batch 的时间会影响服务的响应延时；NLP 场景中的动态序列长度可能引起 shape 的组合爆炸；CV 检测和识别的场景中同样大小的图片可能得到完全不同的个数的输出结果。
为了更加彻底地解决推理应用场景中的动态 shape 问题，PAI-Blade 原生支持BladeDISC 动态 AI 编译引擎。BladeDISC 是基于 MLIR 构建的动态 AI 编译器，它也从 TensorRT、XLA 和 TVM 等框架中汲取了很多经验。但是相比于 ONNX 和 HLO(XLA)， BladeDISC 选择了 MHLO 作为 Tensor-Level IR，因为 MHLO 在 IR 和编译器层面提供了更好的动态 shape 支持。
BladeDISC 首先从 TensorFlow 和 PyTorch 中对支持的算子进行合并圈图，然后将子图转换成为 MHLO 模块。在此之后，BladeDISC 会执行许多编译 pass 来转换和优化 MLIR 模块，包括 placer、shape inference、control flow lowering、fusion、高性能 codegen 等等。
在这些 pass 中间，值得一提的是 fusion 和 codegen。目前 BladeDISC 的 fusion 策略主要会节省 kernel launch 和访存的开销。除此之外，fusion 同样会高性能代码生成创建了更多的机会和空间。PAI-Blade 和 BladeDISC 相关代码已经开源，可以参考 https://github.com/alibaba/BladeDISC。同时，GTC22 S41073-Generalized and Transparent AI Optimization Solutions with AI Compilers from Cloud Service Providers 拥有更加详细的关于 BladeDISC 的技术揭秘。
PAI-Blade 案例介绍
Detecron2 是一个基于 PyTorch 构建的检测分割深度学习算法库。这里选取Detectron2 的主要原因是其框架中提供了一些 SOTA 的检测分割算法网络结构，同时框架拥有很强的自定义和扩展能力。但是，从 Detectron2 中直接导出到 TensorRT 进行优化部署相对较为困难。PAI-Blade 基于 TorchScript 及圈图和 fallback 机制，在解决的导出问题的同时拿到了不错的执行性能，如下图：
Detectron2 上 TorchBlade(fp16) 对比 Torch(fp16) 加速效果(NVIDIA T4, CUDA11.5)
PAI-Blade 特性总结
- Performance: 集成最好的优化后端，提升对动态 Shape 的支持。
- Generality: 良好的通用性，支持主流的机器学习框架、加速硬件。
- Usability and Robustness: 降低工程师学习的成本和系统集成的成本。
PAI-Blade 和 BladeDISC 现已开源， https://github.com/alibaba/BladeDISC，也欢迎大家一起共建。
