[논문 리뷰] MiCS: Near-linear Scaling for Training Gigantic Model on Public Cloud

Zhen Zhang, Shuai Zheng외 6명|arXiv (Cornell University)|2022. 04. 30.

Stochastic Gradient Optimization Techniques인용 수 6

한 줄 요약

MiCS는 스케일 인식 모델 분할, 계층적 통신, 2단계 그래디언트 동기화를 통해 통신 규모를 최소화하여 공용 클라우드에서 거대한 딥러닝 모델을 거의 선형적으로 확장 가능한 훈련 시스템을 제안한다. 1000억 파라미터 모델을 512개 GPU에서 훈련할 때, 2.89배 높은 처리량과 99.4%의 약한 스케일링 효율을 달성하며, 이는 이질적이고 대역폭이 낮은 클라우드 환경에서 DeepSpeed ZeRO를 능가한다.

ABSTRACT

Existing general purpose frameworks for gigantic model training, i.e., dense models with billions of parameters, cannot scale efficiently on cloud environment with various networking conditions due to large communication overheads. In this paper, we propose MiCS, which Minimizes the Communication Scale to bring down communication overhead. Specifically, by decreasing the number of participants in a communication collective, MiCS can utilize heterogeneous network bandwidth, reduce network traffic over slower links, reduce the latency of communications for maintaining high network bandwidth utilization, and amortize expensive global gradient synchronization overhead. Our evaluation on AWS shows that the system throughput of MiCS is up to 2.89$ imes$ that of the state-of-the-art large model training systems. MiCS achieves near-linear scaling efficiency, which is up to 1.27$ imes$ that of DeepSpeed. MiCS allows us to train a proprietary model with 100 billion parameters on 512 GPUs with 99.4% weak-scaling efficiency, and it is able to saturate over 54.5% theoretical computation power of each GPU on a public cloud with less GPU memory and more restricted networks than DGX-A100 clusters.

연구 동기 및 목표

• 공용 클라우드에서 이질적이고 대역폭이 낮은 네트워크를 사용할 때 기존의 거대 모델 훈련 프레임워크에서 발생하는 높은 통신 오버헤드 문제를 해결하기 위해.
• 고속 InfiniBand 인터커넥트가 없는 공용 클라우드 클러스터에서 훈련 처리량과 확장성을 향상시키기 위해.
• 특수 하드웨어나 복잡한 모델 병렬화 없이도 표준 공용 클라우드 GPU 인스턴스에서 거대 모델(예: 1000억 파라미터)을 효율적으로 훈련시킬 수 있도록 하기 위해.
• 대규모 이질적 클러스터에서 통신 비용을 크게 줄이면서도 데이터 병렬화의 단순성과 일반성을 유지하기 위해.
• 제한된 네트워크 조건에서도 고성능 클러스터와 유사한 거의 선형 스케일링 효율을 달성하기 위해.

제안 방법

• 스케일 인식 모델 분할을 도입하여 GPU를 더 작은 고립된 그룹으로 나누며, 각 그룹이 모델 상태의 완전한 복제본을 보유함으로써 빈번한 파라미터 동기화에 참여하는 참가자의 수를 줄인다.
• 계층적 통신 전략을 적용하여 집합 연산을 내부 그룹 수준로 제한하고 그룹 간 통신을 병렬화함으로써 느린 노드 간 링크의 트래픽을 줄인다.
• 2단계 그래디언트 동기화 메커니즘을 설계하여 전역 all-reduce 연산을 피하고, 먼저 그룹 내에서 그래디언트를 집계한 후 그룹 수준 요약을 그룹 간으로 동기화함으로써 전역 동기화 비용을 줄인다.
• GPU 메모리 용량과 네트워크 토폴로지에 기반하여 파artition 그룹을 구성하여 메모리 사용과 통신 효율성을 균형 잡고, 그룹 크기를 조정 가능한 파라미터로 설정한다.
• 기존의 데이터 병렬화 프레임워크(예: ZeRO)를 활용하지만, 통신 패턴을 재구조화하여 노드 내 고대역폭 링크를 활용하고 노드 간 지연과 데이터 양을 줄인다.
• 모델 상태 복제에 최적의 그룹 크기를 결정하기 위한 휴리스틱을 사용하여 메모리 효율성과 통신 병목 현상을 최소화한다.

실험 결과

연구 질문

• RQ1 데이터 병렬 훈련 시스템이 이질적이고 대역폭이 낮은 네트워크를 가진 공용 클라우드 클러스터에서 거의 선형 스케일링을 달성할 수 있는가?
• RQ2 InfiniBand과 같은 전용 하드웨어에 의존하지 않고 거대 모델 훈련에서 통신 오버헤드를 어떻게 줄일 수 있는가?
• RQ3 통신 참가자 수를 줄이는 것이 대규모 GPU 클러스터에서 훈련 처리량과 확장성에 얼마나 기여하는가?
• RQ4 복잡한 모델 또는 파이프라인 병렬화를 피하면서도 높은 효율성과 낮은 통신 비용을 유지할 수 있는가?
• RQ5 제한된 GPU 메모리와 변동하는 네트워크 대역폭 등의 실제 공용 클라우드 제약 조건 하에서 시스템의 성능은 어떠한가?

주요 결과

• MiCS는 AWS에서 DeepSpeed ZeRO보다 최대 2.89배 높은 시스템 처리량을 달성하여 대규모 훈련에서 뚜렷한 성능 향상을 보였다.
• 1000억 파라미터 모델을 512개 GPU에서 훈련할 때 99.4%의 약한 스케일링 효율을 달성하여 거의 선형 확장성을 입증했다.
• 공용 클라우드 인스턴스에서 각 GPU의 이론적 계산 용량의 54.5% 이상을 활용했으며, GPU당 메모리 사용률이 32% 미만이었고, 동일 조건에서 ZeRO를 능가하는 성능을 보였다.
• 계층적 통신 전략이 노드 간 네트워크 트래픽과 지연을 줄여 느린 노드 간 링크에서 더 높은 대역폭 활용도를 가능하게 했다.
• 2단계 그래디언트 동기화 메커니즘이 전역 all-reduce 연산을 그룹 수준 요약으로 제한함으로써 대규모 클러스터에서의 효율성을 향상시켰다.
• 제한된 GPU 메모리와 제한된 네트워크 대역폭 조건에서도 높은 성능을 유지하여 표준 공용 클라우드 배포에 실용적임을 입증했다.