[논문 리뷰] Maximizing Communication Efficiency for Large-scale Training via 0/1 Adam

Yucheng Lu, Conglong Li외 3명|arXiv (Cornell University)|2022. 02. 12.

Stochastic Gradient Optimization Techniques인용 수 4

한 줄 요약

이 논문은 0/1 Adam을 제안하며, 이는 대규모 아담 기반 학습을 위한 통신 효율적인 최적화 방법이다. 이 방법은 오래된 분산과 동등성 추정을 통해 아담 업데이트를 선형화함으로써, 1비트 그래디언트 압축과 로컬 스텝을 동시에 사용할 수 있도록 한다. 이로 인해 128개 GPU에서 통신 볼륨을 최대 87% 감소시키고, 통신 라운드 수를 54% 줄이며, 모델 정확도에 손실 없이 학습 처리량을 2배로 높일 수 있다.

ABSTRACT

1-bit gradient compression and local steps are two representative techniques that enable drastic communication reduction in distributed SGD. Their benefits, however, remain an open question on Adam-based large model pre-training (e.g. BERT and GPT). In this paper, we demonstrate the non-linearity in Adam causes slow convergence even when 1-bit compression or local steps are individually applied. To alleviate this limitation, we propose 0/1 Adam that linearizes each Adam step via approximating its optimizer states using their stale estimates and linear correlation. 0/1 Adam performs an Adam-like step to preserve the adaptivity, while its linearity allows utilizing 1-bit compression and local steps simultaneously for wall-clock time speed up. We provide convergence guarantee for 0/1 Adam on smooth non-convex objectives. On various large-scale benchmarks such as BERT-Base, BERT-Large, GPT-2 pre-training and ImageNet, we demonstrate on up to 128 GPUs that 0/1 Adam is able to reduce up to 87% of data volume, 54% of communication rounds, and achieve up to 2$ imes$ higher training throughput and end-to-end training time reduction compared to the state-of-the-art baseline 1-bit Adam; while enjoying the same statistical convergence speed and end task model accuracy on GLUE dataset and ImageNet validation set.

연구 동기 및 목표

• 비선형성으로 인해 최적화기 상태에서 통신 효율이 저해되는 대규모 아담 기반 학습에서 강력한 1비트 그래디언트 압축과 로컬 스텝을 적용하는 데 도전하는 것.
• 기존의 1비트 아담과 같은 방법들이 전밀도 정밀도의 웜업 단계가 필요하고, 로컬 스텝과 쉽게 통합되지 않는 한계를 극복하는 것.
• 오래된 분산과 동등성 추정을 통해 업데이트 과정을 선형화하여, 아담에서 1비트 압축과 로컬 스텝을 동시에 사용할 수 있도록 하는 것.
• 부드러운 비볼록 목표 함수에 대해 0/1 아담의 이론적 수렴 보장을 제공하는 것.
• BERT, GPT-2, ImageNet 벤치마크에서 모델 정확도를 희생시키지 않고 최신 기술 수준의 학습 속도 향상을 달성하는 것.

제안 방법

• 1비트 아담의 두 단계 웜업을 피하기 위해, 최적화기 상태의 분산을 적응적으로 冻결함으로써 통합된 단일 단계 최적화 방법인 0/1 아담을 제안한다.
• 분산과 동등성의 오래된 추정치를 사용하여 동등성과 파라미터 업데이트를 근사함으로써, 여러 워커 간 아담 단계의 선형 근사를 가능하게 한다.
• 고정된 분산 하에서 동등성이 기울기와 선형적으로 의존함을 활용하여, 오직 동등성 값만을 1비트로 압축함으로써 통신 볼륨을 극적으로 감소시킨다.
• 동일한 오래된 분산과 동등성 추정치를 사용해 워커가 다수의 로컬 업데이트를 수행할 수 있도록 하여 동기화 빈도를 낮춘다.
• 연속된 최적화기 상태 업데이트 간의 선형 상관관계를 활용하여 근사 오차를 제한하고 수렴 안정성을 유지한다.
• 선형 근사와 양자화에서 발생하는 오차를 제한하는 수렴 분석 프레임워크를 도입하여, 부드러운 비볼록 목표 함수에서의 수렴을 증명한다.

실험 결과

연구 질문

• RQ1 아담 기반 학습에서 수렴성이나 정확도를 훼손하지 않고 1비트 그래디언트 압축과 로컬 스텝을 효과적으로 융합할 수 있는가?
• RQ2 아담의 업데이트 규칙에 내재된 비선형성이 대규모 사전 학습에서 강력한 압축과 로컬 스텝의 사용을 본질적으로 방해하는가?
• RQ3 성능을 유지하면서도 1비트 아담의 두 단계 웜업을 대체할 수 있는 통합된 단일 단계 학습 프로세스를 도입할 수 있는가?
• RQ4 1비트 압축과 로컬 업데이트 하에서 선형화된 아담의 이론적 보장은 무엇인가?
• RQ5 0/1 아담를 사용할 경우, BERT, GPT-2, ImageNet 대규모 학습에서 통신 볼륨과 통신 횟수는 얼마나 감소시킬 수 있는가?

주요 결과

• 128개 GPU에서 0/1 아담은 1비트 아담 대비 최대 87%의 통신 볼륨 감소를 기록하여 데이터 전송 비용을 크게 낮춘다.
• 통신 라운드 수를 최대 54% 감소시켜 동기화 오버헤드를 최소화함으로써 학습 속도를 높인다.
• BERT-Base, BERT-Large, GPT-2, ImageNet 벤치마크에서 0/1 아담은 최대 2배 높은 학습 처리량과 종합적인 학습 시간 단축을 달성한다.
• GLUE 및 ImageNet 검증 세트에서 1비트 아담과 동일한 통계적 수렴 속도와 최종 작업 정확도를 유지한다.
• 수렴 분석을 통해 0/1 아담는 유한한 근사 오차와 양자화 조건 하에서 부드러운 비볼록 목표 함수에서 수렴함을 증명한다.
• 0/1 아담 최적화기와 학습 스크립트는 DeepSpeed에 오픈소스화되어 보급 및 재현 가능성을 높였다.