Cerebras、GPT-3を1日で学習できるAIスパコンを発表 - Hot Chips 33

OpenAIの「GPT-3」AI言語モデルはパラメタが1750億個あり、Lambda社のWebページでは、V100 GPUを使う割安な高性能クラウドを使っても、GPT-3の学習には355年かかり、クラウドの使用料は460万ドル掛かると書かれている。

しかし、Hot Chips 33でCerebrasが発表した最大規模のシステムを使えば、GPT-3は1日で学習ができ、Microsoftが提唱している1兆パラメタのモデルでも長めの週末一回で学習ができると見積っているという。1兆パラメタのモデルは人間の脳のパラメタの規模であり、人間の脳レベルのモデルの学習ができるようになったら、何ができるかは興味深い。

図1 CerebrasはAIを進歩させ、将来のAIの能力を変えることを目的として、新しいコンピュータシステムを作るために2016年に設立された。現状、従業員は350人あまりである (このレポートのすべての図はHC33におけるCerebrasの創立者の1人でチーフハードウェアアーキテクトのSean Lie氏の発表資料のコピーである)

図2 Cerebrasの第2世代品となるWSE-2ウェハスケールエンジン。WSE-2はTSMCの7nmプロセスで作られ、300mmウェハ1枚から1個採り。46255mm²のシリコン片で85万AIコアと40GBのオンチップSRAMを搭載する。メモリバンド幅は20PB、ファブリックのバンド幅220Pbit/sである

この図2の写真のチップは直径300mmのウェハで作れる最大サイズのチップ(？　これをチップと呼ぶのは違和感がある)で、これより大きいものは450mmウェハを使うファブが稼働しなければ作れない。それには製造設備を全部450mm対応に変更する必要があり、容易なことではない。したがって、当面は、これより大きなシングルチップは作れない。

今回の発表は、このWafer Scale Engine(WSE)の第2世代品「WSE-2」を搭載した超高速ディープラーニングシステム「CS-2」を最大192台並べると何ができるかという話である。ただし、現在のCS-2は192台を並べるのに適した形にはなっていない。消費電力が大きいので、水冷にして、大きな筐体に何台ものCS-2を詰め込むような実装にする必要があるように思われる。

図3 写真のように、現状のCS-2は12Uのデスクサイドの筐体で、本当に192台を並べようとすると、30Uくらいのラックに詰め込み、水冷するなどの実装の変更が必要と思われる

AI研究の最先端では、たった2年で、1000倍サイズのモデルが出現し、1000倍の計算能力を持つAIエンジンが必要になって来ている。

図4 ニューラルネットの規模は指数関数的に大きくなっている。縦軸は学習に必要な演算量(単位はPFlops-Days)、横軸はモデルのメモリ量(単位はGB)。現在は1750億パラメタのGPT-3が最大であるが、マイクロソフトは1兆のモデルを提案している

なお、指数関数的な増加に勝てるものはなく、192台のCS-2を並べたシステムが作れたとしても、それで間に合うのは一時的で、すぐに、より強力なアクセラレータの開発が必要となる。それはともかく、今作れるのはCS-2を192台並べるシステムである。

図5 Multi-Trillionパラメタモデルは大量のメモリ、大量の計算、大量の通信を必要とする。大量の小規模なノードをつなぐクラスタではメモリ分散と通信の問題が入り混じってシステムが複雑化する。そして、クラスタサイズが大きくなるとさらに問題となる。したがって、CS-2のような強力なノードを少数使うシステムの方が良い

これに対するCerebrasの解は、120兆パラメタのモデルを1台のCS-2に収容することである。後述するように、モデルの重みをストリーミング供給することで、大きなモデルを容易に扱うことができる。そして、192台のCS-2を使うクラスタでも計算速度はほぼリニアにスケールする。

さらにCS-2はSparseな計算を入力の値をチェックして省くハードウェアを持っているので、Sparse性の利用が容易というメリットがある。

図6 これに対するCerebrasの解は、120兆パラメタのモデルを1台のCS-2に収容するという方法である。このやり方で192台のCS-2のクラスタまでほとんどリニアにスケールする。さらに重みなどのSparse性を利用して10倍のスピードアップ。これらのスケールアップに計算法の変更は必要なく、ボタンを押すだけで実現できる

図7 重みはモデルメモリに置き、ストリーミングでComputeに供給。ComputeはTraining Datasetを右側のメモリから受け取るという構造でモデルサイズの影響を切り離し

そして、今回の発表の目玉は「MemoryX」と「SwarmX」である。MemoryXは巨大なモデルを格納するメモリであると同時に各CS-2が計算した傾きを集めて、重みの更新を行う。そして、SwarmXはMemoryXから供給される更新した重みをCS-2に配る。

MemoryXとSwarmXとCS-2とそれぞれの処理を分割して実行することにより、192台のCS-2を使って計算能力を増やしてもAIシステムの学習能力はスケールする。

図8 最大120兆パラメタを1台のCS-2に格納。SwarmXで複数のCS-2の結果を集約してMemoryXへの格納とMemoryXのデータのCS-2への広報などを行う。この方法で192台までのCS-2をサポートできる

重みストリーミング型の実行モデルでは

重みはウェハ外に記憶
重みはウェハにストリーミング供給
アクティベーションはウェハに保持して、重み層ごとに計算を進める

計算した傾きはウェハからMemoryXに吐き出し、重みの更新を行う。

図9 重みストリーミング型実行モデルでは、重みはウェハ外に記憶。計算時にはストリーミングで読み出して演算器に供給。アクティベーションはウェハに保持して計算にともなってアップデートして行く。計算された傾きはウェハから吐き出し、重みのアップデートを行う

図10 重みストリーミングのやり方。MemoryXから重みを供給し、右側からActivationを供給してL1層の傾きを計算。計算を層ごとに繰り返してLossを求める

計算を層ごとに実行するとL1の計算が終わらないとL2の計算が出来ないので、レーテンシが問題となる。そこで、Cerebrasは図11のように、粗粒度と細粒度のパイプラインを組み合わせて使っている。層の間を進んでいくForward/Delta/Gradientと完全にパイプライン化する。層ごとの繰り返しの部分は層間の依存性のない計算であり、計算順序を気にする必要は無い。Fowardパスと重みのアップデートを行う部分のオーバラップは細粒度のパイプライン化で順序を確保する。

図11 レーテンシ問題の解決。粗粒度パイプラインでは、フォワード/デルタ/傾きは完全にパイプライン化。繰り返しの中では層間の依存性は無いので、依存性を気にせず並列に実行できる。細粒度パイプラインではL1傾きとL1重みアップデートは並列に実行できる。また、L1重みアップデートとL1重みストリームアウトは並列に実行できる

巨大モデルを扱うには2つの容量問題がある。1つは巨大モデルをどのように格納するか。もう1つは巨大モデルをどのようにして実行するかである。

図12 巨大モデルを扱うには2つの容量の問題がある

巨大モデルを扱うためにCerebrasはMemoryX Technologyを開発した。図13に示すMemoryXはDRAMとFlashで作られたハイブリッドのストレージで、最大2.4PBの容量を持ち1200億の重みを扱える。

MemoryXは重みのアップデートや最適化のための計算機能も持っている。そして、レーテンシを隠すため、インテリジェントなパイプライン化も行う。

図13 巨大モデルを扱うためにCerebrasはMemoryX Technologyを開発した。MemoryXは最大2.4PBの容量を持ち、120兆の重みを扱える。MemoryXのメモリ容量は計算性能とは独立に決められる

WSEは大容量のローカルメモリを持ちすべてのActivationを格納している。大きなCompute Arrayを持ち、重みのストリームとActivationの積を計算する。十分なメモリバンド幅があるので、フルの演算性能を発揮できる。高バンド幅のインタコネクトを持っているので、部分積の加算をウェハ全域について並列に計算できる。

そして、それらの計算のためにマトリクスのブロッキングやパーティショニングを行う必要はなく、普通に計算を進めれば良い。

図14 ニューラルネットの計算はほとんど全部が行列積の計算。CS-2は100K×100Kの行列積の計算をおこなうことができるので、行列積を複数のデバイスに分割して計算するオーバヘッドが無い

学習は複数のCS-2の上でデータ並列に行われる。重みはMemoryXからすべてのCS-2にブロードキャストされる。傾きは重みと逆方向に更新されていく。結果として、マルチシステムでもシングルステムと同じように実行されていく。

図15 複数のCS-2をSwarmXのファブリックで接続する。扱えるモデルサイズはメモリ容量で決まり、計算能力と無関係に決められる。モデルサイズにかかわらず、同じシステムアーキテクチャ、同じネットワーク実行フロー、同じユーザインタフェースのソフトウェアが使える

同じシステムアーキテクチャのモデルが使えるので、ほぼリニアにスケーリングし、使用するCS-2の台数に逆数に比例して計算時間を短縮できる。

図16 ほぼリニアに性能が向上する。横軸はCS-2の台数(3本しか見えないが、NLPのモデルサイズを示すグラフ)。モデルサイズが小さい場合は性能が飽和するが、モデルサイズが1兆であればほとんど飽和は見られない

しかし、NLPのモデルサイズの拡大は歴史的に継続しており、この拡大が続けば計算能力が不足するので、計算量の拡大を抑える努力も重要である。

図17 横軸はモデルに必要なメモリ量、横軸は学習に必要な計算量。腕力にまかせたやり方では不十分。より少ないFlopsで学習ができるやり方が必要

これまでに発表されている疎計算の研究では10倍余りの効率の改善が示されている。

図18 疎行列であることを利用して、計算量を削減。計算量を半減できるやりかたから、ケースによっては10倍以上の性能になるケースもあり、いろいろと論文が出ている

Cerebrasの計算コアにはDataflow Triggerというハードウェアが付いており、非ゼロのデータが入力された場合だけ計算を行うようになっている。そして、WSEは高バンド幅のメモリを持っているので、ゼロのデータが来ても非ゼロのデータを選択してフルの演算性能を維持することができる。

また、WSEは高バンドのインタコネクトを持っているので、ゼロ入力をスキップしてリダクションを継続できるなど、すべてのタイプの疎計算に対応している。

図19 Cerebrasのアーキテクチャは細粒度のデータコアを持ち、ダイナミックで構造のないスパーシティーに対応できる

行列計算には行列×行列の和を計算するGEMM、行列×ベクトルのGEMV、スカラ×ベクトルのAXPYがある。SparseなGEMMは1つの非ゼロの重みごとに1回のAXPY計算で済む。

図20 Cerebrasの方式はマトリクス×マトリクスだけでなく、マトリクス×ベクトル、定数×ベクトルもフル性能がでる

Sparseな重みはすべてのCS-2にストリーム入力される。計算結果のSparseな傾きは逆方向伝搬に伴って纏められる。Sparseな重みのアップデートはMemoryXで行われる。この計算フローは密行列でも疎行列でも同じである。

図21 MemoryXの中では疎な重みがCS-2にストリームされる。また、疎な傾きが生成される、あるいは疎な重みのアップデートが生成されるケースが出る

WSEはSparseな非ゼロの重みのストリームを受け取る。この部分でゼロの重みが除去される。そして、それぞれのコアはアクティベーションにしたがってAXPYを計算する。

図22 ウェハでのSparse性の利用。ローカルメモリは密なアクティベーションを供給。コアアレイは疎な重みストリームを受け取る。各コアは受け取ったアクティベーションを使ってAXPYを実行。密な場合も疎な場合も同じ処理フロー

図23は非構造のSparsityを利用して、どれだけスピードを上げられるかを示す図である。

図23 Sparse性は学習精度の達成に必要な時間を短縮する。WSEはAXPYをフルスピードで実行し、すべての無構造のSparsityを高速化する。右の折れ線グラフはGPT-3の12K×12KのMatMulのスピードアップ

重みストリーミングでクラスタレベルまでスケールできる。

図24 重みストリーミングはソフトウェアの単純性をクラスタスケールまで拡張。CS-2の40GBのメモリは分割なしに全マトリクスを格納可能、全部、同じマッピングで処理できた

WSE-2はTSMCの7nmプロセスで作られ、300mmウェハから1個採り。46255mm²のシリコン片で85万のAIコアと40GBのオンチップSRAMを搭載する。この贅沢とも言えるハードウェアで通常のAIアクセラレータではモデル分割などが必要な場合でも、WSE 1枚にモデルを収容し、効率的にモデルを実行する。

プログラム開発が容易と言うのを測定するマシンラーニングのベンチマークは無いが、実際の開発を考えれば大きなメリットであると思う。