はじめに

先週は、Hotchips 33 にて、いろいろなチップの情報が公開されました。

Hotchips 33 そのものは、有料なので資料はダウンロードできませんが、Anandtech が実況中継してくれています。プレゼン資料だけでなく、その時に話した内容も文字におこしています。もちろん、Q＆Aもあります。

最後に各セッションの実況ログのページのリンクを張っておきます。

Cerebras Systems の MemoryX

Hotchips 33で、Cerebras Systemsは CS-2 と一緒に使う MemoryX なる装置を発表しました。また、MemoryXと複数のCS-2を接続するためのSwarmXも発表しました。CS-2は今日から販売し、価格は数億円(最初の方)するということです。

Cerebras Systems のセッションの実況中継ログは、下記にあります。タイトルのあるように、Graphcore => Cerebras Systems => SambaNova Systems => Anton3 の順番になっています。

www.anandtech.com

下記がMemoryXの発表時のスライドです。説明のために引用します。MemoryXを使うことにより、120T parameters のモデルを CS-2 にて学習できるようにするようです。

f:id:Vengineer:20210829102217p:plain

下記のスライドもCerebrasのものです。モデルがどんどん大きくなっているのを示しています。説明のために引用します。 Cerebras が WSE として、12インチのシリコンウェハー上に84個のdie相当のものを2次元に並べ、外部にDRAMを接続しない方法で学習するというCS-1 => CS-2 を発表してきましたが、内部SRAMは CS-1で18GB、CS-2で40GBしかありません。40GBというと、NVIDIA A100 のメモリ容量が少ない版、1台と同じです。つまり、40GB、1つのパラメータが4バイトとすると、10Gパラメータまで(実際にはもっと少なくなってしまうと思いまうが)モデルはWSEの中に入れることができるが、それ以上のモデルは入らなくなっちゃうということです。10Gパラメータのモデルというと、Megatron-LM(8.3憶パラメータ)ぐらいまでで、それより大きいモデルは CS-2 には入らないということになっちゃいます。MSFT-1Tは1Tパラメータになり、Megatron-LMの100倍です。

f:id:Vengineer:20210829103120p:plain

そこで、MemoryXが登場します。MemoryXには、4TB ～ 2.4PB のメモリ & Flash Storage が搭載されています。4TBがメモリで、2.6PBがFlash Storageだと思います。 4TBあれば、1Tパラメータのモデルに対応でき、それ以上のモデルに対しては、Flash Storage に置いておき、必要な時に、必要なところのパラメータをメモリにロードすれば、もっと大きなモデルにも対応できるということになります。120Tパラメータまで対応できるということなので、2.4PBと120Tパラメータを比べると1パラメータは20バイトと考えるのか？それとも他に何か理由があるのでしょうか？

下記のスライドも説明のために、引用します。 f:id:Vengineer:20210829104934p:plain

MemoryX と CS-2 の間に、SwarmXなる装置を入れることで、複数のCS-2とMemoryXを接続できるようになり、データパラレルで学習できるようになります。 CS-2が数億円なので、4台のCS-2 + SwarmX + MemoryX で、軽く10億円を超えるのではないでしょうか？ NVIDIA DGX A100 が 20万ドル(A100 40G版)ということなので、50台のDGX A100ぐらいのPodが出来そうです。

MemoryXをどのように使うのか？

MemoryXはどのように使うのか？それは、次のスライドで説明ています(引用します)。今までは、入力データをCS-1/CS-2にストリーミングしていましたが、MemoryXを使う場合は、Weightをストリーミングするモデルを実行するというものです。

f:id:Vengineer:20210829104205p:plain

MemoryXからは Weights を CS-2 に送り、CS-2 からはGradients を送ります。そして、MemoryXでは、Optimizer Compute も行います。どのような Optimizer を使うかはわかりませんが、この部分をMemoryXに入れているのが一つのポイントだと思っています。

次のスライドでは、MemoryXとCS-2とのWeight/GradientのLatencyを解決するためにどのようにしているかを示しています。説明のために、引用します。

CS-2 から MemoryX に L1 Gradient を送り、MemoryX の中で L1 Weight Update を行いつつ、L1 Weight Stream Out ということで、L1 Weight を CS-2 に送り込み、その L1 Weightを使って、CS-2 の中で L1 Foward を行う。このようなフローにすることで、Latency を隠蔽できるということになるようです。

f:id:Vengineer:20210829104611p:plain

SwarmX

MemoryX と CS-2 の間に入る SwarmX は、下図(説明のために引用します)のように、Data parallelを行うための装置です。

Data parallel training across CS-2s
Weights are broadcast to all CS-2s
Gradients are reduced on way back

f:id:Vengineer:20210829105134p:plain

下図(説明のために引用します)にあるよに、192台までCS-2が接続でき、100TパラメータまでのモデルでScaleするというものです。192台で止まっているので、100Tパラメータのモデルは192台でOKだったということになるのでしょうか？あ、MemoryX が 120Tパラメータまでのサポートだからでしょうかね。

SwarmX 経由で 192台ものCS-2 を接続するのって、どのような構成になっているのかがきになります。

f:id:Vengineer:20210829105215p:plain