[논문 리뷰] Automatic Full Compilation of Julia Programs and ML Models to Cloud TPUs

Keno Fischer, Elliot Saba|arXiv (Cornell University)|2018. 10. 23.

Distributed and Parallel Computing Systems참고 문헌 4인용 수 9

한 줄 요약

이 논문은 Julia의 정적 분석 및 타입 추론 기능을 활용하여 XLA 컴파일러를 통해 Google Cloud TPU에 전체 Julia 프로그램—예를 들어 VGG19와 같은 신경망 모델 및 그 기울기—를 자동으로 오프로드하는 컴파일러 기반 방법을 제시한다. 이 방법은 전체 순방향 및 역방향 계산을 단일 융합 TPU 커널으로 컴파일함으로써 100장의 이미지 배치에서 0.23초 대비 52.4초로 227배의 성능 향상을 달성하였으며, 컴파일러나 핵심 패키지의 수정 없이 1000줄 미만의 코드로도 가능하다.

ABSTRACT

Google's Cloud TPUs are a promising new hardware architecture for machine learning workloads. They have powered many of Google's milestone machine learning achievements in recent years. Google has now made TPUs available for general use on their cloud platform and as of very recently has opened them up further to allow use by non-TensorFlow frontends. We describe a method and implementation for offloading suitable sections of Julia programs to TPUs via this new API and the Google XLA compiler. Our method is able to completely fuse the forward pass of a VGG19 model expressed as a Julia program into a single TPU executable to be offloaded to the device. Our method composes well with existing compiler-based automatic differentiation techniques on Julia code, and we are thus able to also automatically obtain the VGG19 backwards pass and similarly offload it to the TPU. Targeting TPUs using our compiler, we are able to evaluate the VGG19 forward pass on a batch of 100 images in 0.23s which compares favorably to the 52.4s required for the original model on the CPU. Our implementation is less than 1000 lines of Julia, with no TPU specific changes made to the core Julia compiler or any other Julia packages.

연구 동기 및 목표

• 일반적인 Julia 프로그램, 특히 머신러닝 모델을 Julia 컴파일러나 핵심 패키지의 수정 없이 Cloud TPU로 전면 자동 컴파일할 수 있도록 하는 것.
• Zygote.jl과의 통합을 통해 엔드투엔드 자동 미분을 지원함으로써 순방향 및 역방향 계산을 모두 TPU로 오프로드할 수 있도록 하는 것.
• 다양한 디스patch, 고차원 함수, 일반화 프로그래밍과 같은 Julia의 고수준 언어 기능이 TPU에서 실행 가능한 XLA IR로 효과적으로 컴파일될 수 있음을 보여주는 것.
• 제어 흐름 구조(예: 반복문, 조건문)를 분석하고 정적 최적화를 적용하여 전체 프로그램 영역을 단일 커널으로 융합함으로써 복잡한 워크로드(예: 학습 루프)를 효율적으로 TPU에서 실행할 수 있도록 하는 것.
• Tracing 기반 또는 프레임워크 전용 접근 방식의 한계를 극복하고, 확장 가능하고 사용자 친화적인 인터페이스를 제공함으로써 Julia 코드를 TPU로 컴파일하는 데에 유용한 방법을 제공하는 것.

제안 방법

• Julia 컴파일러의 정적 분석 및 타입 추론 기능을 활용하여 런타임 트레이싱 없이 Julia 코드를 직접 XLA 중간 표현(IR)으로 컴파일하는 방법.
• Julia의 다형 디스patch 및 일반화 프로그래밍 기능을 활용하여 Julia IR에서 XLA IR로의 변환을 표현하고 최적화함으로써 타입 유형에 관계없이 컴파일할 수 있도록 하는 것.
• 데이터 의존성 분석과 정적 최적화를 적용하여 제어 흐름 요소(예: 반복문, 조건문)를 XLA 연산으로 융합함으로써 오프로드 전에 최적화하는 방법.
• Zygote.jl의 소스 투 소스 자동 미분 기능과 통합하여 순방향 및 역방향 계산을 하나의 융합 TPU 커널로 생성하고 오프로드하는 것.
• XLA 컴파일러의 일반적인 IR을 활용하여 TPU를 타겟으로 하여, TensorFlow 외의 워크로드나 선형 대수 중심의 비-ML 응용 프로그램도 지원할 수 있도록 하는 것.
• CPU와 TPU 간의 데이터 전송을 관리하기 위해 infeed/outfeed 연산을 활용하고, XLA의 타입 제약 조건을 충족시키기 위해 StructsOfArrays를 통한 SoA-to-AoS 변환을 지원하는 것.

실험 결과

연구 질문

• RQ1 제어 흐름과 고수준 추상화를 포함한 전체 Julia 프로그램이 런타임 트레이싱 없이 TPU에서 실행 가능한 XLA 커널로 자동 컴파일될 수 있는가?
• RQ2 Julia의 다형 디스패치 및 일반화 프로그래밍 기능을 얼마나 효과적으로 활용하여 효율적이고 타입에 특화된 XLA 커널을 생성할 수 있는가?
• RQ3 Julia 컴파일러 인프라와 Zygote의 자동 미분 기능의 조합이 딥러닝 모델의 순방향 및 역방향 계산을 모두 TPU로 엔드투엔드 오프로드할 수 있도록 할 수 있는가?
• RQ4 완전히 오프로드된 VGG19 모델의 성능은 CPU 실행 대비 어떻게 되며, Cloud TPU에서 달성 가능한 성능 향상은 어느 정도인가?
• RQ5 XLA를 통한 일반 Julia 코드의 TPU 컴파일 과정에서 디버깅 가능성, 타입 추론 실패, 분산 실행 지원 등의 측면에서 주요 제한 사항은 무엇인가?

주요 결과

• 이 방법은 VGG19 모델의 전체 순방향 계산을 단일 융합 TPU 커널으로 컴파일하여 CPU 실행 대비 227배의 성능 향상을 달성하였다(100장 배치 기준 0.23초 대비 52.4초).
• Zygote.jl의 자동 미분과 XLA의 융합 기능을 활용하여 순방향 및 역방향 계산, 학습 루프까지 엔드투엔드로 오프로드할 수 있도록 지원한다.
• 구현은 Julia 코드로 1000줄 미만이며, 핵심 Julia 컴파일러나 외부 패키지의 수정이 전혀 필요하지 않다.
• 성능 측정이 안정적이었으며, TPU 계산 시간(FluXLA TPU)은 총 월드클록 시간보다 훨씬 일관성 있게 측정되었다.
• 선형 대수 연산이 지배적인 워크로드라면, 이 방법은 머신러닝을 초월한 비-ML Julia 코드로도 일반화 가능함을 시사하며, 더 넓은 응용 가능성을 보여준다.
• 현재 구현은 타입 추론 실패의 디버깅 어려움과 XLA의 분산 프리미티브 지원 부족으로 인해 어려움을 겪고 있으며, 향후 개선이 필요한 핵심 영역로 지목된다.