[論文レビュー] High-performance, Distributed Training of Large-scale Deep Learning Recommendation Models.

Dheevatsa Mudigere, Yuchen Hao+49 more|arXiv (Cornell University)|Apr 12, 2021

Stochastic Gradient Optimization Techniques参考文献 56被引用数 18

ひとこと要約

本論文は、大規模な深層学習推薦モデル（DLRMs）の高性能分散学習を実現するソフトウェare-ハードウェア共同設計システムを提示している。従来のシステムと比較して40倍の高速化を達成した。PyTorchベースのトレーニングスタックとZionEXプラットフォームを統合し、最適化されたデータ/モデル並列処理、階層的埋め込みシャーディング、低精度通信、高帯域幅・低レイテンシのネットワークインfraを活用して、最大12兆パラメータのモデルを学習可能にした。

ABSTRACT

Deep learning recommendation models (DLRMs) are used across many business-critical services at Facebook and are the single largest AI application in terms of infrastructure demand in its data-centers. In this paper we discuss the SW/HW co-designed solution for high-performance distributed training of large-scale DLRMs. We introduce a high-performance scalable software stack based on PyTorch and pair it with the new evolution of Zion platform, namely ZionEX. We demonstrate the capability to train very large DLRMs with up to 12 Trillion parameters and show that we can attain 40X speedup in terms of time to solution over previous systems. We achieve this by (i) designing the ZionEX platform with dedicated scale-out network, provisioned with high bandwidth, optimal topology and efficient transport (ii) implementing an optimized PyTorch-based training stack supporting both model and data parallelism (iii) developing sharding algorithms capable of hierarchical partitioning of the embedding tables along row, column dimensions and load balancing them across multiple workers; (iv) adding high-performance core operators while retaining flexibility to support optimizers with fully deterministic updates (v) leveraging reduced precision communications, multi-level memory hierarchy (HBM+DDR+SSD) and pipelining. Furthermore, we develop and briefly comment on distributed data ingestion and other supporting services that are required for the robust and efficient end-to-end training in production environments.

研究の動機と目的

• 生産環境のデータセンタで増加する大規模な深層学習推薦モデル（DLRMs）のインfraストラクチャ需要に対応すること。
• 数兆パラメータを持つDLRMsの分散学習における性能ボトルネックを克服すること。
• 決定論的な最適化更新を維持しながら、ハードウェアの効率的利用を最大化するスケーラブルかつ高スルーレートのトレーニングシステムを開発すること。
• 生産環境における統合型データインジェストおよび支援サービスを通じて、エンドツーエンドの効率的トレーニングを可能にすること。
• モデルの精度や学習安定性を損なわず、従来のシステムよりも顕著な高速化を達成すること。

提案手法

• 高帯域幅、最適なトポロジ、効率的なトランスポートプロトコルを備えた専用のスケールアウトネットワークを備えたZionEXプラットフォームの設計および導入。
• 分散ワーカー間でのモデル並列処理とデータ並列処理をネイティブにサポートするPyTorchベースのトレーニングスタックの実装。
• 埋め込みテーブルを行方向および列方向に階層的にパーティショニングし、複数のワーカー間で負荷をバランスさせるためのシャーディングアルゴリズムの開発。
• 決定論的な最適化更新を完全に保持する高パフォーマンスコアオペレータの統合により、学習の再現性を確保。
• 低精度通信、マルチレベルメモリハーゲリー（HBM、DDR、SSD）、通信パイプライン化を活用して、レイテンシと帯域幅の圧力を低減。
• 信頼性が高くスケーラブルかつ効率的なエンドツーエンドのトレーニングを生産環境で実現するための分散データインジェストパイプラインおよび支援サービスの構築。

実験結果

リサーチクエスチョン

• RQ1 12兆パラメータに達するDLRMsの分散学習を、生産環境のデータセンタでスケーラブルかつ効率的に行うにはどうすればよいか？
• RQ2 大規模なDLRMsのトレーニングシステムにおいて、従来のシステムと比較して40倍の高速化を達成するには、どのようなシステムレベル最適化が必要か？
• RQ3 埋め込みテーブルのシャーディングをどのように階層的にワーカー間で分割することで、負荷バランスを最適化しつつ通信オーバーヘッドを最小限に抑えられるか？
• RQ4 低精度通信とメモリハーゲリー最適化をどの程度活用すれば、学習安定性を損なわずトレーニングスルーレートを向上できるか？
• RQ5 共同設計されたソフトウェアとハードウェアコンponentsが、大規模モデルの決定論的かつ高性能なトレーニングを実現するために果たす役割は何か？

主な発見

• 本システムは、大規模なDLRMsのトレーニングにおいて、従来のシステムと比較して40倍の高速化を達成した。
• 本プラットフォームは、最大12兆パラメータのDLRMsのトレーニングをサポートしており、最大規模の既知の推薦モデルへのスケーラビリティを示した。
• 埋め込みテーブルの階層的シャーディングにより、分散ワーカー間での効果的な負荷バランスが実現され、通信ボトルネックも最小限に抑えられた。
• 低精度通信とパイプライン化により、トレーニングレイテンシが顕著に低減され、ハードウェア利用効率が向上した。
• HBM、DDR、SSDのメモリレイヤーの統合により、大規模なモデル状態と勾配の効率的処理が可能になった。
• 本システムは、極めて高い最適化を施しても、完全に決定論的な最適化更新を維持しており、学習の再現性を保証している。