LLM向けAIアクセラレータ Cerebras における最適化手法の実証

研究室での位置付け

田浦研究室では、LLMの台頭に伴い、その学習・推論両面での性能最適化（高速化・省メモリ・省IO）にソフトウェア側から取り組んできました。しかし、この類の最適化は、命令セットアーキテクチャからハードウェア構成（メモリサイズ・メモリ帯域幅など）・システムコールへの理解が欠かせません。これらは、元来LLMの台頭以前から、田浦研究室の研究内容のメインストリームを支えてきたものであり、それらのノウハウを昇華させるものです。

Cerebras アーキテクチャ

世間で LLM の学習・推論に用いられるハードウェアとして広く知られているのは、紛れもなく GPU （特に NVIDIA GPU）です。しかし、GPUは Graphical Processing Unit という名前からも分かるように、元々は画像処理用のハードウェアであり、LLMの計算に最適化されたアーキテクチャではありません。近年は、GPU上で、GPUのボトルネックを隠蔽・低減するアルゴリズムが非常にホットなトピックとなってきましたが、GPUのアーキテクチャの根本的な変革が求められてきた面も否めません。

そのような期待を一心に背負って提案されたアーキテクチャ（ハードウェア）が、 Cerebras Wafer Scale Engine (WSE) ¹です。アーキテクチャの優位性を端的に表すエピソードとして、

Chat GPT の応答速度が、GPU上のモデルに比べて20倍ほど高速である

というものが挙げられます。

アーキテクチャの要点をかいつまんで説明すると

GPUの性能ボトルネックである、メモリ帯域幅をカバーするべく、デバイスメモリはSRAM (On-chip memory) のみであり、これをウェハ1枚分44GB確保して使用する
ウェハ上に、約100万のコアが、on-chip interconnect で長方形状の網目を構成する

core2wafer cs-mesh

各コア(PE)は、local SRAM に加えて128bit幅のSIMD計算ユニット、それから他のPEとの通信用のルータを有する

wse3Core

という特徴を持っています。

ソフトウェア的な観点で重要となるのは例えば以下の点です。

コア間の共有メモリが存在せず、アドレス空間も共有されていない
記憶階層が存在しないため、SRAMのサイズがそのまま仮想記憶のサイズになる
配列の要素に対する Iterative な操作において、アドレス計算をサポートするハードウェアが存在し、Tensor Memory Descriptor（フィールドがリッチなポインタのようなイメージ）を通じて効率的なメモリアクセスが可能となる
- 通常のポインタの持つ情報: 先頭アドレス, 型
- Tensor Memory Descriptor の持つ情報: 先頭アドレス, stride, length, 型, その他
コア間の物理通信路には color という抽象化が存在するのに対して、通信を制御するルーターのキューは抽象化が存在せず、プログラマが明示的にハードウェアレベルの制御信号を記述する必要がある

詳細なアーキテクチャはこちら ¹をご覧ください.

研究内容

Cerebras WSE は魅力的なハードウェアである反面、ソフトウェアレベル最適化については、十分な研究が行われているとは言えず（というか未開の領域に等しい）、このハードウェアの優れた能力を真の意味で引き出す点においては、プログラマの能力に大いに依存します。

本プロジェクトでは、LLMにおいて多用される Softmax + Cross-entropy カーネルを題材として、様々なソフトウェア最適化を検討・提唱することを目指します。

exp, log の Cerebras 命令セットを用いた高速アルゴリズムの実装と検証
parallel online normalizer ² ³ が Cerebras アーキテクチャでも高速なアルゴリズムであることを実証
データ依存関係のある複数データのコア間 reduction のスケジューリングのパフォーマンスモデルの提案と実証
- 単純な sum-reduce のパフォーマンスモデルを提案した先行研究 ⁴
コア内計算最適化
- software-managed cache の静的割り当てアルゴリズム
- レジスタ割り付けの最適化アルゴリズム

また、今後は、これらの最適化がプログラマの責任ではなく、コンパイラによって行われるよう、言語処理系の拡張も視野に入れています。