【機械学習ポテンシャル】llmlip.rsの動作をUMAで検証する（WSL2環境）

Posted at 2026-10-31

はじめに

llmlip.rsはGPUに依存せずCPUだけで機械学習原子間ポテンシャル（MLIP）を実行できるように設計された軽量なRust製の推論エンジンで、CPU上でMLIP推論を効率的に実行することを目指して開発されている。

CUDA、PyTorch、専用GPUを必要としないため、H100やB200といった高価なGPUに依存しない選択肢の一つとして位置付けられる。MACEやNequIPだけでなく、UMAやeSENなどの比較的新しい機械学習ポテンシャルにも対応が進められており、汎用的な実行基盤として開発が進められているようである。

公開版の検証はmacOSのArm64環境で実施されており、LinuxやWindowsでの動作はまだ検証されていない。そこで、本稿では llmlip.rs の動作をWSL環境で調べる。機械学習ポテンシャルには、Meta社FairChemチームが開発するUMAを利用する。

マシンの構成は以下の通りだが、今回GPU 2は使用しない。

計算資源	仕様
CPU	AMD Ryzen 7 5800X3D（8コア／16スレッド）
GPU 1	NVIDIA GeForce RTX 5070 Ti（16GB）
GPU 2	NVIDIA GeForce RTX 3070 Ti（8GB）
メモリ	64GB
OS	Windows 11＋WSL2 Ubuntu

llmlip.rsの導入

実行にはRust 1.89以上が必要となる。必要ならRustの導入から始める。

Rustの導入

sudo apt update
sudo apt install -y \
    git curl build-essential pkg-config cmake \
    python3.12-venv

# Rustのインストール（defaultを選択すればよい）
curl https://sh.rustup.rs -sSf | sh

# 環境変数の反映
source "$HOME/.cargo/env"

# ツールチェーンを更新
rustup update stable

rustc --version
cargo --version

記事執筆時点では以下のバージョンがインストールされる。

rustc 1.99.1 (58cf8f9c3 2026-10-13)
cargo 1.99.1 (5f94df478 2026-08-27)

OpenBLAS版のビルド

OpenBLASによる行列演算高速化を取り入れるため、OpenBLASのフラグを有効化してビルドする。結果の項でも示すが、OpenBLASを有効化しない場合は計算効率が悪化し、計算時間が5倍以上に伸びる。

OpenBLASのインストール（未導入の場合）

sudo apt update
sudo apt install -y libopenblas-dev pkg-config

確認と結果

$ ldconfig -p | grep openblas
        libopenblas.so.0 (libc6,x86-64) => /lib/x86_64-linux-gnu/libopenblas.so.0
        libopenblas.so (libc6,x86-64) => /lib/x86_64-linux-gnu/libopenblas.so

続いて、llmlip.rsを取得してビルドする。筆者の環境では3分程度でコンパイルが完了した。

OpenBLAS版のビルド（CPU並列実行を有効化）

git clone https://github.com/mizukami-group/llmlip.rs.git
cd llmlip.rs

RUSTFLAGS="-C target-cpu=native" \
cargo build --locked --release \
    -p llmlip-cli \
    --features parallel,openblas \
    --target-dir target/native-openblas

コンパイル中、ソースコード内の未使用の関数についてwarningが出るが実用上問題は無い。

ビルドを実行する前に RUSTFLAGS="-C target-cpu=native" として実行環境のCPU機能をコンパイラに認識させている（Ryzen 7 5800X3DなどではAVX2やFMAなどをより積極的に利用できる可能性があるため）。
しかし、既存の演算ライブラリにもSIMD最適化が含まれているため、必ずしも大幅に高速化するとは限らない。（筆者の環境では両者に全く違いが見られなかった）

ビルド後、問題なく実行できるか確認する。

実行ファイルの確認

./target/native-openblas/release/llmlip --help

以下のように出力されれば問題ない。

Rust-native fixed-weight equivariant MLIP inference

Usage: llmlip <COMMAND>

Commands:
  predict                   Predict energy and/or forces from an .e3rs model
  finite-difference-forces  Check conservative forces against central finite differences
  benchmark-predict         Benchmark repeated load-once predictions and emit JSON timing stats
  batch-predict             Batch-predict energy/forces/stress for all frames in a multi-frame extXYZ
  validate-model            Validate an .e3rs model without running prediction
  pack-uma-tensors          Pack UMA tensor `.bin` files into a single-file `tensors.pack`
  help                      Print this message or the help of the given subcommand(s)

Options:
  -h, --help  Print help

サンプルモデルで動作確認する

リポジトリにはテスト用の合成ポテンシャル（examples/model.e3rs）が付属している。

./target/native-openblas/release/llmlip validate-model \
    --model examples/model.e3rs

./target/native-openblas/release/llmlip predict \
    --model examples/model.e3rs \
    --structure examples/input.extxyz \
    --energy --forces --json

正常に動作すれば、エネルギーと各原子に働く力がJSON形式で出力される。続いて、ベンチマーク機能を確認する。

./target/native-openblas/release/llmlip benchmark-predict \
    --model examples/model.e3rs \
    --structure examples/input.extxyz \
    --energy --forces \
    --warmup 5 \
    --iterations 100

このコマンドには、平均・中央値・95パーセンタイルなどの統計を出力する機能が実装されている。環境変数が設定されていなければ null が返る。

問題なければ、ウォームアップを含めた105回の推論が実行され、エネルギー・原子間力の計算が正常に完了する。

UMA-S-1.2の導入

今回の実モデルには、過去に検証済みのUMA-S-1.2を採用する。

注意点として、llmlip.rsの本体には学習済みのモデル重みが同梱されていない点が挙げられる。PyTorch形式のチェックポイントを取得し、llmlip.rs専用の.e3rs形式へ変換する必要がある。

Python環境の構築

GPUを搭載していない環境の場合はCPU版のPyTorchをインストールする。GPU版PyTorchをインストールした場合でも、device="cpu"のようにCPUの利用を明示的に指定することが可能。今回はGPU版を利用して検証用の環境を用意する。

llmlipは公式にはFairChem 2.20.0を参照バージョンとして固定しているため、PyTorchは2.8を導入するのが厳密なベンチマークとしては良いが、ここでは速度面が改善された最新のFairChem 2.25.0を導入し、PyTorch 2.13を導入している。

cd ~/llmlip.rs

python3.12 -m venv .venv
source .venv/bin/activate

python -m pip install --upgrade pip

# PyTorch GPU版
python -m pip install --no-cache-dir \
    torch==2.13.0 \
    --index-url https://download.pytorch.org/whl/cu130

# UMA-S-1.2の変換元に対応するFairChem
python -m pip install --no-cache-dir \
    fairchem-core==2.25.0 \
    huggingface_hub

python -m pip check

CPU版の場合はPyTorchのインストールコマンドを以下に差し替える。

PyTorch CPU版

python -m pip install --no-cache-dir \
    torch==2.13.0 \
    --index-url https://download.pytorch.org/whl/cpu

インストール結果の確認

python - <<'PY'
import torch
import importlib.metadata

print("PyTorch:", torch.__version__)
print("CUDA:", torch.version.cuda)
print("FairChem:", importlib.metadata.version("fairchem-core"))
print("CUDA available:", torch.cuda.is_available())

if torch.cuda.is_available():
    for i in range(torch.cuda.device_count()):
        print(f"GPU {i}:", torch.cuda.get_device_name(i))
PY

以下のように出力されれば問題ない。

PyTorch: 2.13.0+cu130
CUDA: 13.0
FairChem: 2.25.0
CUDA available: True
GPU 0: NVIDIA GeForce RTX 5070 Ti
GPU 1: NVIDIA GeForce RTX 3070 Ti

Pythonバインディングのビルド

llmlip.rsには、PythonからRustの計算エンジンを直接呼び出すためのバインディングも実装されている。ASEと組み合わせて構造最適化や分子動力学計算を行う場合に利用できる。

Pythonバインディングはmaturinを使用して構築する。今回はCLI版と同じく、OpenBLASによる行列演算の高速化とRayonによる並列化を有効にする。

cd ~/llmlip.rs

# 既存のPython仮想環境を利用
source .venv/bin/activate

# Python側の依存パッケージ（重複するものも含む）
python -m pip install maturin ase numpy

# OpenBLAS対応Pythonバインディングのビルド
RUSTFLAGS="-C target-cpu=native" \
maturin develop --release --locked \
    --manifest-path crates/llmlip-py/Cargo.toml \
    --features parallel,openblas \
    --target-dir target/native-openblas

--target-dirは既存のCLI版と同じディレクトリを指定している。コンパイルで生成されたファイルを再利用できる場合があり、ビルド時間の短縮につながる。

インストール後は次のように確認する。

python - <<'PY'
from llmlip.ase import LlmlipCalculator

print("llmlip Python binding: OK")
PY

正常にインポートできれば、ASEからLlmlipCalculatorを利用できる。

condaとvenvを併用している環境では、VIRTUAL_ENVとCONDA_PREFIXの両方が設定されるとmaturin developがエラーになる場合がある。ビルド時には使用するPython環境を一つに統一する。

モデルのダウンロード

HuggingFaceのFairChemの公式リポジトリからダウンロード可能。UMAには新しいバージョンも公開されているが、今回のllmlip.rs用の変換ツールには uma-s-1p2.pt を選択する必要がある。

ダウンロードしたチェックポイントファイルは次のディレクトリに格納する。

cd ~/llmlip.rs
mkdir -p weights/checkpoints

Hugging Face CLIでダウンロードする場合

ログインしていない場合は認証が必要。事前にfacebook/UMAのページでモデルの利用条件に同意しておく必要がある。
hf auth login

hf download facebook/UMA \
    checkpoints/uma-s-1p2.pt \
    --revision main \
    --local-dir "$HOME/llmlip.rs/weights"

チェックサムの検証

python converters/convert.py verify \
    uma-s-1p2 \
    weights/checkpoints/uma-s-1p2.pt

以下の値であれば良い。

SHA256：ba5c0d912efa22dc238e5fb1b5b7f66ee2e68c48c1b95b7cfd7fe1da5938398b
ファイルサイズ：2,333,393,167 bytes

利用するUMAモデルのバージョンについて

FairChem公式ではuma-s-1p2を指定した場合に、バグ修正版のuma-s-1p2p1が自動的に参照されるようになっているが、後者は現バージョンのスクリプトでは適切に変換できない。llmlip.rs用の変換スクリプト export_uma_state.py では次の定数が直接指定されている。

MODEL_KEY = "uma-s-1p2"
STATE_SCHEMA = "uma-s-1p2-state-dict-v1"
FAIRCHEM_VERSION = "2.20.0"
CHECKPOINT_SHA256 = "ba5c0d912efa22dc238e5fb1b5b7f66ee2e68c48c1b95b7cfd7fe1da5938398b"
CHECKPOINT_SIZE_BYTES = 2_333_393_167

変換開始時にはチェックポイントのファイルサイズとSHA256を検証しており、新モデルを既存のuma-s-1p2として変換しようとしても、チェックポイントの同一性検証で失敗する。Rust側のローダーもモデル識別子、テンソル構成、状態スキーマなどを検証しているため、現状は古いuma-s-1p2を使うのが無難である。

今後の変換ツールとRustランタイムの対応拡張が期待される。

WarpがCUDAドライバを読み込む経路について

WarpのGPU認識が失敗してGPUの利用でエラーが発生する場合がある。これはCUDAライブラリの参照先を環境変数に指定することで適切に読み込ませて解決可能。

NVIDIA Warpは、NVIDIAが提供するGPUアクセラレーション対応のオープンソースPythonフレームワーク。

環境変数とCUDAライブラリの参照先を確認

# 正しいNVIDIAドライバのディレクトリ（末尾のIDは適切に書き換えること）
DRV="/usr/lib/wsl/drivers/nv_dispig.inf_amd64_cc526ccc3287f5cd"

# 一時的なライブラリディレクトリ
TMP="$HOME/.wsl-cuda-warp"

mkdir -p "$TMP"

# 既存の無効なリンクを上書き
ln -sfn "$DRV/libcuda.so.1.1" "$TMP/libcuda.so.1"
ln -sfn "$DRV/libcuda.so.1.1" "$TMP/libcuda.so"

# リンク先を確認
readlink -e "$TMP/libcuda.so.1"

Warpをテスト

LD_LIBRARY_PATH="$TMP:$DRV" \
python - <<'PY'
import warp as wp

wp.init()
print("CUDA devices:", wp.get_cuda_devices())
PY

以下のように出れば正常に認識できている。

Warp 1.18.0 initialized:
   CUDA Toolkit 13.4, Driver 13.3
   Devices:
     "cpu"      : "x86_64"
     "cuda:0"   : "NVIDIA GeForce RTX 5070 Ti" (16 GiB, sm_120, mempool enabled)
     "cuda:1"   : "NVIDIA GeForce RTX 3070 Ti" (8 GiB, sm_86, mempool enabled)
   CUDA peer access:
     Not supported
   Kernel cache:
     /home/hogehoge/.cache/warp/1.18.0
CUDA devices: ['cuda:0', 'cuda:1']

このようにWarp 1.18.0のCUDA Toolkit 13.4とDriver 13.3の組み合わせでも正常に初期化できる。

llmlip.rs用の重みファイル変換

ptファイルを変換する

cd ~/llmlip.rs

python converters/convert.py convert uma-s-1p2 \
    --checkpoint weights/checkpoints/uma-s-1p2.pt \
    --out out/uma-s-1p2.e3rs \
    --accept-license

末尾に以下のような行が表示されれば成功している。

model validation ok
OK: wrote out/uma-s-1p2.e3rs (validated with .../llmlip)

動作テスト

早速動作テストをしてみよう。周期系としてSi 64原子のバルク結晶、分子系として712原子のシクロパラフェニレンの構造最適化をテストする。（いずれもCPU上で評価）

1点計算

まず、64原子のSi結晶を作成する。既存の.venvにASEが導入されているため、次のコマンドが使える。

Si結晶の生成

cd ~/llmlip.rs
mkdir -p inputs results

python - <<'PY'
from ase.build import bulk
from ase.io import write

atoms = bulk("Si", "diamond", a=5.43, cubic=True)
atoms = atoms.repeat((2, 2, 2))

write("inputs/si64.extxyz", atoms)
print("Atoms:", len(atoms))
PY

続いて、変換したUMA-S-1.2でエネルギーとForceを計算する。

エネルギーとForceの計算

RAYON_NUM_THREADS=8 \
./target/native-openblas/release/llmlip predict \
    --model out/uma-s-1p2.e3rs \
    --structure inputs/si64.extxyz \
    --energy --forces --json \
    > results/uma_si64.json

結果の確認

python - <<'PY'
import json

with open("results/uma_si64.json") as f:
    result = json.load(f)

print("Energy:", result["energy"], "eV")
print("Force count:", len(result["forces"]))
print("First force:", result["forces"][0])
PY

結果

Energy: -346.59921145117494 eV
Force count: 64
First force: [3.712308238590367e-16, 2.8275992658421956e-16, 1.5612511283791264e-17]

計算は問題なく終了している。

構造最適化

最適化処理はllmlip.rs／FairChemともASEのLBFGSを使用した。特に、llmlip.rsではエネルギーとForceを同一構造で同時に要求するラッパーで重複評価を抑制した。

項目	設定
Python / NumPy / ASE	3.12.7 / 2.4.6 / 3.29.0
モデル	UMA-S-1.2（`uma-s-1p2`）
FairChem	CPU、`inference_settings="batch"`
llmlip.rs	`direct`、`float64`、OpenBLASリンクを確認済み
Rust側のスレッド数	`RAYON_NUM_THREADS=8`
BLAS / OpenMP	`OPENBLAS_NUM_THREADS=1`、`OMP_NUM_THREADS=1`
タスク	Si結晶：`omat`、分子系：`omol`

周期系

disorderを含むSi結晶を初期構造として、構造最適化を実施してみる。

Si結晶の検証は次のコマンドで実行した。validate_mlips.pyは今回の比較のために作成した検証スクリプトであり、llmlip.rs本体に付属するものではない。

RAYON_NUM_THREADS=8 \
OPENBLAS_NUM_THREADS=1 \
OMP_NUM_THREADS=1 \
LLMLIP_UMA_F64_DIRECT_CONCAT_ROTATE_MIN_EDGES=1000000 \
python validate_mlips.py si \
  --engine both \
  --llmlip-backend direct \
  --model out/uma-s-1p2.e3rs \
  --fairchem-model weights/checkpoints/uma-s-1p2.pt \
  --fairchem-mode batch \
  --device cpu \
  --repeat 2 --sigma 0.10 --seed 20261101 \
  --fmax 0.03 --steps 150 \
  --out results/si_local_ckpt

ダイヤモンド構造Siのユニットセル（8原子、格子定数5.43 Å）を各方向に2倍し、64原子のスーパーセルを構築した。全方向に周期境界条件を設定し、セルを固定して原子位置だけを最適化した。

初期構造には各原子の直交座標に標準偏差0.10 Åの正規乱数変位を加え、全原子の平均変位を差し引いた。生成した構造の最短Si–Si距離は2.0164 Åだった。両エンジンには同じ初期構造を与えた。

乱数生成にはNumPyのdefault_rngとシード20261101を使用しているため、同じ実装・設定で初期構造を再現可能。

LBFGSの最大ステップ変位は0.04 Å、収束判定は最大Forcefmax < 0.03 eV/Å、ステップ数の上限は150とした。

項目	llmlip.rs	FairChem CPU
初期エネルギー（eV）	-333.724339	-333.724348
初期最大Force（eV/Å）	6.43634	6.60262
最終エネルギー（eV）	-346.597295	-346.597944
エネルギー低下量（eV）	12.872956	12.873596
最終最大Force（eV/Å）	0.02830	0.02262
収束ステップ数	28	26
エネルギー・Forceの評価回数	29	27
最適化時間（秒）	52.65	42.51

図1：64原子Si結晶の構造最適化に伴うエネルギー低下。各エンジンの初期エネルギーを基準とした。

図2：Si結晶の最大Forceの推移。破線は収束基準の0.03 eV/Å。最大Forceは各エンジン自身が出力したForceから計算した。

両者はほぼ同じエネルギー低下を示し、最終的な原子位置のRMSDは周期境界を考慮して0.00392 Åだった。理論上の最適化構造を基準とするRMSDも、llmlip.rsが約0.00600 Å、FairChemが約0.00523 Åと近かった。

処理時間は今回の64原子Si系ではllmlip.rsがFairChemより約24%長かった。1ステップごとの計算時間はそれぞれ約1.82秒と約1.57秒となる（ログ出力などのオーバーヘッドを含む）。

同一のSi初期構造に対するエネルギー差はE(llmlip.rs) − E(FairChem) = +0.00000918 eVで、64原子全体としても極めて小さい。一方、全192成分のForceのRMSEは0.13196 eV/Å、最大成分誤差は0.41748 eV/Åだった。

さらに、Forceがエネルギーの負の座標微分と一致するかを確認した。座標成分を±0.01 Å変位させ、中心差分

$$
F_{i\alpha}^{\mathrm{FD}}=-\frac{E(R_{i\alpha}+h)-E(R_{i\alpha}-h)}{2h},\qquad h=0.01\ \mathrm{Å}
$$

から計算したForceを、エンジンが直接出力したForceと比較した。今回は初期構造の3つの座標成分を抽出して検証した。

対象のForce成分（原子番号は0始まり）	llmlip.rsの絶対誤差（eV/Å）	FairChemの絶対誤差（eV/Å）
32番原子・x	0.297992	0.000454
61番原子・x	0.123230	0.000244
4番原子・x	0.066853	0.000035

判定閾値を0.005 eV/Åとすると、FairChemは3成分とも合格し、llmlip.rsは3成分とも不合格だった。少なくとも今回調べた成分について、llmlip.rsの解析的なForceとエネルギーの有限差分には明確な不整合が認められる。

分子系

続いて、過去の記事「WSL-Ubuntu + RTX 5070 Ti で fairchem-core (UMA) をGPUで動かす」でも扱ったC432H280からなる712原子の孤立分子系について検証した。タスクにはomolを使用し、電荷0、スピン多重度1を指定した。

先に結論から述べると、llmlip.rsによる構造最適化は所定の収束基準に到達しなかった。 そこで、両エンジンの計算結果を同一構造で比較するため、FairChemの最適化軌跡から代表構造を抽出した。

以前FairChem CPU版で実行した構造最適化のStep 0、50、100、150、200、250の同一座標をllmlip.rsで再評価し、エネルギーと全原子のForceを比較した（過去のトラジェクトリは276 stepで収束している）。

FairChem軌跡のStep	エネルギー差（eV）	ForceのRMSE（eV/Å）	Forceの最大成分誤差（eV/Å）
0	-0.092494	0.202420	2.843234
50	+0.039552	0.151500	0.709474
100	+0.048411	0.149465	0.737573
150	+0.061821	0.149018	0.810651
200	+0.070669	0.149005	0.779318
250	+0.073438	0.148404	0.768078

エネルギー差は、llmlip.rsからFairChemの値を引いたものである。712原子系全体に対するエネルギー差は0.1 eV以下に収まるが、ForceのRMSEは最終構造でも約0.148 eV/Å程度残っている。

図3：FairChemの同一構造に対するエネルギー差。0、50、100、150、200、250の6構造を採用した。

ForceのRMSEは、初期構造の0.20242 eV/ÅからStep 50で0.15150 eV/Åまで減少する。しかし、それ以降はStep 250に至っても0.14840 eV/Åであり、構造最適化の進行に対してForceの差が大きく改善していない。Step 250のForceのRMSEを元素別に求めると、C原子で約0.1701 eV/Å、H原子で約0.1066 eV/Åとなり、両元素で誤差が確認された。

図4：FairChemの同一構造に対して計算した、llmlip.rsとFairChemの原子間力のRMSE。全712原子の3方向のForce成分から算出した。

同一構造に対する最大のForceの推移も大きく異なった。FairChemではStep 50の0.19499 eV/ÅからStep 250の0.04910 eV/Åへ減少するのに対し、llmlip.rsではStep 50以降、約0.84～0.93 eV/Åの範囲で推移した。

図5：各Stepの同一構造をそれぞれのエンジンで評価した最大Force。縦軸は対数表示で、破線は0.03 eV/Åの収束基準を示す。

分子系では、FairChemの軌跡上でForceが減少しても、llmlip.rsの評価では大きなForceが残っていた。遷移状態探索、振動解析、分子動力学（MD）など、Forceの正確さに依存する計算では、この不一致を解決する必要がある。

llmlip.rsの実行時にはメモリを30 GB近くを使用する場面があった。数百以上の原子数からなる系を計算する場合は、それなりのメモリ容量が必要になると思われる。

今回はFairChem 2.25.0を比較対象としたため、llmlip.rsの参照バージョンである2.20.0との実装差が結果に影響している可能性がある。ただし、llmlip.rs自身のエネルギーと力の有限差分検証でも不一致が確認されており、これはfairchem-coreのバージョン差だけでは説明できない。モデル変換処理やRust側の勾配計算経路に原因がある可能性があり、詳細な検証が必要である。

まとめ

今回のUMA-S-1.2のCPU検証では、llmlip.rsはSi結晶の構造最適化を完了したものの、エネルギーと原子間力の数値的一致には未解決の課題があることが分かった。712原子からなる大きめの分子系では構造緩和後も大きなForceの誤差が残っており、今回ビルドした実行環境では構造最適化を収束させることができなかった。

計算速度については、64原子Si結晶の構造最適化において、llmlip.rsはFairChemのPyTorch CPU版より約24%長い時間を要した。一方、OpenBLASの導入によってllmlip.rs単体の推論速度が約5.1倍に向上するなど、実装やビルド設定によって性能が大きく改善することも確認できた。

llmlip.rsはリリースからまだ日が浅く、対応するMLIPや各モデルの計算精度については、今後も検証と改善が進むものと思われる。CPU環境で幅広いMLIPを実行できる仕組みは興味深く、機会があればMACE系のMLIPについても同様の検証を行ってみたい。

補足情報

計算時間の内訳（OpenBLAS無し）

CPUベンチマークの計算時間を調べた。

計算時間の測定（8並列、openblas無し）

RAYON_NUM_THREADS=8 \
./target/native/release/llmlip benchmark-predict \
    --model out/uma-s-1p2.e3rs \
    --structure inputs/si64.extxyz \
    --energy --forces \
    --warmup 1 \
    --iterations 5 \
    --profile-stages \
    > results/uma_si64_profile.json

UMA-S-1.2でエネルギーと原子間力を1回計算するのに要した時間は、5回の中央値が10.27秒で合った。（モデル読込時間は5.69秒で、1点計算の時間とは分けて計測している）

計算時間の内訳を解析すると、推論時間の約76%がニューラルネットワークの順方向計算、約20%が原子間力の逆伝播計算に費やされていた。

Stage	所要時間(秒)
SO(2)畳み込み・第1段	4.418
Radial MLP	2.081
SO(2)畳み込み・第2段	1.384
SO(2)第1段の逆伝播	0.931
Radial MLPの逆伝播	0.496
SO(2)第2段の逆伝播	0.328

ボトルネックは明確にSO(2)畳み込みとRadial MLPであった。

また、CPU使用率が約10％と低いことも気になった。調べたところ、SO(2)畳み込みとゲート演算を統合した高速化経路（so2_fused_conv1_gate_conv1/2）が利用されていないことが分かった。llmlip.rsの実装によると、UMAについて、FP64の融合演算経路は単一スレッドの場合や、OpenBLAS・MKLなどのバックエンドが有効な場合に選択できる設計となっている。今回のparallelのみのビルドでは、8スレッドを指定することで却って効率の低い演算経路が選択されている可能性がある。

並列効率について

1スレッドと8スレッドの比較を行ったところ、8スレッドでは1スレッドに対して約1.33倍の処理速度に留まっていた。

指標	1スレッド	8スレッド
平均推論時間	13.664秒	10.230秒
中央値	13.635秒	10.269秒
最短時間	13.635秒	10.132秒
モデル読込時間	11.594秒	5.686秒
checksum	−1386.3968458	−1386.3968458

この場合のビルドには、少なくとも次の改善余地が考えられる。

Radial MLPの行列演算をOpenBLASで高速化する。
融合SO(2)演算を複数スレッド実行でも利用する。
小規模な演算の並列化に伴うオーバーヘッドを減らす。

計算時間の内訳（OpenBLAS有り）

前節では、SO(2)畳み込みとRadial MLPが計算時間の大部分を占めていることが分かった。そこで、行列演算ライブラリのOpenBLASを導入してCPU推論の高速化を試みた。（導入手順の部分ではこちらをデフォルトで推奨）

計算時間の測定（8並列、OpenBLASあり）

RAYON_NUM_THREADS=8 \
OPENBLAS_NUM_THREADS=1 \
OMP_NUM_THREADS=1 \
./target/native-openblas/release/llmlip benchmark-predict \
    --model out/uma-s-1p2.e3rs \
    --structure inputs/si64.extxyz \
    --energy --forces \
    --warmup 5 \
    --iterations 10 \
    > results/uma_si64_openblas.json

先ほどと同じ64原子のSi結晶について、エネルギーと原子間力を計算した。今回はウォームアップ5回、計測10回とし、Rayonのスレッド数を1・4・8・16に変更して比較した。OpenBLAS自身のスレッド数は1に固定した。

8スレッドでの推論時間は、中央値で2.029秒となった。OpenBLASなしの約10.27秒と比較して、約5.1倍の高速化である。この改善には、OpenBLASによる行列積の高速化だけでなく、llmlip.rsの実装でSO(2)畳み込みとゲート演算を統合した融合演算経路が選択されることも寄与していると考えられる。

現行のllmlip.rsでは、SO(2)畳み込みとゲート演算を統合した融合演算経路が実装されている。ただし、FP64計算では、原則として単一スレッド実行時、またはOpenBLAS・MKLなどの対応バックエンドが有効で、融合演算用の重み行列が用意されている場合に選択される。OpenBLASなし・8スレッド実行では融合演算経路が無効だったが、OpenBLASの導入によって利用可能となった。したがって、今回の高速化には行列積の効率化と演算経路の切り替えがともに寄与したと考えられる。

次に、--profile-stagesを追加し、各演算段階の計算時間を詳しく調べた。

演算段階別のプロファイル取得

RAYON_NUM_THREADS=8 \
OPENBLAS_NUM_THREADS=1 \
OMP_NUM_THREADS=1 \
./target/native-openblas/release/llmlip benchmark-predict \
    --model out/uma-s-1p2.e3rs \
    --structure inputs/si64.extxyz \
    --energy --forces \
    --warmup 5 \
    --iterations 10 \
    --profile-stages \
    > results/uma_si64_openblas_profile.json

プロファイルから求めた主要演算の所要時間は以下のとおりである。各数値はウォームアップを含む15回分の合計を15で割った平均値で、4つのネットワークブロックの処理を合算している。

Stage	所要時間（秒）
SO(2)融合演算・順方向	0.762
Radial MLP・順方向	0.190
SO(2)融合演算・逆伝播	0.795
Radial MLP・逆伝播	0.167

SO(2)融合演算は順方向と逆伝播を合わせて約1.56秒を占め、プロファイル付き推論時間の約69%に相当した。Radial MLPも約16%を占めている。OpenBLASによって大幅に高速化されたものの、SO(2)融合演算が依然として最大のボトルネックである。特に、原子間力を計算するための逆伝播の部分で演算の分割方法やBLASの呼び出し方法を改良すれば、さらに高速化できると考えられる。

並列数とLIBXSMMの比較

さらに、小規模行列積に特化したLIBXSMMをOpenBLASと併用した場合についても比較した。

スレッド数	OpenBLAS	OpenBLAS＋LIBXSMM
1	2.838秒	2.979秒
4	2.084秒	2.205秒
8	2.029秒	2.061秒
16	2.099秒	2.119秒

いずれもウォームアップ5回、計測10回の中央値であり、詳細プロファイルの計測処理は含まない。

今回の条件では、OpenBLAS単独・8スレッドが最速だった。4スレッドと8スレッドの差は小さく、16スレッドでは逆に処理時間が増加した。また、LIBXSMMの併用による高速化は確認できなかった。8スレッドではOpenBLAS単独版より約1.6%遅かったが、差が小さいため、測定環境の変動も考慮する必要がある。今回の64原子Si結晶に対しては、OpenBLAS＋Rayon 8スレッドを採用するのが妥当なようである。

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up