[논문 리뷰] Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients

William S. Moses, Valentin Churavy|arXiv (Cornell University)|2020. 10. 04.

Parallel Computing and Optimization Techniques참고 문헌 49인용 수 10

한 줄 요약

이 논문은 최적화된 LLVM 중간 표현(IR)에서 직접 빠른 기울기를 합성함으로써 외부 코드를 재작성하지 않고도 효율적인 다국어 간 자동미분(AD)을 가능하게 하는 고성능 자동미분(AD) 컴파일러 플러그인인 Enzyme를 소개한다. 컴파일러 최적화 이후에 AD를 수행함으로써 Enzyme는 최적화되지 않은 IR에 대한 전통적 AD 대비 기하평균 4.5배의 성능 향상을 달성했으며, 머신러닝 벤치마크에서 최신 기술을 초월한다.

ABSTRACT

Applying differentiable programming techniques and machine learning algorithms to foreign programs requires developers to either rewrite their code in a machine learning framework, or otherwise provide derivatives of the foreign code. This paper presents Enzyme, a high-performance automatic differentiation (AD) compiler plugin for the LLVM compiler framework capable of synthesizing gradients of statically analyzable programs expressed in the LLVM intermediate representation (IR). Enzyme synthesizes gradients for programs written in any language whose compiler targets LLVM IR including C, C++, Fortran, Julia, Rust, Swift, MLIR, etc., thereby providing native AD capabilities in these languages. Unlike traditional source-to-source and operator-overloading tools, Enzyme performs AD on optimized IR. On a machine-learning focused benchmark suite including Microsoft's ADBench, AD on optimized IR achieves a geometric mean speedup of 4.5x over AD on IR before optimization allowing Enzyme to achieve state-of-the-art performance. Packaging Enzyme for PyTorch and TensorFlow provides convenient access to gradients of foreign code with state-of-the art performance, enabling foreign code to be directly incorporated into existing machine learning workflows.

연구 동기 및 목표

• 머신러닝 프레임워크에서 외부 코드를 재작성할 필요 없이 네이티브로 고성능 자동미분을 가능하게 하기 위해.
• 컴파일러 최적화 이전에 AD를 적용할 경우 발생하는 성능 저하 문제를 해결하기 위해.
• 소스 접근 권한 없이도 사전 컴파일된 또는 제3자 라이브러리의 다국어 및 다프레임워크 자동미분을 가능하게 하기 위해.
• 낮은 수준의 최적화된 IR에서 효율적인 AD가 가능하다는 것을 입증하여, AD는 고수준 언어에서만 효과적이라는 가정에 도전하기 위해.
• 과학계산 및 시뮬레이션 코드를 현대적인 머신러닝 워크플로우에 원활하게 통합할 수 있는 길을 제공하기 위해.

제안 방법

• Enzyme는 LLVM 컴파일러 프레임워크 내부에서 작동하는 플러그인으로, 최적화된 LLVM IR에서 직접 역방향 자동미분을 수행한다.
• AD 적용 이전에 LLVM의 전체 최적화 기능(예: 루프 불변 코드 이동)을 활용하여 성능 향상을 유지한다.
• 최적화된 IR의 데이터 플로우 및 제어 플로우를 분석하여 효율적인 수반 코드(Adjoint code)를 합성한다.
• Enzyme는 스택 기반 접근 방식보다 메모리 오버헤드를 줄이기 위해 재귀 테입 구조를 사용한다.
• C, C++, Fortran, Julia, Rust, Swift 등 다양한 언어가 LLVM IR를 통해 통합되는 단일 IR 기반 자동미분 파이프라인을 지원한다.
• PyTorch 및 TensorFlow와 통합되어 기존 머신러닝 워크플로우에서 외부 코드를 직접 미분할 수 있도록 한다.

실험 결과

연구 질문

• RQ1 최적화된 LLVM IR에서 자동미분을 효율적으로 수행할 수 있는가, 즉 최적화되지 않은 소스나 IR이 아닌가?
• RQ2 컴파일러 최적화 이후에 AD를 적용할 경우, 최적화되지 않은 코드에 대한 전통적 AD 대비 기울기 계산 속도가著로 향상되는가?
• RQ3 컴파일러 플러그인 기반의 AD 시스템이 다양한 머신러닝 및 과학계산 워크로드에서 최신 기술 수준의 성능을 달성할 수 있는가?
• RQ4 Enzyme는 소스 수정 없이도 물리 시뮬레이터나 게임 엔진과 같은 외부 코드를 머신러닝 파이프라인에 얼마나 원활하게 통합할 수 있는가?
• RQ5 IR 수준에서 작동하고 GPU 코드 생성과 같은 고급 기능을 지원하는 고성능 다국어 AD 시스템을 구축하는 것은 가능한가?

주요 결과

• Enzyme는 다양한 벤치마크 세트(예: Microsoft의 ADBench 포함)에서 최적화되지 않은 IR에 대한 AD(기준 파이프라인) 대비 기하평균 4.5배의 성능 향상을 달성했다.
• 성능 향상의 주요 원인은 최적화 이후 AD를 수행함으로써, 예를 들어 루프 불변 코드 이동 최적화를 통해 기울기 내부의 중복 계산을 제거할 수 있었기 때문이다.
• ADBench에서 평가된 두 개의 가장 빠른 C/C++ AD 도구보다 Enzyme가 모든 벤치마크에서 뛰어난 성능을 보이며 최신 기술 수준의 성능을 입증했다.
• Enzyme가 사용하는 재귀 테입 구조는 스택 기반 테입 대비 메모리 오버헤드를 줄여 FFT 및 Bruss 벤치마크에서 뛰어난 성능 기여를 했다.
• Enzyme는 소스 변환 기반 AD 도구가 처리할 수 없는 사전 컴파일된 라이브러리 및 헤더 전용 코드를 직접 미분할 수 있다.
• Enzyme는 Euler 및 RK4 해법기, LSTM, GMM과 같은 복잡한 과학계산 워크로드를 성능이 특수 도구와 경쟁하거나 뛰어나게 미분할 수 있었다.