Scalable Construction of Spiking Neural Networks using up to thousands of GPUs

本論文は、マルチGPUクラスターおよびエクサスケールスーパーコンピュータにおける大規模スパイキング神経回路網の構築とシミュレーションのための新規MPIベース手法を提示し、最適化された局所結合およびスパイク交換戦略を通じて複雑な皮質モデルに対する効率的なスケーリングを実証する。

要約

人間の脳をコンピュータ上でシミュレーションすることを想像してみてください。脳は約 860 億個のニューロンからなる巨大な都市であり、各ニューロンは家のようなもので、毎秒何千もの他の家へ微小な電気的な「テキストメッセージ」（スパイクと呼ばれる）を送っています。これをシミュレーションするには、数千枚のグラフィックカード（GPU）が連携して動作するスーパーコンピュータが必要です。

問題は、これらの GPU が島々のような存在だということです。それらは高速ですが、互いに容易に通信できません。ある島から別の島へメッセージを送ろうとすると、「郵便配達員」（通信システム）が行き来しなければならず、それがすべてを遅くしてしまいます。

本論文は、シミュレーション開始前にこれらの接続の「地図」を構築する、はるかに高速な新しい手法を紹介しています。これにより、GPU は交通渋滞に巻き込まれることなくシミュレーションを実行できます。

以下に、彼らがどのように行ったかを簡単に説明します。

1. 旧来の方法：メインランドで地図を構築する

以前、科学者が脳ネットワークをシミュレーションしようとした場合、まず低速な中央コンピュータ（CPU）上で「接続地図」を構築し、その後、この巨大な地図を高速な GPU へコピーする必要がありました。

• 比喩: 大規模なパーティーを主催すると想像してください。旧来の方法では、キッチン（CPU）で紙にすべてのゲストの名前と知り合いをすべて書き出し、その後、各部屋（GPU）へ走り回ってリストのコピーを手渡していました。準備をするだけで非常に時間がかかりました。

2. 新しい方法：部屋の中で地図を構築する

著者たちは、各 GPU が中央コンピュータを待たずに、自らのメモリ内で直接接続地図の「自分自身の部分」を構築する新しい手法を開発しました。

• 比喩: 今度は、リストをキッチンで書く代わりに、各部屋に自分専用のメモ帳があります。パーティーが始まると同時に、各部屋のゲストがそこで知り合いを記録します。キッチンへの往復は不要です。
• 結果: この「オンボード」構築は、旧来の方法よりも10 倍以上高速です。あるテストでは、ネットワークの構築に 12 分近くかかっていたものが、55 秒で完了しました。

3. メッセージを送る 2 つの方法

地図が構築されると、GPU はシミュレーション中に「テキストメッセージ」（スパイク）を交換する必要があります。論文では、ネットワークの構成に応じて、これに対する 2 つの異なる戦略をテストしました。

• 戦略 A：直接電話（ポイント・ツー・ポイント）

  • 仕組み: GPU #1 内のニューロンが GPU #2 内の特定のニューロンと話したい場合、その特定の GPU に直接電話をかけます。
  • 最適: 接続が不均一または特定のネットワーク（ある領域同士は頻繁に話すが、全員とは話さない実際の脳など）に適しています。
  • 論文の主張: 彼らはマウスの視覚皮質（32 の異なる領域）のモデルにこれを使用しました。完全に機能し、新しい地図構築手法が複雑な現実世界の脳構造と互換性があることを証明しました。

• 戦略 B：グループチャット（集合通信）

  • 仕組み: 個人に電話する代わりに、GPU はメッセージを GPU のグループ全体に一度に叫びます。グループ内の全員がその叫びを聞き、メッセージが自分宛かどうかを確認します。
  • 最適: 全員が全員と話をするような、巨大でランダムなネットワーク（バランスの取れた群衆など）に適しています。
  • 論文の主張: 彼らは最大1,024 枚の GPUにスケールアップする巨大な「バランス型ネットワーク」でこれをテストしました。これは多数のグラフィックカードが連携して動作する非常に大きな数です。彼らは、これだけの数のカードがあっても、システムがクラッシュすることなくスムーズにスケールアップすることを示しました。

4. 「メモリレベル」のトリック

GPU には多くのメモリがありますが、無限ではありません。数十億のニューロンに対する接続地図を保存するには、非常に多くのスペースが必要です。

• 比喩: 小さな机（GPU メモリ）と巨大な倉庫（CPU メモリ）を持っていると想像してください。
• 解決策: 著者たちは、4 つの「レベル」の組織化を作成しました。
  • レベル 0: 地図を倉庫（CPU）に置き、必要なものだけを机へ持ち込みます。これは机のスペースを節約しますが、取得に時間がかかります。
  • レベル 3: 机にすべてを詰め込みます。これは最速ですが、より大きな机が必要です。
• 論文の主張: 適切なレベルを選択することで、Leonardo Boosterスーパーコンピュータ（4,096 枚の GPU を搭載）上でシミュレーションを実行でき、さらに今後のJUPITERスーパーコンピュータが2 億 3,000 万ニューロンと 2.5 兆シナプスを持つネットワークをシミュレーションできると予測できることを示しました。これはほぼ人間の皮質のサイズに相当します！

彼らが達成したことのまとめ

• 速度: ネットワーク地図をグラフィックカード上で直接構築することで、脳シミュレーションの「セットアップ」フェーズを 10 倍高速化しました。
• スケーラビリティ: 最大 1,024 枚の GPU を同時に使用しても機能することを証明しました。
• 柔軟性: 通信を処理する 2 つの異なる方法（直接電話対グループチャット）を示し、科学者が特定の脳モデルに最適な方法を選択できるようにしました。
• 将来性: 彼らの手法は、個々のシナプス詳細を備えた完全な人間の脳をシミュレーションできるほど強力になる次世代の「エクサスケール」スーパーコンピュータで動作するように設計されています。

要約すれば、彼らは単にシミュレーションを高速化しただけでなく、レースが始まる前にスーパーコンピュータが交通渋滞に巻き込まれないよう、データのためのより良い「道路システム」を構築したのです。

技術概要

技術概要：最大数千の GPU を用いたスパイクニューラルネットワークのスケーラブルな構築

問題定義

人間の大脳皮質規模での大規模スパイクニューラルネットワーク（SNN）のシミュレーションは、個々のニューロンおよびシナプスに対する膨大なメモリ要件と、サブミリ秒精度でダイナミクスを解くための高処理速度という、2 つの主要な課題に直面している。数千の GPU を搭載した高性能計算（HPC）システムは必要な計算密度を提供するが、既存の GPU ベースのシミュレーションソフトウェアは、計算神経科学のインフラ要件と精度要件を満たしつつ、計算クラスター全体にスケーリングできる能力をまだ実証していない。

大規模ポイントニューロンネットワークの分散シミュレーションにおける特定のボトルネックは、計算クラスター内の異なるノード間でのスパイクの通信である。Digital Brain や GeNN などの従来のアプローチは、個々のシナプス情報を省略するか、単一 GPU 実行に限定されている。さらに、NEST などの従来の CPU ベースのシミュレータは、ラウンドロビン方式によるニューロンの分散と集団通信に依存しており、均質なネットワーク構造を前提とし、生物学的な脳のトポロジー的・空間的な不均一性を活用できていない。NEST GPU はこれらの問題の一部を解決したが、その初期のネットワーク構築は CPU から GPU メモリへのデータ転送に依存しており、動的構築手法は以前は単一 GPU シミュレーションに限定されていた。

手法

本研究は、メッセージパッシングインターフェース（MPI）を用いて、マルチ GPU システム上で大規模 SNN を直接構築・シミュレートするための、新規かつメモリ効率の高い手法を提示する。中核的な革新は、構築フェーズ中のプロセス間通信なしに、ネットワーク構築を GPU メモリ内（オンボード）で完全に実行することにある。

中核アルゴリズム

本手法は、ローカル接続（同じ MPI プロセス内のニューロン）とリモート接続（異なるプロセス間のニューロン）を区別する。

1. 独立構築: 各 MPI プロセスは、他のプロセスと通信することなく、ネットワークの一部を独立して構築する。ローカル接続を作成し、リモート接続のためのデータ構造を準備する。
2. プロキシ表現: リモート接続については、対象プロセス内で「イメージニューロン」（プロキシ）を使用する。これらは、他の MPI ランクに位置するソースニューロンの仮想的な表現である。
3. 通信マップ: アルゴリズムは、スパイクを効率的にルーティングするために、GPU メモリ内に連続した通信マップをインスタンス化する。これらのマップは、ソースランク内のソースニューロンのインデックスと、ターゲットランク内の対応するイメージニューロンのインデックスを関連付ける。
4. 通信方式: このフレームワークは、ユーザーがネットワークアーキテクチャに基づいて選択できる 2 つの MPI 通信モードをサポートする。
   • ポイントツーポイント: 2 つのプロセス間の直接通信を使用する。ニューロンやシナプスの分布が不均一なネットワーク（例：マルチエリアモデル）に最適化されている。スパイクをルーティングするために、特定のマッピング構造 $(R_{\tau,\sigma}, L_{\tau,\sigma})$ とシーケンス $(T, P)$ を利用する。
   • 集団通信: グループベースの通信（例：MPI_Allgather）を使用する。均質な通信ペイロードを持つバランスの取れたネットワークに有利である。複数のプロセスにわたってスパイクのルーティングを管理するために、グループ固有のインデックス配列とホスト配列を採用する。

GPU メモリ最適化

GPU メモリ消費とシミュレーション速度のバランスを取るため、著者は 4 つの**GPU メモリレベル（GML）**を実装した。

• レベル 0: リモート接続マップと接続数は CPU メモリに格納される。
• レベル 1: レベル 0 と同様だが、すべてのソースニューロンが対象プロセスにイメージを持っていると仮定し、実際の使用状況のチェックを回避する（構築が高速だが、メモリ浪費の可能性が高い）。
• レベル 2: マップと接続インデックスは GPU メモリに格納され、接続数はオンザフライで計算される。これがデフォルトレベルである。
• レベル 3: 接続数を含むすべてのデータ構造が GPU メモリに格納され、CPU-GPU 間のデータ転送を最小化するが、その代償として GPU メモリ使用量が増加する。

評価モデル

1. マルチエリアモデル（MAM）: マカクザルの視覚関連 32 領域の生物学的に詳細なモデル（$4.13 \times 10^6$ニューロン、$24.2 \times 10^9$シナプス）。このモデルは複雑で階層的な接続性を特徴とし、ポイントツーポイント通信を用いてシミュレーションされた。
2. スケーラブル・バランスド・ネットワーク: 固定のインデグリー接続性を持つ興奮性および抑制性ニューロンのランダムネットワーク。弱いスケーリング性能を評価するために設計された。このモデルは、最大 1,024 GPU での集団通信を用いてシミュレーションされた。

主要結果

ネットワーク構築性能

「オンボード」GPU 構築手法は、従来の「オフボード」（CPU ベース）アプローチと比較して大幅な高速化を示した。

• MAM シミュレーション: ネットワーク構築時間が、オフボードの686.0 秒からオンボードの55.5 秒に短縮され、12.4 倍の高速化を達成した。
  • ローカル接続の作成で20 倍の高速化。
  • リモート接続の作成で9 倍の高速化。
  • ニューロン/デバイスの作成とシミュレーション準備で、それぞれ350 倍と50 倍の高速化。
• スケーラブル・バランスド・ネットワーク: 本手法は、1,024 GPU（256 ノード）にわたって最大2.304 億ニューロンおよび2.59 兆シナプスのネットワークの構築に成功した。

状態伝播とスケーリング

• MAM: 状態伝播時間（リアルタイムファクターとして測定）は、オフボード版とオンボード版の間でほぼ同等（約 15〜16）であり、構築の最適化がシミュレーションダイナミクスに悪影響を及ぼしていないことを示している。
• バランスド・ネットワーク: システムは 1,024 GPU までの弱いスケーリングを実証した。
  • メモリ効率: GPU メモリレベル 0 を使用することで、NVIDIA A100 GPU（64 GB）のメモリ制限を超えずに 4,096 ノードまでのシミュレーションが可能となった。より高いメモリレベル（2 および 3）は、より高速な構築とシミュレーション速度を提供するが、メモリ制限に達するノード数は低くなる（レベル 3 の場合、約 3,072 ノード）。
  • 性能: バランスド・ネットワークにおけるスパイク記録を無効にすると、状態伝播時間が約 20% 減少した。

検証

新しい構築手法は、従来のオフボード版および CPU ベースの NEST シミュレータに対して検証された。新しいアルゴリズムによる乱数生成シーケンスの変更にもかかわらず、スパイク活動の統計的性質（発火率、スパイク間隔の変動係数、ペアワイズピアソン相関係数）は保持され、シミュレーションの生物学的妥当性が確認された。

意義と主張

本論文は、この研究が個々のシナプス情報を保持しつつ、計算クラスター全体（最大数千の GPU）にスケーリング可能な初の GPU ベース SNN シミュレーションソフトウェアを提供すると主張している。主な貢献は以下の通りである。

1. スケーラブルな構築: GPU メモリ内で直接ネットワーク接続性を構築する新規アルゴリズム。CPU-GPU 転送のボトルネックを排除し、構築フェーズ中の MPI 通信を回避する。
2. 柔軟性: ポイントツーポイントおよび集団 MPI 通信の両方をサポートし、異なるネットワークトポロジー（階層的 vs ランダム/バランスド）への適応を可能にする。
3. エクサスケール対応: 著者は、今後の JUPITER エクサスケールスーパーコンピュータ上で、$2 \times 10^{10}$ニューロンおよび$10^{14}$シナプスのネットワークをシミュレートできることを外挿している。この規模は、個々のシナプス分解能を維持しながら、人間の皮質の接続性に近づいている。
4. 効率性: GML システムによるメモリ使用量の最適化により、既存のハードウェア上でより大規模なネットワークのシミュレーションを可能にする（例：MAM を 32 GPU ではなく 8 GPU に収容）。また、将来のエクサスケールシステムの全容量を活用する道筋を提供する。

著者は、このアプローチが分散シミュレーションにおけるスパイク通信の決定的なボトルネックに対処し、現代の HPC アーキテクチャにおける大規模で生物学的に詳細な神経シミュレーションのための参照プラットフォームとして NEST GPU を確立すると結論付けている。