[論文レビュー] Maximizing Communication Efficiency for Large-scale Training via 0/1 Adam

Yucheng Lu, Conglong Li+3 more|arXiv (Cornell University)|Feb 12, 2022

Stochastic Gradient Optimization Techniques被引用数 4

ひとこと要約

本稿では、1ビット勾配圧縮とローカルステップを同時に利用できる通信効率の高い最適化手法0/1 Adamを提案する。この手法は、古くなった分散とモーメンタムの近似を用いてAdamの更新を線形化することで、1ビット圧縮とローカルステップの併用を可能にし、128 GPU環境で通信量を最大87%削減、通信ラウンド数を54%削減、学習スループットを2倍に向上させつつ、1ビットAdamと同等のモデル精度を維持する。

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.

研究の動機と目的

• 非線形性が最適化状態に与える影響により、通信効率が低下する大規模なAdamベースの学習において、積極的な1ビット勾配圧縮とローカルステップの適用という課題に対処すること。
• 1ビットAdamのような従来手法の限界を克服すること。1ビットAdamはフル精度のウォームアップフェーズを必要とし、ローカルステップと容易に統合できない。
• 分散とモーメンタムの古くなった推定値を用いて更新プロセスを線形化することで、Adamにおける1ビット圧縮とローカルステップの同時利用を可能にすること。
• 滑らかで凸でない目的関数において0/1 Adamの理論的収束保証を提供すること。
• BERT、GPT-2、ImageNetベンチマークにおいて、モデル精度を損なわずに最先端の学習スピードアップを達成すること。

提案手法

• 1ビットAdamの二段階ウォームアップを回避する統合的単一段階最適化手法として0/1 Adamを提案。最適化状態の分散を動的に凍結することで実現する。
• 分散とモーメンタムの古くなった推定値を用いてモーメンタムとパラメータの更新を近似し、ワーカー間でAdamステップの線形近似を可能にする。
• 凍結された分散下でモーメンタムが勾配に対して線形に依存することを活用し、モーメンタム値のみを1ビット圧縮することで、通信量を著しく削減する。
• 同じ古くなった分散とモーメンタム推定値を用いて複数回のローカル更新を許容することで、同期頻度を低減し、ローカルステップを実装する。
• 隣接する最適化状態更新間の線形相関を活用し、近似誤差を制限し、収束安定性を維持する。
• 線形近似と量子化による誤差を制限する収束解析フレームワークを導入し、滑らかで凸でない目的関数における収束を証明する。

実験結果

リサーチクエスチョン

• RQ1 Adamベースの学習において、収束性や精度を損なわずに1ビット勾配圧縮とローカルステップを効果的に統合できるか？
• RQ2 Adamの更新ルールにおける非線形性が、大規模事前学習における積極的圧縮とローカルステップの使用を根本的に行わないか？
• RQ3 2段階ウォームアップを回避しつつ、性能を維持する統合的単一段階トレーニングプロセスを実現できるか？
• RQ4 1ビット圧縮とローカル更新下での線形化されたAdamの理論的保証は何か？
• RQ5 0/1 Adamを用いることで、BERT、GPT-2、ImageNetの大規模トレーニングにおいて、通信量と通信ラウンド数をどの程度削減できるか？

主な発見

• 128 GPU環境において、0/1 Adamは1ビットAdamと比較して通信量を最大87%削減し、データ転送コストを顕著に低減する。
• 通信ラウンド数を最大54%削減し、同期のオーバーヘッドを低減することで、学習を加速する。
• BERT-Base、BERT-Large、GPT-2、ImageNetベンチマークにおいて、0/1 Adamは最大2倍の高い学習スループットとエンド・ツー・エンドのトレーニング時間短縮を達成する。
• GLUEおよびImageNetバリデーションセットにおいて、1ビットAdamと同等の統計的収束速度とタスク精度を維持する。
• 収束解析により、近似誤差と量子化が有界な条件下で、0/1 Adamが滑らかで凸でない目的関数において収束することを証明した。
• 0/1 Adam最適化手法とトレーニングスクリプトはDeepSpeedにオープンソース化され、広範な採用と再現性を可能にした。