[논문 리뷰] Tensor Program Optimization with Probabilistic Programs

Junru Shao, Xiyou Zhou외 8명|arXiv (Cornell University)|2022. 05. 26.

Parallel Computing and Optimization Techniques인용 수 10

한 줄 요약

MetaSchedule는 텐서 프로그램 최적화를 위한 풍부한 검색 공간을 모듈식으로 구성하기 위해 도메인 특화 확률 프로그래밍 언어를 도입한다. 이로써 도메인 전문가가 최적화 기법을 쉽게 확장할 수 있다. MetaSchedule는 하드웨어별 전환 모듈을 학습 기반 검색과 조합함으로써 TVM보다 종단 간 딥러닝 워크로드에서 48% 빠른 추론 성능을 달성한다.

ABSTRACT

Automatic optimization for tensor programs becomes increasingly important as we deploy deep learning in various environments, and efficient optimization relies on a rich search space and effective search. Most existing efforts adopt a search space which lacks the ability to efficiently enable domain experts to grow the search space. This paper introduces MetaSchedule, a domain-specific probabilistic programming language abstraction to construct a rich search space of tensor programs. Our abstraction allows domain experts to analyze the program, and easily propose stochastic choices in a modular way to compose program transformation accordingly. We also build an end-to-end learning-driven framework to find an optimized program for a given search space. Experimental results show that MetaSchedule can cover the search space used in the state-of-the-art tensor program optimization frameworks in a modular way. Additionally, it empowers domain experts to conveniently grow the search space and modularly enhance the system, which brings 48% speedup on end-to-end deep learning workloads.

연구 동기 및 목표

• 자동 텐서 프로그램 최적화에서 검색 공간 구축을 검색 알고리즘에서 분리하기 위해.
• 핵심 프레임워크를 수정하지 않고도 도메인 전문가가 검색 공간을 점진적으로 확장하고 맞춤화할 수 있도록 하기 위해.
• 구축된 검색 공간을 사용해 최적의 텐서 프로그램을 찾는 학습 기반 프레임워크 개발하기 위해.
• 다양한 하드웨어 백엔드에서 최신 기술 수준의 시스템(예: TVM, PyTorch)과 경쟁하거나 그 이상의 성능을 달성하기 위해.
• 모듈식이고 조합 가능한 검색 공간 구축 방식이 실제 딥러닝 워크로드에서 측정 가능한 성능 향상으로 이어지는지 입증하기 위해.

제안 방법

• 프로그램 변환에서의 확률적 선택을 모델링하는 확률 프로그램으로 검색 공간을 표현하기 위해.
• 도메인 전문가가 조합 가능한 모듈식 전환 규칙를 확률 프로그램으로 정의할 수 있도록 하기 위해.
• 확률적 검색 공간을 탐색하고 저지연 시간 텐서 프로그램을 찾기 위해 미분 가능한 검색 알고리즘 사용하기 위해.
• 다양한 전환 모듈(예: 루프 타일링, 벡터화, 텐서화)을 통합된 검색 공간으로 조합하기 위해.
• GPU 텐서 코어와 같은 전용 하드웨어 프리미티브를 활용하기 위해 하드웨어 특화 모듈(예: Use-Tensor-Core) 통합하기 위해.
• 프레임워크 수준의 변경을 최소화하면서 CPU 및 GPU 플랫폼에서의 자동 튜닝을 지원하는 종단 간 시스템 구축하기 위해.

실험 결과

연구 질문

• RQ1 확률 프로그래밍 추상화가 텐서 프로그램 최적화에서 검색 공간 구축을 검색 알고리즘에서 효과적으로 분리하는 데에 적합한가?
• RQ2 제안된 추상화가 TVM 및 PyTorch와 같은 최신 프레임워크에서 사용되는 최적화 기법을 얼마나 잘 커버하는가?
• RQ3 전환 규칙의 모듈식 조합이 최적화 성능 향상에 얼마나 기여하는가?
• RQ4 이 프레임워크는 종단 간 딥러닝 모델에서 기존 시스템과 경쟁하거나 그 이상의 성능을 달성할 수 있는가?
• RQ5 도메인 전문가가 하드웨어 특화 최적화를 얼마나 효율적으로 검색 공간에 확장할 수 있는가?

주요 결과

• MetaSchedule는 CPU 및 GPU에서 다양한 연산자에 대해 TVM(Ansor)와 동등하거나 그 이상의 성능을 달성한다.
• 종단 간 딥러닝 모델에서 MetaSchedule는 TVM 및 PyTorch의 성능을 따라하거나 초월하며, 하드웨어 특화 모듈을 사용할 경우 BERT-Large에서 48%의 성능 향상을 기록한다.
• 두 날 동안 작성한 82줄의 하드웨어 특화 모듈이 TVM(AutoTVM) 대비 48%의 성능 향상을 제공함으로써, 맞춤화의 높은 실용성을 입증한다.
• 전환 모듈을 점진적으로 조합함으로써 일관된 성능 향상이 이루어지며, 이는 모듈성 덕분에 검색 공간의 확장이 스케일러블하다는 것을 보여준다.
• MetaSchedule의 검색 공간는 기존 프레임워크의 최적화 공간을 충분히 커버할 수 있으며, 기반 시스템에 수술적인 변경 없이도 가능하다.
• 이 프레임워크는 도메인 전문가와 학습 기반 검색 간의 효율적인 협업을 가능하게 하여, 새로운 최적화 기법을 통합하는 데에 드는 노력과 시간을 크게 줄인다.