[논문 리뷰] LazyTensor: combining eager execution with domain-specific compilers

Alex Suhan, Davide Libenzi외 5명|arXiv (Cornell University)|2021. 02. 26.

Parallel Computing and Optimization Techniques참고 문헌 28인용 수 6

한 줄 요약

LazyTensor는 도메인 특화 컴파일러(예: XLA)를 통해 텐서 연산을 컴파일함으로써, 정의 기반 실행(Define-by-Run) 머신러닝 프레임워크에서 호스트 언어의 표현력을 완전히 활용할 수 있도록 하면서도 언어 하위집합 문제를 피한다. 이는 연산을 지연 평가되고 조합 가능한 중간 표현(IR) 그래프로 표현함으로써 달성되며, 이 그래프들은 필요할 때만 컴파일된다. 이를 통해 파이토치와 스위프트 포 텐서플로우에서 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)와 텐서 유형을 지원하는 재사용 가능한 다중 언어 구현을 구축한다.
• 지연 컴파일이 실제 ML 워크로드(예: 클라우드 TPU에서의 허깅페이스 트랜스포머)에서 상당한 성능 향상을 달성할 수 있음을 입증한다.

제안 방법

• 지연 평가되고 조합 가능한 IR 그래프로 텐서 연산을 표현하며, 이는 즉각 실행 중에 구성되지만 필요할 때만 컴파일된다.
• 파이썬과 스위프트 등의 다양한 프로그래밍 언어와 파이토치, 스위프트 포 텐서플로우 등의 프레임워크 간에 동일한 IR 그래프 생성 및 컴파일 파이프라인을 공유한다.
• 지정된 도메인 컴파일러(예: XLA)와 통합하기 위해, 지연된 IR을 필요에 따라 컴파일러의 중간 표현(예: HLO)으로 번역한다.
• 동일한 IR 그래프가 다시 나타날 경우 재컴파일을 방지하기 위해 캐싱을 활용한다.
• 제어 흐름, 동적 크기 등 최적화되지 않은 코드가 최적화된 텐서 연산과 공존할 수 있도록, 명시적 장벽인 LazyTensorBarrier()를 통해 실행 경계를 관리한다.
• Tensor 객체를 완료될 때까지 해제되지 않는 프ом프스로 간주함으로써 비동기 실행을 가능하게 한다.

실험 결과

연구 질문

• RQ1 정의 기반 실행 머신러닝 프레임워크는 도메인 특화 컴파일러를 통해 컴파일할 수 있는 동시에 호스트 언어의 표현력을 완전히 유지할 수 있는가?
• RQ2 IR 그래프 구성의 오버헤드를 최소화하여 소규모 배치나 저지연 추론 워크로드를 지원할 수 있는가?
• RQ3 지연 컴파일 시스템의 동일한 구현이 다양한 프로그래밍 언어와 하드웨어 가속기 간에 얼마나 재사용 가능한가?
• RQ4 지연 컴파일이 실제 워크로드에서 전면 추적 또는 그래프 기반 시스템과 비교해 유사하거나 더 뛰어난 성능을 달성할 수 있는가?

주요 결과

• LazyTensor는 파이썬과 스위프트의 기능을 모두 완전히 활용하면서도 XLA로 컴파일함으로써, 추적 기반 시스템에서 발생하는 언어 하위집합 문제를 피했다.
• 동일한 핵심 구현이 파이토치와 스위프트 포 텐서플로우 간에 재사용되어 언어와 프레임워크 간의 이식성을 입증했다.
• 성능 평가 결과, HuggingFace Transformers 라이브러리를 사용한 클라우드 TPU에서 상당한 속도 향상이 확인되어 실제 적용 가능성은 입증되었다.
• LazyTensorBarrier()를 통해 최적화되지 않은 코드와 최적화된 텐서 연산이 공존하는 혼합 워크로드를 지원했다.
• 대부분의 학습 워크로드에서는 그래프 구성 오버헤드가 관리 가능하지만, IR 표현 최적화 및 캐싱을 통해 추가로 감소시킬 수 있다.