MTIA 300、12基のRDMA NICと集団通信エンジンをチップに統合し、PCIeと演算コアの競合を回避

Metaは、自社開発アクセラレータMTIA 300のアーキテクチャに関する詳細をさらに公開した。狙いは汎用GPUで最高の浮動小数点演算スループットを追求することではなく、大規模なレコメンデーションモデルやランキングモデルの通信ボトルネックを解消することにある。この種のモデルでは、パラメータの99%以上がembedding tableに格納される場合があり、学習時には数百基のアクセラレータ間でAllReduce、AllToAll、AllGatherが頻繁に実行される。集団通信を引き続きGPU kernelで処理すると、行列演算と同じコアを奪い合うことになる。

MTIA 300はパッケージ内に2基のネットワークchipletを搭載し、合計12基のカスタム800Gbps RDMA NICを内蔵する。総I/O帯域幅は1.2TB/sで、データがPCIeを経由する必要はない。同じNICリソースを、ラック内の1TB/s scale-upとラック間の200GB/s scale-outに再配分できる。また、チップには独立したmessage engineが16基あり、それぞれがRISC-V制御コア、NICインターフェース、および1サイクル当たり128 bytesを処理するニアメモリ・リダクション・ユニットを統合している。合計リダクション・スループットは2.8TB/sを超える。

関連ソフトウェアのHCCLは、実行時にホストが通信を段階的に駆動する方式ではない。まずcollectiveを、依存関係が明示されたワークキューのサブグラフへコンパイルし、その後message engineに自律実行させる。PyTorchの`c10d`、TorchComms、`torch.compile`を通じて計算グラフに統合できるほか、one-sided communicationとdevice-triggered collectiveの経路も提供する。Metaの報告では、ラック内帯域幅は最大940GB/sに達した。集団通信と大規模GEMMを同時に実行した場合でも、演算スループットの低下は0.5%未満だった。40基のMTIA 300で150Bパラメータの本番レコメンデーションモデルを実行した際、総通信時間は比較対象のGPUクラスタの1/3.9だったという。

ただし、これらの数値は、ベンダー自身のモデル、トポロジー、ソフトウェアスタックを使用して得られた結果である。MTIA 300は216GBのHBM3Eを搭載するが、Metaはチップの消費電力、システム全体のコスト、製造プロセス、比較対象となったGPUの機種を公表していない。また、行列演算が中心となる最先端LLMの学習に適していることも実証されていない。エンジニアリングチームが次に注視すべき点は、HCCLの利用可能性、クロスフレームワーク対応、そして内蔵ネットワーク設計が推論時の小規模メッセージや高並行負荷でも同様の優位性を維持できるかどうかである。