[論文レビュー] Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training

Can Karakus, Rahul Huilgol+8 more|arXiv (Cornell University)|Nov 10, 2021

Topic Modeling参考文献 28被引用数 8

ひとこと要約

本論文では、PyTorchにおける自動パイプライン並列およびテンソル並列を用いた、効率的な大規模モデル学習を可能にする柔軟で汎用的なフレームワーク、Amazon SageMaker Model Parallelismを提示する。このフレームワークは、コード変更なしに非トランスフォーマー、非均質的、およびMoEモデルを含む多様なアーキテクチャをサポートし、ネイティブPyTorchの機能を維持しながら、DeepSpeedと同等のパフォーマンスを達成する。

ABSTRACT

With deep learning models rapidly growing in size, systems-level solutions for large-model training are required. We present Amazon SageMaker model parallelism, a software library that integrates with PyTorch, and enables easy training of large models using model parallelism and other memory-saving features. In contrast to existing solutions, the implementation of the SageMaker library is much more generic and flexible, in that it can automatically partition and run pipeline parallelism over arbitrary model architectures with minimal code change, and also offers a general and extensible framework for tensor parallelism, which supports a wider range of use cases, and is modular enough to be easily applied to new training scripts. The library also preserves the native PyTorch user experience to a much larger degree, supporting module re-use and dynamic graphs, while giving the user full control over the details of the training step. We evaluate performance over GPT-3, RoBERTa, BERT, and neural collaborative filtering, and demonstrate competitive performance over existing solutions.

研究の動機と目的

• 多様なモデルアーキテクチャやユースケースをサポートする一般性と柔軟性に欠ける大規模モデル学習のためのフレームワークの不足に対処すること。
• 新しい学習スクリプトの統合負荷を軽減し、自動的なパーティショニングと最小限のコード変更を可能にすること。
• 動的計算グラフやモジュール再利用などのネイティブPyTorch機能との完全な互換性を維持すること。
• マルチGPUおよびマルチノード環境におけるパイプライン並列およびテンソル並列の両方をスケーラブルかつ拡張可能に実装するシステムを提供すること。
• DeepSpeedなどの既存ソリューションと比較して、競争力のある学習パフォーマンスを提供すること。

提案手法

• フレームワークはモジュールサーバアーキテクチャを採用し、任意のモデルアーキテクチャをデバイス間で自動的かつ負荷分散を考慮したパーティショニングを可能にする。
• 均一な大規模モデル、孤立した大規模コンponentを有するモデル、および混合専門家（MoE）アーキテクチャをサポートする汎用的なテンソル並列フレームワークを導入する。
• 動的に生成された通信リクエストをデバイス間およびノード間で処理する専用のデバイス間通信（D2D）バックエンドを提供する。
• 元のコードを変更せずにネイティブPyTorchモジュールをその分散版にマッピングするための柔軟なAPI（hooks: init_hook, forward_hook, return_hook）を提供する。
• パラメータ作成スコープや分散テンソル操作（例：allgather, allreduce, reduce-scatter）などのユーティリティを提供し、低レベルのテンソル並列制御を可能にする。
• smp.delay_param_initializationというコンテキストマネージャーにより、パラメータの初期化をモデルデプロイまで延期し、CPUメモリ圧力を軽減する。

実験結果

リサーチクエスチョン

• RQ1 任意のモデルアーキテクチャを最小限のコード変更でサポートできる一般用途のモデル並列フレームワークを構築できるか？
• RQ2 テンソル並列を非均一なモデル、大規模な埋め込み層、および混合専門家（MoE）アーキテクチャに対しても一般化できるか？
• RQ3 動的グラフやモジュール再利用などのネイティブPyTorch機能が分散学習システムにおいてどの程度維持できるか？
• RQ4 実際の学習ワークロードにおいて、本フレームワークのパフォーマンスはDeepSpeedなどの既存ソリューションと比較してどの程度か？
• RQ5 自動的な負荷分散パーティショニングが、複雑で非均質的なモデルアーキテクチャに効果的に適用できるか？

主な発見

• 本フレームワークは、GPT-3、RoBERTa、BERT、およびニューラルコラボラティブフィルタリングモデルを、DeepSpeedと同等の学習スループットで効果的に学習可能である。
• 自動パーティショニングにより、手動でのアーキテクチャ解析や再設定を必要とせずに、デバイス間で効果的な負荷分散が実現されている。
• システムはPyTorchの動的計算グラフを完全に維持しており、条件分岐の実行フローやパラメータの再利用もサポートしている。
• D2D通信バックエンドは、動的に生成された通信パターンを効率的に処理し、複雑な学習ワークフローにおけるオーバーヘッドを低減している。
• smp.delay_param_initializationの使用により、CPUメモリが制限されている状況下でも数十億のパラメータを持つモデルの学習が可能である。
• テンソル並列フレームワークは、従来のソリューションよりも広範なユースケースをカバーしており、MoEや非トランスフォーマー型アーキテクチャも含む。