deepmodeling / deepmodeling/deepmd-gnn

[Code scan] Respect model precision without leaking torch default dtype

Open
#140 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
56
Forks
10
Avg merge
15d 8h
Merged PRs (30d)
12

Description

This issue was found by a Codex global scan of the repository at commit 19f9265efe8825a4383e5a13ce128361533b5ac6.

MACE accepts `precision` through config, but direct construction drops it through `**kwargs` and builds the MACE network from the ambient `torch.get_default_dtype()`:

https://github.com/deepmodeling/deepmd-gnn/blob/19f9265efe8825a4383e5a13ce128361533b5ac6/deepmd_gnn/mace.py#L270-L291
https://github.com/deepmodeling/deepmd-gnn/blob/19f9265efe8825a4383e5a13ce128361533b5ac6/deepmd_gnn/mace.py#L296-L315
https://github.com/deepmodeling/deepmd-gnn/blob/19f9265efe8825a4383e5a13ce128361533b5ac6/deepmd_gnn/mace.py#L338-L359

The `get_model()` config path sets the process-global default dtype and never restores it:

https://github.com/deepmodeling/deepmd-gnn/blob/19f9265efe8825a4383e5a13ce128361533b5ac6/deepmd_gnn/mace.py#L1823-L1831

I verified both effects locally:

```text
MaceModel(..., precision="float64") with global float32 builds atomic_energies as torch.float32
MaceModel.get_model(..., precision="float64") leaves torch.get_default_dtype() == torch.float64
```

NequIP has the complementary failure mode: it passes `model_dtype`, but does not temporarily raise PyTorch's default dtype before `model_from_config()`, so `precision="float64"` can fail when the caller's default dtype is still float32.

https://github.com/deepmodeling/deepmd-gnn/blob/19f9265efe8825a4383e5a13ce128361533b5ac6/deepmd_gnn/nequip.py#L258-L299

Local reproduction:

```text
ValueError: Overall default_dtype=float32, but model_dtype=float64 is a higher precision- change default_dtype to float64
```

Suggested fix: validate `precision` explicitly in both constructors, build the underlying MACE/NequIP networks under a temporary dtype context, store the effective precision in serialized params, and always restore the previous global default dtype.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.