[論文レビュー] Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training

Li, Shenggui, Hongxin Liu+6 more|arXiv (Cornell University)|Oct 28, 2021

Topic Modeling参考文献 15被引用数 21

ひとこと要約

Colossal-AI は、データ並列、パイプライン並列、テンソル並列、シーケンス並列に加え、異種トレーニングとZeRO最適化を統合的にサポートすることで、大規模分散学習を効率的に行える統合型ディープラーニングシステムです。動的メモリ管理と柔軟で使いやすいシーケンシャルモデルの並列化を可能にすることで、ベースラインシステム比最大2.76倍のトレーニングスピードアップを達成しています。

ABSTRACT

The success of Transformer models has pushed the deep learning model scale to billions of parameters. Due to the limited memory resource of a single GPU, However, the best practice for choosing the optimal parallel strategy is still lacking, since it requires domain expertise in both deep learning and parallel computing. The Colossal-AI system addressed the above challenge by introducing a unified interface to scale your sequential code of model training to distributed environments. It supports parallel training methods such as data, pipeline, tensor, and sequence parallelism, as well as heterogeneous training methods integrated with zero redundancy optimizer. Compared to the baseline system, Colossal-AI can achieve up to 2.76 times training speedup on large-scale models.

研究の動機と目的

• 1. ビリオンパラメータモデルの大規模分散トレーニングのための統合的で使いやすいシステムの不足に対処すること。
• 2. リサーチャーやエンジニアが最適な並列化戦略を選択する際の専門知識の壁を低減すること。
• 3. シーケンス並列および多次元テンソル並列を含む包括的な並列化手法のサポート。
• 4. 動的シャーディングとオフロードにより、メモリ効率とトレーニングスループットを向上させること。
• 5. アーキテクチャの再実装なしに、シーケンシャルモデルのコードをマルチGPUクラスタにスムーズにスケーリング可能にすること。

提案手法

• 1. Colossal-AI は、最小限のコード変更でシーケンシャルモデルコードを分散環境にスケーリングできる統合APIを提供する。
• 2. データ並列、パイプライン並列、テンソル並列、シーケンス並列を統合し、最適なパフォーマンスを得るための柔軟な組み合わせを可能にする。
• 3. 動的メモリ管理とモジュラー加速技術により、アクティベーションメモリのボトルネックを低減し、スケーラビリティを向上させる多次元テンソル並列を採用する。
• 4. シーケンス並列は、シーケンス次元に沿って入力およびアクティベーションテンソルを分割し、通信とメモリオーバーヘッドを削減する。
• 5. 必要に応じてのみオフロードを行う動的テンソル配置を実装し、GPUメモリ利用効率を向上させる。
• 6. ZeRO-3 と異種トレーニングをサポートし、重複する最適化器状態を最小限に抑え、混合ハードウェア環境でのトレーニングを高速化する。

実験結果

リサーチクエスチョン

• RQ1 1. 統合型ディープラーニングシステムとして、複数の並列化技術を効果的に統合することで、最適なトレーニングパフォーマンスを達成できるか？
• RQ2 2. 1Dテンソル並列と比較して、シーケンス並列はメモリ効率とスケーラビリティにどのような影響を与えるか？
• RQ3 3. 動的シャーディングとオフロードは、低バッチサイズのシナリオにおいてトレーニングスループットとメモリ利用効率を向上させられるか？
• RQ4 4. Megatron-LM や DeepSpeed といった最先端のシステムと比較して、Colossal-AI のトレーニングスピードアップはどの程度か？
• RQ5 5. パイプライン段階数とGPU数の増加に伴い、Colossal-AI はトレーニングスループットをどの程度スケーリングできるか？

主な発見

• 1. Colossal-AI は、大規模モデルにおいてベースラインシステム比最大2.76倍のトレーニングスピードアップを達成した。
• 2. 12GPU環境では、シーケンス並列が1Dテンソル並列と比較して、4.44倍大きなバッチサイズと1.18倍長いシーケンス長を可能にした。
• 3. 4段階のパイプラインを用いた場合、シーケンス並列は1Dテンソル並列と比較して1.55倍速いトレーニングを実現した。
• 4. 低バッチサイズのシナリオでは、Colossal-AI の動的オフロードが不要なCPU転送を回避し、DeepSpeed の静的ポリシーを上回るスループットを達成した。
• 5. 8GPU、13BパラメータのOPTモデル、バッチサイズ32の環境下で、Colossal-AI はDeepSpeed比1.33倍のスピードアップを達成した。
• 6. 線形複雑度のアテンションを用いる場合、GPU数に比例して最大シーケンス長を線形スケーリング可能であり、ドキュメントレベルのテキスト理解を可能にした。