はじめに

AWS Trainium の詳細については、あまり調べていなかったです。わかっていることは、SRAMを scratch pad として使うこと。

今回は、AWS の

にあった、Trainium の詳細について、記録に残します。

以下の図は、上記のサイトから説明のために引用します。

Trainium Device Diagram

下記の Trainium の Block diagram は、

にある下図と、ちょっと違います。

並べてみました。

左側の図では、DMAが32個、CC-Coreが6個と数が出ています。一方、右側の図では、PCIe は Gen5、HBM は HBM2E とあります。

各 NeuronCore-v2 は、SBUFというSRAM経由でHBM(HBM2E)と接続しています。

各 Compute Engine の動作周波数は、下記のようになっています。

SBUF と HBM 間の転送は、16個のDMAで行っています。1つのHBM2Eに対して、16個のDMAということは、疑似チャネルの16 (64bit x 16) に対応しているのだと思います。

HBM2E が2個、付いているのは、各NeuronCore-v2 に対して、1個の HBM2E が接続しているということになります。

他にも色々と書いてあるので、じっくり読むといいかもしれません。