[論文レビュー] Quantized Distributed Training of Large Models with Convergence Guarantees

Ilia Markov, Adrian Vladu+2 more|arXiv (Cornell University)|Feb 5, 2023

Advanced Neural Network Applications被引用数 5

ひとこと要約

本稿では、理論的収束保証を伴う、勾配とモデル重みの両方を完全に量子化可能な通信効率の良い Fully-Sharded Data Parallel (FSDP) の変種 QSDP を提案する。すべての通信フェーズにおいて重みと勾配に不偏量子化を適用することにより、QSDP は FSDP の通信ボトル neck を解消しながらモデルの精度を維持し、1.3Bパラメータまでの GPT モデルで最大 2.2 倍のエンドツーエンドの高速化を達成する。

ABSTRACT

Communication-reduction techniques are a popular way to improve scalability in data-parallel training of deep neural networks (DNNs). The recent emergence of large language models such as GPT has created the need for new approaches to exploit data-parallelism. Among these, fully-sharded data parallel (FSDP) training is highly popular, yet it still encounters scalability bottlenecks. One reason is that applying compression techniques to FSDP is challenging: as the vast majority of the communication involves the model's weights, direct compression alters convergence and leads to accuracy loss. We present QSDP, a variant of FSDP which supports both gradient and weight quantization with theoretical guarantees, is simple to implement and has essentially no overheads. To derive QSDP we prove that a natural modification of SGD achieves convergence even when we only maintain quantized weights, and thus the domain over which we train consists of quantized points and is, therefore, highly non-convex. We validate this approach by training GPT-family models with up to 1.3 billion parameters on a multi-node cluster. Experiments show that QSDP preserves model accuracy, while completely removing the communication bottlenecks of FSDP, providing end-to-end speedups of up to 2.2x.

研究の動機と目的

• 大規模言語モデルの Fully-Sharded Data-Parallel (FSDP) 学習における通信ボトル neck を解消すること。これは、頻繁な全対全通信による重み交換に起因する。
• 収束性やモデル精度を損なわず、FSDP におけるモデル重みと勾配の両方を完全に量子化すること。
• 非凸な量子化点の格子上で学習を行う場合でも、モデル状態の完全な量子化下での SGD の理論的収束保証を提供すること。
• PyTorch などの既存のディープラーニングフレームワークにスムーズに統合可能な、実用的で低オーバーヘッドの実装を設計すること。

提案手法

• すべての全対全通信ステップの前に勾配と重みに量子化を適用する、FSDP の変種である QSDP を提案する。
• 重みに対して不偏量子化演算子を用い、最近接量子化レベルへの確率的丸めを定義することで、期待値にゼロの誤差を保証する。
• 勾配量子化には、Alistarh 他 (2017) が提案した標準的な不偏勾配圧縮技術を適用する。
• 学習プロセスを量子化を射影機構として持つスパース回復問題としてモデル化し、理論的収束解析を可能にする。
• 標準的な滑らかさと勾配の有界性の仮定の下で、QSDP が量子化格子点上の損失関数の最小化点に収束することを示すことにより、収束保証を導出する。
• PyTorch に効率的に実装し、追加のノード内メモリや複雑な誤差補正を必要としない、最小限のメモリおよび計算オーバーヘッドで実現する。

実験結果

リサーチクエスチョン

• RQ1 FSDP におけるモデル重みと勾配の完全な量子化は、収束性や精度を損なわずに行えるか？
• RQ2 偏った勾配推定器を用いても、非凸な量子化重み格子上で学習する場合に、理論的収束保証を達成できるか？
• RQ3 量子化によって FSDP の通信ボトル neck を完全に解消でき、エンドツーエンドの学習高速化を実現できるか？
• RQ4 GPT などの大規模言語モデル（最大 1.3B パラメータ）において、QSDP は実際の性能をどのように発揮するか？
• RQ5 FSDP における分散学習において、量子化精度と収束安定性のトレードオフは何か？

主な発見

• GPT ファミリーのモデル（最大 1.3B パラメータ）を複数ノードクラスタで学習する際、QSDP は標準的な FSDP と比較して最大 2.2 倍のエンドツーエンドの高速化を達成する。
• 評価されたすべてのモデル、特に 1.3B パラメータの GPT モデルにおいて、モデルの精度が完全に保持され、性能の著しい低下は観測されない。
• 標準的な FSDP における通信ボトル neck は、QSDP によって完全に解消され、量子化による帯域幅要件の顕著な低減が達成されている。
• 理論的解析により、標準的な滑らかさと勾配の有界性の仮定の下で、QSDP が量子化格子点上の損失関数の最小化点に収束することを証明した。
• 重みの量子化は、元の重みの期待値と等しくなる不偏確率的丸め演算子により実装されており、誤差の期待値がゼロであることを保証する。
• 実装に伴う追加の計算的・メモリ的オーバーヘッドはほとんどなく、生産環境での学習に実用的である。