はじめに
AWSがAI学習用チップである Trainium を発表したのが、AWS re:Invent 2020 のこと
2022年10月12日の技術ブログ、
によると、チップには HBM2E メモリが2個付いているということがわかります。説明のために下図を引用します。
各チップには、NeuronCore-v2 が2個載っています。1つのNeuronCore-v2が1個のHBM2eを使う感じですかね。
NeronCore-v2 には、
の4つのEngineが搭載されています。GSIMD Engine以外の3つのEngineは、Google TPU v2/v3/v4 に非常に似た構成です。
PCIe は、Gen5です。チップ間接続のために NeuronLink-v2 なるものが4個付いています。
AWS Trainium の実物
CNBC の TECH : How Amazon is racing to catch Microsoft and Google in generative A.I. with custom AWS chips
に Trainium の写真が載っています。下記の写真(複数)を説明のために引用しました。
2つのHBM2E と Trainium die が 1つ載っています。
追記)2023.08.18
Youtube にもアップされました。
によると、メモリは、32GiB of device memory (for storing model state), with 820 GiB/sec of bandwidth とあります。
820GB/s / 2 chips / 1024bit * 8bit = 3.2 Gbps
にようです。HBM2Eのフルスペックで動いているようです。
v1 と v2 の大きな違いは、GPSIMD Engine の有無のようです。
Trainiumは、2D Torus で繋がっているようです。
おわりに
下記は、上記の技術ブログの後半の Pytorch のコードの一部です。torch_xla を使っているようですね。
参考記事
1ノードに、8 x OAM っぽい。下図を説明のために引用します。
下図を見ると、右上に黄色いケーブルが6本あります。Trainium は 2D Torus なので、ノード間の接続に12本のケーブルを出す必要がありそうです。右上は黄色いケーブルですが、左上には赤色のケーブルが見えますね。
