[논문 리뷰] 1-bit Adam: Communication Efficient Large-Scale Training with Adam's Convergence Speed

Hanlin Tang, Shaoduo Gan외 7명|arXiv (Cornell University)|2021. 02. 04.

Advanced Neural Network Applications참고 문헌 39인용 수 15

한 줄 요약

이 논문은 1-bit Adam을 제안하며, 이는 통신 효율적인 최적화 방법으로, 초기 단계에서 Adam을 사용해 전치된 모멘텀 SGD를 압축 단계 동안 조정함으로써 Adam의 수렴 속도와 1-bit 그래디언트 압축을 결합한다. BERT-Large 및 ResNet 학습에서 압축되지 않은 Adam에 비해 통신량을 최대 5배 줄이고, 처리량을 최대 3.3배 향상시키며, 동일한 수렴 속도와 정확도를 유지한다.

ABSTRACT

Scalable training of large models (like BERT and GPT-3) requires careful optimization rooted in model design, architecture, and system capabilities. From a system standpoint, communication has become a major bottleneck, especially on commodity systems with standard TCP interconnects that offer limited network bandwidth. Communication compression is an important technique to reduce training time on such systems. One of the most effective methods is error-compensated compression, which offers robust convergence speed even under 1-bit compression. However, state-of-the-art error compensation techniques only work with basic optimizers like SGD and momentum SGD, which are linearly dependent on the gradients. They do not work with non-linear gradient-based optimizers like Adam, which offer state-of-the-art convergence efficiency and accuracy for models like BERT. In this paper, we propose 1-bit Adam that reduces the communication volume by up to $5 imes$, offers much better scalability, and provides the same convergence speed as uncompressed Adam. Our key finding is that Adam's variance (non-linear term) becomes stable (after a warmup phase) and can be used as a fixed precondition for the rest of the training (compression phase). Experiments on up to 256 GPUs show that 1-bit Adam enables up to $3.3 imes$ higher throughput for BERT-Large pre-training and up to $2.9 imes$ higher throughput for SQuAD fine-tuning. In addition, we provide theoretical analysis for our proposed work.

연구 동기 및 목표

• 제한된 네트워크 대역폭을 가진 일반 하드웨어에서 대규모 분산 학습의 통신 병목 현상을 해결한다.
• BERT와 같은 최신 모델에 필수적인 비선형 최적화 방법인 Adam과 오차 보정 압축 간의 상충 관계를 해결한다.
• 통신량을 줄이기 위해 강력한 1-bit 그래디언트 압축을 실현하면서도 Adam의 빠른 수렴 속도를 유지한다.
• 워밍업 단계에서는 Adam의 적응형 학습률을 활용하고, 주 학습 단계에서는 압축된 모멘텀 SGD를 적용하는 하이브리드 최적화 전략을 설계한다.
• 표준 TCP 인터커넥트를 사용하여 BERT 및 ResNet과 같은 대규모 모델에서 높은 확장성과 종단 간 처리량 향상을 달성한다.

제안 방법

• 두 단계 최적화 프레임워크를 제안한다: 먼저 전체 정밀도 Adam을 사용해 워밍업 단계(일般적으로 학습 스텝의 20%)를 거치고, 이후 나머지 기간 동안 1-bit 압축 모멘텀 SGD로 전환한다.
• Adam의 분산(비선형 항)이 학습 초반에 안정화됨을 관찰하여, 이후 모멘텀 SGD의 전치자로 기능할 수 있음을 활용한다.
• 오차 보정 압축을 기울기 대신 모멘텀 항에 적용함으로써, 1-bit 압축 하에서의 수렴 안정성을 확보한다.
• 감소된 통신량을 효율적으로 처리하고 고확장성을 달성하기 위해 MPI 기반의 커스터마이징된 집합적 통신 프리미티브를 설계한다.
• 1-bit Adam 최적화 방법과 통신 백엔드를 DeepSpeed에 통합하여 실사용 환경에 구현하고, 구현 코드를 오픈소스로 공개한다.

실험 결과

연구 질문

• RQ1 Adam과 같은 비선형 최적화 방법에 대해 오차 보정 압축을 효과적으로 적용할 수 있는가? 이는 BERT와 같은 현대 모델 학습에 필수적인 요소이다.
• RQ2 Adam의 비선형 분산 항이 학습 초반에 안정적으로 유지되어, 이후 압축 최적화의 전치자로 기능할 수 있는가?
• RQ3 Adam과 압축된 모멘텀 SGD를 조합한 하이브리드 최적화 방법이 전체 정밀도 Adam과 동일한 수렴 속도를 달성하면서도 통신량을 줄일 수 있는가?
• RQ4 256개 GPU에서 BERT 및 ResNet과 같은 대규모 모델에 대해 이러한 하이브리드 최적화 방법의 확장성과 종단 간 처리량 향상 수준은 어떠한가?
• RQ5 1-bit Adam은 단순한 Adam 기울기의 1-bit 압축과 비교해 성능가능성이 높은가? 또한 최종 모델 정확도를 유지하는가?

주요 결과

• 1-bit Adam은 압축되지 않은 Adam에 비해 통신량을 최대 5배 줄였으며, 수렴 속도나 최종 모델 정확도에 영향을 주지 않았다.
• 64개 GPU에서 배치 크기 4K로 BERT-Large 사전학습을 수행한 결과, 1-bit Adam은 51.5시간에 학습을 완료하여 기준 Adam(174.3시간) 대비 3.4배 빠른 속도를 기록했다.
• 256개 GPU에서 1-bit Adam은 압축되지 않은 Adam 대비 최대 3.3배 높은 종단 간 처리량을 달성했으며, 이더넷 환경에서 인피니밴드 대비 더 뛰어난 확장성을 보였다.
• BERT 사전학습, SQuAD 미세조정, CIFAR-10에서의 ResNet-18, DCGAN 학습 등에서 1-bit Adam은 전체 정밀도 Adam과 동일한 샘플 단위 수렴 속도를 유지했다.
• 10Gbps 이더넷 환경에서 1-bit Adam은 100Gbps 인피니밴드에서의 Adam과 유사한 처리량을 달성하여 뛰어난 하드웨어 효율성을 입증했다.
• 실험 결과, 단순한 Adam 기울기의 1-bit 압축은 수렴하지 못하는 반면, 전치된 모멘텀 SGD를 적용한 1-bit Adam은 뛰어난 성능을 유지를 해냈다.