5nmの限界に早くもたどり着いてしまったWSE-3 Hot Chips 2024で注目を浴びたオモシロCPU

　Hot Chips第6弾は、CerebrasのWSE-3を取り上げる。Cerebrasは連載572回で紹介したが、Celebrasは2019年のHotChips 31でWSEを発表して以来、HotChipsの常連となって毎年のようになにか発表している。

　もちろんIBMやインテル/AMDのように毎年多数のチップを発表しているメーカーならこれは珍しくないのだが、WSEシリーズしかないメーカーでこういうのは珍しい(単に発表の場をスポンサー枠で購入しているだけ、というのは口の悪い言い方かもしれないが)。

　ちなみに同社は2018年にTraining Session 3(Accelerating Training in the Cloud)も実施しており、これまで入れると7年連続での登場になる。

85万コアを有するウエハースケール計算エンジンWSE-2

　HotChips 32から今年のHot Chips 2024までの間の発表を簡単にまとめると、まず2020年のHotChips 32は"Software Co-design for the First Wafer-Scale Processor (and Beyond)"と題して、初代WSEことWSE-1上でどうソフトウェアが動作するのか、あるいは逆にソフトウェアでどうWSE-1の上の動作を制御できるのかの説明があった。詳細は割愛するが、この講演の最後で第2世代製品に関するプレビューが示された。

WSE-1が1.2兆トランジスタ、40万コアなので、第2世代製品は倍以上である

　翌2021年のHotChips 33ではWSE-2ことWSE-2が発表された。チップの面積そのものはWSE-1と同じく46,225mm²であり、一方でトランジスタ数やコアだけでなく、オンチップメモリーは18GB→40GB、メモリー帯域は9PB/秒→20PB/秒といずれも2.2倍ほどに膨れ上がっている。要するにプロセスをTSMC 16FFからTSMC N7に変更しただけで、これだけのコアやメモリー容量向上が実現した格好だ。

WSE-2。チップサイズは、もともと300mmウェハーのギリギリを攻めていたので、300mmウェハーを使う限りこれ以上大きくはならない

　WSE-1からWSE-2の変更はこれだけでなく、新たにSwarmXおよびMemoryXという名称のシステムが提供開始され、これで最大192台のWSE-2でクラスターを構築可能になったとしている。

このスライドではSwarmXがスイッチ6台で実装されているように描かれているが、実際はもう少し大きくなると思われる

　このクラスターの目的は、大規模ネットワークへの対応である。オンチップメモリーで40GBというのはある意味驚異的ではあるのだが(なにせこれは全部SRAMである)、絶対的な量という意味では全然足りていない。

　そもそもNVIDIAのA100ですら40GB版に加えて80GB版をリリース。H200は141GB、GB200ではHBMだけで384GBを搭載する。AMDもInstinct MI300Xで192GBを搭載したが、次のInstinct MI325Xでは288GBに増量することを発表している。LLMの規模がどんどん大きくなることでメモリー不足になるのは明白である。

　これはCerebrasも認識しており、そのための対策として提供されたのがチップの外にメモリープールを置き、これをインターコネクトで接続するという上の画像の構成である。

この時点ですでに1T(1兆)パラメーターのMSFT-1Tが出ており、10TBを超えるメモリーと10EFlopsの処理能力が必要、とみなされていた。なんというか縦軸と横軸がどちらも対数表現というのが結構イヤな話である

　このMemoryXを併用する場合、重みデータはすべてMemoryXに置かれる形になり、アクティベーションのみがWSE-2のSRAMに格納される形になる。またこの重みのアップデートはMemoryX内のプロセッサーで行なうため、WSE-2には負荷をかけずにできるとする。

この場合、MemoryXからWSE-2への重み(Weight)の伝達は、間に入るSwarmXがブロードキャストで実行されるので、仮に192台のWSE-2があってもMemoryXからの読み出しは1回で済む

　ちなみにこのMemoryXは、2～192台までのWSE-2に対応し、4TB～2.4PBまでのメモリー搭載量であるとしている。

当然2.4PBものDRAMを実装できるわけもなく、"DRAM and flash hybrid storage"であるとしている

　4TBの場合で200B個、2.4PBなら120T個のパラメーターを格納できるとしているほか、内部動作をパイプライン化することでレイテンシーを遮蔽できるとしているが、このあたりの詳細は語られていない。

　全体でクラスターをいくつまで論理的に分割できるのかは明示されていない。分割可能であることはわかっているが、その1つの論理的なクラスターに所属するWSE-2は、すべて同じネットワークを動かし、重みも共通して持つ形になる。

　2つ上の画像で"Execute one layer at a time"という言い方をしているあたり、ネットワークのある層を複数のWSE-2で分割して処理する形であり、その際にMemoryXはSwarmX経由で、その層の重みを対象のWSE-2にブロードキャストするようだ。

　この方式のメリットは、WSE-2の台数にほぼ比例する形で性能がスケールすることだとしている。従来のGPUやAIプロセッサーの場合、結局重みをおのおののローカルメモリーで保持するこになるため、最初にメモリーネックになる。結果、多数のAIプロセッサーやGPUをクラスターにしても、性能が生かしきれない(煩雑にホストとの間で重みのやり取りが発生する)。ところがMemoryX+SwarmXを組み合わせると、この問題が一気に解決するこになる。

　もちろんメリットばかりではない。MemoryXを組み合わせると、1層の処理ごとにMemoryXから重みを受信する必要があるので、カード1枚当たりの性能はオンチップSRAMだけを使った場合に比べれば落ちることは避けられない。

　もっとも先に書いたようにWSE-2でも40GB「しか」メモリーはないので、大規模LLMを稼働させようとするとホストから煩雑に重みを受け取る必要があり、それよりはMemoryX+SmarmX経由で重みを受け取る方が速い上にスケールする、というわけだ。

WSE-2は事前に処理が必要なGPUと違い 200倍の帯域が利用できて効率が10倍良い

　2022年のHotChips 34では"Cerebras Architecture Deep Dive:First Look Inside the HW/SW Co-Design for Deep Learning"と題し、WSE-2の内部構造の説明が行なわれた。そもそもWSE-1の時も非常にラフな構造の説明しかなかった。まずWSE-2では85万個存在する、個々のコア(Cerebras用語ではPE:Processor Element)が下の画像だ。

Standard Cellで製造、というのも昨今ではすごい。後述するが、SRAMは48KBを8バンクに分割しており、個々のバンクは6KBということになる

　1.1GHz駆動と動作周波数は控えめだが、ピークでも30mWと控えめではある。もっともこれが85万個あるので、本当に全コアがフル稼働するとピークでは25.5KWもの消費電力になる。もっとも前ページのウェハー写真でもわかるように、実際にはこの85万個のPEが84個のダイに分割されており、個々のダイの消費電力は300Wほどになるため、妥当と言えば妥当な数字である。

　ここのPEの内部が下の画像で、6ステージのパイプライン構造のインオーダー構造の演算器であるが、FP16を1サイクルで4演算できるようになっている。ここで言う演算はFMAC(Multiply and ACcumulate:積和演算)で、なので見かけ上は1サイクルで8演算になる。

PEの内部構造。テンソル演算でも性能は変わらない。単にデータの持ち方を変えるだけで、演算そのものは4つの16bit FMACで行なうからだ

　このPEはデータフローの形で動作する、というのは前回の記事でも説明した通りで、なにもしなくてもSparsityを実現できることになる。このあたりが、事前に処理が必要なGPUとの違いであり、結果効率が10倍良いとするのは大げさではあるかもしれないが、嘘ではないだろう。

テンソルの場合、1回の演算に複数サイクルを要する(FMACが4つしかないから)。そこでマイクロスレッディングを利用して、演算順序の制御やプライオリティの管理などができるというのはなかなかおもしろい

　SRAMブロックの構成が下の画像で、6KBのバンク×8が用意され、かつSWキャッシュが256B搭載される仕組みになっている。SWキャッシュがスクラッチパッドなどではなくキャッシュなのは、これは外部から書き込みするだけで、PE自身にはキャッシュを変更できないためと思われる。

SW-Managed cacheとはあるが、そのSWはどこで動いているものか？　というと、PEというかダイというか、全体を管理する上位のコントローラーだろうと思われる

　ここの表現を見ると、特にテンソル演算の場合など、同じ処理をずーっと繰り返す、DSP的な動き方になる。PE自身でこれを制御するためには、分岐予測やBTB的なものを用意して、処理が終わったら次のコードに移行するのではなくまた元に戻って繰り返すことをハンドリングする必要があるが、PEにそれを持たせるのは無駄が多いと判断したのだと思われる。

　結果的に、DSP的にぶん回すような用途で、GPUと比較して200倍の帯域が利用できるとしているが、この数字が正しいのかどうかは判断がつきかねる。

　ちなみにその帯域の話で言えば、処理内容に応じて要求される帯域にはけっこう差があるのは知られた話である。

おのおののバンクは32bit幅なので、つまり4バンクからの同時読み出しと、2バンクへの書き込みを1サイクルでできる計算だ

　例えばNVIDIAのBlackwellは20PFlops演算性能と8TB/秒のメモリー帯域とされているわけだが、ということは2万TFlopsと8TB/秒だから、演算当たりのメモリー帯域は0.0004Bytes/Flopsという計算になる。

　実際には内部の2次キャッシュの帯域はもう少し大きいが、容量は小さいのですぐに使い切ってしまい、結局HBMアクセス待ちになる。これがWSE-2では16bitのFMAC×4に対して、64bit/サイクルの読み込み×2 + 64bit/サイクルの書き込み×1で、3Bytes/Flopsになる計算で、圧倒的に有利としている。

　WSE-2ではこのPEを1万156個集積したダイを84個つなぐ形で構成される。おのおののPEの外には5ポートのルーター(うち4ポートは隣接ノードに、1つはPEに接続)が設けられ、これで2次元メッシュを持つ構造だ。バス幅は32bitで、ただしデータは16bit(残りはIndex)なので、PEをまたいでのメモリー参照などはあまり現実的ではない。あくまでPEは自身のSRAMだけを対象に演算すると思われる。

実際には1つのダイには66×154で1万164個のPEがあり、これが84ダイで85万3776個のPEがあるはずだが、当然欠陥があるのでそこは無効化する形になる。つまり冗長PEが3776個あると考えるのが妥当だろう。そもそもダイあたりのPEが8つほど少ないのも、冗長PEのためと思われる

　続く2023年のHotChipsでは、クラスターに関する詳細が示された。まずMemoryX、2021年には「2～192台までのWSE-2に対応し、4TB～2.4PBまでのメモリー搭載」と発表されていたが、2023年には多少アップデートされた。

1台のMemoryXは、1TBメモリー/500TBフラッシュで、1台のWSE-2にはこれを12台接続可能。64台なら768台接続できる、ということらしい。やはりシャーシ1個で2.4PBのストレージは無理があったようだ

　このスペックからすると、どうも1ノードのMemoryXは1Uのシャーシに収まる規模になり、あとは顧客のニーズというか、どの規模のメモリーを必要とするかで最大12台まで接続できるようになっている模様だ(Dedicated interfaces to WSE-2 and *other MemX*とあるあたり、MemoryX同士での相互接続も可能になっているようだ)。

　下の画像はCerebrasが提供するAndromedaという1 Exaflopsのシステムの写真だが、左から9ラック目に12本装着されているのがMemoryXではないかと思われる。

SwarmXは上側であろうか？　この一列だけではSwarmX(やEPYC接続用スイッチ)が収まらず、別のラックに収まっている可能性もある

　ちなみにこのシステムは他にEPYC Gen3を合計1万8176コア搭載しており、左から1～8および9・10・12ラックに収められた1Uシャーシ(合計132枚)と、11・14～16ラック目の2Uシャーシ(合計12枚)には、全部64コアの2ソケットEPYCが搭載されているものと考えられる。

　同様にSwarmXもシャーシが複数必要である。4×WSE-2で12ノードだから単純計算すればWSE-2 1つあたり3ノードのSwarmXが必要になる計算だが、これはノード数が増えるほど大規模になる。

SwarmXもシャーシが複数必要だ。例えばBroadcomのTomahawk 5なら、800GbEで64本、400GbEで128本、200GbEなら256本のリンクを1つのチップで接続できるから、64台のWSE-2でもフラットなネットワークを構成可能である。ただこれは汎用品であることを考えると、そのまま使うことはできないだろう

　WSE-2が64ノードだと3 level tree(おそらくFat Treeを構成するのだろう)が必要で、合計252ノードなどになるのは少し多すぎる。独自のプロトコルを実装している関係で、既存のスイッチが使えない結果がこの有様で、このあたりは今後スイッチベンダーと協業して、もっと大規模かつ高速なスイッチになるかもしれない。

128bit SIMDに変更されたWSE-3 DGX-H100と比較しても5倍以上高速

　やっと2024年のHotChipsでの発表に移る。今回Cerebrasは第3世代のWSE-3を発表した。ただTSMC N7→N5ではあまり微細化の効果が得られていない。トランジスタ数こそ2.6兆個→4兆個で、倍増とは言わないまでも54%程増えているが、PEの数は85万個→90万個で5.8%増。搭載メモリー量は40GB→44GBだから10%増と、ほんのわずかでしかない。

WSE-3。トランジスタ数が大幅に増えているというのは、PUの内部が大幅に変わったことが主要因だろう

　なぜPEの増加量とメモリーの増加量が一致していないかは後で説明するとして、7nm→5nmではロジックは微細化の恩恵を多少は受けるが、SRAMの微細化はほとんどない、という実例のような結果になってしまった。もちろんそれでも動作速度向上(あるいは同じ動作周波数で消費電力削減)の効果は見込めるため、意味がないわけではないのだが。

　さてWSE-3では、PUの内部が大幅に変更になった。WSE-2までは16bit FMAC×4の構成だったのが、WSE-3では128bit SIMDに変更され、また8bitでは16個のデータを同時に処理できるようになった。演算性能だけで言えば、16bitで従来の2倍、8bitなら4倍になる計算である。

SRAM増加量がPU増加量より多いのは、一つはCacheが256Bytes→512Bytesになったこと、それとDSRが44→48になったこと。それとGPR/DSRともに従来は64bit幅で済んだのが今回128bit幅に拡張されたこと、の3つが関係していると思われる。逆にSRAMそのものには変更はなさそうだ

　ただ、演算性能が2倍なり4倍になったりするのであれば、それに合わせてメッシュネットワークの高速化もしないと追い付かない気もするのだが、それに関する説明は今回なかった。変更がないのか、変更はあるけどまだ未公開なのかは不明である。

　性能という意味で示されたのが下の画像である。Llama 3.1-70Bの推論が、Cloudベースのソリューションは元より、DGX-H100と比較しても5倍以上高速としている。

WSE-3の性能。それでもGroqが半分以上の性能を残しているのはさすがだと思うが。ちなみにLlama3.1-8Bの場合、Groqは750Token/s、WSE-3は1800Token/sだそうだ(DGX-H100は不明)

　ちなみにこの理由に関してのCerebrasの説明が下の画像だ。要するにHBMの帯域がネックになっているから、ということだ。

どんなモデルであっても、だいたい200トークン/秒あたりでメモリーの限界に達する、としている

　そうなると、WSE-3とH100を比較するとメモリー帯域は7000倍と言っているわりに性能は5倍にしかならないのか？　という疑問が出てくるわけだが、これは実際には大規模ネットワークを稼働させると煩雑にMemoryXから重みを読み出す必要がある。層ごとにこの読み出しが入るわけで、これがボトルネックになっているものと思われる。

　それにしてもWSE-3はもう5nmの限界にたどり着いてしまった感がある。これ、このまま3nmに持っていっても性能が大して上がらずにコストだけが跳ね上がりそうだ。WSEの欠点は少ないメモリー量で、これを補うためには次世代ではロジックウェハーの下にSRAMウェハーを3D接続するなどしてローカルのSRAM量を増やすとともに、SwarmX/MemoryXとの帯域を大幅に引き上げる必要がある。

　ただ現状でもSwarmXが明らかに能力不足の感がある。これをスイッチメーカーと協業するなどして改善できれば、もう少し性能の伸びが期待できそうではある。