[論文レビュー] ZeRO++: Extremely Efficient Collective Communication for Giant Model Training

Guanhua Wang, Heyang Qin+7 more|arXiv (Cornell University)|Jun 16, 2023

Advanced Neural Network Applications被引用数 15

ひとこと要約

ZeRO++ は、ZeRO のノード間データ移動を最大 4 倍削減する3つの通信削減技術（qwZ、hpZ、qgZ）を導入し、巨大モデル学習のスループットを向上させる。

ABSTRACT

Zero Redundancy Optimizer (ZeRO) has been used to train a wide range of large language models on massive GPUs clusters due to its ease of use, efficiency, and good scalability. However, when training on low-bandwidth clusters, or at scale which forces batch size per GPU to be small, ZeRO's effective throughput is limited because of high communication volume from gathering weights in forward pass, backward pass, and averaging gradients. This paper introduces three communication volume reduction techniques, which we collectively refer to as ZeRO++, targeting each of the communication collectives in ZeRO. First is block-quantization based all-gather. Second is data remapping that trades-off communication for more memory. Third is a novel all-to-all based quantized gradient averaging paradigm as replacement of reduce-scatter collective, which preserves accuracy despite communicating low precision data. Collectively, ZeRO++ reduces communication volume of ZeRO by 4x, enabling up to 2.16x better throughput at 384 GPU scale.

研究の動機と目的

• 限られたノード間帯域幅のマルチ-GPUクラスターで、非常に大規模な言語モデルの効率的な学習を動機づける。
• 前方のウェイト収集、後方のウェイト収集、勾配還元全体にわたる ZeRO の通信オーバーヘッドを低減する。
• 3D並列への再設計や収束精度の犠牲を伴わず、より高いスループットとスケーラビリティを実現する。

提案手法

• 3 つの通信最適化を導入する：qwZ（前方全収集のブロックベースのINT8 重み量子化）、hpZ（ノード内の二次的な重み分割でノード間後方全収集を削除）、qgZ（テンソルスライスの再配置を伴う All-to-All ベースの INT4 勾配還元）。
• 精度を保持するためブロックベースの量子化を使用し、量子化・非量子化・演算子融合の高性能 CUDA カーネルを実装する。
• ノード内通信を先に、ノード間通信を後に行う2段階の階層的 All-to-All 勾配還元を採用し、ノード間トラフィックを最小化する。
• 待ち時間を隠すため計算と通信を重ね合わせ、メモリトラフィックを最小化するためカーネルを融合する。

Figure 1. Large scale training throughput are constrained by network bandwidth and batch size per GPU

実験結果

リサーチクエスチョン

• RQ1 非常に大規模なモデルでも学習精度を犠牲にせず、ZeRO のデバイス間通信をどのように削減できるか？
• RQ2 ZeRO-3 で前方・後方の通信と勾配集約を最小化するためには、どの量子化と分割技術の組み合わせが有効か？
• RQ3 低精度データを用いた all-to-all に基づく新しい勾配還元パラダイムは、通信量を削減しつつ収束を維持できるか？

主な発見

通信	フォワード全収集	バックワード全収集	バックワードリデュース-スキャッター
ZeRO-3	M	M	M
ZeRO++	0.5M	0	0.25M

• ZeRO++ は、トレーニング1反復あたりのノード間通信量を 3M から 0.75M に削減し、スループット向上を実現。
• 10B〜138B パラメータのモデルで ZeRO ベースラインに対し、エンドツーエンドのスループットを最大 2.4 倍改善。
• GPT-3 風モデルで最大384GPUのスケーラビリティを実証し、持続的ピークスループットの45%以上を達成。
• ZeRO ベースラインと同等の収束と学習精度を維持。
• 低帯域設定でも ZeRO++ は高帯域ベース設定と同等のスループットを達成。

Figure 2. Illustration & example of block based quantization vs. baseline