[논문 리뷰] Quantized Distributed Training of Large Models with Convergence Guarantees

Ilia Markov, Adrian Vladu외 2명|arXiv (Cornell University)|2023. 02. 05.

Advanced Neural Network Applications인용 수 5

한 줄 요약

이 논문은 이론적 수렴 보장을 제공하는 통신 효율적인 Fully-Sharded Data Parallel (FSDP) 훈련의 변종인 QSDP를 제안한다. QSDP는 기울기와 모델 가중치를 모두 완전히 양자화할 수 있도록 하며, 모든-to-all 통신 중에 비편향 양자화를 적용함으로써 FSDP의 통신 병목 현상을 제거한다. 이로 인해 모델 정확도를 유지하면서도 GPT 모델(최대 13억 파라미터)에서 종합적인 속도 향상 최대 2.2배를 달성한다.

ABSTRACT

Communication-reduction techniques are a popular way to improve scalability in data-parallel training of deep neural networks (DNNs). The recent emergence of large language models such as GPT has created the need for new approaches to exploit data-parallelism. Among these, fully-sharded data parallel (FSDP) training is highly popular, yet it still encounters scalability bottlenecks. One reason is that applying compression techniques to FSDP is challenging: as the vast majority of the communication involves the model's weights, direct compression alters convergence and leads to accuracy loss. We present QSDP, a variant of FSDP which supports both gradient and weight quantization with theoretical guarantees, is simple to implement and has essentially no overheads. To derive QSDP we prove that a natural modification of SGD achieves convergence even when we only maintain quantized weights, and thus the domain over which we train consists of quantized points and is, therefore, highly non-convex. We validate this approach by training GPT-family models with up to 1.3 billion parameters on a multi-node cluster. Experiments show that QSDP preserves model accuracy, while completely removing the communication bottlenecks of FSDP, providing end-to-end speedups of up to 2.2x.

연구 동기 및 목표

• 대규모 언어 모델의 Fully-Sharded Data-Parallel (FSDP) 훈련에서 발생하는 빈번한 all-to-all 가중치 교환으로 인한 통신 병목 현상을 해결하기 위해.
• 모델 가중치와 기울기를 모두 양자화하면서도 수렴성이나 모델 정확도를 손상시키지 않고 FSDP에서 가능하게 하기 위해.
• 비볼록인 양자화된 점들로 이루어진 격자 위에서 훈련함에도 불구하고, 전체 모델 상태의 완전한 양자화 하에 SGD의 이론적 수렴 보장을 제공하기 위해.
• 기존 딥러닝 프레임워크(예: PyTorch)에 원활하게 통합될 수 있는 실용적이고 저비용의 구현을 설계하기 위해.

제안 방법

• 모든-to-all 통신 단계 이전에 기울기와 가중치에 모두 양자화를 적용하는 FSDP의 변형인 QSDP를 제안한다.
• 가중치에 대해 비편향 양자화 연산자를 사용하며, 이는 근접한 양자화 수준로의 확률적 반올림으로 정의되며, 기대값에서 편향 오차가 0이 되도록 보장한다.
• 기울기 양자화에 표준적인 비편향 기울기 압축 기법(예: Alistarh 등, 2017)을 적용한다.
• 훈련 과정을 양자화를 투영 수단으로 사용하는 희박 복구 문제로 모델링함으로써 이론적 수렴 분석을 가능하게 한다.
• 표준적인 부드러움과 유한 기울기 가정 하에, QSDP가 양자화된 격자점들 위에서 손실 함수의 최소화점을 향해 수렴함을 보여줌으로써 수렴 보장을 도출한다.
• PyTorch에 효율적으로 구현하여 메모리 및 계산 오버헤드를 최소화하였으며, 노드당 추가 메모리나 복잡한 오차 보정이 필요하지 않다.

실험 결과

연구 질문

• RQ1 FSDP에서 모델 가중치와 기울기를 완전히 양자화해도 모델 수렴성이나 정확도가 저하되지 않을 수 있는가?
• RQ2 편향된 기울기 추정기를 사용하고도 비볼록인 양자화된 가중치 격자 위에서 훈련할 경우, 이론적 수렴 보장을 확보할 수 있는가?
• RQ3 양자화를 통해 FSDP의 통신 병목 현상을 완전히 제거할 수 있으며, 종단 간 훈련 속도 향상을 유지할 수 있는가?
• RQ4 QSDP는 최대 13억 파라미터를 가진 대규모 언어 모델인 GPT와 같은 모델에서 실제로 어떻게 성능을 발휘하는가?
• RQ5 FSDP 환경에서 분산 훈련 시, 양자화 정밀도와 수렴 안정성 사이의 상호 교환 관계는 어떠한가?

주요 결과

• QSDP는 최대 13억 파라미터를 가진 GPT 계열 모델을 훈련할 때, 표준 FSDP 대비 멀티노드 클러스터에서 종합적인 속도 향상 최대 2.2배를 달성한다.
• 모델 정확도가 평가된 모든 모델에서 그대로 유지되며, 13억 파라미터 GPT 모델에서도 성능 저하가 측정되지 않는다.
• 표준 FSDP의 통신 병목 현상은 QSDP에 의해 완전히 제거되며, 양자화로 인해 대역폭 요구 사항이 크게 감소하기 때문이다.
• 이론적 분석을 통해 QSDP가 표준 부드러움과 유한 기울기 가정 하에 양자화된 격자점들 위에서 손실 함수의 최소화점을 향해 수렴함을 증명하였다.
• 가중치의 양자화는 기대값이 원래 가중치와 일치하도록 보장하는 비편향 확률적 반올림 연산자로 구현되었다.
• 구현 과정에서 실질적인 추가 계산 또는 메모리 오버헤드가 발생하지 않아, 생산 규모의 훈련에 실용적이다.