[논문 리뷰] FusionStitching: Boosting Memory Intensive Computations for Deep Learning Workloads

Zhen Zheng, Pengzhan Zhao외 7명|arXiv (Cornell University)|2020. 09. 23.

Advanced Neural Network Applications참고 문헌 46인용 수 12

한 줄 요약

FusionStitching는 메모리 집약적인 딥러닝 워크로드를 최적화하기 위해 복잡한 데이터 의존성과 비균일한 병렬성을 가진 연산자들을 대규모 GPU 커널으로 융합하는 JIT 딥러닝 컴파일러이다. 공유 메모리와 레지스터 셔플을 통해 중간 값 재사용을 가능하게 하고, 외부 메모리 액세스와 컨텍스트 스위치 오버헤드를 줄이며, 최신 기술 대비 최대 2.21배의 성능 향상을 기록했으며 평균 1.45배 향상되었다.

ABSTRACT

We show in this work that memory intensive computations can result in severe performance problems due to off-chip memory access and CPU-GPU context switch overheads in a wide range of deep learning models. For this problem, current just-in-time (JIT) kernel fusion and code generation techniques have limitations, such as rough fusion plan exploration strategies and limited code generation ability. We propose FusionStitching, a deep learning compiler capable of fusing memory intensive operators, with varied data dependencies and non-homogeneous parallelism, into large GPU kernels to reduce global memory access and context switch overhead automatically. FusionStitching widens the range of operation combinations that fusion can target beyond previous JIT works by introducing data reuse of intermediate values. It explores large fusion spaces to decide optimal fusion plans with considerations of memory access costs, kernel calls and resource usage constraints. FusionStitching tunes the optimal stitching scheme with a domain-specific cost model efficiently. Experimental results show that FusionStitching can reach up to 2.21x speedup compared to state-of-the-art, with 1.45x on average. Besides these experimental results, we integrated our approach into a compiler product and deployed it onto a production cluster for AI workloads with thousands of GPUs. The system has been in operation for more than 4 months and saves 7,000 GPU hours on average for approximately 30,000 tasks per month.

연구 동기 및 목표

• 메모리 집약적인 딥러닝 연산자에서 발생하는 외부 메모리 액세스와 CPU-GPU 컨텍스트 스위치로 인한 성능 저하 문제를 해결한다.
• 기존의 JIT 융합 기술의 한계, 즉 좁은 융합 공간과 데이터 재사용 지원 부족 문제를 극복한다.
• 다양한 데이터 의존성과 병렬성 구조를 가진 메모리 집약적인 연산자(예: 원소별 연산, 전치, 감소)의 복잡하고 다양한 조합을 효율적으로 융합할 수 있도록 한다.
• 최적의 융합 계획 선택과 고성능 커널 생성을 위한 확장 가능한 도메인 특화 비용 모델 및 튜닝 프레임워크를 개발한다.
• 생산용 AI 클러스터 환경에 통합하고 검증하여 실제 워크로드에서의 성능 향상과 자원 절감 효과를 입증한다.

제안 방법

• 메모리 집약적인 연산자 간의 복잡한 의존성과 비균일한 병렬성을 고려한 다양한 융합 패턴을 모델링하기 위해 네 가지 스티치 스킴 추상화를 도입한다.
• 공유 메모리와 레지스터 셔플을 활용해 스레드 간 중간 값 재사용을 가능하게 하여 재계산 오버헤드를 줄인다.
• 메모리 액세스 비용, 커널 런처 비용, 자원 사용량을 추정함으로써 융합 공간 탐색을 안내하는 이중 계층의 도메인 특화 비용 모델을 설계한다.
• 큰 융합 공간에서 비용 기반의 지능적인 탐색을 수행하는 융합 탐색기(Explorer)를 구현하여 탐욕적 국소 최적화를 피한다.
• 융합 계획을 최적화된 GPU 커널로 매핑하고 성능을 향상시키기 위한 스케줄 튜닝을 수행하는 코드 생성기(Generator)를 개발한다.
• 동적 형상 지원과 생산 환경 배포를 위해 FusionStitching을 XLA 프레임워크에 통합한다.

실험 결과

연구 질문

• RQ1 메모리 집약적인 딥러닝 연산자에 대한 JIT 융합에서 중간 데이터 재사용을 체계적으로 어떻게 활용할 수 있는가?
• RQ2 복잡하고 동적인 워크로드에서 메모리 액세스, 커널 런처 오버헤드, 자원 사용량 간 최적의 트레이드오프를 달성하기 위해 어떤 융합 패턴과 비용 모델이 필요한가?
• RQ3 고도로 발전한 탐색 및 코드 생성 기능을 갖춘 JIT 융합 프레임워크가 실제 딥러닝 모델에서 룰 기반 또는 탐욕적 접근 방식을 능가할 수 있는가?
• RQ4 중간 값 재사용과 片상 메모리 활용이 외부 메모리 트래픽과 컨텍스트 스위치 오버헤드를 얼마나 줄일 수 있는가?
• RQ5 동적 형상과 높은 연산자 다양성을 가진 대규모 생산 워크로드에서 제안된 융합 프레임워크의 효과성은 어느 정도인가?

주요 결과

• FusionStitching는 실제 워크로드의 딥러닝 모델에서 최신 기술 대비 최대 2.21배의 성능 향상과 평균 1.45배의 성능 향상을 기록했다.
• 다양한 메모리 집약적인 연산자들을 단일 최적화 커널으로 융합함으로써 GPU 커널 런처 오버헤드와 외부 메모리 액세스를 줄였다.
• 공유 메모리와 레지스터 셔플을 통해 중간 값 재사용을 가능하게 하여, XLA와 같은 기존 JIT 도구에서 발생하는 중복 계산 문제를 방지했다.
• 프레임워크는 생산용 XLA 기반 컴파일러 스택에 성공적으로 통합되었으며, 수천 대의 GPU를 포함한 클러스터에서 4개월 이상 운영되었다.
• 배포로 인해 30,000개의 작업에서 월 평균 약 7,000 GPU 시간을 절감하여 실질적인 효율성 향상을 입증했다.
• 도메인 특화 비용 모델은 동적 형상과 복잡한 연산자 조합에서도 낮은 튜닝 오버헤드로 효과적인 융합 계획 선택을 가능하게 했다.