マジック: メモリプーリングと分散 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-Torus トポロジーに新しい空間を開くことです。結局のところ、スイッチのレイテンシも節約され、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