deepmodeling / deepmodeling/deepmd-kit
bug(pt-expt): compiled DPA2 lower miscompiles descriptors that depend on atype_ext
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
## Summary
Making a DPA2/DPA3 or shared descriptor path depend on `atype_ext` (e.g. by masking `em`/`diff`/`sw` with `atype_ext >= 0` inside `EnvMat.call`) produces **wrong forces** when the model is run through the **compiled pt_expt DPA2 lower**.
## Reproduction context
This was hit while developing deepmd/dpmodel/utils/env_mat.py for #5833 (commit 6d67a1257). The naive fix of re-masking the zeroed virtual-center outputs using `atype_ext` made the descriptor read `atype_ext`, and the pt_expt DPA2 compiled (`torch.compile` / graph-lowered `call_lower`) path returned incorrect forces; forces were correct in the eager/dense path. This is why #5833 chose to neutralize the `davg`/`dstd` normalization instead of re-masking, keeping `EnvMat.call` independent of `atype_ext`.
## Impact / risk
- A future developer who adds an `atype_ext` dependency to a descriptor for a reasonable reason could silently reintroduce wrong forces when a compiled DPA2 model happens to use that descriptor.
- Nothing in CI currently pins this behavior; there is no regression test or assertion.
## Suggested follow-ups
- Characterize the pt_expt DPA2 lower miscompilation (which transformation of the `atype_ext`-dependent branch breaks) and add a regression test on the compiled lower.
- At minimum, add a comment/assertion documenting that `EnvMat.call` must remain `atype_ext`-independent.
Related: #5833
Contributor guide
Assessment
This issue has not been assessed yet.