Skip to content

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

Open

Description

@njzjz-bot

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions