Skip to content

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

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

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