Parmを使ってAIのトレーニングスピードを改善する

AIの大規模モデルのトレーニングが研究者たちの大きな焦点になってる。こういう大きなモデルを扱う一般的なアプローチの一つがMixture-of-Experts（MoE）レイヤーを使うこと。これにより、計算コストを大きく上げることなくモデルのサイズを増やせるんだ。でも、モデルが大きくなるにつれて、トレーニング中に必要なコミュニケーションも増えてくる。特に、たくさんのGPUやTPUを使ってると、トレーニングプロセスが遅くなることがあるんだ。

これを克服するために、研究者たちはタスクの処理とデータの通信の仕方をいろいろと工夫してるんだけど、複数の並列セットアップを使うことで発生する通信コストにはまだ苦労してる。この文では、MoEモデルのトレーニング速度を向上させつつ、通信コストを抑えるために設計された新しいシステム「Parm」を探るよ。

#大モデルの課題

AIモデルが大きくなるに連れて、トレーニングに必要な計算リソースも増加する。例えば、数十億のパラメータを持つモデルのトレーニングには、ものすごい計算量が必要なんだ。ここでMoEレイヤーが役立つ。これにより、追加の計算コストを最小限に抑えつつモデルを大きくできる。MoEレイヤーは、特定のタスクのためにモデルの一部（エキスパートと呼ばれる部分）だけをアクティブにすることで機能する。この選択的なアクティベーションによって、全体の計算ニーズが低く抑えられる。

でも、MoEを使ってても、複数のGPUを使うときの通信時間がトレーニング全体の時間の大部分を占めることがある。GPU間の通信がうまく管理されていないと、タスクの分担が遅れちゃう。だから、タスクのスケジューリングやデータ転送の管理をもっとうまくやる必要があるんだ。

#Mixture-of-Expertsレイヤーの理解

MoEレイヤーは、異なる作業負荷を処理するために設計された複数のエキスパートで構成されてる。このシステムは、どのエキスパートがどのデータの部分を処理するかを決めるゲーティング関数を使う。結果として、処理されるときには、特定の時点でアクティブなのはごく一部のエキスパートだけ。これによって、計算時間を二倍にすることなくモデルをスケールアップできるんだ。

各エキスパートは通常独立して動いていて、小さなニューラルネットワークだと考えられる。サイズは小さいけど、組み合わせると強力なモデルができる。でも、ゲーティングメカニズムによってエキスパート間で負荷が不均一になることもある。そのため、エキスパートが同時に処理できるタスクの数に制限をかけて、よりバランスの取れた負荷を確保してるんだ。

#並列処理の種類

大規模モデルのトレーニングには、いくつかの異なる並列処理が使われてる：

1. データ並列処理（DP）：データを複数のデバイスに分け、各デバイスで同じモデルを使う。各デバイスがデータの一部を同時に処理する。

2. モデル並列処理（MP）：モデル自体を異なるデバイスに分割する。モデルの一部が別々のGPUで処理される。

3. パイプライン並列処理（PP）：モデルをステージに分けて、異なる部分を同時にデータが流れるようにする。

4. エキスパート並列処理（EP）：エキスパートを複数のデバイスに配置する。一つのデバイスにすべてのエキスパートを持てない場合、作業がいくつかのデバイスに分配される。

5. エキスパートシャーディング並列処理（ESP）：一つのエキスパートを小さな部分に分解し、複数のデバイスで実行できるようにする。

MoEモデルは通常、効率のためにEPとESPを一緒に使うけど、トレーニング中の通信オーバーヘッドにはまだ課題があるんだ。

#パフォーマンスのボトルネック

大規模MoEモデルのトレーニング速度は、GPU間の通信にかかる時間によって大きく影響される。研究によると、通信は高性能GPUクラスター上のMoEレイヤーのトレーニング時間の最大60%を占めることがある。この問題は、複数の並列セットアップを同時に使うと悪化するんだ。

トレーニングパフォーマンスを改善する方法を探る中で、研究者たちは様々なアプローチを取ってきた。一部はデバイス間の作業負荷をバランスよく分散させるアルゴリズムに焦点を当て、他は通信方法の効率を改善することに取り組んできた。でも、過去の多くの解決策は主にEPに関連する通信コストに焦点を当てていたけど、MPやESPに伴う他の通信の側面を見落としてたんだ。

#Parmシステム

これらの課題に応えるために、Parmが作られた。これは、MP、EP、ESPを一緒に使うときにタスク管理を改善するための新しい通信スケジュールを2つ提案してる。このスケジュールの核心アイデアは、通信タスクを戦略的に配置することによって不必要な計算と通信を減らすことだよ。

#Parmの主な特徴

1. 通信量の削減：Parmはデバイス間でデータがどれだけ頻繁に送受信されるかを管理することに焦点を当てて、全体の通信を減らす。

2. 重複タスクの排除：同じ計算が異なるデバイスで何度も行われないように、タスクを整理する。

3. タスクの重なり：Parmは異なるタイプの通信が互いに重なることを許可する。これは利用可能な帯域幅をフル活用するために重要。

専用のスケジュールを導入することで、Parmはタスクをより効率的に実行できるようになり、トレーニング時間を短縮するんだ。

#Parmの動き方

Parmは特定のタスクの間、MPを一時的に無効にしてパフォーマンスを最適化する。二つのスケジュールが実装されていて、一つはMoEレイヤーのゲーティング操作の前、もう一つは結合操作の前だ。これにより、タスクの重複を防ぎ、情報の流れをよりスムーズにする。

#スケジュールの概要

• 最初のスケジュール：ゲーティング関数の前にMPを無効にして、結合操作の後に再びアクティブにする。これによりタスクの効率的な分担が可能になる。

• 二つ目のスケジュール：ゲーティング関数の後にMPをオフにして、結合プロセスの前に再びアクティブにする。これは特定の条件に応じた柔軟性を提供する。

各スケジュールは、さまざまな設定に応じた明確な利点があり、適切なものを選ぶことでトレーニング時間を最小限に抑えることができる。

#実験結果

Parmの効果を検証するために、8-GPUサーバーと32-GPUクラスターを使って広範なテストが行われた。結果は、ParmがDeepSpeed-MoEのような既存のシステムを大きく上回り、多様なMoE設定で1.13倍から5.77倍のスピードアップを達成することが示された。

#パフォーマンスの指標

さまざまなセットアップにおいて、Parmは全体のトレーニング時間を削減し、GPT-2やBERTに基づく実世界モデルもトレーニングできた。実験は、Parmのスケジュール戦略が複数のシナリオで一貫して速いトレーニングパフォーマンスを生み出すことを確認した。

#結論

Parmの誕生は、大規模AIモデルのトレーニングにおける重要な前進を示してる。タスクと通信のスケジューリングを改善することで、効率と速度において大きな成果を上げてる。より大きくて複雑なAIモデルの需要が高まる中で、Parmのようなシステムは、これらのモデルを効果的にトレーニングするための障害を克服するのに重要なんだ。