[논문 리뷰] APMSqueeze: A Communication Efficient Adam-Preconditioned Momentum SGD Algorithm

Hanlin Tang, Shaoduo Gan외 5명|arXiv (Cornell University)|2020. 08. 26.

Vehicle License Plate Recognition참고 문헌 35인용 수 4

한 줄 요약

APMSqueeze는 BERT 및 ResNet-18와 같은 대규모 모델의 분산 학습에서 발생하는 통신 병목 현상을 해결하기 위해, 첫 번째 단계에서 Adam을 사용해 모멘텀 SGD를 전처리한 후, 오차 보정 기반의 기울기 압축을 적용하여 통신량을 최대 97%까지 줄이는 통신 효율적인 최적화 알고리즘입니다 (32배 감소). 이로 인해 BERT 및 ResNet-18에서 종료 시점까지의 속도를 최대 10배 향상시키며, 수렴성과 정확도를 손상시키지 않습니다.

ABSTRACT

Adam is the important optimization algorithm to guarantee efficiency and accuracy for training many important tasks such as BERT and ImageNet. However, Adam is generally not compatible with information (gradient) compression technology. Therefore, the communication usually becomes the bottleneck for parallelizing Adam. In this paper, we propose a communication efficient {\bf A}DAM {\bf p}reconditioned {\bf M}omentum SGD algorithm-- named APMSqueeze-- through an error compensated method compressing gradients. The proposed algorithm achieves a similar convergence efficiency to Adam in term of epochs, but significantly reduces the running time per epoch. In terms of end-to-end performance (including the full-precision pre-condition step), APMSqueeze is able to provide {sometimes by up to $2-10 imes$ speed-up depending on network bandwidth.} We also conduct theoretical analysis on the convergence and efficiency.

연구 동기 및 목표

• BERT 및 ResNet-18와 같은 대규모 모델의 분산 학습에서 발생하는 통신 병목 현상을 해결하기 위해, Adam이 필수적이지만 기존 압축 기법과 호환되지 않는 문제를 해결합니다.
• 기존 방법에서 기울기 압축이 Adam의 비선형 업데이트 규칙과 충돌하여 수렴성이 깨지는 문제를 해결합니다.
• 통신량을 극적으로 줄이면서도 Adam의 수렴 속도와 최종 모델 정확도를 유지하는 통신 효율적인 알고리즘을 개발합니다.
• 제한된 네트워크 대역폭 조건에서도 최신 기술 수준의 모델을 확장 가능하고 고성능으로 분산 학습할 수 있도록, 새로운 전처리 및 압축 파이프라인을 제공합니다.

제안 방법

• 압축 학습으로 전환하기 전에, 일부 에포크 동안 전체 정밀도의 분산 Adam 최적화를 통해 모멘텀 SGD를 전처리합니다.
• 전처리 단계 이후의 모멘텀 업데이트 방향에 대해 오차 보정 기반 기울기 압축을 적용하여 수렴성이 유지되도록 합니다.
• 양자화와 희소화를 모두 지원하는 무작위 압축 연산자 $\bm{C}_{\omega}(\cdot)$ 를 사용하며, 압축 손실을 보정하기 위해 오차 보정을 적용합니다.
• 전처리 단계 동안 Adam의 모멘텀 벡터 $\bm{m}_t$ 와 적응형 학습률 $\bm{v}_t$ 를 유지한 후, 주 학습 단계에서는 $\bm{v}_t$ 를 冻결하고 $\bm{m}_t$ 를 압축합니다.
• 하이브리드 학습 전략을 도입합니다: 전체 학습의 15% 동안은 전체 정밀도의 Adam을 사용하고, 나머지 85% 동안은 오차 보정이 적용된 압축된 모멘텀 SGD를 사용합니다.
• 이론적 분석을 통해, 표준 가정 하에 압축된 알고리즘이 압축되지 않은 버전과 동일한 渐近 수렴 속도를 달성함을 입증합니다.

실험 결과

연구 질문

• RQ1 기울기 압축을 Adam 기반 학습에 효과적으로 적용할 수 있을까? 이로 인해 수렴성이나 모델 정확도가 떨어지지 않을까?
• RQ2 단지 몇 에포크의 Adam을 통해 모멘텀 SGD를 전처리하면, 이후 학습 단계에서 안정적이고 효율적인 통신 압축이 가능할까?
• RQ3 Adam으로 전처리된 모멘텀 SGD 알고리즘에 대해, 압축된 경우의 이론적 수렴 보장은 무엇일까?
• RQ4 BERT 및 ResNet-18와 같은 대규모 모델에서 이 방법을 사용하면 통신 오버헤드 감소와 종료 시점까지의 속도 향상은 어느 정도 달성될 수 있을까?
• RQ5 이러한 방법은 다양한 네트워크 대역폭과 압축 기법에 대해 강건한가?

주요 결과

• APMSqueeze는 BERT-Base 및 BERT-Large에서 전체 정밀도의 Adam과 동일한 수렴 특성과 최종 정확도를 유지하면서 통신량을 최대 97%까지 줄였습니다 (32배 압축에 해당).
• 128개 GPU에서 APMSqueeze는 에포크당 런타임에서 최대 8배의 속도 향상을 기록했으며, 1Gbps 네트워크 조건에서 종료 시점까지의 속도 향상은 최대 10배였습니다.
• 10Gbps 대역폭 조건에서는 2배의 종료 시점 속도 향상을 달성하여, 다양한 네트워크 환경에서 뛰어난 확장성을 입증했습니다.
• BERT-Base, BERT-Large, CIFAR-10에서의 ResNet-18를 포함한 모든 평가된 작업에서, APMSqueeze는 Adam과 동일한 학습 손실과 테스트 정확도를 유지했습니다.
• 이론적 분석을 통해 APMSqueeze가 압축되지 않은 기준과 동일한 渐近 수렴 속도를 확보했으며, 워커 수에 비례하여 선형 속도 향상을 달성함을 확인했습니다.
• APMSqueeze는 기울기 압축과 Adam 유사 최적화기를 함께 사용할 수 있는 통신 효율적인 분산 알고리즘으로서, BERT 수준의 복잡한 모델 학습을 성공적으로 지원하는 최초의 알고리즘입니다.