deepmodeling / deepmodeling/deepmd-gnn
[Code scan] MACE-OFF loader fails official checkpoints with symmetric contraction shape mismatches
- Dominant language
- Python
- Stars
- 56
- Forks
- 10
- Avg merge
- 14d 19h
- Merged PRs (30d)
- 14
Description
This issue was found by a Codex global scan of the repository at commit 19f9265efe8825a4383e5a13ce128361533b5ac6.
The conservative MACE-OFF loader reconstructs a DeePMD-GNN `MaceModel` from inferred checkpoint attributes and then tries to load the original checkpoint state dict into the reconstructed submodel:
https://github.com/deepmodeling/deepmd-gnn/blob/19f9265efe8825a4383e5a13ce128361533b5ac6/deepmd_gnn/mace_off.py#L249-L267
https://github.com/deepmodeling/deepmd-gnn/blob/19f9265efe8825a4383e5a13ce128361533b5ac6/deepmd_gnn/mace_off.py#L358-L385
https://github.com/deepmodeling/deepmd-gnn/blob/19f9265efe8825a4383e5a13ce128361533b5ac6/deepmd_gnn/mace_network.py#L125-L148
On this checkout, `python -m pytest -q` fails the MACE-OFF tests for `off23_small` with shape mismatches such as:
```text
size mismatch for products.0.symmetric_contractions.contractions.0.weights_max: copying a param with shape torch.Size([10, 23, 96]) from checkpoint, the shape in current model is torch.Size([10, 8, 96])
size mismatch for products.0.symmetric_contractions.contractions.0.U_matrix_3: copying a param with shape torch.Size([16, 16, 16, 23]) from checkpoint, the shape in current model is torch.Size([16, 16, 16, 8])
```
This means selected official MACE-OFF checkpoints can be downloaded but not loaded or converted through the advertised helper path.
Suggested fix: infer and preserve the checkpoint-side MACE options that control symmetric contraction basis sizes, or avoid constructing an intermediate scripted `ScaleShiftMACE` with incompatible defaults before assigning the trusted original checkpoint model. Add a focused regression test for `off23_small` load/convert.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start in deepmd_gnn/mace_off.py at the checkpoint reconstruction and loading paths, then inspect deepmd_gnn/mace_network.py for the MACE options used by the intermediate model. Run the MACE-OFF tests for off23_small and trace the checkpoint-side contraction sizes. Done means the official off23_small checkpoint loads and converts successfully with a focused regression test.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 50/100