ウエハースケールエンジンにSimulated Annealingを分散並列実装しCS-2実機で動作確認しました

皆さんこんにちは！

本記事では、Cerebras Wafer Scale Engine (WSE)にSimulated Annealing(SA)を分散並列で実装する方法を紹介します。 本記事の内容は前回の記事 の続きの話となります。

本記事では前回の記事で行ったウエハースケール上での並列化をさらに発展させ、より分散並列な形でSimulated Annealingを実装しました。 さらに、東京エレクトロンデバイス社様からCerebras CS-2実機を貸与していただくことが出来たため、実機での評価を行いました。

まだ前回の記事を読んでいない方は、そちらから読むことをおすすめします。 もし前回の記事をお読みになった方も、再度読んでみるとよりこの記事を理解しやすくなると思います。

簡単なおさらい

Cerebras ウエハースケールエンジンとは

Cerebras ウエハースケールエンジンとは、Cerebras Systems社が開発したウエハ一枚を一つのチップとして使い切るプロセッサです。 その巨大さとデータフローアーキテクチャから、唯一無二のプロセッサとして機械学習やHPCなどのワークロードで使用されています。

チップ上にコアが２次元メッシュをなす形で分散しており、チップ内でも事実上の分散並列コンピューティングを行うことが出来ます。 各コア(PE)は演算ユニットや通信ユニットを備えているほか、独立したメモリ空間を持つ小さなスクラッチパッドメモリを持っています。

Simulated annealingとは

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

SAは様々な組合せ最適化問題を解く可能性を秘めています。 前回に引き続き、今回もスピン変数を0,1で表現したQUBOモデルを対象として実装を行うことにしました。

QUBOモデルは次のような形で定義されます。

\[f(s)=-\sum_{i\le j}Q_{ij}s_{i}s_{j}\]

ここで、\(Q\)はパラメータ行列、\(s\)は0か1の値を持つバイナリ変数です。

この\(f(s)\)が最も小さくなる\(s\)の組み合わせを探索して見つけることが、SAの目的となります。

前回の記事での並列化方法

前回の記事 では、複数のPEにおいて乱数のシード値を変えて各々が独立したアニーリングを解き、最後に最もよい結果を集めるという方針で実装を行いました。

そのため、各PEがこの式の\(f(s)\)を最小化する\(s\)を独立に探索していき、最後に最も良い結果を持つPEが結果となる形です。

この実装はすべてのPEがほぼ同じ動作をするため実装が単純であり、WSEには85万ものPEがあるので、高い並列性を期待することが出来ます。

前回の記事での並列化方法の問題点

前回の実装は高い並列性を期待できる一方で、1つ問題があります。メモリの容量です。

WSEは1PE当たり40KiB程度のスクラッチパッドメモリ容量を持ちます。 前回の記事における実装方針の場合、この式における\(Q\)及び\(s\)は複製され、すべてのPEが同じデータを持ちます。

そのため、\(Q\)および\(s\)が合わせて40KiBを超えるサイズにような大きさの問題は、前回の実装方針では解くことが出来ません。 本記事では、この問題を解決するために複数のPEが一つの\(Q\)を分散して持つ形で実装を行います。

今回の並列化方針と実装内容について

\(Q\)を分散してPEに配置する概略

図に示すように、\(Q\)を複数のPEを跨いで配置します。 この図では、6×6の\(Q\)を3×3のPEで分散しています。各PEは2×2のサイズの\(Q\)の一部を担当するという事になります。

これにより、各PEが必要となるメモリの量は1/9となり、1つのPEがすべての\(Q\)を担当するよりもはるかに大きいサイズの\(Q\)を扱うことが出来るようになります。 WSEには85万個ほどのPEがあるので、仮に全部使うと膨大なサイズの問題を解くことが出来ます。ワクワクしますね？

\(Q\)を分散してPEでSAを行う

\(Q\)をPEに分散したので、Simulated Annealingを行うためには複数のPEが協力して計算する必要があります。 また、Simulated Annealingは\(s\)のどれか一つを確率的にフリップさせることで探索を行うので、どのPEがフリップを行うのか？という問題を解決する必要があります。

Simulated Annealingの計算順序は次の通りになります。

1. \(s\)のフリップ位置を乱数で決める
2. フリップした際の差分エネルギーを\(Q\)と\(s\)から計算する
3. 差分エネルギーと温度をもとに、次の状態に遷移するか乱数で決める
4. 温度を少し下げて1に戻る

これを\(Q\)を分散した状態で行うと、次のような処理手順となります。

1. ホストからWSEに\(Q\)を転送する
2. \(s\)のフリップ位置を乱数に基づいて決定し、フリップ位置をブロードキャストする
3. 現在の\(s\)の状況をブロードキャストする
4. 差分エネルギーを計算し、x方向にリデュースする
5. x方向に集まった差分エネルギーをy方向にリデュースする
6. フリップを行い温度を少し下げて、結果をブロードキャストする
7. 最大イテレーション回数に至るまで2.に戻る
8. WSEからホストへエネルギーと\(s\)を転送する

ここで、「ブロードキャスト」とは同じデータをコピーしながらPEへと転送していくことを指し、「リデュース」とはなんらかの演算を行いながらデータを集めることを指します。

では一つ一つのステップを順番に見ていきます。

1. ホストからWSEに\(Q\)を転送する

ホスト側から用意した\(Q\)をWSEへと転送を行います。 各PEはそれぞれ自身が担当する位置の\(Q\)を受け取ります。

2. \(s\)のフリップ位置を乱数に基づいて決定し、フリップ位置をブロードキャストする

\(s\)のフリップ位置を決定します。\(Q\)を分散しているので、\(s\)も分散した状態になっています。 ここでの処理は3ステップに分かれています。

1. 前イテレーションの結果を反映させるために横方向にデータを受け渡し、\(s\)を最新の状態にします
2. 左上のPEにおいて次はどの\(s\)をフリップさせるのか位置を決定します
3. 左上PEで決定したフリップ位置を横方向にブロードキャストします

実際に有効なフリップ位置を持っているのは左上だけなので、緑色で示されている最上行だけが正しいフリップ位置を持っている状態になります。

3. 現在の\(s\)の状況をブロードキャストする

前ステップでは最新状態の\(s\)と次のフリップ位置を持つのは最上行のPEだけでした。 そこで、最新状態の\(s\)と次のフリップ位置をを縦方向にブロードキャストします。

これにより、すべてのPEが最新状態の\(s\)と次のフリップ位置を持っている状態になります。

4. 差分エネルギーを計算し、x方向にリデュースする

さて、最新状態の\(s\)とフリップ位置が得られたので、そこをフリップした場合の差分エネルギーを計算しましょう。 計算が行われるのはフリップ位置が自分の担当領域だった場合のPEだけです。

今回の図の例だと、flip_indexが3の場合の例なので、真ん中の列の3つPEだけが差分エネルギーを計算していることになります。 他のPEは担当していなので、0を返します。

この結果を左側のPEへとリデュースを行います。

5. x方向に集まった差分エネルギーをy方向にリデュースする

前のステップで左側のPEへとリデュースを行ったので、\(Q\)の各行に対応する差分エネルギーを左側のPEが持っています。 これをさらにリデュースしたものが、求めたい差分エネルギーです。

そこで、今度は縦方向にリデュースを行い左上のPEへと差分エネルギーを集めます。 これで差分エネルギーが求まりました。

6. フリップを行い温度を少し下げて、結果をブロードキャストする

差分エネルギーを求められたので、その差分エネルギー及び温度から次の状態に遷移するかどうかを最終決定します。 もちろん温度やエネルギーによっては状態を遷移しない場合もあります。

7. 最大イテレーション回数に至るまで2.に戻る

このように2.から7.の処理をあらかじめ決められた最大イテレーション回数に達するまで繰り返していきます。 この時、温度を徐々に下げていくことでエネルギーが高い方向へフリップする確率を下げていきます。

8. WSEからホストへエネルギーと\(s\)を転送する

最大イテレーション回数に達したら、その時点のエネルギーと\(s\)をホストへと転送します。

2.の時点で、最新状態の\(s\)を持っているのは最上行のPEですので、ここから\(s\)を送信します。 また、エネルギーに関しては左上のPEが持っているので、ここからエネルギーを送信します。

以上で、\(Q\)を分散したSimulated AnnealingをWSE上で行うことが出来ました。

シミュレータでの評価および結果

シミュレータでの実行結果

以上の実装を、まずはシミュレータで試してみましょう。 プログラムは GitHub上で公開しています。 SDKはv1.1.0を使用しました。

環境構築を簡単にするために、Dockerfileおよびビルド・実行スクリプトを付随しています。 今回はこれを使います。 実行には正解データを生成するための無料Fixstars Amplifyトークンが必要です。 まだお持ちでない方は、こちらのページから会員登録を行い、無料トークンを取得してください。

Dockerが使える環境でダウンロードしてきたSDKをCerebras-SDK.tar.gzという名前でcerebras_sa/cerebras_docker内に配置した後、build.shおよびexec_cerebras.shを行います。

すると、次のような結果が得られます。

+ dirname /home/cerebras/cerebras_sa/src/commands.sh
+ cd /home/cerebras/cerebras_sa/src
+ pwd
+ SCRIPT_DIR=/home/cerebras/cerebras_sa/src
+ cd /home/cerebras/cerebras_sa/src
+ command -v sdk_debug_shell
+ command -v sdk_debug_shell
+ dirname /home/cerebras/cs_sdk/sdk_debug_shell
+ sif_path=/home/cerebras/cs_sdk/sdk-cbcore-202404122356-8-a1162951.sif
+ width=2
+ height=3
+ singularity instance list
+ grep cerebras
+ singularity instance start /home/cerebras/cs_sdk/sdk-cbcore-202404122356-8-a1162951.sif cerebras
INFO:    Instance stats will not be available - system configuration does not support cgroup management.
INFO:    instance started successfully
+ singularity exec instance://cerebras python -c import amplify
+ singularity exec instance://cerebras pip install amplify

... 途中省略...

fab w,h = 9,5
Kernel x,y w,h = 4,1 2,3
memcpy x,y w,h = 1,1 7,3
Q=array([[-0.992232  ,  0.2096471 , -0.9569606 , -0.5456669 , -0.16260374,
         0.48462176],
       [ 0.2096471 ,  0.90178144, -0.83920753, -0.7692129 ,  0.80470073,
         0.9533495 ],
       [-0.9569606 , -0.83920753,  0.33332357, -0.8376797 ,  0.38912448,
        -0.5340135 ],
       [-0.5456669 , -0.7692129 , -0.8376797 ,  0.52016926, -0.1768382 ,
         0.31264699],
       [-0.16260374,  0.80470073,  0.38912448, -0.1768382 , -0.16835177,
        -0.6802061 ],
       [ 0.48462176,  0.9533495 , -0.5340135 ,  0.31264699, -0.6802061 ,
         0.5482536 ]], dtype=float32)
energy_function=Poly(0.209647104144096 q_0 q_1 - 0.956960618495941 q_0 q_2 - 0.545666873455048 q_0 q_3 - 0.162603735923767 q_0 q_4 + 0.48462176322937 q_0 q_5 - 0.839207530021667 q_1 q_2 - 0.769212901592255 q_1 q_3 + 0.80470073223114 q_1 q_4 + 0.953349471092224 q_1 q_5 - 0.83767968416214 q_2 q_3 + 0.389124482870102 q_2 q_4 - 0.534013509750366 q_2 q_5 - 0.176838204264641 q_3 q_4 + 0.312646985054016 q_3 q_5 - 0.680206120014191 q_4 q_5 - 0.992232024669647 q_0 + 0.901781439781189 q_1 + 0.333323568105698 q_2 + 0.520169258117676 q_3 - 0.168351769447327 q_4 + 0.54825359582901 q_5)
s_amplify=array([1, 1, 1, 1, 0, 0], dtype=int32)
e_amplify=-2.9760382622480392
5.781631834s
best_s=array([1, 1, 1, 1, 0, 0], dtype=int32)
min_energy=-2.9760382622480392
OK

ここで、e_amplifyはFixstars Amplifyで得られた結果です。best_sが今回の実装をシミュレータで実行した結果になります。

これにより、一致した結果が得られていることが分かると思います。

実行時間のブレイクダウン

プログラムを作って動いたら、やはりどこに時間がかかっているのか調べるのが我々の本業です。

プログラム中に実行時間を測定するコードを挿入し、時間計測を行います。 このコードはGitHubのsrc/trace以下で公開しています。

このブレイクダウン図は縦軸にPE、横軸に実行時間（サイクル）を示しています。

この図から、PE[0,0]の緑色の部分が実行時間を大きく占めており、ボトルネックとなっていることが分かると思います。

ここは乱数を計算する部分で、動作説明の「2. \(s\)のフリップ位置を乱数に基づいて決定し、フリップ位置をブロードキャストする」に対応します。

この乱数計算の高速化結果については、次以降の記事にてまとめてお知らせしようと思います。楽しみにお待ちください！

実機での評価および結果

次に、実機での評価を行います。 今回、CS-2実機を東京エレクトロンデバイス様より貸与いただきました。

もし実機をお持ちの方は、GitHubで公開しているコードを実機で動作させることが出来ます。 ただし、大きすぎる問題サイズはFixstars Amplifyでは解くことが出来ないので、コメントアウトする必要があることに注意してください。 ここでは動作手順は省略します。

実行すると実行結果が得られます。パラメータを調節することで\(Q\)=4096×4096の問題がおおよそ4分程度で解けています。

※以下の出力は開発中のバージョンのため、GitHubで公開しているものと違います。

Q=array([[ 0.18473563,  0.07461128, -0.88200694, ...,  0.1582715 ,
        -0.37397817, -0.04449175],
       [ 0.07461128, -0.10853962, -0.7037109 , ..., -0.96829927,
         0.35810596,  0.9692767 ],
       [-0.88200694, -0.7037109 , -0.389984  , ..., -0.22095379,
        -0.26828793,  0.319895  ],
       ...,
       [ 0.1582715 , -0.96829927, -0.22095379, ...,  0.36359495,
         0.07861157,  0.82223004],
       [-0.37397817,  0.35810596, -0.26828793, ...,  0.07861157,
        -0.68019027, -0.15861152],
       [-0.04449175,  0.9692767 ,  0.319895  , ...,  0.82223004,
        -0.15861152,  0.44415537]], dtype=float32)
params={'Num': '4096', 'block_height': '64', 'block_width': '32', 'max_iters': '4194304', 'time_constant': '597980', 'log_init_temperature': '33901', 'MEMCPYH2D_DATA_1_ID': '0', 'MEMCPYD2H_DATA_1_ID': '1'}
runner.load 172.114181294s
runner.run 8.009724081s
memcpy_h2d 0.880734969s
memcpy_d2h 64.882085465s
memcpy_d2h 0.262949031s
runner.stop 0.000385073s
total 246.150059913s
best_s=array([0, 1, 1, ..., 1, 1, 1], dtype=int32)
min_energy=-43873.973
e_amplify   = -43873.381, [0 1 1 ... 1 1 1]
best_s  = -43873.973, [0 1 1 ... 1 1 1]

Fixstars Amplifyとほぼ同じ解が246秒程度で得られています。

今回の場合だとFixstars Amplifyは100ms程度で解いているので、アルゴリズムの最適化なども含めてまだまだ製品として動いている最適化されたGPU向けの実装には遠く及びません。

さらに、もっと大きな問題を解いてみましょう。

WSEのPEをすべて使い切るサイズは、\(Q\)=83328×83328のサイズです。 これを解いてみましょう。

CS-2に搭載されたWSE-2のスペック上は85万PEですが、ホストのやり取り等にいくつか必要なため実際に計算で使うのは992×744=738’048コアになります。 これにより、PEあたりだと82×112のサイズの\(Q\)を担当することになります。

※以下の出力は開発中のバージョンのため、GitHubで公開しているものと違います。

params={'Num': '83328', 'block_height': '992', 'block_width': '744', 'grid_height': '1', 'grid_width': '1', 'max_iters': '67108864', 'time_constant': '4000000', 'log_init_temperature': '34079', 'MEMCPYH2D_DATA_1_ID': '0', 'MEMCPYD2H_DATA_1_ID': '1'}
runner.load 215.652606623s
runner.run 8.067431309s
memcpy_h2d 887.698785822s
memcpy_d2h 2036.899615608s
memcpy_d2h 0.267256359s
runner.stop 0.000545655s
total 3148.586241376s
best_s=array([1, 1, 1, ..., 1, 1, 1], dtype=int32)
min_energy=-4111129.2
opt_s   =   0.000, [0 0 0 ... 0 0 0]
best_s  = -4111129.250, [1 1 1 ... 1 1 1]

およそ1時間程度で結果が得られていることが分かります。 このぐらい大きな問題をSAで解こうとした例はあまりないのではないでしょうか？ WSEの大きさを感じますね。

おわりに

今回の記事では、\(Q\)を複数のPEに分散して解く実装について述べました。

また、シミュレータで実行して実行時間のブレイクダウンを見たところ、乱数生成に時間がかかっていることが分かりました。

さらに、CS-2実機で動作させたところ、4096×4096の問題は246秒程度で解けました。 また、WSE上の最大サイズは83328×83328で、1時間弱で解くことが出来ました。

まだまだ発展途上であるためSimulated Annealingのパラメータ次第では全く解が求まらない場合もありますが、大きなサイズの\(Q\)をWSE上で解けるようになったのは大きな前進です。

今後の展望として、乱数生成部分の高速化をしたり、さらにこの\(Q\)を分散したPEの塊を複数用意して並列に解いたり、マルチテンパリング実装を行ったりしていく予定です。 これらはコードの公開とともにブログ上で解説していく予定なので、ぜひ楽しみにしていてください！

謝辞

今回CS-2実機を東京エレクトロンデバイス様より貸与いただきました。

この場を借りてお礼申し上げます。

開発に携わっているFixstarsメンバー (アルファベット順)

このプロジェクトはFixstarsの計算機好きな有志が自己研鑽の一部として、業務時間外（たまに時間内）で好き勝手に開発しています。

• Hikaru Takayashiki
• Hiroki Nishimoto
• Keisuke Kimura
• Naoki Yoshifuji
• Toru Fukaya
• Yoshiki Imaizumi
• Yuki Ito

また、株式会社フィックスターズでは一緒に働く仲間を募集しています。 WSEのみならず様々なプロセッサでのプログラミング・高速化に興味がある方は、ぜひ採用ページよりご応募ください。

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