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

Open
  • Is something wrong?
#5,953 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 3 days

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
45/100
Issue type
Bug
Clarity
Mostly clear
Activity status
Quiet
Tech stack
python, pytorch
Domain
compilers, machine-learning, testing

Research direction

Start in deepmd/dpmodel/utils/env_mat.py and trace EnvMat.call into the pt_expt DPA2 lower and call_lower path used by torch.compile. Reproduce the atype_ext-dependent descriptor case and compare compiled forces with the eager/dense path. Done means the lower behavior is characterized and covered by a regression test, or the atype_ext-independence constraint is documented with an assertion.

Written by the indexing model from the issue text.

Description

bug

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

Dominant language
Python
Stars
2.1k
Forks
663
Avg merge
3d 11h
Merged PRs (30d)
25

Getting set up

Open in Codespaces

Starts the project's dev container in your browser, under your own GitHub account.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from deepmodeling/deepmd-kit

All issues in deepmodeling/deepmd-kit

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.