This issue comes from a Codex global scan of deepmodeling/deepmd-kit at commit 73de44b1f94471b2e3bdb6b11f57b34d7bc791bb.
Problem
The experimental PyTorch compiled-training wrapper builds neighbor lists with distinguish_types=False unconditionally:
|
def forward( |
|
self, |
|
coord: torch.Tensor, |
|
atype: torch.Tensor, |
|
box: torch.Tensor | None = None, |
|
fparam: torch.Tensor | None = None, |
|
aparam: torch.Tensor | None = None, |
|
do_atomic_virial: bool = False, |
|
charge_spin: torch.Tensor | None = None, |
|
) -> dict[str, torch.Tensor]: |
|
from deepmd.dpmodel.utils.nlist import ( |
|
build_neighbor_list, |
|
extend_coord_with_ghosts, |
|
) |
|
from deepmd.dpmodel.utils.region import ( |
|
normalize_coord, |
|
) |
|
|
|
nframes, nloc = atype.shape[:2] |
|
rcut = self.original_model.get_rcut() |
|
sel = self.original_model.get_sel() |
|
|
|
# coord extension + nlist (data-dependent, run in eager) |
|
coord_3d = coord.detach().reshape(nframes, nloc, 3) |
|
box_flat = box.detach().reshape(nframes, 9) if box is not None else None |
|
|
|
if box_flat is not None: |
|
coord_norm = normalize_coord(coord_3d, box_flat.reshape(nframes, 3, 3)) |
|
else: |
|
coord_norm = coord_3d |
|
|
|
ext_coord, ext_atype, mapping = extend_coord_with_ghosts( |
|
coord_norm, atype, box_flat, rcut |
|
) |
|
nlist = build_neighbor_list( |
|
ext_coord, |
|
ext_atype, |
|
nloc, |
|
rcut, |
|
sel, |
|
distinguish_types=False, |
|
) |
For non-mixed-type descriptors, sel is per atom type. A combined nearest-sum(sel) list can discard a farther but still required neighbor of a specific type before the lower model later formats/splits the list by type.
The inference path handles this with distinguish_types=not mixed_types:
|
if cells is not None: |
|
box_input = cells.reshape(nframes, 3, 3) |
|
coord_normalized = normalize_coord(coords, box_input) |
|
else: |
|
coord_normalized = coords |
|
|
|
extended_coord, extended_atype, mapping = extend_coord_with_ghosts( |
|
coord_normalized, |
|
atom_types, |
|
cells, |
|
rcut, |
|
) |
|
nlist = build_neighbor_list( |
|
extended_coord, |
|
extended_atype, |
|
natoms, |
|
rcut, |
|
sel, |
|
distinguish_types=not mixed_types, |
|
) |
The export/sample-input path uses the same distinction:
|
nlist = build_neighbor_list( |
|
extended_coord, |
|
extended_atype, |
|
nloc, |
|
rcut, |
|
sel, |
|
distinguish_types=not mixed_types, |
|
) |
Impact
torch.compile training can train on different local environments than eager training or inference for ordinary multi-type, non-mixed descriptors. This can silently change losses and gradients instead of only changing execution mode.
Suggested fix
Derive mixed_types from the wrapped model and call build_neighbor_list(..., distinguish_types=not mixed_types) in _CompiledModel.forward. Add a regression test where sel is type-specific and a same-type neighbor lies farther than the combined nearest-neighbor cutoff would otherwise keep.
This issue comes from a Codex global scan of
deepmodeling/deepmd-kitat commit73de44b1f94471b2e3bdb6b11f57b34d7bc791bb.Problem
The experimental PyTorch compiled-training wrapper builds neighbor lists with
distinguish_types=Falseunconditionally:deepmd-kit/deepmd/pt_expt/train/training.py
Lines 562 to 603 in 73de44b
For non-mixed-type descriptors,
selis per atom type. A combined nearest-sum(sel)list can discard a farther but still required neighbor of a specific type before the lower model later formats/splits the list by type.The inference path handles this with
distinguish_types=not mixed_types:deepmd-kit/deepmd/pt_expt/infer/deep_eval.py
Lines 923 to 942 in 73de44b
The export/sample-input path uses the same distinction:
deepmd-kit/deepmd/pt_expt/utils/serialization.py
Lines 254 to 261 in 73de44b
Impact
torch.compiletraining can train on different local environments than eager training or inference for ordinary multi-type, non-mixed descriptors. This can silently change losses and gradients instead of only changing execution mode.Suggested fix
Derive
mixed_typesfrom the wrapped model and callbuild_neighbor_list(..., distinguish_types=not mixed_types)in_CompiledModel.forward. Add a regression test whereselis type-specific and a same-type neighbor lies farther than the combined nearest-neighbor cutoff would otherwise keep.