マジック: メモリプーリングと分散 AI クラスターの最適化

[[429309]]

分散機械学習が登場した理由は非常に単純です。一方で、トレーニングに利用できるデータはますます増えており、他方ではモデル自体の規模もますます大きくなっているため、複数のマシンが必要になります。 RoCE などの通信プロトコルが自然に使用されており、これが実際に nVidia が Mellanox を買収したい根本的な理由です。並列処理には主に 2 つの方法があります。

データの並列性は簡単に説明できます。主に、トレーニング サンプルを保存し、複数のマシン間で混合サンプルを転送する方法について説明します。基本的に、誰もがこれらの問題を解決するために SSD と分散ストレージを使用しています。もちろん、メモリプーリングも必要です。

もう 1 つの問題はモデルの並列性です。単一のワーカーノードにデータを保存できない場合は、モデル自体を分割する必要があります。分散トレーニングの各ラウンドが完了したら、パラメータを同期する必要があります。通常、各モデルに対応するパラメータを合計し、平均値を求めます。この通信は AllReduce と呼ばれます。

当初は集中型パラメータ サーバーが使用されていましたが、すぐにクラスター全体のボトルネックになっていることが判明しました。次にリングトポロジーのAll-Reduce

nVidia に関しては、NVLink の帯域幅を拡大し、NCCL を迅速に反復して AllReduce 問題を解決するよう努めていますが、これらは単一のマシン上または非常に密に結合されたクラスター内でのみ利用可能です。一方、ホスト間の通信には、スーパーコンピュータで非常に一般的な RDMA ROCE が自然に選択されます。

しかし、それでも AllReduce の遅延はトレーニング クラスター全体の規模に大きく影響します。

Allreduce アルゴリズムの概要については、Goose Factory の次の概要を参照してください。

テンセントJizhiチームがAllReduceアルゴリズムの歴史と開発を共有[1]

もう一つの仕事: EFLOPS

アリババがHPCA2020で論文を発表

Alibaba も、PCIe の輻輳、内部スケジューリングの輻輳、ネットワーク カードの輻輳など、この問題を明確に認識しています。

解決策は簡単です。お金がたくさんあるので、1 つの GPU に 1 つのネットワーク カードを装備できます。次に、ネットワーク ポートを増やし、スイッチング ネットワークを 2 セットの Fat-Tree に変更します。

アリの記事には結論がある

つまり、HDRM を使用しても、帯域幅は 40Gbps しかありません。 NetDAM の結論を述べます。100Gbps は簡単にフル活用でき、1 台の Alveo U55N で 200Gbps をフル活用できます。ここに一文があります: 私の中で虎がバラの匂いを嗅ぎます。

NetDAMはAllReduceを実装する

まず、AllReduce アルゴリズムと対応するトポロジについては説明しません。一定の帯域幅の条件下では、制約は通信遅延と計算遅延です。 RoCE を使用すると、マシンからの読み取りと書き込みは 1 回 PCIe を経由する必要があるため、この問題の根本的な解決策はメモリを前に置くことであり、遅延は削減されることになりますよね?

通信遅延が短縮されたので、計算遅延とそれを従来の方法で追加する方法を見てみましょう。

コンピューティング領域では、CPU は AVX512 を使用して何ができるでしょうか?キャッシュ遅延ジッタを制御するのは困難です。 GPU にはさらに 1 つのメモリ コピーがあります。 GPU-Direct を直接使用する場合でも、PCIe を経由する必要があります。そのため、100Gbps の回線速度で動作できないのは当然です。

多数の ALU をネットワーク カード上に直接配置します。パケットが受信されると、そのパケットはまだ SRAM バッファ内に残ります。このとき、パケット ヘッダー内の NetDAM 命令に従って、複数の ALU が同時にローカル DRAM をロードし、対応する SRAM に追加することができます。追加後、IP ヘッダーを変更した後、パケット全体が直接転送されます。この方法では、9000B メッセージは 2048 個の float32 を伝送でき、これは AVX (32*2048) SIMD-Add に相当します。だから当然、私の CPU はあなたの CPU よりも高速です。また、追加する場合は DRAM ストアはなく、最後のホップのみが格納されます。どれくらい節約できるでしょうか？

さらに、AI トレーニング シナリオでは、NetDAM ALU を介して直接フィルター処理できる方法が多数あります。例えば、SIMD内の2048個のfloat32の半分以上が0の場合、次のように単純に使用できます。

AllGather 段階、つまり計算されたデータを再度配布する必要がある場合、RoCE のマルチキャストは単なる概念であり、NetDAM はこの段階でイーサネット マルチキャストまたはブロードキャスト機能を最大限に活用できます。もちろん、NetDAM 間で特定のパケット損失の再送信を完了することもできます。 FPGA は Seq 損失を検出すると、ソースへの READ メッセージを直接生成し、CPU が関与する必要はまったくありません。具体的な内容については、明日の混雑制御のお話の際に詳しくお話しさせていただきます。

重要なのはそれだけではありません。セグメント ルーティング ヘッドも搭載されており、原子爆弾のように連鎖反応を起こすことができます。

連鎖反応のもう 1 つの特徴は、3D トーラス トポロジーに新しい空間を開くことです。結局のところ、スイッチのレイテンシも節約され、RingAllreduce はインキャストを考慮せずに全帯域幅で実行するために使用できます。美しいでしょう？素晴らしいと思いませんか？

Google TPU クラスターが Torus-Ring を使用し、一部のスーパーコンピューターが 6D-Torus を使用する理由をご存知ですか?実際、それはスケーラビリティのためです。ノンブロッキングは非常にコストがかかり、クラスター サイズを一時的に拡張または縮小するには、FatTree を形成するために追加のデバイスを追加する必要があります。インキャストも制御が難しいです。 Torus はブロックされていますが、通信モードを通じてブロックを回避できます。

Fat-Tree データセンターを使用している場合でも、トラフィック エンジニアリングを行うための Ruta ソリューションを用意しています。 PortRankと比較すると、混雑していない場所に行くのがより簡単で直感的です。明日は混雑緩和の手配をさせていただきます。

NetDAMはメモリプールを実装する

NetDAM は標準の UDP プロトコルです。 NetDAM はホストから独立して展開できるため、非常に大きなメモリ プールを形成できます。

したがって、通常のホスト ユーザー モードでは特別な開発キットは必要なく、UDP ソケットでメモリ クラスター全体を直接制御できます。かっこいいでしょう？

ホストに独自の NetDAM カードがある場合は、より高度な操作を行って、パーティション化されたグローバル アドレス空間 (PGAS) について学習できます。このシナリオでは、スイッチ チップを MMU に変換し、外部に仮想 IP アドレスと UDP ポートを提供し、大規模な仮想化プールを形成して内部トポロジを隠すことができます。各 netDAM メッセージによってアクセスされるメモリ アドレスは、スイッチング チップによってテーブル参照を通じて最終的な NetDAM に変換されます。この場合、スイッチ MMU はインターリーブ アドレッシングを使用して、ローカル メモリの過熱の問題を解決することもできます...

分散型 AI トレーニング クラスターの観点から続けると、メモリ プールの需要は主にトレーニング データ セットの分散と混同のためであり、もう 1 つはパラメーターと勾配の更新のためです。そこで今回は、HotChip の Cerebras が Memory-X キットを提供します。

コンピューティングタスクでは、MemoryXはオプティマイザーも追加します

結論 NetDAM でもこの機能を実現できます:)

プレビュー....EFLOPS について話した後は、HPCC について話しましょうか?遅延が決定論的である場合、バッファの深さのみを考慮する必要があるため、アルゴリズムはより単純になります。

参照

[1] Tencent JizhiチームはAllReduceアルゴリズムの歴史を共有しました：

https://zhuanlan.zhihu.com/p/79030485