[논문 리뷰] Data Movement Is All You Need: A Case Study on Optimizing Transformers

Andrei Ivanov, Nikoli Dryden외 3명|arXiv (Cornell University)|2020. 06. 30.

Advanced Neural Network Applications참고 문헌 96인용 수 24

한 줄 요약

이 논문은 트랜스포머 학습에서 데이터 이동을 주요 성능 저하 원인으로 규명하며, 계산 성능 향상에도 불구하고 애뮬러의 법칙에 따라 학습 과정이 메모리 기반으로 전환되었음을 입증한다. 체계적인 데이터플로우 분석과 레이아웃 변환을 통해 전역적으로 데이터 레이아웃과 이동을 최적화함으로써, 저자들은 데이터 이동을 최대 22.91% 감소시키고, 상태기반 프레임워크 대비 BERT 인코더 레이어에서 1.30×, 전체 BERT 모델에서 1.19×의 성능 향상을 달성한다.

ABSTRACT

Transformers are one of the most important machine learning workloads today. Training one is a very compute-intensive task, often taking days or weeks, and significant attention has been given to optimizing transformers. Despite this, existing implementations do not efficiently utilize GPUs. We find that data movement is the key bottleneck when training. Due to Amdahl's Law and massive improvements in compute performance, training has now become memory-bound. Further, existing frameworks use suboptimal data layouts. Using these insights, we present a recipe for globally optimizing data movement in transformers. We reduce data movement by up to 22.91% and overall achieve a 1.30x performance improvement over state-of-the-art frameworks when training a BERT encoder layer and 1.19x for the entire BERT. Our approach is applicable more broadly to optimizing deep neural networks, and offers insight into how to tackle emerging performance bottlenecks.

연구 동기 및 목표

• 트랜스포머 모델 학습에서의 주요 성능 저하 원인을 규명하는 것.
• 고도로 최적화된 커널에도 불구하고 기존 딥러닝 프레임워크가 GPU 계산 자원을 효율적으로 활용하지 못하는 이유를 분석하는 것.
• 데이터 이동과 비효율적인 데이터 레이아웃이 학습 성능 저하에 미치는 영향을 조사하는 것.
• 전체 계산 그래프에 걸쳐 데이터 이동을 체계적으로 감소시키는 전역 최적화 전략을 개발하는 것.
• 데이터 이동 최적화가 알고리즘적 또는 하드웨어 수준의 개선만으로도 달성할 수 있는 성능 향상보다 뛰어나다는 것을 입증하는 것.

제안 방법

• 저자들은 상태 기반 데이터플로우 다중그래프(SDFG)로 트랜스포머 계산을 모델링하여, 다양한 연산 간의 데이터 이동을 체계적으로 분석할 수 있도록 한다.
• 메모리 기반 연산자, 특히 다중 헤드 어텐션에서 데이터 이동에 37%의 런타임이 소요됨을 규명하기 위해 세밀한 데이터플로우 분석을 수행한다.
• 이 방법은 전역적인 레이아웃 변환 전략을 도입하여, 다양한 연산 간의 데이터 레이아웃을 재정렬하고 재구조화하여 중복된 데이터 이동을 최소화한다.
• 이 접근법은 다중 수준의 분할 해상도에서 근본적이고 세밀한 레이아웃 및 융합 최적화를 가능하게 하는 파arametric하고 평탄한 데이터플로우 표현을 사용한다.
• 저자들은 상태기반 SDFG 기반 컴파일러 스택을 사용하여 성능을 평가하였으며, PyTorch, XLA 기반 텐서플로우, cuDNN와의 성능 비교를 수행하였다.
• 메모리 액세스 오버헤드를 감소시키기 위해 레이아웃 변환을 적용하여, 행렬-행렬 곱셈(MMM) 성능을 최대 52% 향상시켰다.

실험 결과

연구 질문

• RQ1 현대 트랜스포머 모델 학습에서의 주요 성능 저하 원인은 무엇인가?
• RQ2 고도로 최적화된 커널에도 불구하고 기존 딥러닝 프레임워크가 GPU 계산 자원을 효율적으로 활용하지 못하는 이유는 무엇인가?
• RQ3 데이터 레이아웃은 트랜스포머에서 데이터 이동과 전체 학습 성능에 어떤 영향을 미치는가?
• RQ4 국소적 또는 커널 수준 최적화보다 전역적이고 종단 간 데이터 이동 최적화가 더 뛰어난 성능을 낼 수 있는가?
• RQ5 데이터 이동 감소는 BERT 및 그 구성 요소에서 학습 스루풋을 얼마나 향상시킬 수 있는가?

주요 결과

• 트랜스포머 학습 반복에서 데이터 이동이 37%의 런타임을 차지하고 있음에도 불구하고, 텐서 축소 연산이 전체 산술 연산의 99% 이상을 차지하고 있음.
• 전역적 레이아웃 및 융합 최적화를 통해 데이터 이동을 최대 22.91% 감소시켜 성능 향상에 기여함.
• 제안된 방법은 상태기반 프레임워크 대비 단일 BERT 인코더 레이어에서 1.30×, 전체 BERT 모델에서 1.19×의 성능 향상을 달성함.
• 성능 향상의 주요 원인은 최적화된 데이터 레이아웃으로, 이는 행렬-행렬 곱셈(MMM) 성능을 최대 52% 향상시킴.
• 애뮬러의 법칙과 계산 대비 메모리 대역폭의 비례적 증가로 인해 학습 과정이 메모리 기반으로 전환되었음을 확인함.
• 이 방법은 트랜스포머를 초월해 다른 딥 뉴럴 네트워크의 데이터 이동 최적화를 위한 일반화 가능한 프레임워크를 제공함.