分散型機械学習通信の最適化 - 融合計算-集合演算を用いて

Core Concepts

GPUの大規模並列性と GPU主導の通信を活用して、依存する集合通信と計算を融合することで、通信の潜在時間を隠蔽し、ピークネットワーク帯域幅の需要を削減する。

Abstract

本論文では、分散型機械学習モデルにおける集合通信の課題に取り組む。大規模な機械学習モデルを分散システムで実行するために、並列化手法が開発されているが、その結果生じる重要な通信(重み更新、層間の活性化交換など)は、独立した計算がないため隠蔽が困難となり、システムの規模が大きくなるにつれボトルネックとなっている。 提案手法では、GPUの大規模並列性とGPU主導の通信を活用して、依存する集合通信と計算を融合する。具体的には、スレッドブロック/ワークグループ単位で自己完結型のGPUカーネルを開発し、計算結果を即座にリモートGPUに通信する。他のワークグループは並行して計算を行い、ALUの利用率を維持する。このような細粒度の重複により、ピークネットワーク帯域幅の需要が低減され、アプリケーション全体にわたって通信が分散される。さらに、スケールアップ通信では、1つのGPUで計算された結果を直接ピアGPUのメモリに書き込むことで、中間バッファリングと複製を排除する。 提案手法を3つのプロトタイプ融合演算子(埋め込み+All-to-All、GEMV+AllReduce、GEMM+All-to-All)に適用し、深層学習推奨モデル(DLRM)、Transformer、Mixture of Experts(MoE)アーキテクチャにおける顕著な通信オーバーヘッドに取り組む。さらに、提案手融合演算子をPyTorchの新しい演算子として公開し、Tritonフレームワークにも拡張することで、本手法を広く採用可能にする。 評価の結果、提案手法は通信と計算の重複を効果的に行い、従来の集合ライブラリベースのアプローチよりも実行時間を短縮できることを示した。スケールアップのGEMV+AllReduceとGEMM+All-to-Allは13%(最大22%)と12%(最大20%)の実行時間短縮を達成し、スケールアウトの埋め込み+All-to-Allは内部ノードで20%、ノード間で31%の実行時間短縮を実現した。大規模なスケールアウトシミュレーションでは、提案手法がDLRMの実行時間を128ノードシステムで21%短縮することを示した。

Stats

機械学習モデルのサイズは2018年から2022年の間に5桁増加した。 分散型機械学習モデルの通信(重み更新、層間の活性化交換など)は、独立した計算がないため隠蔽が困難で、システムの規模が大きくなるにつれボトルネックとなっている。 提案手法の埋め込み+All-to-Allカーネルは、内部ノードで20%、ノード間で31%の実行時間短縮を実現した。 提案手法のGEMV+AllReduceとGEMM+All-to-Allは、それぞれ13%(最大22%)と12%(最大20%)の実行時間短縮を達成した。 大規模なスケールアウトシミュレーションでは、提案手法がDLRMの実行時間を128ノードシステムで21%短縮した。

Quotes

"機械学習は現在、分類(不正検出、医療診断、顔認識)、パターン分析(製品推奨、株価予測)、コンテンツ生成(コード生成、チャットボット、画像/動画生成)など、幅広い分野で使用されている。" "大規模な機械学習モデルは、より複雑な問題に取り組むために増大しており、その容量と計算要件を満たすための分散システムの開発が促進されている。" "分散型機械学習モデルの通信(重み更新、層間の活性化交換など)は、独立した計算がないため隠蔽が困難で、システムの規模が大きくなるにつれボトルネックとなっている。"

Key Insights Distilled From

Optimizing Distributed ML Communication with Fused Computation-Collective Operations

by Kishore Punn... at arxiv.org 04-24-2024

https://arxiv.org/pdf/2305.06942.pdf

Deeper Inquiries

分散型機械学習の通信ボトルネックを解決するための代替アプローチはあるか?

提案された手法は、GPU-initiated intra-kernel communicationを活用して通信と依存する計算を細かい粒度で重ねることで、通信の遅延を隠すことを目的としています。このアプローチは、従来のCPU-initiated networkingに比べてより細かい粒度での通信と計算のオーバーラップを可能にし、通信の遅延を効果的に隠すことができます。他の代替アプローチとしては、通信と計算をさらに細かい単位に分割して重ねる方法や、特定のハードウェアトポロジーに合わせた最適な通信スケジュールを提案する方法などが考えられます。

提案手法の適用範囲を拡大するためには、どのようなハードウェアの変更や新しい並列化手法が必要か

提案手法の適用範囲を拡大するためには、どのようなハードウェアの変更や新しい並列化手法が必要か? 提案手法を他の分散型アプリケーションに適用するためには、より高度なGPU-initiated networking機能や通信プリミティブが必要となります。また、ハードウェアレベルでの最適化や新しい並列化手法の導入が考えられます。例えば、より効率的なGPU間通信を可能にする新しいネットワークトポロジーや、より高速なデータ転送を実現する新しい通信プロトコルの導入が考えられます。さらに、GPU間のデータ転送や通信のオーバーヘッドを最小限に抑えるためのハードウェアの改良や新しい並列化アルゴリズムの開発が重要です。

提案手法を他の分散型アプリケーションに適用することで、どのような新しい洞察が得られるか

提案手法を他の分散型アプリケーションに適用することで、どのような新しい洞察が得られるか? 提案手法を他の分散型アプリケーションに適用することで、通信と計算のオーバーラップによる性能向上や通信の効率化に関する新しい洞察が得られるでしょう。特定のアプリケーションにおいて、通信と計算の間に生じるボトルネックを特定し、それらを効果的に隠すための最適な手法やアルゴリズムを開発することが可能です。さらに、異なる分散型アプリケーションに対して提案手法を適用することで、通信と計算の最適なバランスや並列化戦略に関する新たな知見を得ることができるでしょう。