[论文解读] FusionStitching: Boosting Memory Intensive Computations for Deep Learning Workloads

Zhen Zheng, Pengzhan Zhao+7 more|arXiv (Cornell University)|Sep 23, 2020

Advanced Neural Network Applications参考文献 46被引用 12

一句话总结

FusionStitching 是一种 JIT 深度学习编译器，通过将具有复杂数据依赖性和非均匀并行性的内存密集型深度学习算子融合为大型 GPU 内核，优化内存密集型深度学习工作负载。它通过共享内存和寄存器重排实现中间值复用，减少访主存访问和上下文切换开销，相较于最先进的技术最高可实现 2.21× 的加速，平均提升 1.45×。

ABSTRACT

We show in this work that memory intensive computations can result in severe performance problems due to off-chip memory access and CPU-GPU context switch overheads in a wide range of deep learning models. For this problem, current just-in-time (JIT) kernel fusion and code generation techniques have limitations, such as rough fusion plan exploration strategies and limited code generation ability. We propose FusionStitching, a deep learning compiler capable of fusing memory intensive operators, with varied data dependencies and non-homogeneous parallelism, into large GPU kernels to reduce global memory access and context switch overhead automatically. FusionStitching widens the range of operation combinations that fusion can target beyond previous JIT works by introducing data reuse of intermediate values. It explores large fusion spaces to decide optimal fusion plans with considerations of memory access costs, kernel calls and resource usage constraints. FusionStitching tunes the optimal stitching scheme with a domain-specific cost model efficiently. Experimental results show that FusionStitching can reach up to 2.21x speedup compared to state-of-the-art, with 1.45x on average. Besides these experimental results, we integrated our approach into a compiler product and deployed it onto a production cluster for AI workloads with thousands of GPUs. The system has been in operation for more than 4 months and saves 7,000 GPU hours on average for approximately 30,000 tasks per month.

研究动机与目标

• 解决内存密集型深度学习算子中由访主存访问和 CPU-GPU 上下文切换引起的性能瓶颈。
• 克服现有 JIT 融合技术的局限性，如融合空间狭窄和缺乏数据复用支持。
• 实现对复杂、多样的内存密集型算子组合（例如逐元素运算、转置、规约）的高效融合，这些算子具有不同的数据依赖性和并行性。
• 设计可扩展的领域特定成本模型与调优框架，以实现最优融合计划选择和高性能内核生成。
• 在生产级人工智能集群环境中集成并验证系统，以展示真实场景下的性能提升与资源节约。

提出的方法

• 引入四种缝合方案抽象，以建模涉及内存密集型算子、复杂依赖关系和非均匀并行性的多样化融合模式。
• 通过片上内存（共享内存）和寄存器重排，在线程间实现中间值复用，减少重复计算开销。
• 设计两层领域特定成本模型，通过估算内存访问成本、内核启动成本和资源使用情况，指导融合空间探索。
• 实现一个融合探索器，对大规模融合空间执行智能、基于成本的搜索，避免陷入贪婪的局部最优。
• 开发代码生成器，将融合计划映射为经过调度调优的优化 GPU 内核。
• 将 FusionStitching 集成到 XLA 框架中，以支持动态形状和生产环境部署。

实验结果

研究问题

• RQ1 如何系统性地在内存密集型深度学习算子的 JIT 融合中利用中间数据复用？
• RQ2 哪些融合模式和成本模型能够在复杂、动态的工作负载中实现内存访问、内核启动开销和资源使用之间的最优权衡？
• RQ3 具备先进探索与代码生成能力的 JIT 融合框架，是否能在真实世界深度学习模型中超越基于规则或贪婪的方法？
• RQ4 中间值复用和片上内存利用率在多大程度上可减少访主存流量和上下文切换开销？
• RQ5 所提出的融合框架在具有动态形状和高算子多样性的生产规模人工智能工作负载中效果如何？

主要发现

• FusionStitching 在真实世界深度学习模型上相较于最先进的 JIT 优化技术，最高实现 2.21× 的加速，平均提升 1.45×。
• 该系统通过将多种内存密集型算子融合为单个高度优化的内核，减少了 GPU 内核启动开销和访主存访问。
• 通过共享内存和寄存器重排实现中间值复用，FusionStitching 避免了现有 JIT 工具（如 XLA）中普遍存在的冗余计算问题。
• 该框架成功集成到基于生产级 XLA 的编译器栈中，并已在包含数千个 GPU 的集群中稳定运行超过四个月。
• 部署后每月为 30,000 个任务节省约 7,000 个 GPU 小时，显著提升了真实场景下的效率。
• 领域特定成本模型即使在处理动态形状和复杂算子组合时，也能以极低的调优开销实现有效的融合计划选择。