From b39a8a1397c399fdf3d050f7f397139c468ab523 Mon Sep 17 00:00:00 2001 From: Owen Price Skelly <21372141+OwenPriceSkelly@users.noreply.github.com> Date: Tue, 29 Sep 2026 11:36:47 -0500 Subject: [PATCH] orbmol configs: name the arch for the :custom smoke-test leg The weights= leg borrows the first fetched family checkpoint (declaration order: orbmol-v1-conservative) and reloads its weights through setup_from_path, whose default arch is orbmol-v2. The v2 energy head has 257 input columns against v1's 256, so the leg failed on Frontier with 'size mismatch for heads.energy.mlp.NN-0.weight' (2026-08-24 run). A VERIFY_KWARGS entry keyed by the :custom id selects the matching loader for the leg only; users still pass arch= for their own weights. Co-Authored-By: Claude Fable 5.1 --- sample_model_configurations/amd_configs/orbmol.py | 10 ++++++++++ sample_model_configurations/nvidia_configs/orbmol.py | 10 ++++++++++ 2 files changed, 20 insertions(+) diff --git a/sample_model_configurations/amd_configs/orbmol.py b/sample_model_configurations/amd_configs/orbmol.py index d1b9223..8a24ff3 100644 --- a/sample_model_configurations/amd_configs/orbmol.py +++ b/sample_model_configurations/amd_configs/orbmol.py @@ -57,6 +57,16 @@ "orbmol:custom": None, } +# The smoke-test's weights= leg borrows the first fetched family checkpoint +# (CHECKPOINTS order → orbmol-v1-conservative) and reloads its weights through +# setup_from_path, whose default arch is orbmol-v2. The v2 energy head has 257 +# input columns (CoulombModule adds one) against v1's 256, so the default arch +# fails with a state_dict size mismatch (Frontier, 2026-08-24). Name the +# matching architecture for the leg; your own weights still take arch= per call. +VERIFY_KWARGS = { + "orbmol:custom": {"arch": "orbmol-v1-conservative"}, +} + def _default_weights_url(load_fn) -> str: """The upstream URL baked into the loader's ``weights_path`` default.""" diff --git a/sample_model_configurations/nvidia_configs/orbmol.py b/sample_model_configurations/nvidia_configs/orbmol.py index 67847e1..4c57bee 100644 --- a/sample_model_configurations/nvidia_configs/orbmol.py +++ b/sample_model_configurations/nvidia_configs/orbmol.py @@ -49,6 +49,16 @@ "orbmol:custom": None, } +# The smoke-test's weights= leg borrows the first fetched family checkpoint +# (CHECKPOINTS order → orbmol-v1-conservative) and reloads its weights through +# setup_from_path, whose default arch is orbmol-v2. The v2 energy head has 257 +# input columns (CoulombModule adds one) against v1's 256, so the default arch +# fails with a state_dict size mismatch (Frontier, 2026-08-24). Name the +# matching architecture for the leg; your own weights still take arch= per call. +VERIFY_KWARGS = { + "orbmol:custom": {"arch": "orbmol-v1-conservative"}, +} + def _default_weights_url(load_fn) -> str: """The upstream URL baked into the loader's ``weights_path`` default."""