Magnum.np.distributed: Accelerating Finite Difference Micromagnetic Simulations with Multiple GPUs

本論文では、PyTorch Distributed上に構築された初のPythonネイティブなマルチGPUマイクロマグネティック・フレームワークであるMagnum.np.distributedを紹介するが、これは減磁界の計算を大幅に加速させ（8つのGPUで最大7.0倍の高速化を実現）、より大規模かつ複雑なスピントロニクス系のシミュレーションを可能にするものである。

要約

あなたは、コンピュータチップの中で小さな磁石がどのように振る舞うかをシミュレーションしようとしていると想像してください。これを行うために、科学者たちは磁石を数百万個の小さなレゴブロック（「メッシュ」と呼ばれます）に分解し、各ブロックが隣り合うブロックにどのように押し引きするかを計算します。これはマイクロマグネティック・シミュレーションと呼ばれます。

長い間、これらのシミュレーションは、たった一人の人間で山を動かそうとするようなものでした。強力なグラフィックスカード（GPU）の助けを借りても、ほとんどのソフトウェアは一度に一つのカードしか使用できませんでした。それは高速でしたが、最も大きく複雑な磁気のパズルを解くには、十分な速さではありませんでした。

この論文は、Magnum.np.distributedと呼ばれる新しいツールを紹介しています。これは、一人の人間が山を動かすことから、全員が自分専用のスーパーパワーを持ったシャベルを持つ建設現場の作業員チームへとアップグレードすることだと考えてください。

その仕組みを、シンプルな概念に分解して説明します。

1. 「チームワーク」の問題

昔は、もし8つのグラフィックスカードを使いたいと思ったら、それらが互いに通信できるように、複雑で困難なコード（C++やCUDAなど）を書かなければなりませんでした。これは、8人が異なる言語を話し、異なる制服を着て家を建てようとしているようなものです。

新しいMagnum.np.distributedフレームワークは、現代のデータサイエンスの「英語」であるPythonの言語を話します。これは、複数のGPUが即座に互いに通信できるようにするために、PyTorch Distributedというライブラリを使用しています。

• 比喩： リレーレースを想像してください。旧システムでは、バトン（データ）は遅くて手動の受け渡しを経なければなりませんでした。この新しいシステムでは、ランナー（GPU）は高速な光ファイバーのトンネル（NVLink）で接続されており、バトンをほぼ瞬時に受け渡すことができます。

2. 仕事の分割方法

シミュレーションのグリッドは、パンの塊のようにスライスされます。

• ローカルタスク： 一部の計算は、特定のパンのスライスだけに依存しています。これらは簡単です。各GPUは独立して自分のスライスを行います。これは、8人が部屋の8つの異なる壁を塗っているようなもので、彼らは互いに会話する必要はありません。
• 「ハロー（Halo）」問題： 一部の計算（「交換場」など）は、隣のブロックが何をしているかを知る必要があります。例えば、あなたが壁の端を塗っている場合、隣の壁がどのような色であるかを知る必要があります。
  • 解決策： システムは、各スライスの端に「ハロー」（バッファゾーン）を作成します。GPUは、この端のデータを隣接するGPUと交換します。
  • 落とし穴： 論文では、これが非常にうまく機能する一方で、コンピューターの脳（CPU）がGPUにタスクの開始を指示するまでの時間が、この端データの「受け渡し」を時として遅らせることがあると指摘されています。これは、ランナーは速いのに、コーチが「行け！」と叫ぶのが遅すぎるようなものです。

3. 大きな挑戦： 「消磁（Demagnetization）」場

シミュレーションの中で最も難しい部分は、消磁場を計算することです。これは、磁石全体のすべてのレゴブロックが、他のすべてのブロックの影響を受けるという「グローバル」な計算です。

• 比喩： 部屋の中にいる人々の中で、全員が同時に全員に対してメッセージを叫ぶ必要がある状況を想像してください。
• 結果： この新しいフレームワークは、FFT（高速フーリエ変換）と呼ばれる数学的なトリックを使用して、この問題を処理します。これにより、データを並列処理できるようにシャッフルします。
  • スピード： 8つの強力なGPUが超高速のトンネル（NVLink）で接続されたシステムでは、この特定のタスクは、単一のGPUを使用した場合よりも7倍速くなりました。ほぼ完璧な線形加速を実現しました。

4. 一般的なコンピュータ（CPU）については？

誰もが高性能な8枚のGPUを備えたクラスターを持っているわけではありません。著者らは、標準的なコンピュータプロセッサ（CPU）上でもこれをテストしました。

• 問題： CPUには異なるメモリゾーン（NUMA）があります。プログラムが「遠い」メモリゾーンからデータを取得すると、速度が低下します。
• 解決策： 彼らは、プログラムを自身の「ローカル」なメモリ近傍に留まらせるNUMAピニングという手法を使用しました。
• 結果： 強力なデュアルソケットCPUにおいて、この手法はピニングなしの場合よりも6.8倍速くなりました。高性能なGPUと比較すると依然として約10倍遅いものの、高価なGPUへのアクセスを持たない研究者でも、以前よりずっと速くこれらの複雑なシミュレーションを実行できることを意味しています。

5. 実世界のテスト： 「迷路」磁石

それが機能することを証明するために、彼らは高度なメモリデバイスに使用される実際の磁気スタック（Pt/Gd/Co/Ni）をシミュレートしました。

• セットアップ： 2,360万個のセルを持つグリッド。
• 結果：
  • 1 GPU： 50.6時間。
  • 4 GPU： 8.4時間。
  • 加速率： 6倍高速化。
• 予想以上に優れていた理由： 仕事が4つのGPUに分割されたとき、データがGPUの小さく高速な内部メモリ（キャッシュ）により適合したため、プロセス全体がさらにスムーズに進行しました。

まとめ

この論文は、科学者が複数のGPUを使用して微小な磁石をシミュレートすることを可能にする、最初のPythonネイティブなツールを提示しています。

• コーディングの悪夢なし： C++の専門家である必要はありません。Pythonを書くだけです。
• 圧倒的なスピード： 8つのGPUで7倍速くすることができます。
• 多用途： 高性能なGPUクラスターでも、適切な設定を行えば標準的なコンピュータのCPUでも動作します。

これにより、研究者はより大規模で複雑な磁気システムを、わずかな時間でシミュレートできるようになり、次世代のコンピュータメモリやロジックであるスピントロニクスデバイスをより迅速に設計することが可能になります。

技術概要

技術要約: Magnum.np.distributed

問題提起

マイクロマグネティック・シミュレーションは、ナノ磁性学やスピントロニクスの研究において極めて重要であるが、既存の高性能ソルバーはアクセシビリティとスケーラビリティに関して大きな制限に直面している。Mumax3やPythonネイティブのmagnum.npのようなツールはGPU加速を利用しているものの、これらはシングルデバイスの計算に限定されている。この制約は、現代のデバイス設計に求められる、より大規模で複雑なシステムのシミュレーションを妨げている。さらに、広く使用されているソルバーの多くは、非Pythonスクリプト言語（GoやTCLなど）に依存しているか、あるいはC++/CUDAコードをソースからコンパイルする必要があり、これらがインストール、クロスプラットフォームの互換性、およびPythonベースの解析ワークフローへのシームレスな統合における障壁となっている。Borisのような一部のC++/CUDAベースのシミュレータはマルチGPU操作をサポートしているが、Pythonネイティブのフレームワークのようなインストールの容易さやプラットフォームに依存しない設計を欠いている。

手法

著者らは、既存のmagnum.npソルバーをPyTorch Distributedによって拡張することにより、初のPythonネイティブ・マルチGPUマイクロマグネティック・フレームワークであるmagnum.np.distributedを提示する。この実装は、プラットフォームに依存しない設計を維持するためにベンダー固有のCUDA最適化を避けつつ、PyTorchのJust-In-Time (JIT) コンパイル（TorchDynamoおよびTorchInductor経由）とオートグラッド（autograd）機能を活用している。

コア実装の詳細

1. 領域分割（Domain Decomposition）: シミュレーションメッシュは、x軸に沿って連続したスラブ（slab）に分割され、各プロセスランクに1つのスラブが割り当てられる。
2. ローカル場とグローバル場:
   • ローカル場: スロンツェウスキー・スピン注入トルクや一軸異方性などの項は、ローカルデータに対して作用し、プロセス間通信なしに分散スラブに直接適用される。
   • ハロー交換（Halo Exchange）: ハイゼンベルク交換相互作用やジャロシンスキー・守谷相互作用（DMI）のような非局所的な相互作用のために、フレームワークはハロー領域（スラブ境界における2つの余剰セル）を実装している。通信は、バッチ化された非ブロッキングのポイント・ツー・ポイント転送によって処理される。重要な点として、これらのハロー交換は、TorchDynamoの最適化問題を回避するために、JITコンパイルのスコープの外側に配置されており、計算のみがコンパイル対象となっている。
3. 脱磁場（Demagnetization Field）: 最も計算負荷の高いグローバルな相互作用として、脱磁場はNewellの公式を用いた離散畳み込みによって計算される。実装では、分散3D高速フーリエ変換（FFT）戦略を利用している：
   • メッシュはx軸に沿って分解される。
   • y方向およびz方向の2D FFTが実行される。
   • **all-to-all転置（all-to-all transpose）**により、データをy軸に沿ったベクトルへと再分配する。
   • x方向の1D FFTを実行するためにゼロパディングが行われる。
   • 脱磁カーネルとの点別乗算が行われる。
   • 逆パスがこれらのステップを反転させる。
   • 注記： 著者らは、コードの可読性を維持するため、また現在のJITコンパイラが複素数演算を単一のカーネルコールに融合できないため、他のソルバー（Borisなど）で使用されている通信時の精度を半分にする機能を削除した。
4. ソルバーの統合: フレームワークは、時間積分用のRunge-Kutta-Fehlberg法（RKF45）と、エネルギー最小化のためのBarzilai-Borromweinアルゴリズムをサポートしている。これらのソルバーは、最大誤差や収束基準のためのグローバルなリダクションを行うように修正されており、すべてのランクが同じタイムステップで進行するか、同時に終了することを保証している。

バックエンドサポート

フレームワークは、NVLink/NVSwitchを介して接続されたNVIDIA GPU向けのNCCLバックエンドと、CPUベースの計算向けのMPIバック後をターゲットとしている。MPIバックエンドは、CPUのパフォーマンスを最適化するためにNUMA（Non-Uniform Memory Access）ピニングを具体的にサポートしている。

主な貢献

• 初のPythonネイティブ・マルチGPUフレームワーク: C++やCUDAのコンパイルを必要とせずにmagnum.npを拡張し、インストールの容易さとPythonエコシステムの互換性を維持したまま、マルチGPUをサポートする。
• 分散FFTアルゴリズム: 分散スラブ間で効率的にall-to-all転置とゼロパディングを処理する、脱磁場のためのスケーラブルな3D FFTアルゴリズムを実装した。
• ハイブリッドMPI+OpenMP CPUサポート: PyTorch MPIバックエンドを介してNUMAピニングを用いたCPUクラスターを効果的に利用できることを示し、GPUリソースが制約された環境における実行可能な代替案を提供した。
• 最小限のコード移行: 非分散版から分散版への移行には、約8行のコード変更（分散ソルバーのインポート、プロセスグループの初期化、および標準的なメッシュ/状態オブジェクトの分散型への置き換え）のみが必要であることを示した。

結果

検証

フレームワークは、muMAG標準問題4（パーマロイ磁石における磁気反転）およびDMIとドメイン壁ピンニングを含む問題を用いて検証された。分散シミュレーション（2、4、8ランクを使用）の結果は、元の非分散版magnum.npおよびMumax3（単精度）の結果と一致しており、分散LLG-RKF45ソルバー、脱磁、および交換場の実装の正当性を確認した。

パフォーマンス・ベンチマーク

• マルチGPUスケーリング（NVLink/NVSwitch）:
  • 脱磁場: 8枚のNVIDIA H100 HBM3 GPU（NVSwitch接続）において、7.0倍の高速化を達成した。スケーリングはほぼ線形であり、異なるメモリ帯域世代（HBM3 vs HBM2e）間で一貫していた。これは、性能が生の帯域幅ではなく、アルゴリズムの通信対計算比によって支配されていることを示している。
  • 交換/DMI場: ポイント・ツー・ポイント通信におけるカーネルディスパッチのレイテンシにより、小さな問題サイズではスケーリングが制限された。意味のある高速化（8 GPUで4.2倍など）は、より大きなメッシュ（>8Mセル）においてのみ観察された。
  • ローカル場（STT）: インターランク通信を必要としないため、ほぼ理想的な高速化を達成した。
• マルチノード・スケーリング:
  • 2ノードクラスター（計8 GPU）を400 Gbps InfiniBandで接続したテストでは、特に脱磁場において、シングルノードNVSwitchシステムと比較して大幅な低下（72%の劣化）が見られた。これは、インターコネクトのボトルネック（InfiniBand vs NVSwitch）と、all-to-all転置中に転送される大量のデータに起因している。
• CPUパフォーマンス（NUMAピニング）:
  • デュアルソケットIntel Xeonノードにおいて、MPIバックエンドによるNUMAピニングを有効にすることで、$10^6$セルの問題に対する脱磁場の計算時間を6.8倍短縮した（204.0 msから29.8 msへ）。
  • 単一のH100 GPUよりは依然として10〜15倍遅いものの、この性能により、GPUへのアクセス権を持たない研究者にとってもCPUベースのシミュレーションが実行可能なものとなっている。
• 実世界への応用:
  • 2360万セルを持つ界面DMI誘起多層膜（Pt/Gd/Co/Ni）のシミュレーションにおいて、フレームワークは単一GPUと比較して4枚のGPUで6.0倍の高速化を達成した。この超線形スケーリングは、ワーキングセットが複数のデバイスに分散されることで、GPUキャッシュの利用効率が向上したことによるものである。

重要性と主張

本論文は、magnum.np.distributedが、高性能なマルチGPUコンピューティングと、Pythonネイティブの科学ソフトウェアに固有の使いやすさとの間の溝を埋めることに成功したと主張している。C++/CUDAのコンパイルを不要にし、GPUおよびCPU（NUMA最適化を含む）の両方の実行のための統一されたインターフェースを提供することで、高忠実度のマイクロマグネティック・シミュレーションへの参入障壁を下げている。

著者らは、ターンアラウンドタイムの短縮が、研究者がより大規模で複雑なナノ磁性システムを探索することを可能にし、それによって新しいスピントロニクスデバイスの設計サイクルを加速させることを強調している。本研究は、将来の高帯域幅インターコネクト（例：NVIDIA NVL72）の恩恵を受ける位置付けにある一方で、現在はシングルノード・マルチGPUおよびNUMA対応CPUクラスターに対する堅牢なソリューションを提供している。さらなる採用と発展を促進するため、ソースコードは公開されている。