arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2603.24239cs.PLcs.AIcs.LG
收藏

DVM: 一种用于动态张量计算的字节码虚拟机方法

DVM: A Bytecode Virtual Machine Approach for Dynamic Tensor Computation

Huawei(华为);HKUST(香港科技大学);HKUST(GZ)(香港科技大学(广州));Hunan University(湖南大学)

Jingzhi Fang, Xiong Gao, Renwei Zhang, Zichun Ye, Lei Chen, Jie Zhao, Chengnuo Huang, Hui Xu, Xuefeng Jin

AI总结:

本文提出DVM,通过字节码虚拟机实现实时编译,提升动态张量计算的效率和编译速度,相比其他框架在效率和编译时间上有显著优势。

摘要:

动态是AI计算中的常见特性,例如模型中的动态张量形状和动态控制流。由于编译时间长,现有运行时编译损害了模型效率,而离线编译者要么因长编译时间和设备内存占用无法覆盖所有动态模型的执行实例,要么牺牲优化机会以换取易用性。本文重新审视动态模型运行时编译的可行性,发现关键在于加速编译或隐藏编译开销。为此,我们提出实时编译器DVM。在DVM中,我们基于字节码虚拟机设计了运行时操作符编译器,以对每个动态操作符实例进行有效且高效的编译。具体来说,而不是将程序编译成机器代码,我们将在CPU上将操作符程序编码为字节码,并将其解码为虚拟指令,直接在NPU上执行。基于运行时操作符编译器,我们进一步提出操作符融合器,对静态图进行基于符号推导的融合,并对动态图进行运行时融合。两种基于模式和堆叠的融合方式均被支持,以增加融合机会。对操作符、子图和模型的评估表明,与TorchInductor、PyTorch-eager和MindSpore-graph-O0相比,在操作符/模型效率方面提高了11.77倍,在最大编译时间方面快了五个数量级。

英文摘要:

Dynamism is common in AI computation, e.g., the dynamic tensor shapes and the dynamic control flows in models. Due to the long compilation time, existing runtime compilation damages the model efficiency, while the offline compilers either suffer from the long compilation time and device memory footprint to cover all the possible execution instances of a dynamic model, or sacrifice optimization opportunities for usability. In this paper, we rethink the feasibility of runtime compilation for dynamic models and identify that the key for it to work is to speed up the compilation or hide the compilation overhead. To do this, we propose a real-time compiler, DVM. In DVM, we design a runtime operator compiler based on a bytecode virtual machine to perform effective and efficient compilation for each dynamic operator instance given its input. Specifically, instead of compiling programs into machine code, we encode the operator program into bytecode on the CPU and decode the bytecode into virtual instructions for direct execution on the NPU. Based on the runtime operator compiler, we further propose an operator fuser, which performs symbol-deduction-based fusion on static graphs and runtime fusion on dynamic graphs. Both pattern- and stacking-based fusion are supported to increase fusion opportunities. Evaluation on operators, subgraphs, and models shows that, compared with TorchInductor, PyTorch-eager and MindSpore-graph-O0, we are up to 11.77$\times$ better in terms of the operator/model efficiency and up to 5 orders of magnitude faster in terms of the maximum compilation time.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑