NVIDIA ALCHEMI Toolkit でカスタム分子動力学パイプラインを構築する実践ガイド

なぜ今、GPUネイティブな分子シミュレーションなのか

計算化学の分野では長らく、精度と速度のトレードオフが課題でした。DFT（密度汎関数理論）に代表されるab initio手法は高い精度を提供するものの、計算コストが高く、数百原子程度の系に制限されます。一方、古典力場は高速ですが、複雑な結合解離や遷移状態解析に必要な化学精度が不足します。

機械学習原子間ポテンシャル（MLIP） はこのギャップを埋める技術として登場し、量子精度を古典速度で実現します。しかし、ソフトウェアエコシステムが新たなボトルネックとなっていました。MLIPモデル自体はGPU上で動作するものの、周辺のシミュレーション基盤はCPU中心のレガシーコードに依存しているケースが多かったのです。

NVIDIA ALCHEMI Toolkitは、こうしたボトルネックを解消するために設計された GPUアクセラレーテッド・シミュレーション・ビルディングブロック のコレクションです。PyTorchネイティブ環境で、近接リスト構築、構造最適化、分子動力学（MD）積分ステップをすべてGPU上で処理することで、ホスト-デバイス間のメモリ転送オーバーヘッドを排除します。

参考: 本記事はNVIDIA Developer Blogの公式資料を元に再構成しました。詳細はこちらをご参照ください。

コア機能1: カスタマイズ可能なバッチシミュレーションワークフロー

ALCHEMI Toolkitの最大の特徴は GPUネイティブのバッチダイナミクスエンジン です。単一のMLIPモデルがすべての化学環境に完璧であることは稀で、特に非局所的な長距離相互作用を扱う際には顕著です。

Toolkitは、研究者がモジュール式の化学/材料科学カーネルと深層学習モデルを組み合わせて カスタムシミュレーションワークフロー を構成することを可能にします。以下は、FIRE2構造最適化の後にVelocity Verlet動力学を実行する基本的な例です。

from ase import Atoms
from nvalchemi.data import AtomicData, Batch
from nvalchemi.dynamics import ConvergenceHook
from nvalchemi.dynamics.optimizers import FIRE2
from nvalchemi.dynamics.integrator import VelocityVerlet

# 1. バッチ原子構造データをGPUに直接ロード
atomic_data = [AtomicData.from_atoms(Atoms(...), device="cuda") for _ in range(16)]
batch = Batch.from_data_list(atomic_data)

# 2. MLIPモデルを設定 (MACE, TensorNetなど)
mlip = ...

# 3. 収束条件: force norm 0.05 eV/Å以下、max 0.1 eV/Å以下
conv_criteria = ConvergenceHook(
    criteria=[
        {"key": "forces", "threshold": 0.05, "reduce_op": "norm"},
        {"key": "forces", "threshold": 0.1, "reduce_op": "max"}
    ]
)

# 4. 最適化器と積分器を定義
optimizer = FIRE2(mlip, convergence_hook=conv_criteria, n_steps=200)
velverlet = VelocityVerlet(mlip, n_steps=1000)

単一GPUでパイプラインを実行

FusedStageクラスを使うと、2つのダイナミクスオブジェクトを加算演算子で結合し、torch.compileでラップできます。

fused = optimizer + velverlet

with fused:
    # 200ステップ最適化 + 1000ステップMDを自動実行
    fused.run(batch)

マルチGPU分散実行

パイプ演算子（|）を使うと、各ステップを異なるGPUに割り当てられます。

pipeline = optimizer | velverlet
# FIRE2はGPU 0、Velocity VerletはGPU 1で実行

with pipeline:
    pipeline.run(batch)

さらに、16台のGPUを使って8台で構造最適化、残り8台でLangevin動力学を実行することも可能です。

from torch import distributed as dist
from torch.utils.data.distributed import DistributedSampler
from nvalchemi.data.datapipes import Dataset, DataLoader

dist.initialize_process_group()

dataset = Dataset(...)
data_sampler = DistributedSampler(dataset, num_replicas=dist.get_world_size(), rank=dist.get_rank())
loader = DataLoader(dataset, batch_size=128, sampler=data_sampler, use_stream=True)

# 8 rankが最適化、8 rankがLangevin動力学
optimizers = [FIRE2(mlip, ..., next_rank=index + 8) for index in range(8)]
dynamics = [Langevin(mlip, ..., prior_rank=index) for index in range(8)]
pipeline = DistributedPipeline(
    {index: stage for index, stage in enumerate(optimizers + dynamics)}
)

with pipeline:
    for batch in loader:
        pipeline.run(batch)

この抽象化により、ユーザーは複雑な分散システムの詳細を意識せずにスケールアップできます。

コア機能2: 独自のダイナミクスクラスを構築する

ALCHEMI Toolkitはモジュラーアーキテクチャを提供し、完全に新しいダイナミクスクラスをゼロから構築することを可能にします。以下は、シミュレーテッドアニーリング（焼きなまし）フックを含むVelocity Verletの実装例です。

from enum import Enum
import torch
from nvalchemi.data import Batch
from nvalchemi.dynamics.base import BaseDynamics, DynamicsStage
from nvalchemi.hooks import Hook, HookContext

# カスタムフック: MDステップごとに温度を徐々に下げる
class MySimulatedAnnealer(Hook):
    def __init__(self, t_start: float, t_end: float, cooldown_steps: int, frequency: int):
        self.frequency = frequency
        self.t_start = t_start
        self.t_end = t_end
        self.cooldown_steps = cooldown_steps
        self.stage = DynamicsStage.BEFORE_STEP
        self.decay = (t_end / t_start) ** (1.0 / cooldown_steps)
        self._current_temp = t_start

    def __call__(self, ctx: HookContext, stage: Enum) -> None:
        dynamics = ctx.workflow
        dynamics.target_temperature = max(
            dynamics.target_temperature * self.decay, self.t_end
        )

# カスタムVelocity Verlet (温度制御付き)
class VelocityVerlet(BaseDynamics):
    __needs_keys__ = {"energies", "forces", "masses", "velocities"}
    __provides_keys__ = {"positions"}

    def __init__(self, model, n_steps, dt=1.0, target_temperature=300.0, tau=10.0, hooks=None):
        super().__init__(model=model, n_steps=n_steps, hooks=hooks)
        self.dt = dt
        self.target_temperature = target_temperature
        self.tau = tau
        self._prev_accelerations = None

    def pre_update(self, batch: Batch) -> None:
        with torch.no_grad():
            accelerations = batch.forces / batch.masses
            self._prev_accelerations = accelerations.clone()
            batch.positions.add_(batch.velocities * self.dt + 0.5 * accelerations * self.dt**2.0)

    def post_update(self, batch: Batch) -> None:
        with torch.no_grad():
            new_accelerations = batch.forces / batch.masses
            batch.velocities.add_(0.5 * (self._prev_accelerations + new_accelerations) * self.dt)
            # 温度スケーリング (Berendsenサーモスタット)
            ke_per_atom = 0.5 * batch.masses * (batch.velocities**2).sum(dim=-1, keepdim=True)
            total_ke = scatter_add_(...)  # システムごとの運動エネルギー合計
            current_temp = 2.0 * total_ke / (batch.num_atoms * 3.0)
            ratio = self.target_temperature / current_temp
            lam = torch.sqrt(torch.tensor(1.0 + (self.dt / self.tau) * (ratio - 1.0))).clamp(min=0.8, max=1.2)
            batch.velocities.mul_(lam)

# 使用例: 900Kから300Kまで1000ステップかけて冷却
my_vv = VelocityVerlet(
    mlip,
    n_steps=1000,
    hooks=[MySimulatedAnnealer(t_start=900.0, t_end=300.0, cooldown_steps=10, frequency=100)]
)

コア機能3: モデルラッパー

自身の事前学習モデルをALCHEMIパイプラインに接続するには、BaseModelMixinを継承してadapt_inputとadapt_outputメソッドを実装するだけです。

from beartype import beartype
from nvalchemi._typing import ModelOutputs
from nvalchemi.models.base import BaseModelMixin, ModelConfig, NeighborConfig

class BestMLIPWrapper(nn.Module, BaseModelMixin):
    def __init__(self, model):
        super().__init__()
        self.model = model
        self.model_config = ModelConfig(
            outputs=frozenset({"energy", "forces"}),
            required_inputs=frozenset({"positions", "atomic_numbers"}),
            neighbor_config=NeighborConfig(cutoff=5.0, format="coo")
        )

    def adapt_input(self, data: Batch, **kwargs) -> dict:
        return {
            "atom_numbers": data.atomic_numbers,
            "coords": data.positions
        }

    def adapt_output(self, model_output, data: Batch) -> ModelOutputs:
        output = ModelOutputs()
        output["energies"] = model_output["energies"]
        if "forces" in self.model_config.active_outputs:
            output["forces"] = model_output["forces"]
        return output

    @beartype
    def forward(self, data: Batch, **kwargs) -> ModelOutputs:
        model_inputs = self.adapt_input(data, **kwargs)
        model_outputs = self.model(**model_inputs)
        return self.adapt_output(model_outputs, data)

コア機能4: 高度なデータ管理

従来、CPUとGPU間のデータ移動はAI駆動型発見の主要なボトルネックでした。ALCHEMI ToolkitはAtomicDataとBatchオブジェクトを通じてデータをGPU上に常駐させることでこの問題を解決します。

from nvalchemi import AtomicData, Batch
from nvalchemi import data
from ase.build import slab

# ASE Atomsオブジェクト -> AtomicData (即座にGPUへ)
atoms = slab(...)
data = AtomicData.from_atoms(atoms, device="cuda")

# 複数構造を1つのバッチに
batch = Batch.from_data_list([data, data, data])

# Zarrベースのデータセットに書き込み/読み込み
writer = data.AtomicDataZarrWriter("atom_dataset.zarr")
writer.write(batch)

reader = data.AtomicDataZarrReader("atom_dataset.zarr")
dataset = data.Dataset(reader, device="cuda", num_workers=4)
dataloader = data.DataLoader(dataset, batch_size=16)
for batch in dataloader:
    # ここでバッチ処理
    pass

この技術の限界と注意点

1. モデル依存性: ALCHEMI ToolkitはMLIPモデルのラッパーとして機能するため、MLIP自体の精度がシミュレーション品質を決定します。DFTレベルの精度が必要な場合は、MACEやTensorNetなど検証済みのモデルを使用すべきです。
2. メモリ制限: バッチサイズが大きくなりすぎるとGPUメモリが不足する可能性があります。batch_sizeを調整するか、勾配チェックポインティング手法を併用することを推奨します。
3. OSサポート: 公式サポートはLinuxが主力で、macOSは限定的に動作します。Windowsは現時点では公式サポートされていません。
4. GPU要件: CUDA Compute Capability 7.0以上（RTX 20xx世代以降）が必要です。古いGPUクラスタを使用する研究室は事前に互換性を確認してください。

次のステップ

• ALCHEMI ToolkitのGitHubリポジトリをクローンし、提供されているJupyterノートブックを実行してみてください。
• MACE、TensorNet、AIMNet2など様々なMLIPモデルを自身のデータに適用し、バッチシミュレーションのパフォーマンスを測定することをお勧めします。
• 分子動力学結果の解析には、nvalchemi.hooksモジュールの多彩な分析フックを活用してみてください。

合わせて読みたい記事

• エヌビディア ルービン発表前に準備完了したAzureのAIデータセンター戦略
• CSSハイライト擬似要素完全ガイド search-textからアクセシビリティまで