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 6d67a12). 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
Summary
Making a DPA2/DPA3 or shared descriptor path depend on
atype_ext(e.g. by maskingem/diff/swwithatype_ext >= 0insideEnvMat.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 6d67a12). The naive fix of re-masking the zeroed virtual-center outputs using
atype_extmade the descriptor readatype_ext, and the pt_expt DPA2 compiled (torch.compile/ graph-loweredcall_lower) path returned incorrect forces; forces were correct in the eager/dense path. This is why #5833 chose to neutralize thedavg/dstdnormalization instead of re-masking, keepingEnvMat.callindependent ofatype_ext.Impact / risk
atype_extdependency to a descriptor for a reasonable reason could silently reintroduce wrong forces when a compiled DPA2 model happens to use that descriptor.Suggested follow-ups
atype_ext-dependent branch breaks) and add a regression test on the compiled lower.EnvMat.callmust remainatype_ext-independent.Related: #5833