[논문 리뷰] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Samyam Rajbhandari, Jeff Rasley외 2명|arXiv (Cornell University)|2019. 10. 04.

Advanced Neural Network Applications참고 문헌 21인용 수 72

한 줄 요약

ZeRO는 메모리 효율적인 옵티마이저 상태, 그래디언트, 파라미터의 파티셔닝(여기에 잔여 메모리 최적화 포함)을 도입하여 데이터 병렬성으로 1조 파라미터 규모의 모델 학습을 가능하게 하고 대규모 속도 향상을 달성하며 대중화된 대형 모델 학습을 가능하게 한다.

ABSTRACT

Large deep learning models offer significant accuracy gains, but training billions to trillions of parameters is challenging. Existing solutions such as data and model parallelisms exhibit fundamental limitations to fit these models into limited device memory, while obtaining computation, communication and development efficiency. We develop a novel solution, Zero Redundancy Optimizer (ZeRO), to optimize memory, vastly improving training speed while increasing the model size that can be efficiently trained. ZeRO eliminates memory redundancies in data- and model-parallel training while retaining low communication volume and high computational granularity, allowing us to scale the model size proportional to the number of devices with sustained high efficiency. Our analysis on memory requirements and communication volume demonstrates: ZeRO has the potential to scale beyond 1 Trillion parameters using today's hardware. We implement and evaluate ZeRO: it trains large models of over 100B parameter with super-linear speedup on 400 GPUs, achieving throughput of 15 Petaflops. This represents an 8x increase in model size and 10x increase in achievable performance over state-of-the-art. In terms of usability, ZeRO can train large models of up to 13B parameters (e.g., larger than Megatron GPT 8.3B and T5 11B) without requiring model parallelism which is harder for scientists to apply. Last but not the least, researchers have used the system breakthroughs of ZeRO to create the world's largest language model (Turing-NLG, 17B parameters) with record breaking accuracy.

연구 동기 및 목표

• 트릴리언(10의12제곱) 파라미터 모델을 학습해야 할 필요성을 동기화하고 기존의 데이터-모델-파이프라인 병렬성에서의 메모리 병목 현상을 식별한다.
• 효율성을 유지하면서 모델 상태와 잔여 메모리의 중복을 제거하도록 ZeRO를 제안한다.
• 가용 하드웨어로 모델 크기를 확장하기 위한 세 가지 파티셔닝 단계(옵티마이저 상태, 그래디언트, 파라미터)와 잔여 메모리 최적화를 개발·분석한다.
• 실용적인 학습 개선을 입증하고 MP와의 통합을 보여주며 대형 모델 학습의 대중화를 위한 오픈 소스 도구를 제공한다.

제안 방법

• 데이터 병렬 프로세스(P_os) 간에 옵티마이저 상태를 분할하여 각 랭크가 상태의 1/N_d를 보유하도록 한다.
• 각 파라미터 파티션(P_g)에 대응하는 그래디언트를 분할하고 파티션 경계에서 버켓화된 reduce-scatter를 수행한다.
• 파라미터를 데이터 병렬 프로세스(P_p) 간에 분할하고 비국부 파라미터의 필요 시 브로드캐스트를 통해 순전파/역전파에 사용한다.
• ZeRO-DP를 ZeRO-R과 결합하여 활성화도 분할하고 잔여 메모리(체크포인트, 버퍼, 단편화)를 관리하며 선택적 CPU 오프로딩(P_a, P_a+cpu)을 포함한다.
• 연산과 메모리 절약을 분리하면서 통신을 최소화하는 동적 통신 스케줄을 사용하고 통신 오버헤드를 분석한다( P_p 사용 시 최대 1.5배).
• 유리할 때 MP와의 결합을 지원하되 다수의 경우 ZeRO-DP가 MP에 필적하거나 이를 능가하고 광범위한 모델 리팩토링 없이도 학습을 가능하게 함을 보여준다.

실험 결과

연구 질문

• RQ1 ZeRO가 모델 상태를 분할하여 DP 효율성을 해치지 않으면서 각 디바이스의 메모리를 크게 감소시킬 수 있을까?
• RQ2 데이터 병렬도 N_d의 함수로 옵티마이저 상태, 그래디언트, 파라미터에 대해 얼마나 많은 메모리가 절약될 수 있을까?
• RQ3 잔여 메모리(활성화, 버퍼, 단편화)를 어떻게 관리하여 아주 큰 모델을 더 가능하게 할 수 있을까?
• RQ4 ZeRO-DP와 ZeRO-R을 적용했을 때 통신량과 확장성에 미치는 결과적 영향은 무엇인가?
• RQ5 ZeRO가 수백억~조 파라미터 수준의 학습을 어떻게 가능하게 하며 실제 하드웨어에서의 실용적 속도 향상은 어떠한가?

주요 결과

DP 차수	7.5B 모델 (GB) P_os	7.5B 모델 (GB) P_os+g	7.5B 모델 (GB) P_os+g+p	128B 모델 (GB) P_os	128B 모델 (GB) P_os+g	128B 모델 (GB) P_os+g+p	1T 모델 (GB) P_os	1T 모델 (GB) P_os+g	1T 모델 (GB) P_os+g+p
1	120	120	120	2048	2048	2048	16000	16000	16000
4	52.5	41.3	30	896	704	512	7000	5500	4000
16	35.6	21.6	7.5	608	368	128	4750	2875	1000
64	31.4	16.6	1.88	536	284	32	4187	2218	250
256	30.4	15.4	0.47	518	263	8	4046	2054	62.5
1024	30.1	15.1	0.12	513	257	2	4011	2013	15.6

• ZeRO는 전통적인 DP/MP 한계보다 훨씬 더 큰 모델을 데이터 병렬 학습으로 적합하게 해주며 파라미터와 그래디언트의 메모리 감소는 N_d에 비례하는 수준까지, 옵티마이저 상태의 감소는 단계에 따라 4배~8배까지 가능하다.
• 세 가지 단계(P_os, P_g, P_p)로 각 디바이스의 메모리 풋프린트를 크게 줄여 대형 GPU 클러스터에서 트릴리언 파라미터 모델을 가능하게 한다.
• ZeRO-DP를 64방 DP로 구성하면 7.5B 모델을 훨씬 낮은 메모리로 학습할 수 있으며 P_os+g+p를 사용하면 7.5B 모델의 메모리가 대략 1.88 GB로 축소되어 훨씬 더 큰 규모를 가능하게 한다.
• ZeRO-R은 잔여 활성화 메모리를 추가로 감소시키고 효율성을 위해 상수 크기 버퍼를 사용하며 조각화로 인한 메모리 초과를 완화하기 위해 즉시 메모리 단편화 제거를 수행한다.
• 구현 결과는 100B 파라미터 모델을 400개의 Nvidia V100 GPU에서 학습하고 GPU당 38 TFlops 이상, 총 성능이 15 페타플롭스를 넘는 것을 보여주며, 당시 기준으로 최대 모델 크기가 최대 8배 커지고 달성 가능한 성능이 약 10배 더 높아진다고 주장한다.
• ZeRO는 대형 모델 학습의 민주화를 위해 DeepSpeed 라이브러리의 일부로 공개되었으며 더 나아가 추가 분할(P_os+g+p)로 1조 파라미터까지 확장하는 것을 목표로 한다.