コンテンツにスキップ

ローカルLLMはAllegro/NeQUIPの学習設定YAMLを作れるのか：Nemotronでmetadata-only preflightを検証

これは性能ベンチマークではありません。A100x4上のローカルLLMで、Allegro / NeQUIPの学習設定YAMLをどこまで安全に作り、実行前に確認できるかを調べた機能検証です。

結論から言うと、NemotronはAllegro / NeQUIPのtraining YAML候補を生成できました。既存SIF内で torch / nequip / allegro のimport、A100x4のCUDA可視性、YAML parse、NeQUIP reference configとの比較、Hydra config表示まで確認できています。

ただし、trainingはまだしていません。dataset読込、model構築、checkpoint/model load、torch.load、inference、LAMMPS連携も未実施です。ここを明確に分けることが、今回のいちばん重要な点です。

何が通って、何が未確認か

項目 状態
SIF内A100x4可視性 確認済み
torch / nequip / allegro import 確認済み
NemotronによるYAML候補生成 成功
YAML静的チェック PASS
SIF内 yaml.safe_load PASS
SIF内NeQUIP minimal configとの比較 実施
isolated config.yaml で nequip-train --help exit code 0
nequip-train --cfg job / --cfg all exit code 0
--resolve 系 OmegaConf型エラー
dataset読込 未実施
model load / torch.load 未実施
training 未実施
inference / evaluate / deploy 未実施
LAMMPS連携 未実施

なぜ第1弾QE/LAMMPSと違う進め方にしたか

第1弾では、Quantum ESPRESSOとLAMMPSでローカルNemotronによる入力生成、実行、意図的エラー、ログからの修正、再実行成功まで確認しました。小さなQE入力やLAMMPS入力は、短いcanary runで実行検証できます。

Allegro / NeQUIPは事情が違います。training YAMLを実行すると、dataset読込、データ統計、model構築、checkpoint処理、GPU-heavyなtrainingへ進む可能性があります。つまり、「YAMLがそれらしく見える」だけでは実行してよいとは言えません。

そのため今回は、runtimeへ進む前のmetadata-only preflightとして設計しました。

検証環境

検証はA100 80GB x4のローカル環境で行いました。コンテナ実行にはApptainer 1.4.5を使い、既存のAllegro / NeQUIP import-readiness SIFを利用しました。SIFの詳細パスやraw proofpackは社内証跡に保管しています。

項目 結果
GPU NVIDIA A100 80GB x4
Container runtime Apptainer 1.4.5
torch 2.6.0+cu124
nequip 0.17.1
allegro 0.8.2
numpy 2.2.2
ase 3.28.0
e3nn 0.6.0
lmdb 2.2.0
h5py 未導入
torch.cuda.is_available True
torch.cuda.device_count 4

Phase 7: SIF / import readiness

まず、既存SIFをmetadata-onlyで確認しました。apptainer exec --nv によりcontainer内からA100 4枚が見え、torch, nequip, allegro のimportが通りました。

CLIについては、nequip-train が存在しました。一方、このSIFでは nequip-evaluate, nequip-deploy, allegro-train, allegro-deploy, allegro-evaluate は見つかりませんでした。

この時点では、trainingに入る可能性があるCLI操作は実行していません。

Phase 8: NemotronによるYAML候補生成

ローカルNemotronに、Allegro / NeQUIP styleのtraining YAML候補を生成させました。要求した条件は次の通りです。

  • dataset pathはplaceholder: ./data/example.extxyz
  • chemical speciesはSi
  • h5/hdf5を避ける
  • checkpoint/model pathを書かない
  • command execution文を書かない
  • 出力はYAMLのみ

生成されたYAMLは、静的チェックとSIF内 yaml.safe_load を通りました。これは「YAMLとして読める」ことの確認であり、training可能性の確認ではありません。

Phase 8B / 9: reference比較とYAML補正

次に、既存runbookやassetからreference config候補を探し、Nemotron生成YAMLと比較しました。この段階で、実行用configとしては怪しい可能性があるキーが見つかりました。

  • hidden_neurals
  • envelope_cutoff
  • num_spherical
  • num_irreps
  • num_neurals

そのため、元のYAMLをそのまま実行へ進めず、reference-aligned YAMLとして作り直しました。ここでも実行はせず、静的チェックと yaml.safe_load に限定しています。

Phase 10-13: Hydra / nequip-train metadata確認

nequip-train --help は、configなしでは普通のhelp表示になりませんでした。Hydraの MissingConfigException が出て、primary config config が見つからないという結果でした。

SIF内のPython metadataを確認すると、nequip-train のentrypointは nequip.scripts.train:main で、Hydraはcurrent working directoryから config を探す構造でした。

@hydra.main(version_base=None, config_path=os.getcwd(), config_name="config")

HydraはPythonアプリケーションの設定を実行前にcomposeする仕組みです。そのため、helpやconfig表示でも、current working directoryのconfig contextに依存します。

そこで、SIFに含まれるNeQUIPの最小テスト用YAMLを抽出して比較しました。

  • minimal_aspirin.yaml
  • minimal_toy_emt.yaml

これらを参考に、data, trainer, training_module, run, _target_, loss, val_metrics などを持つSIF-reference-aligned YAMLを作成しました。

このYAMLをisolated workdirの config.yaml として配置すると、metadata-onlyの範囲で次が通りました。

Command Result
nequip-train --help exit code 0
nequip-train --cfg job exit code 0
nequip-train --cfg all exit code 0
nequip-train --cfg job --resolve exit code 1
nequip-train --cfg all --resolve exit code 1

--resolve 系はOmegaConfの型エラーでした。

omegaconf.errors.UnsupportedValueType: Value 'dict' is not a supported primitive type

これは悪い知らせでもありますが、良い検出でもあります。datasetやtrainingへ入る前に、config resolver上の未解決課題を発見できたからです。

今回の失敗と教訓

今回の教訓は、QE/LAMMPSより慎重な段階分けが必要だということです。

  • nequip-train --help でもconfig contextがないと失敗する
  • yaml.safe_load はtraining安全性を保証しない
  • --cfg が通ってもdataset読込やmodel構築が安全とは限らない
  • --resolve でOmegaConf型エラーが出るならruntimeへ進む前に止めるべき
  • .pth, .pt, .nequip.pth はmetadata確認に留め、loadしない
  • torch.load やpickle展開は明示的に禁止して進めるべき

training smokeへ進む条件

次にtraining smokeへ進む場合、少なくとも次を決める必要があります。

条件 内容
Dataset 極小、自己作成または明確に利用可能なextxyz
Config trusted referenceに寄せたYAML
Command exact commandを事前承認
Timeout 短いtimeoutを設定
Monitoring nvidia-smi とprocess監視
Output outputs配下へ隔離
Model load 原則禁止、必要なら別承認
Stop condition dataset path逸脱、model load、GPU負荷異常、timeoutなど

ライセンスとprovenance

NemotronはNVIDIAのmodel cardおよびlicense条件を確認して利用する必要があります。NeQUIP、Allegro、PyTorch、Apptainerにもそれぞれのlicenseがあります。本記事では検証手順と結果を説明するだけで、モデル、dataset、SIF、第三者ソフトウェアを再配布しません。

将来training smokeへ進む場合、datasetとmodel artifactのprovenance/license確認が必須です。

まとめ

Allegro / NeQUIPについては、metadata-only preflightとしては成功です。ローカルNemotronはYAML候補を生成でき、静的チェック、reference比較、Hydra config表示まで進めました。

ただし、これはtraining成功ではありません。QE/LAMMPSでは小さな入力をすぐ実行できましたが、ML interatomic potentialのtraining configでは、schema、dataset、model、GPU負荷、license/provenanceの安全ゲートが先に必要です。

次は、この記事を第2弾として整理したうえで、別フェーズとして極小dataset + 厳格監視付きtraining smokeを設計するのが妥当です。H200 NVLやRTX PRO 6000 Blackwellでの比較、GPU Start Consoleへのworkflow card化も、その後の展開候補です。