deepmodeling / deepmodeling/deepmd-kit

bug(pt-expt): compiled DPA2 lower miscompiles descriptors that depend on atype_ext

Open
#5,953 0 comments 0 reactions 0 assignees View on GitHub
bug
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.