[논문 리뷰] ZeRO++: Extremely Efficient Collective Communication for Giant Model Training

Guanhua Wang, Heyang Qin외 7명|arXiv (Cornell University)|2023. 06. 16.

Advanced Neural Network Applications인용 수 15

한 줄 요약

ZeRO++은 세 가지 통신 감소 기술(qwZ, hpZ, qgZ)을 도입하여 ZeRO의 노드 간 데이터 이동을 최대 4배까지 줄이고 거대 모델 학습의 처리량을 높입니다.

ABSTRACT

Zero Redundancy Optimizer (ZeRO) has been used to train a wide range of large language models on massive GPUs clusters due to its ease of use, efficiency, and good scalability. However, when training on low-bandwidth clusters, or at scale which forces batch size per GPU to be small, ZeRO's effective throughput is limited because of high communication volume from gathering weights in forward pass, backward pass, and averaging gradients. This paper introduces three communication volume reduction techniques, which we collectively refer to as ZeRO++, targeting each of the communication collectives in ZeRO. First is block-quantization based all-gather. Second is data remapping that trades-off communication for more memory. Third is a novel all-to-all based quantized gradient averaging paradigm as replacement of reduce-scatter collective, which preserves accuracy despite communicating low precision data. Collectively, ZeRO++ reduces communication volume of ZeRO by 4x, enabling up to 2.16x better throughput at 384 GPU scale.

연구 동기 및 목표

• 대규모 언어 모델을 다중-GPU 클러스터에서 제한된 노드 간 대역폭으로 효율적으로 학습하도록 동기를 부여한다.
• 전방 가중치 수집, 역방향 가중치 수집, 그래디언트 축소에 걸친 ZeRO의 통신 오버헤드를 줄인다.
• 3D 병렬화로의 리팩토링이나 수렴 정확도의 희생 없이 더 높은 처리량과 확장성을 가능하게 한다.

제안 방법

• 세 가지 통신 최적화 기법 도입: qwZ(전방 all-gather를 위한 블록 기반 INT8 가중치 양자화), hpZ(노드 내 보조 가중치 분할로 노드 간 역방향 all-gather 제거), qgZ(텐서 조각 재배치를 통한 all-to-all 기반 INT4 그래디언트 축소).
• 정확도 유지를 위한 블록 기반 양자화 사용 및 양자화, 디양자화, 연산자 융합을 위한 고성능 CUDA 커널 구현.
• 노드 내 우선 통신 후 노드 간 통신으로 상호 all-to-all 그래디언트 축소를 두 단계로 채택하여 노드 간 트래픽 최소화.
• 지연을 숨기기 위해 계산과 통신을 중첩하고 메모리 트래픽을 최소화하기 위해 커널 융합을 수행한다.

Figure 1. Large scale training throughput are constrained by network bandwidth and batch size per GPU

실험 결과

연구 질문

• RQ1 매우 큰 모델에 대해 training 정확도를 희생하지 않으면서 ZeRO의 디바이스 간 통신을 어떻게 감소시킬 수 있는가?
• RQ2 ZeRO-3에서 전방 및 역방향 통신과 그래디언트 집계를 최소화하기 위해 어떤 양자화 및 분할 기법 조합이 효과적인가?
• RQ3 낮은 정밀도의 데이터를 기반으로 모든-대-모아(all-to-all) 그래디언트 축소에 의해 수렴을 보존하면서 통신 볼륨을 낮출 수 있는 새로운 축적(paradigm)인가?

주요 결과

통신	전방 all-gather	역방향 all-gather	역방향 reduce-scatter
ZeRO-3	M	M	M
ZeRO++	0.5M	0	0.25M

• ZeRO++는 노드 간 통신 볼륨을 한 학습 반복당 3M에서 0.75M으로 줄여 처리량을 향상시킨다.
• 10B에서 138B 매개변수 모델 전반에 걸쳐 ZeRO 기준선 대비 엔드-투-엔드 처리량을 최대 2.4배 향상시킨다.
• GPT-3 유사 모델로 최대 384개의 GPU까지 확장 가능함을 보여주고, 지속적 피크 처리량의 45% 이상을 달성한다.
• 모델 수렴 및 학습 정확도가 ZeRO 기준선과 비교해 유지된다.
• 저대역대 설정에서 ZeRO++는 고대역폭 기준 구성과 유사한 처리량을 달성한다.

Figure 2. Illustration & example of block based quantization vs. baseline