[BUG] State_dict error when packaging/compiling fine-tuned models

[kavanase]

When trying to nequip-compile or nequip-package a checkpoint from a fine-tuning run (which used ModelFromPackage, and where that package file came from an earlier nequip version), I am getting this error message:

(pytorch_2.7) Perlmutter: > nequip-package build ckpts/nequip_MPA/*-0.ckpt test.nequip.zip
[2025-11-16 10:19:41,024][nequip.utils.versions.package_versions][INFO] - [rank: 0] Version Information:
[2025-11-16 10:19:41,025][nequip.utils.versions.package_versions][INFO] - [rank: 0] torch 2.7.0+cu128
[2025-11-16 10:19:41,025][nequip.utils.versions.package_versions][INFO] - [rank: 0] e3nn 0.5.6
[2025-11-16 10:19:41,025][nequip.utils.versions.package_versions][INFO] - [rank: 0] nequip 0.15.0
[2025-11-16 10:19:41,025][nequip.utils.versions.package_versions][INFO] - [rank: 0] talaria 0.1.0
[2025-11-16 10:19:41,025][nequip.utils.versions.package_versions][INFO] - [rank: 0] allegro 0.7.1
[2025-11-16 10:19:41,047][nequip.scripts.package][INFO] - [rank: 0] Building `eager` model for packaging ...
[2025-11-16 10:19:41,047][nequip.model.saved_models.load_utils][INFO] - [rank: 0] Loading model from ckpts/nequip_MPA/weighted_metric_0.0533-0.ckpt ...
[2025-11-16 10:19:41,047][nequip.model.saved_models.checkpoint][INFO] - [rank: 0] Loading model from checkpoint file: ckpts/nequip_MPA/weighted_metric_0.0533-0.ckpt ...
[2025-11-16 10:19:42,976][nequip.model.saved_models.package][INFO] - [rank: 0] Loading model from package file: XL_rtol_0.1_netF_l4_32_WM_0.0690-11.nequip.zip ...
Traceback (most recent call last):
  File "/global/homes/k/kavanase/miniconda3/envs/pytorch_2.7/bin/nequip-package", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/global/u2/k/kavanase/Packages/nequip/nequip/scripts/package.py", line 208, in main
    eager_model, data = load_saved_model(
                        ^^^^^^^^^^^^^^^^^
  File "/global/u2/k/kavanase/Packages/nequip/nequip/model/saved_models/load_utils.py", line 126, in load_saved_model
    model = ModelFromCheckpoint(
            ^^^^^^^^^^^^^^^^^^^^
  File "/global/u2/k/kavanase/Packages/nequip/nequip/model/saved_models/checkpoint.py", line 78, in ModelFromCheckpoint
    lightning_module = training_module.load_from_checkpoint(checkpoint_path)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/global/homes/k/kavanase/miniconda3/envs/pytorch_2.7/lib/python3.12/site-packages/lightning/pytorch/utilities/model_helpers.py", line 125, in wrapper
    return self.method(cls, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/global/homes/k/kavanase/miniconda3/envs/pytorch_2.7/lib/python3.12/site-packages/lightning/pytorch/core/module.py", line 1581, in load_from_checkpoint
    loaded = _load_from_checkpoint(
             ^^^^^^^^^^^^^^^^^^^^^^
  File "/global/homes/k/kavanase/miniconda3/envs/pytorch_2.7/lib/python3.12/site-packages/lightning/pytorch/core/saving.py", line 91, in _load_from_checkpoint
    model = _load_state(cls, checkpoint, strict=strict, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/global/homes/k/kavanase/miniconda3/envs/pytorch_2.7/lib/python3.12/site-packages/lightning/pytorch/core/saving.py", line 187, in _load_state
    keys = obj.load_state_dict(checkpoint["state_dict"], strict=strict)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/global/homes/k/kavanase/miniconda3/envs/pytorch_2.7/lib/python3.12/site-packages/torch/nn/modules/module.py", line 2593, in load_state_dict
    raise RuntimeError(
RuntimeError: Error(s) in loading state_dict for EMALightningModule:
	Missing key(s) in state_dict: "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant0", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant1", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant2", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant3", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant4", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant5", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant6", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant7", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant8", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant9", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant10", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant11", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant12", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant0", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant1", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant2", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant3", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant4", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant5", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant6", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant7", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant8", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant9", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant10", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant11", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant12", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant0", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant1", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant2", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant3", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant4", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant5", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant6", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant7", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant8", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant9", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant10", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant11", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant12", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant0", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant1", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant2", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant3", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant4", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant5", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant6", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant7", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant8", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant9", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant10", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant11", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant12".

As mentioned, the original model package file is from an earlier nequip version (still 0.15.0, but before some recent commits which changed model structure). It trained fine with ModelFromPackage, but compilation/packaging of the output checkpoints from this fine-tuning run are failing here.

Environment:

• OS: Linux
• See output above for package versions

[cw-tan]

I believe the state dict entries refers to the buffers from the e3nn TP codegen: https://github.com/e3nn/e3nn/blob/d7661462f07773b78e9dd1141520d41ab33aed1c/e3nn/o3/_tensor_product/_tensor_product.py#L400

Why the error is manifesting as such now, I am unsure. This might take a while.

[kavanase]

I ran into this same issue again with another run. I'm surprised it happens when attempting to package/compile, when training is fine, but I guess it's related to some of the loading/querying that occurs when packaging/compiling, where it tries to load something from the original model package?

[kavanase]

Just to confirm as a sanity check; I checked that directly compiling the packaged model (from the earlier nequip version) works fine. It's just compiling/packaging the checkpoint from a training run with the packaged model that causes this

[kavanase]

Noting that a user (@savya10) has reported this same issue to me, from fine-tuning the Nequip.net packaged FM and then trying to compile the checkpoint

[cw-tan]

Does this still occur with nequip v0.16.1? I'm wondering if b80c50b changed anything.

[kavanase]

Yep I just checked and the issue still happens with nequip v0.16.1 (when using this version to attempt the packaging/compilation, of the checkpoint and package file from 0.15.0. Do you mean if I use nequip v0.16.1 for fine-tuning training from the package file, and use it for the packaging/compilation?

[cw-tan]

I would like to check the case when the original (compiled) pre-training and packaging (and everything that happens subsequently, i.e. fine-tuning workflow) happened with nequip v0.16.1. I'll also check independently to figure out if different training configurations affect it.

[kavanase]

Ah ok! For that I'll need to put on some new big training runs (if we want this for the FMs), ofc can test with some smaller datasets if needed.

Is there a manual fix we could do, with our FMs, that fixes their forward compatibility?

[kavanase]

Just to note, I have confirmed that this issue still happens when taking a model package from 0.15.0 (e.g. our foundation models), then using 0.16.1 to fine-tune (with ModelFromPackage) and then trying to compile/package.

I will try testing the other case of everything with 0.16.1 – though this should be perfectly fine, as it was for 0.15.0 etc? (It's just the standard fine-tuning workflow, with everything the same version the whole way through – I guess the question will be more for later version/model updates, if we can use a previous version package file for fine-tuning without these issues... which I guess so, as I don't think we had this issue before, so hopefully just something specific to the changes from 0.15.0 to 0.16.0 🙏 ).

I think it would really save a lot of headaches (and GPU time) if we could somehow manually edit the 0.15.0 FM package files somehow, to make them forward-compatible with >=0.16.0.

[kavanase]

Ok, surprisingly (to me – but at least this maybe makes this an easier (while more severe) problem to solve), taking a model trained on 0.16.1, packaged, then fine-tuned with ModelFromPackage, and then trying to package/compile a ckpt, I get the same error – where all these steps were performed with the same 0.16.1 version.

Output:

(pytorch_2.8.0) FasRC: FromPackage_FineTune_then_Compile_Package_Test > nequip-package build outputs/2025-12-24/09-26-28/*3.ckpt test.nequip.zip
[2025-12-24 09:33:50,760][nequip.utils.versions.package_versions][INFO] - [rank: 0] Version Information:
[2025-12-24 09:33:50,760][nequip.utils.versions.package_versions][INFO] - [rank: 0] torch 2.8.0+cu129
[2025-12-24 09:33:50,760][nequip.utils.versions.package_versions][INFO] - [rank: 0] e3nn 0.5.9
[2025-12-24 09:33:50,760][nequip.utils.versions.package_versions][INFO] - [rank: 0] nequip 0.16.1
[2025-12-24 09:33:50,760][nequip.utils.versions.package_versions][INFO] - [rank: 0] allegro 0.7.1
[2025-12-24 09:33:50,760][nequip.utils.versions.package_versions][INFO] - [rank: 0] talaria 0.1.0
[2025-12-24 09:33:50,907][nequip.scripts.package][INFO] - [rank: 0] Building `eager` model for packaging ...
[2025-12-24 09:33:50,907][nequip.model.saved_models.load_utils][INFO] - [rank: 0] Loading model from outputs/2025-12-24/09-26-28/weighted_metric_0.2126-3.ckpt ...
[2025-12-24 09:33:50,908][nequip.model.saved_models.checkpoint][INFO] - [rank: 0] Loading model from checkpoint file: outputs/2025-12-24/09-26-28/weighted_metric_0.2126-3.ckpt ...
[2025-12-24 09:33:51,306][nequip.model.saved_models.package][INFO] - [rank: 0] Loading model from package file: NequIP_OMol_4M_No_ZBL_TE_ChgSpin_16_WM_0.0782-143.nequip.zip ...
Traceback (most recent call last):
  File "/n/home03/skavanagh/miniconda3/envs/pytorch_2.8.0/bin/nequip-package", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/n/home03/skavanagh/miniconda3/envs/pytorch_2.8.0/lib/python3.12/site-packages/nequip/scripts/package.py", line 208, in main
    eager_model, data = load_saved_model(
                        ^^^^^^^^^^^^^^^^^
  File "/n/home03/skavanagh/miniconda3/envs/pytorch_2.8.0/lib/python3.12/site-packages/nequip/model/saved_models/load_utils.py", line 126, in load_saved_model
    model = ModelFromCheckpoint(
            ^^^^^^^^^^^^^^^^^^^^
  File "/n/home03/skavanagh/miniconda3/envs/pytorch_2.8.0/lib/python3.12/site-packages/nequip/model/saved_models/checkpoint.py", line 78, in ModelFromCheckpoint
    lightning_module = training_module.load_from_checkpoint(checkpoint_path)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/n/home03/skavanagh/miniconda3/envs/pytorch_2.8.0/lib/python3.12/site-packages/lightning/pytorch/utilities/model_helpers.py", line 125, in wrapper
    return self.method(cls, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/n/home03/skavanagh/miniconda3/envs/pytorch_2.8.0/lib/python3.12/site-packages/lightning/pytorch/core/module.py", line 1611, in load_from_checkpoint
    loaded = _load_from_checkpoint(
             ^^^^^^^^^^^^^^^^^^^^^^
  File "/n/home03/skavanagh/miniconda3/envs/pytorch_2.8.0/lib/python3.12/site-packages/lightning/pytorch/core/saving.py", line 91, in _load_from_checkpoint
    model = _load_state(cls, checkpoint, strict=strict, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/n/home03/skavanagh/miniconda3/envs/pytorch_2.8.0/lib/python3.12/site-packages/lightning/pytorch/core/saving.py", line 187, in _load_state
    keys = obj.load_state_dict(checkpoint["state_dict"], strict=strict)  # type: ignore[arg-type]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/n/home03/skavanagh/miniconda3/envs/pytorch_2.8.0/lib/python3.12/site-packages/torch/nn/modules/module.py", line 2624, in load_state_dict
    raise RuntimeError(
RuntimeError: Error(s) in loading state_dict for EMALightningModule:
	Missing key(s) in state_dict: "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant0", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant1", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant2", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant3", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant4", "model.sole_model.model.func.layer1_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant5", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant0", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant1", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant2", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant3", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant4", "model.sole_model.model.func.layer2_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant5", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant0", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant1", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant2", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant3", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant4", "model.sole_model.model.func.layer3_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant5", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant0", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant1", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant2", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant3", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant4", "model.sole_model.model.func.layer4_convnet.conv.tp_scatter.tp._compiled_main_left_right._tensor_constant5".

So it seems it's actually not a version mismatch issue between packages/checkpoints, but an issue specific to trying to package/compile a checkpoint from a ModelFromPackage run (while packaging/compiling checkpoints from ModelFromCheckpoint runs, which essentially requires the same version to be used throughout, works fine).