[论文解读] LazyTensor: combining eager execution with domain-specific compilers

Alex Suhan, Davide Libenzi+5 more|arXiv (Cornell University)|Feb 26, 2021

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

一句话总结

LazyTensor 使定义运行型机器学习框架能够充分利用宿主语言的表达能力，同时通过领域专用编译器（如 XLA）编译张量操作，从而避免语言子集问题。它通过将操作表示为惰性、可组合的中间表示（IR）图来实现这一点，这些图仅在需要时才进行编译，从而在 PyTorch 和 Swift for TensorFlow 中均实现 TPU 和 GPU 上的高性能，并且跨语言和硬件共享实现代码。

ABSTRACT

Domain-specific optimizing compilers have demonstrated significant performance and portability benefits, but require programs to be represented in their specialized IRs. Existing frontends to these compilers suffer from the "language subset problem" where some host language features are unsupported in the subset of the user's program that interacts with the domain-specific compiler. By contrast, define-by-run ML frameworks-colloquially called "eager" mode-are popular due to their ease of use and expressivity, where the full power of the host programming language can be used. LazyTensor is a technique to target domain specific compilers without sacrificing define-by-run ergonomics. Initially developed to support PyTorch on Cloud TPUs, the technique, along with a substantially shared implementation, has been used by Swift for TensorFlow across CPUs, GPUs, and TPUs, demonstrating the generality of the approach across (1) Tensor implementations, (2) hardware accelerators, and (3) programming languages.

研究动机与目标

• 解决领域专用编译器中的语言子集问题，即用户程序必须符合受限的中间表示（IR）表达能力。
• 使定义运行型机器学习框架能够利用高性能、硬件优化的编译器，同时不牺牲调试能力或语言灵活性。
• 构建一个可重用的跨语言实现，支持多种硬件加速器（CPU、GPU、TPU）和张量类型。
• 证明惰性编译可在真实生产级机器学习工作负载（如 Cloud TPU 上的 HuggingFace Transformers）中实现显著的性能提升。

提出的方法

• 将张量操作表示为惰性、可组合的中间表示（IR）图，这些图在即时执行期间构建，但仅在需要时才进行编译。
• 在不同编程语言（Python、Swift）和框架（PyTorch、Swift for TensorFlow）之间共享 IR 图构建与编译管道的实现。
• 通过将惰性 IR 图即时翻译为编译器的中间表示（如 HLO），与领域专用编译器（如 XLA）集成。
• 利用缓存机制，避免在再次遇到相同 IR 图时重复编译，从而降低开销。
• 通过显式屏障（如 LazyTensorBarrier()）支持混合工作负载，允许非优化代码（如控制流、动态形状）与优化的张量操作共存。
• 通过将张量对象视为承诺（promises）来实现异步执行，仅在计算完成后才解析。

实验结果

研究问题

• RQ1 定义运行型机器学习框架是否能够在仍通过领域专用编译器进行编译的同时，保留宿主语言的完整表达能力？
• RQ2 如何最小化 IR 图构建的开销，以支持小批量或低延迟推理工作负载？
• RQ3 单个惰性编译系统实现可在多大程度上跨不同编程语言和硬件加速器复用？
• RQ4 惰性编译在真实世界机器学习工作负载上是否能实现与完整追踪或基于图的系统相当甚至更优的性能？

主要发现

• LazyTensor 在将张量程序编译为 XLA 的同时，成功实现了对 Python 和 Swift 特性的完整支持，避免了基于追踪系统固有的语言子集问题。
• 相同的底层实现被复用于 PyTorch 和 Swift for TensorFlow，证明了其在不同语言和框架间的可移植性。
• 性能评估显示，在 Cloud TPU 上使用 HuggingFace Transformers 库时实现了显著的速度提升，证明了其在真实场景中的适用性。
• 该系统通过 LazyTensorBarrier() 支持混合工作负载，能够管理非优化代码与优化张量操作之间的执行边界。
• 图构建开销对大多数训练工作负载而言是可接受的，但可通过优化 IR 表示和缓存机制进一步降低。