Apache TVM 是一个用于 CPU、GPU 和机器学习加速器的开源机器学习编译器框架。TVM 支持 TensorFlow、Pytorch、MXNet、ONNX 等几乎所有的主流框架，目标是优化机器学习模型让其高效运行在不同的硬件平台上。TVM 提供了深度学习模型编译、优化和部署的端到端解决方案，支持从模型定义到部署的全流程自动化。
近日，TVM 社区举办了 TVMCon2023 会议。会议上，腾讯 BlazerML 深度学习编译器团队发表了题为《TVM at Tencent》的演讲，主要介绍了 BlazerML 的相关工作。BlazerML 是基于开源深度学习编译器（TVM/XLA）深度定制的端到端推理/训练加速方案，致力于降低深度学习编译器使用门槛，帮助编译优化技术在业务中落地。
本次分享主要介绍了 BlazerML 加速强化学习推理的方案以及 CPU 上高性能卷积算子的优化。以下为分享内容的具体介绍，完整视频见文末。
强化学习推理加速
强化学习是通过为成功和失败的行为提供奖励和惩罚来提高 AI 性能的一种方法。下图展示了游戏 AI 中的 RL 训练框架。其中 Learner 部署在 GPU 集群上，负责训练模型并接收 Actor 样本；而 Actor 则部署在 CPU 集群上，加载实时策略进行样本生成。样本数据的数量对 RL 训练的有效性非常重要。因此，需要高效的推理框架以提高样本生成速度。
下图展示了 BlazerML 加速进行强化学习推理的完整工作流程。BlazerML 的目标是从 Learner 训练的模型经过自动调优实时加速，发送给 Actor 进程所在的所有节点进行推理抽样。为此，BlazerML 基于 TVM 进行了一系列优化。
第一，TVM 拥有自动调优模型和算子的能力，但自动调优和编译的耗时非常长。BlazerML 使用一个自动调优数据库保存调优后的子图 Schedule 以加速 TVM 的自动调优流程。在对模型的自动调优流程中，对于已经调优过的子图或算子，通过搜索数据库，BlazerML 可以快速找到目标硬件设备上的最佳性能调度。
第二，TVM 在构造模型时需要提供固定的权重，一旦权重改变，就需要重新构建。然而，在强化学习过程中同步模型时，只有模型的权重值会发生变化，而结构保持不变。BlazerML 记录了 TVM 自动调优的 Schedule 在每个权重上的变换，直接操作原始权重。
第三，在强化学习推理中，服务器上运行多个智能体的多个模型。使用 TVM 部署可能会引起 DRAM 限制问题，极大地影响训练中样本的生成。BlazerML 提供了共享内存版本的 TVM 模型，以便多个智能体可以共享权重。使用共享内存可以为 TVM 模型部署提供显著的好处，包括减少内存使用，提高性能，简化代码。使用共享内存后，服务器 QPS 可以达到原始值的 2.3 倍，同时内存使用也更小，服务器成本减半。
下图展示了游戏模型中推理的比较结果。相比社区版 TVM，模型的性能提升了超过 1.3 倍，Actor 的推理时间占总时间的 30％至 40％，节省了 6％至 16％的机器资源。
高性能卷积算子优化
卷积算子是 CNN 网络中的常见算子，对卷积算子的优化一直是深度学习编译器的热门问题。解决卷积问题的算法有非常多，常见的有 DirectConv，Im2Col，Winograd，FFT。BlazerML 主要针对 Winograd 算法实现的卷积算子进行了调优。
Winograd 算法的主要流程包含 4 部分：输入变换，权重变换，核心矩阵乘，以及输出变换。Winograd 算法先将数据张量切分为若干个数据块，接着对每个数据块做输入变换，得到张量 V；对权重张量做权重变换得到张量 U。U 和 V 通过核心矩阵乘得到张量 M，再经过输出变换得到结果张量 O。考虑到输入输出变换的计算特殊性，是和常数矩阵进行连续两次矩阵乘，BlazerML 做了针对性的优化，降低了变换计算的计算量，并降低了数据排布变换的开销。
在输入变换中，数据块的大小是可选的，常见的大小有 6，4，2 等。仅考虑核心矩阵乘计算部分大小为 6 的数据块有显著更高的加速比，但实际实现中输入输出变换的计算量也有很大影响。因此，对于有些形状，大小为 4 的数据块会有更高的加速比。BlazerML 根据对理论复杂度的分析，增加了对数据块大小选择的逻辑，在部分形状上性能取得了显著的提升。
对计算模板调优完成后，BlazerML 基于 TVM 对模板进行了自动调优，并在 x86 和 arm 平台上分别进行了性能测试。具体性能如下图所示。
Intel(R) Xeon(R) Platinum 8255C CPU @ 2.50GHz THREAD_NUM=1
Ampere(R) Altra(R) Neoverse-N1 @ 2.80GHz THREAD_NUM=1
未来计划：Cost Model 优化
在使用 TVM 对 Winograd 模板进行自动调优时，多次调优结果的差距很大。如下图所示，与 DirectConv 不同，Winograd 模板多次自动调优的结果存在非常大的波动。
这导致常常需要多次自动调优才能获得较好的 Schedule。在探索这个问题时定位到，在 TVM 这种 Cost Model Guide Search 中，Cost Model 的准确度对结果影响很大。因此，对比了在 DirectConv 模板和 Winograd 模板的 Tuning 过程中 Cost model 部分指标的差异。如下图所示，Winograd 模板的误差更大，拟合度更低。
我们已经初步定位到目前 TVM 使用的 Cost Model 设计的 Bug，正在寻找更好的 Cost Model 方案。未来，BlazerML 也将提供更多高性能算子的支持。
相关文章
