付録4: 並列実行のベンチマーク

ブックキーピングと並列実行ベンチマーク#

LightPFPのExample modelとして提供しているモデル examplesv1alni3s を使用して、複数のジョブをjoblibライブラリを用いて並列に実行した際のパフォーマンスを比較しました。
本章で使用したベンチマークのコードの概要は以下の通りです。

import ase
from ase.build import bulk
from light_pfp_client.estimator import Estimator
from light_pfp_client.ase_calculator import ASECalculator

from ase import units
from ase.md.langevin import Langevin

from joblib import Parallel, delayed
import time


def run_md(atoms: ase.Atoms, book_keeping: bool):
    estimator = Estimator(model_id="examplesv1alni3s", book_keeping=book_keeping)
    calculator = ASECalculator(estimator)
    atoms.calc = calculator
    dyn = Langevin(atoms, 5 * units.fs, temperature_K=300.0, friction=0.002)
    dyn.run(20)


# Plot the retuned value varying n_atoms_per, n_jobs and book_keeping
# n_atoms_per defines the number of atoms in the system
# n_jobs defines the number of parallel jobs to run
def benchmark_parallel_md(n_atoms_per: int, n_job: int, book_keeping: bool):
    atoms_list = [bulk("Al") * (n_atoms_per, n_atoms_per, n_atoms_per) for _ in range(n_job)]
    start = time.time()

    Parallel(n_jobs=n_job, backend="threading")(
        delayed(run_md)(atoms, book_keeping) for atoms in atoms_list
    )

    end = time.time()
    elapsed = end - start
    return elapsed

まず一つ目のベンチマークでは、並列実行するジョブの数を1から10まで変化させ、全てのジョブが完了するまでにかかる時間を比較しました。
（なお、並列実行のセクションで述べたように現在LightPFPでは10個までの推論に必要な情報のみをサーバーに保存することができることに注意してください。）
また、各ジョブで使用する構造の原子数も変化させ、その影響を調査しました。

最初の図はブックキーピングを使用しない場合の結果を示し、次の図はブックキーピングを使用した場合の結果を示しています。
横軸がジョブで使用する構造中の原子数(100原子から200,000原子程度まで)、縦軸が全てのジョブの完了までに要する時間であり、異なるジョブの数の結果については色で区別しています。
なお、縦軸の縮尺は統一されています。
最初の画像がブックキーピングを使用しない場合の結果を示し、次の画像がブックキーピングを使用した場合の結果を示しています。

この図から、並列実行の数が10以下の場合には、ブックキーピングを使用することで常にパフォーマンスが向上している様子が見て取れます。
また、ブックキーピングを使用した場合も使用していない場合も、ジョブの数に対してほぼ線形に全てのジョブが完了するまでに要している時間が増加していることがわかります。

この結果から、並列実行するジョブの数が10以下であれば、ブックキーピングを使用することでパフォーマンスが向上することが期待できることがわかります。

次に同様のベンチマークを、ジョブの数を8から15まで変化させ、全てのジョブが完了するまでにかかる時間を比較しました。
縦軸、横軸ともに先ほどのベンチマークと同様であり、縦軸の縮尺も統一されています。

この図から、並列実行の数が10を超える場合、ブックキーピングを使用することでパフォーマンスが低下していることが見られます。
これはブックキーピングのセクションで述べたように、10以上のジョブを並列で実行すると、隣接情報がうまく再利用されないが隣接情報の再作成・削除のオーバーヘッドが発生し、パフォーマンスの低下が生じるためです。

また、ブックキーピングを使用しない場合には、10以上のジョブを並列実行してもパフォーマンスの劣化は見られないことも分かります。
これはブックキーピングを使用しない場合は、推論に必要なモデルなどの情報のロードや削除にかかるオーバーヘッドが比較的少ないからだと考えられます。
しかし、ブックキーピングを有効活用した方が各ジョブが高速に完了することが期待されるため、10個以上のジョブを実行する際は、10個以下のジョブを並列でブックキーピングで実行し、
それらが完了してから次の10個のジョブを実行するなどの方法で実行することが最もパフォーマンスが高いことが期待されます。

最後に、並列実行およびブックキーピングの計算可能な原子数に関する影響を調査しました。
このベンチマークでは、ジョブの数を1から15まで、そして原子数を500,000程度から1,000,000程度まで変化させ、各ジョブ数ごとにどのサイズの構造までシミュレーションできるのかについて計測しました。
横軸と縦軸の意味は先ほどのベンチマークと同様であり、ブックキーピングを有効にした際の結果が示されています。
グラフが途中で途切れているのは、メモリが不足して計算が中断されたことを示しています。

この図から、並列で実行されているジョブの数が増えると計算可能な原子数が減少することがわかります。
これはブックキーピングのセクションで述べたように、ブックキーピングを使用すると隣接情報の保存のためにメモリが占有され、推論に割り当て可能なメモリのサイズが小さくなるためです。