[論文レビュー] 1-bit Adam: Communication Efficient Large-Scale Training with Adam's Convergence Speed

Hanlin Tang, Shaoduo Gan+7 more|arXiv (Cornell University)|Feb 4, 2021

Advanced Neural Network Applications参考文献 39被引用数 15

ひとこと要約

本論文では、1ビットAdamを提案する。これは、1ビット勾配圧縮とAdamの収束速度を組み合わせた通信効率の高い最適化手法であり、暖 warmed-up 階段でAdamを用いて勾配圧縮段階におけるモーメンタムSGDの前処理を実行する。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.

研究の動機と目的

• 通信帯域が限られたコンmodityハードウェアにおける大規模分散学習の通信ボトル neck 問題に対処すること。
• BERTのような最先端モデルに不可欠な非線形最適化手法（例：Adam）と誤差補償圧縮の不適合性を克服すること。
• 通信量を削減するための積極的な1ビット勾配圧縮を実現しながら、Adamの高速収束速度を維持すること。
• 暖 warmed-up 段階でAdamの適応的学習率を活用し、本番学習段階で圧縮付きモーメンタムSGDを適用するハイブリッド最適化戦略を設計すること。
• 標準的なTCPインタコネクトを用いて、BERT や ResNet などの大規模モデルにおいて高いスケーラビリティとエンドツーエンドスループットの向上を達成すること。

提案手法

• 2段階の最適化フレームワークを提案：まず、通常の訓練ステップの20%程度を想定して、全精度Adamを用いた暖 warmed-up 段階を実行し、その後、残りの期間は1ビット圧縮付きモーメンタムSGDに切り替える。
• Adamの分散（非線形項）が訓練の初期段階で安定することに着目し、これにより後続の圧縮付きモーメンタムSGDの前処理として機能できることを活用する。
• 誤差補償圧縮を勾配ではなく、SGDのモーメンタム項に適用することで、1ビット圧縮下でも収束の安定性を保証する。
• 通信量が削減された状況に適した効率的なスケーラビリティを実現するため、独自のMPIベースの集約通信プリミティブを設計する。
• DeepSpeedに1ビットAdam最適化手法と通信バックエンドを統合し、本番環境でのデプロイを可能にするとともに、実装をオープンソース化する。

実験結果

リサーチクエスチョン

• RQ1 誤差補償圧縮は、BERTのような現代のモデルに不可欠な非線形最適化手法（例：Adam）に効果的に適用可能か？
• RQ2 Adamの非線形分散項は訓練の初期段階で安定するか？ その安定性により、圧縮最適化の前処理としての役割を果たせるか？
• RQ3 Adamと圧縮付きモーメンタムSGDを組み合わせたハイブリッド最適化手法は、全精度Adamと同等の収束速度を達成しながら通信量を削減できるか？
• RQ4 256 GPUで大規模モデル（例：BERT や ResNet）に対して、そのようなハイブリッド最適化手法のスケーラビリティとエンドツーエンドスループットの向上はどの程度か？
• RQ5 1ビットAdamは、Adam勾配の単純な1ビット圧縮と比較して性能が優れており、最終的なモデル精度を維持できるか？

主な発見

• 1ビットAdamは、非圧縮Adamと比較して通信量を最大5倍削減し、収束速度や最終的なモデル精度に影響を及げない。
• 64 GPUでバッチサイズ4KでBERT-Largeの事前学習を実行した場合、1ビットAdamは174.3時間のベースラインAdamと比較して51.5時間で完了し、3.4倍の高速化を達成した。
• 256 GPUで最大3.3倍のエンドツーエンドスループット向上を達成し、イーサネット環境でもInfiniBandよりも優れたスケーラビリティを示した。
• BERT事前学習、SQuAD微調整、CIFAR-10におけるResNet-18、DCGAN学習の全テストケースにおいて、1ビットAdamは全精度Adamと同等のサンプル単位の収束速度を維持した。
• 10Gbpsイーサネット環境でも、1ビットAdamは100Gbps InfiniBandで動作するAdamと同等のスループットを達成し、ハードウェア効率が優れていることを示した。
• 実験の結果、Adam勾配の単純な1ビット圧縮では収束が失敗する一方、前処理付きモーメンタムSGDを用いた1ビットAdamは、安定した性能を維持した。