大規模な言語モデルの開発には、時間と GPU リソースへの多大な投資が必要であり、それは直接的にコストの増加につながります。モデルが大きくなるほど、これらの課題は顕著になります。 最近、Yandex は新しいソリューションを導入しました。YaFSDP 、GPU リソースの消費とトレーニング時間を大幅に削減することで、LLM トレーニングに革命を起こすと期待されるオープンソース ツールです。700 億のパラメータを持つモデルを含む事前トレーニング シナリオでは、YaFSDP を使用すると約 150 個の GPU のリソースを節約できます。これは、仮想 GPU プロバイダーまたはプラットフォームによって異なりますが、毎月約 50 万ドルから 150 万ドルの節約に相当します。 は YandexはYaFSDPを一般公開しました。 。 GitHub 分散型LLMトレーニングの課題 複数の GPU にわたる LLM のトレーニングには、非効率性とメモリ消費量の増加につながる複雑な操作が伴います。主な問題の 1 つは、GPU 間で大量のデータを送受信する必要があることです。たとえば、一般的な all_reduce 操作では、ネットワーク パラメータの 2 倍の量の勾配データを通信する必要があります。Llama 70B モデルの場合、これは反復ごとに 280 GB のデータを転送することを意味します。 さらに、重み、勾配、およびオプティマイザーの状態が GPU 間で複製されるため、膨大なメモリ負荷が発生します。Llama 70B モデルと Adam オプティマイザーは 1 TB を超えるメモリを必要とし、ほとんどの GPU の標準的な 80 GB のメモリ容量をはるかに超えています。この冗長性によりトレーニング プロセスが大幅に遅くなり、中程度のサイズのモデルでさえ GPU メモリに収めるのが困難になることがよくあります。 YaFSDPの紹介 Yandex の YaFSDP は、これらの課題に対して非常に効果的なソリューションを提供します。メモリ消費の最適化と通信ボトルネックの解消に重点を置くことで、YaFSDP は LLM トレーニングの効率を高めます。個々のパラメータではなくレイヤーをシャーディングすることで、効率的な通信を維持し、冗長な操作を回避します。さらに、YaFSDP は必要なすべてのデータにバッファを事前割り当てし、Torch アロケータが非効率性を引き起こさないようにします。 YaFSDP は中間重みと勾配に 2 つのバッファーを利用して動作し、奇数レイヤーは 1 つのバッファーを使用し、偶数レイヤーはもう 1 つのバッファーを使用します。 異なるレイヤーからの重みは同じメモリに保存されます。レイヤーの構造が同じであれば、それらは常に同一になります。レイヤー X が必要なときに、バッファーにレイヤー X の重みが含まれていることを確認することが重要です。すべてのパラメーターは、バッファー内の対応するメモリ チャンクに保存されます。 メモリ消費量 トレーニング中、メモリを主に消費するのは、重み、勾配、オプティマイザーの状態、バッファ、アクティベーションです。YaFSDP は、これらの要素の保存方法とアクセス方法を最適化することで、メモリ消費を大幅に削減します。 : これらはプロセスの数に依存し、プロセスの数が増えるにつれてメモリ消費量はゼロに近づく傾向があります。これらのコンポーネントを GPU 間でシャーディングすることにより、YaFSDP は重複を最小限に抑え、メモリ使用量を削減します。 重み、勾配、およびオプティマイザーの状態 一定量のメモリを消費し、計算中に中間値を保存します。 バッファは モデルのサイズと GPU ごとに処理されるトークンの数によって異なります。 アクティベーションは、 アクティベーションチェックポイント アクティベーション チェックポイントは、フォワード パス中に必要なアクティベーションのみを保存し、バックワード パス中にそれらを再計算する手法です。これにより、必須データのみが保存されるため、メモリ フットプリントが大幅に削減されます。たとえば、バッチ サイズが 8192 トークンの Llama 2 70B モデルをトレーニングする場合、アクティベーション ストレージを 110 GB 以上からわずか 5 GB に削減できます。 ただし、このアプローチでは追加の計算オーバーヘッドが発生しますが、YaFSDP ではメモリの最適化により一部のレイヤーでアクティベーション チェックポイントを使用しないことでこれを回避できます。 コミュニケーションの最適化 YaFSDP は、必要なときだけデータを転送し、通信と計算をオーバーラップさせることで、GPU 通信の効率を向上させます。CUDA ストリームを利用して、同時計算と通信を効率的に管理します。 このツールは、計算ストリームと通信ストリームの 2 つのストリームを使用します。イベントはこれらのストリームを同期し、デッドロックが発生することなく操作が正しい順序で実行されるようにします。 3 番目のレイヤーのフォワード パスは、all_gather 操作が完了するまで開始されません (条件 1)。同様に、3 番目のレイヤーの all_gather 操作は、同じバッファーを使用する 1 番目のレイヤーのフォワード パスが完了するまで開始されません (条件 2)。この方式にはサイクルがないため、デッドロックは発生しません。 実験結果とパフォーマンスの向上 YaFSDP の実装により、トレーニング効率が著しく向上しました。700 億のパラメータを持つモデルを使用した事前トレーニング シナリオでは、YaFSDP は約 150 個の GPU のリソースを節約できました。これは、仮想 GPU プロバイダーまたはプラットフォームに応じて、毎月 50 万ドルから 150 万ドルの大幅なコスト削減につながります。 YaFSDP は、FSDP などの既存の方法と比較してトレーニング時間を最大 26% 短縮し、メモリ使用量を最適化するため、より大規模なモデルをより効率的にトレーニングできるようになります。 YandexはYaFSDPを一般公開しました。 ML エンジニアはこのツールを活用して、LLM トレーニング プロセスの効率を高めることができます。Yandex は YaFSDP をオープンソース化することで、AI コミュニティにおけるイノベーションとコラボレーションを促進し、開発者がより迅速かつコスト効率よくモデルをトレーニングできるようにすることを目指しています。 GitHub YaFSDP は、LLM トレーニングにおける大きな進歩を表しています。メモリ消費と通信の非効率性という重大な課題に対処することで、大規模言語モデルのトレーニングをより高速かつ効率的に行うことができます。