🌡️

ウエハースケールエンジン向けSimulated Annealingを複数タイルによる並列化で実装しました

hikaru.takayashiki

2024/10/02に公開

皆さんこんにちは！
本記事では、Cerebras Wafer Scale Engine (WSE)にSimulated Annealing(SA)を並列分散で実装する方法を紹介します。
本記事の内容はこれらの記事の続きとなります。
ウエハースケール計算エンジンWSE-2においてSimulated-Annealingを実装しました
ウエハースケールエンジンにSimulated Annealingを分散並列実装し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の目的となります。

 これまでの実装ウエハースケールエンジンにSimulated Annealingを分散並列実装しCS-2実機で動作確認しました までの実装で、一つの大きな行列Qに対して並列分散グループ（この記事ではタイルと呼びます）を構築し、並列分散によるSAを行いました。
この並列化により、一つのPEに収まり切らない大きなQについて並列化を行うことが出来るようになりました。

 前回の記事での並列化方法の問題点前回の記事では、一つのPEに収まり切らない大きなQを並列化しました。 しかし大きいといえど、最適化問題はたかだか数十PE分で収まるQがほとんどで、WSE上にある数万のPEをすべて必要とするような大きな問題は稀です。
実機でこのような数十PE分で収まるQの問題を解く場合、使っていない残りのPEに関しては余っている状態になり、WSEの膨大な計算資源を活かすことが出来ません。
そこで本記事では、WSEの計算資源を活かすために、比較的小さなQをWSE上で複数同時に解く実装を紹介します。

 今回の並列化方針と実装内容について
 Qを複数タイルで解く並列化の概略実は、最初の記事 の段階では、複数のPEが同じQについて別の乱数シードを用いて解くことが出来ていました。
しかし、前の記事 において、一旦複数のPEが同じQに関して並列化して解く形に変更しました。
そこで、今回はこの2つの方法を組み合わせたような方針を取ります。
前の記事 の場合のようにQを分散させたタイルを、複数用意します。
各タイルでは異なる乱数シードにより独立して探索を行うことで、より効率の良い解の探索が出来る可能性を高めます。
ここから、3x3のサイズのタイルを2x2で並べた図で説明していきます。

 タイルを複数用意してSAを行う
図に示すように、同じQを持つタイルを複数用意します。 この複数のタイルが同じQについて問題を異なる乱数シードで解くようにします。
各タイル内で行われている処理については、前の記事 と同一ですので参考にしてください。
この実装に変える際に、以下の二つの新しい実装が必要になります。
スタート時に、すべてのタイルに同じQをブロードキャストする処理
終了時に、すべてのタイルの中から最も低いエネルギーとその時のsの組み合わせの回収をする処理
順に説明していきます。

 すべてのタイルに同じQをブロードキャストする処理
図に示すように、ホストからデータを各タイルに渡します。 この処理は、ブロードキャストにより行います。 WSE上のすべてのPEに同一のQのデータを流していき、流れてきたデータがもし自分の担当するQのデータ領域ならPE内のメモリに保存します。
各PEが担当するQのデータが集まり次第、SAのイテレーションを開始します。
だんだん分散並列らしくなってきましたね。

 すべてのタイルの中から最も低いエネルギーとその時のsの組み合わせの回収をする処理SAが規定のイテレーションを終えた後、複数のタイルがそれぞれ異なるエネルギー状態を持っている異なります。
この中で、最も小さいエネルギーを持つ組み合わせsが、SAにより求めた値であるとします。
最も小さいエネルギーを持つ組み合わせsを複数タイルから回収する方針は二つ考えられます。
すべてのタイルからエネルギーと組み合わせsをホスト側に送り、ホスト側で判断する
まずタイル同士で比較を行い、最も小さいエネルギーを持つタイルからホストに組み合わせを送る
どちらも善し悪しがあります。WSEとホスト間の通信にはコストがかかるので、タイルの数が多い場合は2．の方が良さそうです。
今回は、せっかくなのでWSEの中で行う 2.による実装になります。
図に、最大イテレーション回の計算が終了した時のPEの状態を示します。
各タイルの左上のPEは、そのタイルにおいて計算した最小エネルギーを持っている状態になっています。また、最上行のPEはその時の組み合わせを分散して持っています。
まず、すべてのタイルの左上PE同士のエネルギーを比較します。
各タイルの左上PEのエネルギーとそのPEの位置を示す構造データを、WSEの左上のPEに集めます。まずは左に比較しながら送り、そのあと上に比較しながら送ります。
これにより、WSEの左上PEに最小エネルギーとその出どころとなるタイルの位置の情報を集めることが出来ました。これをホスト側に送ります。
最後に、最小エネルギーの出どころとなるタイルの最上行から、組み合わせ状態sを転送します。
これにより、ホストは最終的な値を得ることが出来ます。

 評価
 挙動の確認評価は解の良さなどではなく、コードの挙動に関してを見るために実機で行います。
今回実装したコードは前回同様GitHubにて公開するので、ご覧ください。
まず、PEを3x3に並べたタイルをさらに3x3に並べた場合の、各PEのタイムラインを示します。
実行時間のプロファイルには各PEでの経過時間を測定するために、測定担当のPEをタイル上部に追加することで行います。プロファイル方法についてはまた後の記事で紹介します。
大きい画像なので、右クリックして別タブで開いて拡大するのをおすすめします。
このタイムラインでは、3x3のタイルの左上に位置するPEがボトルネックとなっていることがわかります。
前回の記事でもふれたように、乱数を計算しフリップする位置を決める処理をタイル内の左上だけが行うためです。
この図より、ひとまず複数のタイルが用意できていることがわかります。また、それぞれが独立したSAが行われています。

 タイル数と実行時間の変化まず気になるのは、タイル数を変えると実行時間がどのように変化するのかです。
WSEでは750×994までPEが使えるため、3x3のPEを詰め込んでいくと最大で250x331のタイルを並べることが出来ます。
タイル数を変化させた場合の実行時間の変化を図に示します。
どうも、タイル数を変化させるとWSEの起動にかかる時間が変化するようです。 タイル数が多いほど複数のPEにデータを流す必要があるからだと考えられます。
起動時間を除いた場合の実行時間の変化は次のようになります。
各タイルが行う処理は並列で行われるので、時間はほぼ変化していません。

 イテレーション数と実行時間の変化次に、イテレーション数を変化させた場合の実行時間の変化見てみましょう。 タイルのサイズは3x3で、そのタイルを3x3で用意した場合にイテレーション数を変化させると実行時間がどのように変化するのかを示します。
イテレーション数を変化させた場合の実行時間の変化を図に示します。 縦軸が指数時間になっていることに注意する必要があります。
実機においてプログラムを実行する場合、ジョブ投入から起動までに80秒ほど時間がかかるため、イテレーション数が小さい場合はほとんどその時間がすべてです。
そのため、かなり大きいイテレーション数を行わないと起動時間以上の計算を行うことが出来ません。
起動時間を除いてどこの部分に時間がかかっているのかを図に示します。
起動時間を除くと、ウエハ内でデータを転送している時間が大きな割合を占めます。 イテレーション数が大きくなるにつれて、データ転送時間に比べて計算している時間が長くなります。

 おわりに今回の記事では、タイルを複数用意することで並列にSAを解く実装について述べました。
また、実機で動作確認を行い、複数のタイルで同じ問題を異なる乱数シードで解きました。
これにより、大きなQの問題を解くだけでなく、小さなQについてもWSEのポテンシャルを活かしてSAをすることが出来そうです。
今後の展望としては、プロファイラを作った部分についての紹介や、マルチテンパリングの実装などに取り組む予定です。 これらもコードの公開と共にブログ上で説明していく予定なので、ぜひ楽しみにしていてください！

 謝辞今回CS-2実機を東京エレクトロンデバイス様より貸与いただきました。
この場を借りてお礼申し上げます。

 開発に関わったメンバー (アルファベット順)このプロジェクトは計算機好きな有志が自己研鑽の一部として、業務時間外（たまに時間内）で好き勝手に開発しています。
Hikaru Takayashiki
Hiroki Nishimoto
Keisuke Kimura
Naoki Yoshifuji
Toru Fukaya
Yoshiki Imaizumi
Yuki Ito
また、株式会社フィックスターズでは一緒に働く仲間を募集しています。 WSEのみならず様々なプロセッサでのプログラミング・高速化に興味がある方は、ぜひ採用ページよりご応募ください。