[论文解读] A Tensor Compiler for Unified Machine Learning Prediction Serving

Supun Nakandala, Karla Saur+5 more|arXiv (Cornell University)|Oct 9, 2020

Parallel Computing and Optimization Techniques参考文献 28被引用 10

一句话总结

Hummingbird 是一个张量编译器，可将传统机器学习模型和特征工程算子编译为一组统一的张量操作，从而通过现有的深度学习框架在 CPU 和 GPU 上实现高效执行。与手工优化的内核相比，它实现了具有竞争力或更优的性能，同时降低了基础设施复杂性，并实现了端到端的硬件加速。

ABSTRACT

Machine Learning (ML) adoption in the enterprise requires simpler and more efficient software infrastructure---the bespoke solutions typical in large web companies are simply untenable. Model scoring, the process of obtaining predictions from a trained model over new data, is a primary contributor to infrastructure complexity and cost as models are trained once but used many times. In this paper we propose HUMMINGBIRD, a novel approach to model scoring, which compiles featurization operators and traditional ML models (e.g., decision trees) into a small set of tensor operations. This approach inherently reduces infrastructure complexity and directly leverages existing investments in Neural Network compilers and runtimes to generate efficient computations for both CPU and hardware accelerators. Our performance results are intriguing: despite replacing imperative computations (e.g., tree traversals) with tensor computation abstractions, HUMMINGBIRD is competitive and often outperforms hand-crafted kernels on micro-benchmarks on both CPU and GPU, while enabling seamless end-to-end acceleration of ML pipelines. We have released HUMMINGBIRD as open source.

研究动机与目标

• 为解决在企业环境中部署传统机器学习模型日益增长的复杂性和成本问题。
• 将分散的特征工程和模型算子统一到单一的基于张量的抽象下，避免因支持多个框架和硬件目标而产生的 N×M 复杂性爆炸。
• 通过利用现有的深度神经网络（DNN）编译器和运行时，实现高效、可移植且经过优化的传统机器学习推理。
• 证明高级张量抽象在模型评分方面可以超越手工优化的命令式实现。

提出的方法

• 将传统机器学习模型（如决策树、线性模型）和特征工程算子（如分词、归一化）编译为一组最小化的核心张量操作。
• 采用两级优化流水线：基于算子组合的逻辑优化，随后通过 DNN 运行时后端进行物理优化。
• 为树模型集成提出新颖的推理策略，包括 GEMM 和 TT，根据模型特征选择最高效的策略。
• 利用现有的 DNN 框架（如 PyTorch、TVM）及其硬件优化后端（CPU、GPU）来执行编译后的模型。
• 与推理运行时集成，以实现自动硬件感知的代码生成和优化。
• 支持密集和稀疏张量操作，并可通过 TACO 等技术扩展未来稀疏优化能力。

实验结果

研究问题

• RQ1 传统机器学习模型和特征工程算子能否被有效编译为统一的张量计算抽象？
• RQ2 这种基于张量的编译能否在 CPU 和 GPU 上实现与手工优化的命令式内核相当或更优的性能？
• RQ3 利用 DNN 编译器和运行时是否能降低基础设施复杂性，并为传统机器学习实现无缝硬件加速？
• RQ4 编译后的张量模型在不同模型类型、输入大小和硬件平台上的性能表现如何？

主要发现

• 与 scikit-learn 相比，Hummingbird 在最快流水线中实现了 1200× 的速度提升，而最慢的流水线仅因输入过小或稀疏而出现 60× 的性能下降。
• 在 CPU 上，2,317 个评估流水线中有 60% 实现了加速，其中 27% 因稀疏性或计算负载过低而出现性能下降。
• 在 GPU 加速下，73% 的流水线实现了加速，最快流水线相比 scikit-learn 实现了 1,000× 的性能提升。
• 尽管使用了高级张量抽象，该系统在微基准测试中仍优于手工编写的 C++ 和 CUDA 内核。
• 该方法实现了跨多种硬件（包括 CPU 和 GPU）的完整机器学习流水线端到端加速，且工程投入极少。
• 该系统具备可扩展性，并能无需额外工作即受益于未来 DNN 框架的改进。