ウエハースケール計算エンジンWSE-2においてSimulated-Annealingを実装しました

2024年4月25日

皆さんは大規模計算のアーキテクチャといえば何を思い浮かべますか？

現在世界中のAI の学習の中心を担っている GPU、スーパーコンピュータ『富岳』、Preferred Networks 社が発表している『MN-Core』と人により様々だと思います。

世はまさに AI 戦国時代、世界中の半導体企業がより高速で、より省電力で、より扱いやすいアーキテクチャの開発に切磋琢磨しています。

そのような中でもチップ一つを製造するのにウエハを一枚使い切ってしまうというプロセッサが 2021年8月に発表されました。

Cerebras Systems 社 の 第二世代 Wafer-Scale Engine (WSE-2)です。1

おおよそ20cm x 20cmのサイズ一枚で一つのチップです。※写真はTED AI Labの許諾の元掲載

さらに、Cerebras Systems 社の WSE 向けにプログラミングする開発環境『Cerebras SDK』が一般公開が開始されました。

今回はそのWSE-2へ、Fixstars の有志でSimulated Annealingを実装しました。

この記事はその解説記事です。

アーキテクチャ・開発環境の紹介

実装内容の説明の前に、Cerebras WSEについて簡単に紹介します。

  • 実装対象アーキテクチャ：WSE-2
  • 開発環境：Cerebras SDK

チップ 概要

ウエハースケールプロセッサ(WSE)ということで、とにかく大きいことが特徴です。

Cerebras が公開している Whitepaper をもとに、 NVIDIA の H100 との比較すると次のようになります。

WSE-2 H100 Cerebras Advantage
Chip size 46,225 mm^2 826 mm^2 56x
Cores 850,000 16,896 FP32 + 528 Tensor 48x
On chip memory 40 GB 0.05 GB 800x
Memory bandwidth 20 PB/s 0.003 PB/s 6666x
Fabric bandwidth 220 Pb/s 0.0576 Pb/s 3819x

チップサイズが56倍あり、その分コアがたくさんあります。 各コアがそれぞれのオンチップメモリを持っており、メッシュ状に接続されています。 そのため、巨大な分散メモリシステムがチップ内に形成されています。

アーキテクチャ概要

WSE-2 はデータフローアーキテクチャと呼ばれるアーキテクチャです。

以下のように二次元に Processing Element (以下、PE ) が配置されています。

各コアは2次元メッシュ状に並んでおり、隣り合ったのコアのみ通信が可能です。

この PE 一つひとつに 48 KB のメモリと演算器が積まれており、『あるコアで計算した値を隣のコアに渡す』という処理が可能です。

このように複数のコアを同時に動かすことで高い並列性やパイプライン性を実現しています。

Cerebras SDK 概要

WSE-2 上で任意の処理を実行するために Cerebras Systems 社が用意している、 Cerebras SDK について紹介します。

WSE-2を搭載したシステムであるCS-2は、制御を行うホスト部とWSE-2部に分かれています。 ホスト部とWSE-2部で使用するプログラミング言語は次のようになっています。

プログラミング対象 言語
ホスト Python
WSE-2 CSL (Zig言語ベースの専用プログラミング言語)

CSL で書かれたプログラムをコンパイルし、ホスト部のプログラムでロードし、WSE-2 でプログラムを実行します。

WSE-2の実行モデル図
CS-2のWSE-2部 ※写真はTED AI Labの許諾の元掲載
CS-2を搭載したサーバー ※写真はTED AI Labの許諾の元掲載

CSL でのプログラミング

CSL での WSE-2 のプログラミングには、color, route の 2つの概念があります。これにタスクベースのプログラミングモデルを組み合わせることで処理を記述します。

概念 概要
color
  • data に付随できる情報
  • 各 PE は『何色のデータとして別のPEに送るかを決める機能』、 『ある色のデータを受信する機能』、 『ある色のデータが来るたびにタスクを実行する機能』などを持つ
route
  • rx(受信先), tx(送信先) で構成される
  • データをどこから受け取り、どこに送信するかを決定する
  • color とPEの位置ごとに異なるrouteを指定できる
  • {North, West, South, East} + RAMP が指定可能

colorによりデータに色を付け、routeによりどの通信経路からデータが来てどの通信経路へデータを送るのかを制御します。 ここで、RAMPは受け取ったPE自身のメモリへの転送を指します。

さらに、タスクには2種類あります。

task の種類 概要
local_task 他の task からの activate によって実行されるタスク
data_task 指定した color の data が入力された際に実行されるタスク

タスクはCSL上の関数とは異なり、呼び出し元には戻りません。

local_taskは処理をタスクとして記述することで、非同期な操作が完了した時に特定のタスクを実行するといったことが可能になります。

また、data_taskを用いることで、『ある色のデータが流れてきた際に、流れてきたデータの総和を取る』といった処理が可能となります。

Simulated Annealing について

SA は NP困難 である組合せ最適化問題に対して、ヒューリスティックに近似解を求めるアルゴリズムです。

基本的にはエネルギーが減少する向きに状態を遷移させますが、一定確率でエネルギーを増加させる向きにも遷移させるため、局所解に陥りにくいという特徴があります。

アニーリングによる解空間探索のイメージ図

詳しくは Fixstars Amplify が公開している記事をご参照ください。

Simulated AnnealingのWSE-2への実装方針

今回はスピン変数を0,1で表現したQUBOモデルを対象として実装を行うことにしました。 QUBO – Wikipedia

Simulated Annealingを並列化する方法は様々あると思います。 今回は最も単純な方法として、複数のPEで乱数のシード値を変えて各々が独立にアニーリングをした後に、各PEの中から最も良い結果を集めてホストに送るという方針をとります。

実装は、以下のようなフローとなります。

  1. 入力をホストからのデバイスに送る
  2. 入力を (0,0) の PE から範囲内の各 PE に配る（Broadcast）
  3. 各 PE で乱数のシード値を変えて SA を実行 (SA の実行中は PE 間での通信は行わない)
  4. 最小のエネルギーを得るための Reductionを行う
  5. 最もエネルギーが低い状態とその時のエネルギーをデバイスからホストに送る

入力をホストからデバイスに送る

ホストからデバイスへの転送は、SDKで用意されているストリーミングを使用します。

ストリーミングにより、左上の(0,0)のPEへ入力を送信します。

入力を (0,0) の PE から範囲内の各 PE に配る

入力を(x,y)=(0,0)から縦方向に(0,h-1)までデータを配った後、横方向にデータを広げていくことで計算を行うPEに入力データを配ります。

次の図はその様子を示しています。

入力を4x4のPEにブロードキャストする様子のイメージ図

各 PE で乱数のシード値を変えて Simulated Annealing を実行

各PEでSimulated Annealingを行います。

ここでは、各PEが並列でSimulated Annealingによる探索計算を繰り返している状態になります。

最小のエネルギーを得るための Reductionを行う

最終的に最小のエネルギーを得るために、各PEは『送られてきたデータと自分のデータを比較し、より小さいエネルギーを持つデータを指定方向に流す』ことで右下に最小のエネルギーを集めます。

まずは横方向にこの処理を行い、最終的に一番右端にいるPEが縦方向に処理を行うことで実現しています。

次の図はその様子を示しています。

エネルギーを4x4のPEから右下へ集約する様子のイメージ図

最もエネルギーが低い状態とその時のエネルギーをデバイスからホストに送る

デバイスからホストへの転送も、SDKで用意されているストリーミングを使用します。

ストリーミングにより、右下の(w-1,h-1)のPEからホストへ最終的なデータを転送します。

実行および評価結果

実装したSimulated AnnealingをWSE-2の実機で試したいところですが、シミュレータを使用して実行してみます。 SDKはv1.0.0を使用しました。

サイズは4×8として、N=6の問題を解いてみましょう。入力は乱数で与えます。

シミュレータはSDKに同梱されており、環境を設定して実行をすると次のような結果となります。

+ width=4
+ height=8
+ cslc ./layout.csl --fabric-dims=11,10 --fabric-offsets=4,1 --params=Num:6,width:4,height:8,max_iters:16,T:1000 --params=MEMCPYH2D_DATA_1_ID:0,MEMCPYD2H_DATA_1_ID:1 -o out --memcpy --channels 1
INFO: Using SIF: /home/Owner/Cerebras-SDK-1.0.0-202311111408-11-4ef8698c/cbcore_sdk-202311111408-10-4a54bce5.sif
INFO: User's specified CSL_IMPORT_PATH=
NOTE: CSL_IMPORT_PATH accepts colon separated list of paths generated by 'realpath <path>'
compile successful
+ cs_python run.py --name out
INFO: Using SIF: /home/Owner/Cerebras-SDK-1.0.0-202311111408-11-4ef8698c/cbcore_sdk-202311111408-10-4a54bce5.sif
Reading file out/out.json
Reading file out/west/out.json
Reading file out/east/out.json
fab w,h = 11,10
Kernel x,y w,h = 4,1 4,8
memcpy x,y w,h = 1,1 9,8
inputs = [ 0.5479121  -0.12224312  0.71719587  0.39473605 -0.8116453   0.9512447
  0.5222794   0.5721286  -0.74377275 -0.09922812 -0.25840396  0.85353
  0.28773025  0.64552325 -0.1131716  -0.5455226   0.10916957 -0.8723655
  0.65526235  0.2633288   0.5161755 ]
18.179776396s
best_s = [0 1 0 1 0 1]
energy = -1.3816108703613281

この結果は、同様のアルゴリズムをCPUへ実装したものと一致します。

今後の課題と展望

今後も有志メンバーで時間を作りながら開発を進めていきます。 現在考えている展望は次のようなものとなります。

  • 入力行列Q を複数の PE で分散して同じ問題を解く
  • 入力行列Q を分散した複数の PE ブロックをさらに複数用意し、その中で最も良い結果を得る
  • 擬似乱数の改良
  • 出力結果の既存 SA との比較評価

本記事について

本記事にて解説している内容は、日本でCerebras Systems 社の販売代理店である、東京エレクトロンデバイス様が開催したハッカソンにおいて発表した内容です。

ハッカソンを開催し、実機での動作確認をしていただけたことに感謝申し上げます。

ハッカソン説明会ページ

付録

本記事で紹介した実装内容は、弊社GitHub上でOSSとして公開しています。 興味がある方はぜひ見てください。

GitHub Link

  1. 2024年3月13日にさらに最新のWSE-3が発表されました。↩︎

About Author

hikaru.takayashiki

Leave a Comment

メールアドレスが公開されることはありません。 ※ が付いている欄は必須項目です

このサイトはスパムを低減するために Akismet を使っています。コメントデータの処理方法の詳細はこちらをご覧ください。

Recent Comments

Social Media