[논문 리뷰] ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning

Samyam Rajbhandari, Olatunji Ruwase외 3명|arXiv (Cornell University)|2021. 04. 16.

Advanced Neural Network Applications참고 문헌 30인용 수 57

한 줄 요약

ZeRO-Infinity는 모델 상태와 활성화를 CPU/NVMe로 오프로드하고 메모리 중심 타일링으로 제한된 GPU 자원에서 트릴리언-파라미터 모델을 훈련시키며 높은 처리량과 접근성을 달성한다.

ABSTRACT

In the last three years, the largest dense deep learning models have grown over 1000x to reach hundreds of billions of parameters, while the GPU memory has only grown by 5x (16 GB to 80 GB). Therefore, the growth in model scale has been supported primarily though system innovations that allow large models to fit in the aggregate GPU memory of multiple GPUs. However, we are getting close to the GPU memory wall. It requires 800 NVIDIA V100 GPUs just to fit a trillion parameter model for training, and such clusters are simply out of reach for most data scientists. In addition, training models at that scale requires complex combinations of parallelism techniques that puts a big burden on the data scientists to refactor their model. In this paper we present ZeRO-Infinity, a novel heterogeneous system technology that leverages GPU, CPU, and NVMe memory to allow for unprecedented model scale on limited resources without requiring model code refactoring. At the same time it achieves excellent training throughput and scalability, unencumbered by the limited CPU or NVMe bandwidth. ZeRO-Infinity can fit models with tens and even hundreds of trillions of parameters for training on current generation GPU clusters. It can be used to fine-tune trillion parameter models on a single NVIDIA DGX-2 node, making large models more accessible. In terms of training throughput and scalability, it sustains over 25 petaflops on 512 NVIDIA V100 GPUs(40% of peak), while also demonstrating super linear scalability. An open source implementation of ZeRO-Infinity is available through DeepSpeed, a deep learning optimization library that makes distributed training easy, efficient, and effective.

연구 동기 및 목표

• 극대 규모 DL에서의 GPU 메모리 벽 과제를 설명하고 리팩토링 없이 도달 가능한 모델 규모를 증가시키는 시스템을 고무하는 동기를 제시한다.
• 기존 클러스터에 거대한 모델을 맞추기 위한 ZeRO-Infinity의 이질적 메모리 오프로드 및 타일링 기법을 도입한다.
• 트릴리언-파라미터 규모를 포함하여 매우 큰 모델에서의 훈련 효율성과 확장성을 입증한다.
• 모델 리팩토링이나 복잡한 병렬처리의 필요성을 줄이거나 제거하는 사용성 개선을 평가한다.
• 오픈 소스 구현 세부 정보를 제공하고 향후 하드웨어 설계에 대한 시사점을 논의한다.

제안 방법

• GPU, CPU, NVMe 메모리를 동시에 활용하는 Infinity 오프로드 엔진을 제안한다.
• 작업 메모리 감소를 위해 대형 연산자를 타일로 분해하는 메모리 중심 타일링을 도입한다.
• 장치 간의 집합적 메모리 대역폭을 활용하기 위한 대역폭 중심 분할을 정의한다.
• 이종 메모리 간의 통신과 계산을 겹치며 숨기는 오버랩 중심 설계를 개발한다.
• 임의의 아키텍처를 훈련하는 동안 모델 리팩토링을 피하도록 사용 편의성에 영감을 받은 구현을 제공한다.
• Transformer 기반 모델로 대형 모델 훈련에 필요한 메모리와 대역폭 요구사항을 특징화한다.

실험 결과

연구 질문

• RQ1 실용적인 GPU 클러스터에서 수백 조 파라미터의 모델을 훈련하기 위해 GPU 메모리 벽을 어떻게 허물 수 있는가?
• RQ2 CPU와 NVMe 메모리를 활용하여 거대한 GPU 클러스터가 없는 연구자들에게 대형 모델을 접근 가능하게 만들 수 있는가?
• RQ3 모델 코드를 리팩토링하거나 복잡한 3D 병렬처리를 채택하지 않고도 트릴리언-파라미터 모델을 훈련하는 것이 가능한가?
• RQ4 이질적 메모리로 지속 가능한 효율적 훈련을 위해 필요한 메모리와 대역폭 특성은 무엇인가?
• RQ5 현재 하드웨어에서 ZeRO-Infinity의 처리량, 확장성 및 접근성 측면의 성능은 어떠한가?

주요 결과

• ZeRO-Infinity는 32 trillion 파라미터 모델을 32 NVIDIA DGX-2 노드(512 GPUs)에서 훈련시킬 수 있다.
• 동일 하드웨어에서 25 PFLOPS 이상의 처리량을 지속한다(피크의 40%).
• 트릴리언-파라미터 모델에 대해 초선형 확장성을 보여준다.
• 모델 병렬화나 리팩토링 없이 단일 DGX-2 노드에서 최대 트릴리언 파라미터에 이르는 미세 조정을 가능하게 한다.
• 메모리 중심 타일링은 연산자당 작업 메모리를 줄여 모델 병렬화 없이도 매우 큰 계층을 적합하게 한다.
• ZeRO-Infinity의 오픈 소스 구현은 DeepSpeed에서 사용할 수 있다.