HUOXIU

LLM分散トレーニングの6つの主要テクノロジーの紹介

編集者注:この記事は分散型ニューラル ネットワークのトレーニング手法に焦点を当てており、著者は大規模モデルのトレーニングにおけるハードウェア利用効率を向上させる革新的な方法を体系的に詳しく説明しています。

この記事では、6 つの主要な分散トレーニング手法に焦点を当てています。

  1. データ並列トレーニング: データのミニバッチを複数のワーカーに分散することで、並列勾配計算と効率的なトレーニングを実現します。
  2. Butterfly All-Reduce: 革新的なデータ分割および集約方法により、通信コストを効果的に削減します。
  3. ゴシップベースの平均化: システムのフォールト トレランスとスケーラビリティを向上させる分散型通信戦略。
  4. Moshpit 勾配降下法: ワーカーが小規模で独立したグループ内で勾配平均化を実行できるようにし、トレーニングのフォールト トレランスを強化します。
  5. DiLoCo: ローカル パラメータ更新とグローバル パラメータ更新を組み合わせて、収束速度とシステム パフォーマンスのバランスをとる革新的な内部および外部最適化アルゴリズムです。
  6. SWARM: 異機種ハードウェア環境でのリソース割り当てを最適化するために、動的なタスク割り当てとフォールト トレランス メカニズムを導入します。

著者 | ロバート・ランゲ

編纂者:岳陽

人工知能技術の進歩に伴い、大規模ニューラルネットワーク(大規模言語モデルを含む)の学習の重要性がますます高まっています。これらのモデルの規模と複雑さは、学習コストとエネルギー消費量の増加につながるだけでなく、ハードウェアの利用効率の向上も急務となっています。これらの課題に対処するため、研究者やエンジニアは分散型・非中央集権型の学習手法を模索しています。本稿では、データ並列学習手法やゴシップベース平均化といった様々な分散学習手法について解説し、これらの手法がモデルの学習効率を最適化しながら、この分野の高まる需要に応える方法を示します。

多数の小型 GPU を組み込んだ、シンプルな日本風に描かれた GPU クラスターの図 (OpenAI の Dallé-3 API を使用して生成)。

01 データ並列トレーニング技術、完全縮約操作、ノード同期

データ並列学習技術は、複数のワーカーノードにデータのミニバッチを分散させることで効率的な学習を実現します。このアプローチは、複数のワーカーが並列に勾配を計算できるため学習プロセスを高速化するだけでなく、単一のデバイスよりも大きなバッチサイズを処理できるようになります。すべてのワーカー間でモデルの更新を同期させるために、完全な縮約演算を採用しています。この演算は、すべてのワーカーからの勾配を集約して平均化し、モデルを均一に更新することで、分散システム全体でモデルの一貫性を確保します。

以下は、PyTorch を使用して Python でこのプロセスを示す簡単な例です。

完全縮約に加えて、パラメータサーバーを使用するアプローチもあります。このアーキテクチャでは、中央サーバーが勾配情報の収集とオプティマイザーの状態監視を担当します。これにより同期プロセスは簡素化されますが、単一障害点のリスクがあり、システムパフォーマンスのボトルネックとなる可能性があります。

分散学習において、Hogwild(Recht et al., 2011)[1]はもう一つのよく知られた手法である。これは非同期方式を用いてモデルパラメータを更新し、全ての計算ノードを同期させることなく実行できる。この手法は教師あり学習だけでなく、非同期アクタークリティックアルゴリズム(A3C、Mnih et al., 2016)[2]などの強化学習(RL)シナリオにも適用できる。A3Cでは、複数のエージェントが同時に環境と相互作用し、自身の経験に基づいて同じモデルを非同期的に更新することができる。これにより、リソース利用効率が向上するだけでなく、複数のエージェントの異なる経験を通じて収束速度が加速され、複雑な環境でのパフォーマンスが向上する。

データ並列トレーニング手法に加えて、モデル並列やパイプライン並列といった並列トレーニング手法も存在します(詳細はLlian Wengのブログ[3]を参照)。モデル並列化では、モデルを複数のコンピューティングデバイスに分割し、モデルの異なる部分を同時に処理することができます。これは、単一のデバイスでは処理できない超大規模モデルに特に有効です。パイプライン並列化では、モデルを複数のステージに分割し、各ミニバッチのデータをこれらのステージで順次処理します。これにより、並列計算と通信が可能になり、トレーニング全体の効率とスループットが向上します。これらの手法は互いに補完し合い、大規模なトレーニングシナリオにおけるリソース利用を共同で最適化します。

02 バタフライオールリデュース

バタフライ・オールリデュース(Butterfly All-Reduce)手法(Zhao and Canny, 2013)は、従来の完全リダクション手法が直面する課題を効果的に解決します。この手法では、参加ノード(合計N個)がローカルデータをN個の部分に分割します。次に、i番目のノードが他のすべてのノードからi番目の部分のデータを集め、集約した後、各ノードに均等に分配します。

このアプローチは通信負荷を大幅に軽減し、システムのスケーラビリティを向上させます。分散学習において、いわゆる「ワールドサイズ」とは、学習に参加するプロセスまたはデバイスの総数を指します。このパラメータは、ノード間でデータがどのように集約され、同期されるかを決定する上で重要な役割を果たします。

以下は、Butterfly All-Reduce テクノロジーの概念的な実装例です。

このコードは、分散システムで同期を維持しながら、Butterfly all-reduce 手法によって並列処理の利点を効果的に活用する方法を示しています。

バタフライオールリデュース法は、従来のフルリデュース法と比較して通信コストを大幅に削減し、優れたスケーラビリティを提供するという利点があり、大規模分散システムに適しています。しかし、この手法にはいくつかの欠点もあります。例えば、実装が比較的複雑であり、通信ネットワークのトポロジやネットワーク状況によってパフォーマンスが影響を受ける可能性があります。さらに、参加ノードのいずれかに障害が発生すると、システム全体の同期プロセスに影響を及ぼす可能性があります。

特定のアプリケーションシナリオ、特に連合学習においては、学習プロセスは不安定なネットワーク帯域幅や信頼性の低いワーカーノードに適応できる必要があります。連合学習は、機密性の高いプライバシーデータを保持する複数の独立した参加ノードが関与するため、特に複雑です。このような状況では、モデル学習の信頼性を確保するために、堅牢な戦略を採用する必要があります。次に、収束速度とシステムのフォールトトレランスのバランスをとることを目的としたいくつかの手法を検討します。

03 ゴシップベースの平均化

ゴシップベース平均化(Boyd et al., 2005)は、参加ノードが疎な通信ネットワークを構築する分散通信戦略です。各ノードは定期的に隣接ノードからパラメータを取得し、自身のローカルパラメータと組み合わせます。このアプローチはパラメータサーバの通信負荷を軽減しますが、各ノードが異なるローカルパラメータを使用して計算を行う可能性があることも意味します。

ゴシップベースの平均化の収束特性は、通信ネットワークのアーキテクチャに大きく影響されます。以下は、ゴシップベースの平均化の実装例です。

ゴシップベースの平均化には、次の利点があります。

  • 通信ボトルネックの軽減: 集中型パラメータ サーバーが不要なため、ゴシップの平均化によって通信の混雑が大幅に軽減され、パラメータの更新がより効率的になります。
  • スケーラビリティ: このアプローチは分散型であるため、スケーラビリティの点で優れており、過度のオーバーヘッドを発生させることなく、参加ノード数の増加を簡単に処理できます。
  • フォールトトレランス:分散設計により、システムのフォールトトレランスが向上します。1つのワーカーに障害が発生しても、トレーニングプロセス全体が中断されることはありません。他のワーカーは引き続き通信を行い、パラメータを更新できます。

ただし、この方法にはいくつかの欠点がある可能性があることにも注意する必要があります。

  • 収束速度の低下: 集中型の更新方法と比較すると、ゴシップ平均化ではパラメータの集約が頻繁ではなく、各ワーカーが比較的最近のデータに基づいて計算を実行する必要があるため、収束速度が遅くなる可能性があります。
  • パラメータ更新の不一致: 各ノードが異なるローカルパラメータを使用するため、パラメータ更新に不一致が生じる可能性があり、それが収束の安定性と速度に影響します。
  • 通信グラフへの依存性:ゴシップ平均化の有効性は、通信グラフの構造に大きく依存します。グラフの接続性が低い場合や、グラフ構造が不均衡な場合、アルゴリズム全体のパフォーマンスに悪影響を与える可能性があります。

要約すると、分散型パラメータ更新方法であるゴシップベースの平均化は大きな可能性を秘めていますが、実際のアプリケーションでは、特定のトレーニングシナリオに基づいてその利点と欠点を比較検討する必要があります。

04 モッシュピット・グラディエント・ディセント

Moshpit勾配降下法(Ryabinin et al., 2021)[4]は、分散学習の概念をさらに発展させ、ワーカーが小規模かつ独立したグループ内で勾配平均化を実行できるようにします。この設計により、参加ノードに問題が発生した場合でも、影響はそのグループに限定されるため、学習プロセス全体の耐障害性が向上し、グローバルな学習の中断を回避できます。

これらのグループの動的な形成は、トレーニングの有効性を確保するために不可欠です。グループ構造を最適化することで、この手法は収束に必要なステップ数を大幅に削減します。これは、ワーカーが小規模なチーム内で勾配情報をより効率的に交換・更新できるためです。この適応的なグループ化戦略は、既存のリソースをより有効に活用し、さまざまなネットワーク環境において優れたパフォーマンスを実現するのに役立ちます。

以下は、モッシュピット勾配降下法を実装するための概念フレームワークです。

モッシュピット勾配降下法の利点は次のとおりです。

  • フォールト トレランス: 個々のワーカーの障害はそのグループにのみ影響し、トレーニング プロセス全体には影響しません。一方、他のグループは通常どおりトレーニングを続行できます。
  • リソース利用効率: 小規模グループ内で更新することで、ネットワーク状況や作業者の可用性の変化に柔軟に対応でき、トレーニングの効率が向上します。
  • 通信負荷の軽減:通信がグループ内に限定されるため、全体の通信量が軽減され、特に帯域幅が制限されている場合に有効です。

ただし、この方法にもいくつかの欠点があります。

  • 収束の課題: グループ構造が絶えず変化すると、パラメータの更新に矛盾が生じ、トレーニングの収束と安定性に課題が生じる可能性があります。
  • 管理の複雑さの増大:グループを動的に管理・調整することは、トレーニングプロセスの複雑さを間違いなく増大させます。最適なグループ構成を見つけるには、より洗練されたメカニズムを開発する必要があります。
  • スケーラビリティの問題: グループを小さくするとシステムのフォールト トレランスが向上しますが、効果的な管理がなければ、大規模なトレーニング シナリオではこのアプローチのスケーラビリティが制限される可能性があります。

要約すると、モッシュピット勾配降下法は、新しい分散学習手法として大きな可能性を秘めています。フォールトトレランスとリソース利用効率という利点と、収束の課題や実装の複雑さとの間で、絶妙なバランスを実現しています。

05 DiLoCo: 内外最適化

DiLoCo (Douillard et al., 2023) [5] は、分散学習の効率を向上させるために設計された革新的な内部外部最適化アルゴリズムを導入しました。このアルゴリズムでは、各コンピューティングノードは、内部最適化フェーズ中にローカルAdamWオプティマイザーを用いて複数のパラメータ更新を実行します。この設計により、ノードは他のノードとリアルタイムで同期することなく、ローカルデータに基づいて独立してパラメータを最適化できます。一定回数(通常約500回)のローカル更新が完了すると、外部最適化フェーズが開始され、その時点ですべてのノードの擬似勾配(ローカル更新の結果の合計)が同期されます。

このアプローチは、ローカル更新とグローバル更新の利点を巧みに組み合わせることで、収束を加速し、トレーニングパフォーマンスを向上させる可能性があります。DiLoCoは、ノードがグローバルモデルと同期する前にパラメータをローカルに最適化できるようにすることで、両方の更新戦略の長所を活用します。

以下は、DiLoCo 更新プロセスの概念的な説明です。

DiLoCoはもともとGoogle DeepMindによって実装されましたが、現在では新興企業PrimeIntellectがこのアプローチを再現することに成功しています。OpenDiLoCo (Jaghouar et al., 2024)[6]はGitHub[7]で公開され、Hivemindライブラリ[8]を用いて10億パラメータのモデルを学習しました。最近、PrimeIntellectは独自のカスタムインフラストラクチャ[9]を立ち上げました。これには、カスタムのall-reduceアルゴリズムや通信プロトコルなど、多くのエンジニアリングイノベーションが含まれています。同社は現在、Intellect-1[10]と呼ばれる100億パラメータのモデルを学習中です。この実験の結果は、既存モデルのブレークスルーに大きな影響を与えると信じています。現在、大規模モデルの学習は依然として集中型のコンピューティングリソースに依存しています。しかし将来的には、誰もが次世代の主要な基礎モデルの構築に貢献できるようになるかもしれません。

06 SWARM: フォールトトレランスと動的タスク割り当て

SWARMアルゴリズム(Ryabinin et al., 2023)[11]は、各ワーカーノードがトレーニングプロセスの後半段階で出力を他のワーカーノードに送信できる、新しい分散トレーニング手法を導入しています。この柔軟なタスク割り当て手法により、強力なコンピューティング能力を持つデバイスがより多くのタスクを処理できるようになり、多様なハードウェア環境において最適なリソース割り当てを実現します。この戦略は、コンピューティングリソースの変動が大きいシナリオで特に効果的であり、よりバランスの取れたワークロードを実現し、アイドル時間を削減できます。

SWARMアルゴリズムは、ワーカーノードに障害が発生した場合、障害が発生したノードから正常に動作している他のノードにタスクを迅速に転送することで、フォールトトレランスを発揮します。このメカニズムは、トレーニングプロセスの継続性を維持し、予期せぬ中断の影響を効果的に軽減し、処理能力のタイムリーな回復を保証するために不可欠です。ワーカーノード間の通信パスはランダムかつ動的に調整されるため、アルゴリズムはネットワーク状況やノードの状態の変化に応じてリアルタイムに適応できます。

この適応型通信方式は、データフローを効率化するだけでなく、トレーニングプロセスの安定性も向上させます。以下は、SWARM通信の実装例を簡略化したものです。

この例では、アクティブなワーカーノードはそれぞれ、情報伝達の対象として隣接ノードをランダムに選択します。この分散通信モデルは、現在のシステム状態にリアルタイムで適応できます。動的なタスク割り当てと強力なフォールトトレランスを備えたSWARMアルゴリズムは、大規模な機械学習シナリオにおける分散学習の効率と信頼性を大幅に向上させます。

07 結論

分散学習は、大規模ニューラルネットワークを効率的に学習するための強力な基盤を提供します。データ並列学習手法、バタフライオールリデュース、ゴシップベースの平均化などの手法を活用することで、実践者は様々な環境におけるモデル学習の課題に取り組むことができます。これらの手法を習得することは、大規模AIシステムのパフォーマンスを最適化しようとするすべての人にとって不可欠です。この分野の研究が深まるにつれ、これらの手法を理解することが、分散学習の潜在能力を最大限に引き出す鍵となるでしょう。この記事では、すべての分散学習手法と最新の研究成果を網羅するのではなく、大まかな概要を提供するにとどめています。そのため、読者の皆様には、より技術的な詳細についてはご自身で調査することをお勧めします。

読んでくれてありがとう!

このブログを楽しんで、新しいことを学んでいただければ幸いです。

著者について

ロバート・ランゲ

ベルリン工科大学でディープラーニングの博士号を取得。Sakana.AIで研究員。✍️ Google DeepMindインターンシップ2回経験

終わり

今週のインタラクティブコンテンツ🍻

❓分散トレーニングにおいて、最大の技術的ボトルネックは何だと思いますか?通信オーバーヘッド、収束速度、システムのフォールトトレランス、それとも何か他のものでしょうか?

🔗記事内のリンク🔗

[1]https://proceedings.neurips.cc/paper_files/paper/2011/file/218a0aefd1d1a4be65601cc6ddc1520e-Paper.pdf

[2]https://arxiv.org/abs/1602.01783

[3]https://lilianweng.github.io/posts/2021-09-25-train-large/

[4]https://openreview.net/pdf?id=cwWfDHYpb1z

[5]https://arxiv.org/abs/2311.08105

[6]https://arxiv.org/abs/2407.07852

[7]https://github.com/PrimeIntellect-ai/OpenDiLoCo

[8]https://github.com/learning-at-home/hivemind

[9]https://github.com/PrimeIntellect-ai/prime

[10]https://www.primeintellect.ai/blog/intellect-1

[11]https://arxiv.org/abs/2301.11913

オリジナルリンク:

https://towardsdatascience.com/分散分散型ニューラルネットワークのトレーニング-入門-21e5e961fce1