[논문 리뷰] ZeRO-Offload: Democratizing Billion-Scale Model Training

Jie Ren, Samyam Rajbhandari외 6명|arXiv (Cornell University)|2021. 01. 18.

Advanced Neural Network Applications참고 문헌 27인용 수 61

한 줄 요약

ZeRO-Offload는 그래디언트, 옵티마이저 상태, 및 옵티마이저 계산을 CPU로 오프로드하여 단일 GPU에서 최대 13B 파라미터 규모의 모델을 학습시키고 확장 가능한 성능과 PyTorch와의 원활한 통합을 제공합니다.

ABSTRACT

Large-scale model training has been a playing ground for a limited few requiring complex model refactoring and access to prohibitively expensive GPU clusters. ZeRO-Offload changes the large model training landscape by making large model training accessible to nearly everyone. It can train models with over 13 billion parameters on a single GPU, a 10x increase in size compared to popular framework such as PyTorch, and it does so without requiring any model change from the data scientists or sacrificing computational efficiency. ZeRO-Offload enables large model training by offloading data and compute to CPU. To preserve compute efficiency, it is designed to minimize the data movement to/from GPU, and reduce CPU compute time while maximizing memory savings on GPU. As a result, ZeRO-Offload can achieve 40 TFlops/GPU on a single NVIDIA V100 GPU for 10B parameter model compared to 30TF using PyTorch alone for a 1.4B parameter model, the largest that can be trained without running out of memory. ZeRO-Offload is also designed to scale on multiple-GPUs when available, offering near linear speedup on up to 128 GPUs. Additionally, it can work together with model parallelism to train models with over 70 billion parameters on a single DGX-2 box, a 4.5x increase in model size compared to using model parallelism alone. By combining compute and memory efficiency with ease-of-use, ZeRO-Offload democratizes large-scale model training making it accessible to even data scientists with access to just a single GPU.

연구 동기 및 목표

• 접근 가능한 수십억 파라미터 모델 훈련의 필요성과 하드웨어 장벽 감소를 동기화합니다.
• GPU 메모리 절약을 극대화하면서 CPU 계산 및 CPU-GPU 통신을 최소화하는 독특한 오프로드 전략을 제안합니다.
• ZeRO-기반 데이터 병렬화와 가능한 모델 병렬화를 사용하여 단일 GPU에서 최대 128 GPU까지의 확장성을 입증합니다.
• 빠른 Adam을 활용한 최적화된 CPU 실행과 처리량을 보존하는 스케줄링 기법을 제공합니다.

제안 방법

• 모델 훈련을 데이터 흐름 그래프로 표현하고 CPU와 GPU 사이에서 계산, 통신, 메모리 절감을 최적화하도록 분할합니다.
• 오프로드 전략은 GPU에 FP16 매개변수를 남겨두고 FP16 그래디언트 및 FP32 옵티마이저 상태를 CPU로 오프로드하며 매개변수 업데이트를 CPU에서 수행합니다.
• 분할 분석은 최소 CPU 계산 및 최소 CPU-GPU 통신을 보여주며 최적의 오프로드 결정을 도출합니다.
• 오프로드를 ZeRO-Stage-2 데이터 병렬화와 결합하여 128 GPUs까지 거의 선형 스케일링을 달성하고 더 큰 모델에 대해 모델 병렬화와의 호환성을 제공합니다.
• SIMD, 루프 언롤링, 다중 스레딩을 사용하는 최적화된 CPU Adam 구현을 사용하고 CPU와 GPU 작업을 중첩시키기 위한 한 단계 지연 매개변수 업데이트를 도입합니다.
• 지연 매개변수 업데이트(DPU) 옵션은 CPU 업데이트와 GPU 계산을 겹치게 하여 처리량을 개선하고 정확도 손실 없이 성능을 향상시킵니다.

실험 결과

연구 질문

• RQ1 ZeRO-Offload가 단일 GPU 메모리에 맞지 않는 더 큰 모델의 학습을 효율 손실 없이 가능하게 할 수 있는가?
• RQ2 CPU로의 오프로드가 ZeRO 데이터 병렬화와 어떻게 상호 작용하여 다수의 GPU에 걸쳐 확장되나?
• RQ3 단일/다중 GPU에서 수십억 파라미터 모델을 학습할 때Practical한 처리량과 메모리 절감은 얼마나 되는가?
• RQ4 최적화된 CPU Adam과 DPU가 성능을 향상시키면서 모델 수렴을 유지하는가?

주요 결과

• 단일 NVIDIA V100 GPU에서 40 TFLOPS의 성능으로 최대 13B 파라미터의 모델 학습이 가능하며, 오프로드 없이도 1.2B까지 가능하다는 결과가 나타납니다.
• ZeRO-기반 데이터 병렬화와 결합했을 때 최대 128 GPUs에서 거의 선형 확장성을 보입니다.
• DGX-2 노드에서 ZeRO-Offload와 모델 병렬화를 결합하면 최대 70B 파라미터까지 학습 가능.
• CPU-Adam 최적화는 표준 PyTorch Adam 구현 대비 6배 이상 속도 향상을 제공하며, 지연 매개변수 업데이트를 통해 전체 처리량이 최대 1.5배 향상됩니다.
• 이 접근법은 모델 리팩토링 없이도 최소한의 CPU 계산 오버헤드와 제한된 통신으로 10배 더 큰 모델의 학습을 가능하게 합니다.
• 오프로드 전략은 CPU 계산 및 CPU-GPU 통신을 최소화하면서 GPU 메모리 절감을 극대화하는 정의된 목표 하에서 독창적이고 최적입니다.