一句话：提出 EVT 编译器，用图级融合优化与 ILP 分区自动生成灵活 epilogue，加速深度学习训练 1.26~3.1 倍。
深度学习编译器在推理上已取得显著加速，但在训练负载上受限：训练计算图包含归一化、损失函数、归约等难以融合的算子；前向与反向算子之间的附加边使寻找最优且可行的融合分区变得困难；现有编译器既无法生成在现代 GPU 上达到最先进性能的 kernel，也难以支持多样化的融合模式。本文提出 Epilogue Visitor Tree（EVT）编译器：通过新的图级编译 pass 挖掘隐藏的融合与优化机会，设计基于整数线性规划的分区器，在复杂的前后向联合图中高效求解最优可行分区，并提出 Epilogue Visitor 抽象与 EVT 算子编译器，自动生成灵活的 epilogue 并与 CUTLASS 等最先进库的高性能主循环实现集成。在跨领域的多种训练负载上，EVT 取得 1.26~3.1× 的加速。
As deep learning models become increasingly complex, the deep learning compilers are critical for enhancing the system efficiency and unlocking hidden optimization opportunities. Although excellent speedups have been achieved in inference workloads, existing compilers face significant limitations in training. Firstly, the training computation graph involves intricate operations challenging to fuse, such as normalization, loss functions, and reductions, which limit optimization opportunities like kernel fusion. Secondly, the training graph's additional edges connecting forward and backward operators pose challenges in finding optimal and feasible partitions for kernel fusion. More importantly, existing compilers cannot either generate kernels with state-of-the-art performance on modern GPUs or accommodate diverse fusion patterns. In this paper, we introduce Epilogue Visitor Tree (EVT), a novel compiler that overcomes these limitations. EVT employs novel graph-level compilation passes to unlock hidden fusion and optimization opportunities. It also incorporates a novel integer linear programming-based partitioner that efficiently solves the optimal and feasible partitions in complex joint forward-backward graphs. Moreover, we present the Epilogue Visitor Abstraction and introduce the EVT operator compiler that automatically generates flexible epilogues that can be integrated with high-performance main loop implementations from CUTLASS and other SOTA libraries. EVT is evaluated on diverse training workloads across domains and achieves 1.26~3.1× speedup.
@inproceedings{ChenKCKWD024,
title = {EVT: Accelerating Deep Learning Training with Epilogue Visitor Tree},
author = {Zhaodong Chen and Andrew Kerr and Richard Cai and Jack Kosaian and Haicheng Wu and Yufei Ding and Yuan Xie},
booktitle = {ASPLOS},
year = {2024},
doi = {10.1145/3620666.3651369}
}