Cerebrasのウェハスケールエンジン(Wafer Scale Engine:WSE)は直径300mmのシリコンウェハをまるまる1枚を使う超巨大LSIのAIエンジンである。Hot Chips 32では、このWSEの使い方に関する考え方、そして次世代のWSEに関する話が語られた。

システムとしての利用が進むCerebrasのWSE

このチップ(?)は1.2Trillionトランジスタを集積している。NVIDIAのA100 GPUが54.2Billionトランジスタであるので、その20倍以上のトランジスタを集積しているわけである。そして、そのトランジスタを使って、WSEは18GBのオンチップメモリとAI演算に最適化した40万個のコアをチップ上に作っている。

  • Cerebras

    CerebrasのWafer Scale Engineは1.2Trillionトランジスタと、NVIDIAのA100 GPUの20倍以上のトランジスタを集積している (このレポートのすべての図は、Hot Chips 32におけるCerebrasのSean Lie氏の発表スライドをキャプチャしたものである)

CerebrasはすでにWSEチップを搭載したCS-1システムを販売している。CS-1は1チップのシステムであるが、Cerebrasはクラスタシステム級の性能を持っていると言っている。すでにアルゴンヌ国立研究所で使用されており、ピッツバーグスパコンセンターからも2台のCS-1を受注しているほか、一般の企業にも導入されているようであるが、それらの顧客は明らかにされていない。

  • Cerebras

    CS-1はクラスタシステム級のディープラーニング性能を発揮する

WSEのプログラミング方法

WSEのコアはマシンラーニングに最適化された拡張命令を持ち、データフローアーキテクチャで入力が入ってきた時しか演算を行わないようになっている。メモリはコアごとに分散されており、メモリとロジックが混然一体となった造りになっている。そして、高バンド幅、低レーテンシのネットワークを持ち、ユーザの指定したトポロジの接続を行うことができる。

  • Cerebras

    CerebrasのコアはFMACのデータパスとレジスタ、SRAMメモリを持ち、ファブリックにつながっている。演算器はML演算に最適化された拡張がなされている

WSEのプログラミングは、TensorFlow、PyTorchやCaféなどのフレームワークで記述されたニューラルネットを、まず、LAIRと呼ぶ中間表現に変換し、それをカーネルグラフとのマッチングを調べてWSEでの実行方法を決定する。そして、各部の実行プランをリンクして、全体の実行モジュールを作ってWSEで実行するという流れになる。

  • Cerebras

    CerebrasのWSEはKernel Graphと呼ぶグラフを実行する。そのため、まず、フレームワークで記述されたニューラルネットをLAIR表現に変換し、それをKernel Graphに変換する。Cerebrasのグラフコンパイラは、DNNグラフを自動的にコンパイルしてくれる

Cerebrasのコンパイラはフレームワークの演算とグラフカーネルの一致を調べて、フレームワークの記述をカーネルに変換する。カーネルとしては汎用の自動生成したものと、人手で最適化したカーネルが用意されている。

  • Cerebras

    Cerebrasのコンパイラはフレームワークの記述とカーネルのライブラリとの一致を調べ、一致していればカーネル表現に置き換える。一致するものが無い場合は、自動生成したカーネルを使うものと思われる

次の図はコンパイラの最適化の例を示すもので、1つの層の処理を3×3から、12×12までの面積と性能の異なる5種の候補で、自動的にトレードオフを検討した例である。

  • Cerebras

    ニューラルネットの1つの層で面積を3×3から12×12まで自動的に振って、面積、性能のトレードオフを検討した例

そして、それぞれの層をWSE上のコアに配置、配線して作り、カーネルを実現する。このプロセスは、LSIの配置配線と同様のプロセスを実行している。

  • Cerebras

    カーネルの面積、性能のトレードオフが決まると、WSE上の配置や配線が行なわれる

データ並列でもモデル並列でも使うことができるWSE

WSE上での並列性の利用法は柔軟で、1つのアルゴリズムでモデル並列とデータ並列の両方の並列性が利用されることもある。

データ並列は同じモデルをDevice1 、Device2のように複数作り、それぞれのデバイスは入力データの異なる部分を処理する。モデル並列は大きなモデルを複数のパーツに分け、それらのパーツを並列に実行する。

WSEはたくさんのコアがあるので、モデル並列でも高い実行並列度が得られる。一方、GPUなどはバッチサイズを大きくして高いデータ並列の処理はできるが、並列に作れるモデルの数は制約があるので、高いモデル並列を実現するのは難しい。

  • Cerebras

    ニューラルネットの並列実行には、データ並列とモデル並列というやり方が考えられる。WSEは多くのコアを持ち、多数のモデルを並列に実行できるが、集積度の小さいGPUなどでは、モデルは1つでデータ並列の実行になる。WSEでは必要に応じて、両方の並列度を利用する柔軟な並列処理ができる

モデル並列の場合は、ネットワークのすべての層をパイプライン的に並列に実行させる。すべての層を並列に実行するので、高い性能が得られる。ただし、このためには高いバンド幅を持つインタコネクトが必要である。

モデル並列では、複数のモデルの間で重みの値を同期させるオーバヘッドは生じない。そのため、小さなバッチサイズでの実行でも良く、直線的なスケーリングが得られる。右側のグラフは4レイヤのBERTの性能をコア数を横軸にとってプロットしたもので、1つのモデルに使うコア数は100Kで、4コアまで並列にすることができている。

BERTの相対性能は100Kコアの時を1.0として、400Kコア、4モデルの場合には3.6倍の性能となっている。

  • Cerebras

    モデル並列での性能スケーリングを示す。モデル並列は、各層をパイプラインで並列実行することにより、性能を上げている。このグラフは、4層のBERTの性能を示しており、横軸は使用コア数で、1モデルで100Kコアを必要としている。性能は1モデルの場合を1.0として、4モデル並列での性能は3.6程度となっている

データ並列の場合は、各層のレプリカを作り、異なるデータでレプリカを並列に実行する。レプリカの数が増えると相対性能は上がるが、使用コア数が少ない場合は、レプリカのデータの入れ替えのオーバヘッドが影響しているのか、性能は飽和する傾向にある。

  • Cerebras

    データ並列の場合は、レプリカを作って異なるデータで並列に実行する。並列実行により性能は上がるが、コア数が少ない場合は、演算能力が限られるので実行性能は飽和する

WSEは小さな粒度のSparsityも利用可能

例えば非線形の関数としてReLUを使うと、入力値が負の場合は出力がゼロになる。このように、ニューラルネットではゼロが多く出てくる。掛け算の入力がゼロの場合、積和演算は何もしないのと同じ結果になり、演算量を減らせる。このため、CerebrasのWSEでは、ゼロの入力となる疎な行列のケースを有効に利用するようにしている。

  • Cerebras

    ニューラルネットではゼロの値が出てくることが多い。入力がゼロの疎な行列を掛ける場合は、実質的な演算量が減るので、WSEでは積極的に利用するようにしている

Cerebrasのコアはファブリックからのデータの値によって命令の実行の可否を判断する。ファブリックへのデータの送り出し元はゼロのデータを省略し、受け手は不必要な演算を飛ばしてしまう。

WSEのコアは、独立の命令を実行しているので、細かい粒度で不必要な実行を飛ばして効率を上げることができる。

  • Cerebras

    WSEのコアはファブリックのデータによって、命令の実行をコントロールできるようになっている。このため、ゼロを掛ける積和演算などは、実質、スキップされる

また、モデルを多少変形して疎な行列にして、計算を減らすという手もある。やり方としては、ReLUのようなゼロを作り出す関数を使う、あるいは小さな値はゼロとみなしてしまうというような方法である。

右の棒グラフは34層のフルコネクトのネットワークとスパース化を行ったネットワークを比較したもので、左端の密行列では、青のGPUとオレンジのCS-1で差は無いが、その右のペアの棒グラフのようにReLUを使うと、CS-1では1.7倍程度に性能が上がる。さらに値の小さい項を無視して25%疎にすると性能向上は約2倍になり、50%疎にすると2.4倍程度の性能となる。

ただし、ここでは計算性能の向上だけの比較で、スパース化に伴う精度への影響は比べられていない。

  • Cerebras

    ReLUのような入力が負の場合に出力がゼロになる関数を使ったり、値が小さい要素を無視してゼロにするなどの方法で、ゼロの値を増やし、行列を疎にすることができる。34層のフルコネクトの密行列の場合、ReLUを使うことにより演算性能は1.7倍に向上し、ReLUと50%のスパース化を行うと、性能は2.4倍になる

自然言語処理のBERTの場合、自然に疎行列にはならないが、フォワードやバックワードの計算過程でゼロ要素を増やすことはできる。次の図に示すように、精度の劣化を最小にする形で50%のゼロを追加することにより、計算性能を1.5倍に引き上げることができた。

  • Cerebras

    自然言語処理のBERTは自然に疎行列にはならないが、ReLUの使用や小さい値の無視でゼロを増やすことはできる。精度の低下を最小に抑えて、50%のゼロの追加で性能を1.5倍に引き上げることができた

CS-1は最も強力な単一ノードのマシンラーニングのシステムである。CS-1は細粒度の疎な計算の削減を利用できる柔軟性を持ち、適応的にダイナミックに計算を行うことができる。

  • Cerebras

    CS-1は最も強力な単一ノードのMLエンジンである。そして、CS-1は柔軟性を備え、動的な計算が可能となっており、次世代のモデルや計算テクニックの利用を可能にしている

開発が進む第2世代Wafer Scale Engine

そして、Hot Chips 32における発表の最後、Cerabrasは第2世代のWafer Scale Engineを発表した。

第1世代は16nmプロセスであったが、第2世代はTSMCの7nmプロセスを採用した。チップの大きさは第1世代と同じ300mmウェハからの1個採りであり、AI最適化されたコアを850K個搭載している。搭載トランジスタ数は2.6Trillionとなっている。

第1世代は400Kコア、1.2Trillionトランジスタであったので、第2世代は2倍強のコア数、トランジスタ数になっている。

  • Cerebras

    第2世代のWafer Scale EngineはTSMCの7nmプロセスで作られ、第1世代の2倍強にあたる850Kコアを集積する

なお、第2世代WSEを作るにあたって、第1世代のWSEの開発と使用の経験からの改良のフィードバックが加えられたことは想像に難くないが、今回の発表では、どのような改良が加えられたのかは明らかにされなかった。