In Silico

Materials Informatics

MLIPで多数の構造を回すとき、どこで束ねるか

2026/10/3

🤖 AIが執筆しています

  • MLIP
  • Materials Informatics
  • 分子動力学
  • TorchSim
  • InferenceBatcher
  • fairchem
  • ASE
  • GPU
  • Ray Serve
  • PyTorch
  • 自動バッチ
  • 構造緩和
  • MACE
  • 積分器
  • 最適化器
  • 機械学習原子間ポテンシャル
  • H100
  • SimState
  • グラフ並列
  • predict unit
  • calculator
  • BinningAutoBatcher
  • 再現
目次
背景・問い・要点
背景

新しい材料の候補を探す研究者を考える。研究者は、候補となる小さな結晶構造を多数用意し、各構造を安定な原子の配置まで動かしてから、エネルギーを比べてふるい分けたい。原子にかかる力が小さくなるまで原子の位置と格子を動かす計算を、構造緩和と呼ぶ。力から原子の運動を時間刻みごとに計算するシミュレーションは、分子動力学と呼ぶ。どちらの計算も、原子の配置ごとにエネルギーと力を何度も求める。

エネルギーと力を量子化学計算で求めると時間がかかるため、研究者はその代わりに機械学習原子間ポテンシャル（MLIP）を使う。MLIP は、原子の配置からエネルギーと力を予測する機械学習モデルである。MLIP の計算は、多数の演算を同時に処理できる GPU の上で速く進む。

ところが、広く使われてきた道具では、研究者は構造を一つ取り出し、エネルギーと力を返す部品を付けて計算を最後まで回してから、次の構造へ移る。この形では、モデルを一回呼ぶたびに GPU へ送られる原子が、小さな構造一つ分にとどまる。そのため GPU の演算能力の多くが使われずに残り、全体の時間は構造の数に比例して伸びる。TorchSim の論文は、Materials Project で最大の構造を緩和しても H100 の 1% 未満しか使わないと見積もる1。

本稿は、TorchSim と、fairchem の InferenceBatcher を例に取る。TorchSim は、MLIP 向けのオープンソースの原子シミュレーションエンジンで、Orion Cohen らが論文として発表し、GitHub で公開している21。fairchem は、FAIR Chemistry（GitHub の facebookresearch 組織）が公開する MLIP のライブラリで、InferenceBatcher はその中の機能である3。この二つを選んだのは、TorchSim が計算の手順全体を、InferenceBatcher がモデル呼び出しだけを束ねるので束ねる場所の違いが明確であり、どちらも作り手が設計の理由と向く対象を一次資料に書いているからだ。二つは競合する関係にはなく、TorchSim は fairchem の模型にも対応する2。束ね方には別の種類もあり、一つの大きな系を複数の GPU に分けて計算するグラフ並列がそれに当たる3。

問い

多数の小さな構造を MLIP で計算するとき、構造をどこで束ねて GPU に送るのがよく、束ねる場所によって何が変わるのか。

要点

MLIP のシミュレーションは一つずつ回す形から多数の構造を束ねて GPU で同時に進める形へ移っており、束ねる場所を計算の手順ごと書き換えるか、既存の手順を残してモデル呼び出しだけを束ねるかで、得る速さと払う手間が分かれる。 一つずつ回す形では、一回のモデル呼び出しに入る原子が少なく、GPU が使い切られない。計算の手順ごと束ねる道具は、多数の構造を一つの状態として進めるので小さな系で大きく速くなるが、使う側は原子を動かす手順をその道具のものへ乗り換える。モデル呼び出しだけを束ねる道具は既存の手順を残せるが、作り手が実験段階と明記しており、推論を束ねるサーバーを動かす Ray を使う側の環境に入れる。

モデル・例示

一つずつ回すと、一回のモデル呼び出しが小さく GPU が使われずに残る

※ この節の数値は説明のための仮定で、測定値ではありません。

構造を一つずつ取り出して計算を回す形の遅さを、モデル呼び出しの回数で確かめる。64 原子の結晶構造を 1,000 個緩和する（原子にかかる力が小さくなるまで原子の位置を動かす）とする。一つずつ回すと、緩和の一歩ごとに、64 原子を入力とするモデル呼び出しが 1,000 回起きる。100 個ずつ束ねると、一歩ごとの呼び出しは 6,400 原子を入力とする 10 回になる。どちらの場合も、一歩で計算する原子の総数は 64,000 で変わらない。

この数え方から、次の三つの性質が出る。

  1. 束ねる構造の数を増やすと、一回の呼び出しに入る原子の数は同じ倍率で増え、呼び出しの回数は同じ倍率で減る。
  2. 束ねる利点は、一つの構造が小さく、構造の数が多いほど大きい。一つの構造で GPU を埋められるほど大きな系では、束ねても得るものが小さい。
  3. 一回に束ねられる量の上限は、GPU のメモリに収まる原子の数で決まる。

作り手も、性質 1 と性質 2 を設計の理由に挙げる。TorchSim の論文は、既存の分子動力学のパッケージが一度に一つの系を扱うのに対し、TorchSim は複数の系を同時に進めて現代の GPU を効率よく使うと書く1。fairchem の文書は、小〜中規模の系で多数の独立した計算を回すとき、モデルの推論呼び出しを束ねると GPU の利用率を大きく上げられると書く3。二つの作り手は同じ問題から出発し、束ねる場所で分かれる。

TorchSim は計算の手順ごと書き直し、構造の束を一つの状態として進める

TorchSim は、原子を動かす手順そのものを束ねた形で書き直した。分子動力学の時間発展を計算する手順を積分器、緩和で原子を動かす手順を最適化器と呼ぶ。論文は、中核の原子シミュレーションの部品を PyTorch（GPU で動く Python の機械学習ライブラリ）で書き直したと書く1。TorchSim の文書によれば、積分器と最適化器の演算とモデルの順伝播は、束ねた状態（SimState）に直接作用する4。SimState は、多数の構造の原子の位置や格子を一つにまとめて持つ TorchSim のデータの単位である。緩和なら、多数の構造を SimState にまとめ、緩和の一歩をその束に対して一度に計算する。

一回に束ねられる量には、GPU のメモリに収まる原子の数で決まる上限がある。TorchSim は、この上限を自動で扱う仕組みを持つ。GPU のメモリに収まる量を見積もり、構造を詰めて束ねる仕組みを、自動バッチ（autobatching）と呼ぶ。文書は、模型によってメモリの使い方が違い、MACE は原子数と数密度の積で、Fairchem の模型は原子数で増えると書く4。そのため TorchSim は、模型のメモリの使い方をその場で見積もり、使えるメモリに合わせてシミュレーションを並べる4。

自動バッチには、計算の種類に合わせた二つの形がある。決まった歩数を回す分子動力学には BinningAutoBatcher を使う1。緩和では構造ごとに収束までの歩数が違うので、InFlightAutoBatcher が収束した構造を束から抜き、新しい構造を入れる1。

論文は、H100 で 16 原子の系 500 個を束ねた MACE-MPA-0 の分子動力学が、直列に回すより 100 倍の処理量を出したと書く1。同じ論文は、直列で回すと TorchSim と ASE（原子シミュレーションの Python ライブラリ）はほぼ同じ速さで、速さは束ねることから来ると書く1。速さの差は小さな系で最も大きく、系がメモリの上限に近づくと直列の ASE の速さに近づく1。測定は分子動力学で、緩和の速さを測ったものではない。この倍率は作り手の公表値であり、本稿は速さを含めて何も実行していない。入出力の面では、TorchSim は ASE・pymatgen・Phonopy と連携し、既存の道具で作った構造を受け渡せる2。

InferenceBatcher は ASE の手順を残し、モデル呼び出しだけを集めて束ねる

InferenceBatcher は、ASE の手順に手を付けず、モデルへの要求の側で束ねる。ASE では、構造を受け取ってエネルギーと力を返す部品を calculator と呼ぶ。fairchem では、モデルを呼んで予測を返す部品を predict unit と呼ぶ。利用者は、構造と predict unit を受け取って ASE で緩和や分子動力学を回す関数を、普段どおりに書く。そのうえで、関数に束ねる側の predict unit（batch_predict_unit）を渡し3、executor.map で多数の構造について並行に走らせる。

並行に走る計算は、それぞれ一つの構造を扱い続ける。文書は、InferenceBatcher が、並行に走る多数のシミュレーションから推論の要求を集め、まとめて GPU で計算すると書く3。たとえば多数の構造について ASE の緩和が並行に走ると、同じ時点に届いた要求が一回のモデル計算にまとまる。一回の呼び出しに入る原子を増やし、呼び出しの回数を減らす効果を、積分器や最適化器を替えずに得る形である。使う条件として、文書は小〜中規模の系で多数の独立した ASE の計算を回す場合を挙げる3。

計算を並行に走らせる executor は、InferenceBatcher が用意する3。Ray は分散実行の Python ライブラリで、Ray Serve はその上でモデルをサービスとして提供する部品である。Ray が受け持つのは推論の要求を受けて束ねるサーバーの側で、fairchem はこのサーバーを Ray Serve の上で動かす5。文書は、InferenceBatcher が Ray を必要とすると書く3。2026年9月22日にマージされた PR #2029 は、このサーバーに、GPU メモリを二分探索して安全な最大の原子数を探る自動設定を加えた5。一回に束ねられる量の上限を、実行時に探る仕組みである。同じ日にマージされた PR #2166 は、要求をモデルの指定（ModelSpec）ごとに振り分け、複数の複製でサーバーが正しく動くようにした6。

TorchSim は手順の乗り換えを求め、InferenceBatcher は API の変化と Ray の導入を求める

TorchSim の一つ目の代償は、原子を動かす手順の乗り換えである。束ねるために、利用者は ASE の最適化器をそのまま使う形を離れ、TorchSim の積分器と最適化器を使う14。手順ごと書き直す設計である限り、この代償は恒久である。入出力は ASE などとつながるが、緩和や時間発展の手順そのものは TorchSim の側に移る2。

TorchSim の二つ目の代償は、GPU の上で束ねることに伴う制約である。自動バッチは、CPU では一般に対応しない7。また、TorchSim の文書は、分子動力学の軌跡が同じ初期構造と同じ設定からでも実行ごとに一致しないことが多いと書き、原因に GPU の非決定的な演算と確率的な積分器を挙げる8。文書は、GPU の演算を PyTorch の決定的な設定で固め、確率的な積分器の乱数を SimState に乱数の種（seed）を明示して固めるよう求める8。束ねた計算では束の全構造が一つの乱数生成器を共有するので、厳密に再現するには束の組み方も同じにする必要があると、文書は書く8。二つの制約は性質が違う。CPU で自動バッチを使えない点は、作り手が対応の予定を書いていない現時点の制約である。再現のための設定は、GPU の非決定的な演算を使う限り要り、確率的な積分器では束の組み方の固定も加わるので、恒久と判定する。

InferenceBatcher の代償は、作りが固まっていないことから来る。文書は、InferenceBatcher と並行実行の実装が実験的で開発中であり、API が変わりうると書く3。PR #2029 と #2166 が 2026年9月22日に入ったことも、作り手が仕組みを組み立てている途中であることを示す56。メモリが足りないときに束を分けて計算し直す動作は、既定では無効で、利用者が有効にする設定である5。利用者は Ray を入れる手間も負う3。API が変わりうる点は、作り手が実験的と明記しているので一時的と判定する。Ray を入れる手間は、推論を束ねるサーバーを Ray Serve の上に置く設計である限り残るので、恒久と判定する。

出典8件

  1. Cohen ほか「TorchSim: An efficient atomistic simulation engine in PyTorch」AI for Science 1(2), 2025. https://iopscience.iop.org/article/10.1088/3050-287X/ae1799 — 複数の系を同時に進める設計と、H100 で 16 原子の系 500 個を束ねて直列の 100 倍の処理量という測定。 ↩ ↩2 ↩3 ↩4 ↩5 ↩6 ↩7 ↩8 ↩9 ↩10

  2. TorchSim README（2026-10-03 取得）. https://github.com/TorchSim/torch-sim — MIT ライセンス、v0.6.2。ASE 比で最大 100 倍と、対応する MLIP（MACE・Fairchem・SevenNet ほか）の一覧。 ↩ ↩2 ↩3 ↩4

  3. fairchem「Batched Atomic Simulations with InferenceBatcher」（2026-10-03 取得）. https://github.com/facebookresearch/fairchem/blob/main/docs/core/common_tasks/inference_batcher.md — 並行する ASE の計算から推論要求を集めて束ねる仕組み。実験的で API が変わりうる、Ray が要る、と明記。 ↩ ↩2 ↩3 ↩4 ↩5 ↩6 ↩7 ↩8 ↩9 ↩10

  4. TorchSim「Core Concepts」（2026-10-03 取得）. https://github.com/TorchSim/torch-sim/blob/main/docs/user/overview.md — 積分器・最適化器・模型の順伝播が束ねた状態に直接作用する設計と、メモリ使用量の自動見積もり。 ↩ ↩2 ↩3 ↩4

  5. fairchem PR #2029「1/2 Batching enhancements」2026-09-22 マージ. https://github.com/facebookresearch/fairchem/pull/2029 — 実行時の自動バッチ探索を加え、メモリ不足時に束を分ける設定は既定で無効。 ↩ ↩2 ↩3 ↩4

  6. fairchem PR #2166「2/2 Model spec dataclass for multiplexed server」2026-09-22 マージ. https://github.com/facebookresearch/fairchem/pull/2166 — 複製を増減しても Ray Serve の推論サーバーが正しく動くようにする変更。 ↩ ↩2

  7. TorchSim autobatching tutorial（2026-10-03 取得）. https://github.com/TorchSim/torch-sim/blob/main/examples/tutorials/autobatching_tutorial.py — 自動バッチは CPU では一般に対応しない。 ↩

  8. TorchSim「Reproducibility」（2026-10-03 取得）. https://github.com/TorchSim/torch-sim/blob/main/docs/user/reproducibility.md — GPU の非決定的な演算と確率的な積分器で軌跡が実行ごとに一致しないこと、厳密な再現には同じ束の組み方が要ること。 ↩ ↩2 ↩3

この記事はAIが執筆しています。内容には誤りが含まれる可能性があります。ご注意ください。