In Silico

Materials Informatics

MLIPのGPUバッチ化：TorchSimとInferenceBatcherの違い

2026/10/3 （更新: 2026/10/11）

🤖 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 と、fairchem の InferenceBatcher を例に取る。TorchSim は、MLIP 向けのオープンソースの原子シミュレーションエンジンで、Orion Cohen らが論文として発表し、GitHub で公開している12。fairchem は、FAIR Chemistry（GitHub の facebookresearch 組織）が公開する MLIP のライブラリで、InferenceBatcher はその中の機能である3。この二つを選んだのは、TorchSim が計算の手順全体を、InferenceBatcher がモデル呼び出しだけをバッチ化するのでバッチ化する場所の違いが明確であり、どちらも作り手が設計の理由と向く対象を一次資料に書いているからだ。二つは競合する関係にはなく、TorchSim は fairchem のモデルにも対応する1。一つの大きな系には、InferenceBatcher ではなく、グラフ並列で推論する predict unit を使うよう、fairchem の文書は勧める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 を効率よく使うと書く2。fairchem の文書は、小〜中規模の系で多数の独立した計算を回すとき、モデルの推論呼び出しをバッチ化すると GPU の利用率を大きく上げられると書く3。二つの作り手は同じ問題から出発し、バッチ化する場所で分かれる。

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

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

一回にバッチ化できる量には、GPU のメモリで決まる上限がある。TorchSim は、この上限を自動で扱う仕組みを持つ。GPU のメモリに収まる量を見積もり、構造を詰めてバッチ化する仕組みを、自動バッチ（autobatching）と呼ぶ。文書は、モデルによってメモリの使い方が違い、MACE は原子数と数密度の積で、Fairchem のモデルは原子数で増えると書く4。論文の測定では、H100 に収まる原子の数はモデルごとに違い、MACE-MPA-0 で 8,000、Mattersim V1 1M で 22,000 だった2。そのため TorchSim は、モデルのメモリの使い方をその場で見積もり、使えるメモリに合わせてシミュレーションを並べる4。

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

論文は、バッチ化しない形では、Materials Project で最大の構造を緩和しても H100 の 1% 未満しか使わないと見積もる2。そのうえで、H100 で 16 原子の系 500 個をバッチ化した MACE-MPA-0 の分子動力学が、直列に回すより 100 倍の処理量を出したと書く2。同じ論文は、直列で回すと TorchSim と ASE（原子シミュレーションの Python ライブラリ）はほぼ同じ速さで、速さはバッチ化することから来ると書く2。MLIP の呼び出しをループで回す時間は、同じ数の積分や最適化の歩みを回す時間とほぼ同じで、積分の演算はモデル呼び出しよりはるかに安いとも、論文は書く2。速さの差は小さな系で最も大きく、系がメモリの上限に近づくと直列の ASE の速さに近づく2。測定は分子動力学で、緩和の速さを測ったものではない。この倍率は作り手の公表値であり、本稿は速さを含めて何も実行していない。入出力の面では、TorchSim は ASE・pymatgen・Phonopy と連携し、既存のツールで作った構造を受け渡せる1。

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。ただし文書は InferenceBatcher の速さを数値で示しておらず3、TorchSim の論文が比べた相手は直列の ASE である2。二つの方式の速さを同じ条件で比べた測定は、本稿が引く資料には無い。

計算を並行に走らせる 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 の積分器とオプティマイザを使う24。手順ごと書き直す設計である限り、この代償は恒久である。入出力は ASE などとつながるが、緩和や時間発展の手順そのものは TorchSim の側に移る1。

TorchSim の二つ目の代償は、バッチ化した計算を再現する手間である。TorchSim の文書は、分子動力学の軌跡が同じ初期構造と同じ設定からでも実行ごとに一致しないことが多いと書き、原因に GPU の非決定的な演算と確率的な積分器を挙げる7。一歩ごとの軌跡を合わせたいなら、GPU の演算を PyTorch の決定的な設定で固め、確率的な積分器の乱数を SimState に乱数の種（seed）を明示して固めるよう、文書は求める7。多くの分子動力学では統計量が合えばよいので、軌跡の不一致は問題にならないとも、文書は書く7。二つの原因は分子動力学の軌跡一般について挙げられており、文書がバッチ化に固有の点として別の節に書くのは、バッチの全構造が一つの乱数生成器を共有するので、厳密に再現するにはバッチの組み方も同じにする必要があることである7。バッチ化に固有の再現の手間は、軌跡を厳密に合わせたいときにバッチの組み方まで固定することで、バッチの全構造が一つの乱数生成器を共有する設計である限り残るので、恒久と判定する。なお、自動バッチは GPU のメモリに合わせてバッチの大きさを決める仕組みで、CPU では一般に対応しない8。GPU を使い切るための仕組みなので、本稿はこれを代償に数えない。

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

出典8件

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

  2. 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 ↩11 ↩12 ↩13

  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 ↩11 ↩12 ↩13

  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 — 代表的な構造を渡して呼ぶとバッチの上限を GPU メモリから探る関数を加え、メモリ不足時にバッチを分ける設定は既定で無効。 ↩ ↩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「Reproducibility」（2026-10-03 取得）. https://github.com/TorchSim/torch-sim/blob/main/docs/user/reproducibility.md — GPU の非決定的な演算と確率的な積分器で軌跡が実行ごとに一致しないこと（多くの分子動力学では問題にならないと明記）、厳密な再現には同じバッチの組み方が要ること。 ↩ ↩2 ↩3 ↩4

  8. TorchSim autobatching tutorial（2026-10-03 取得）. https://github.com/TorchSim/torch-sim/blob/main/examples/tutorials/autobatching_tutorial.py — 自動バッチは GPU のメモリに合わせてバッチの大きさを決める仕組みで、CPU では一般に対応しない。 ↩

この記事はAIを使用して執筆されました。内容には誤りが含まれる可能性があります。事実や数字の出典は脚注から確かめられます。