diff --git a/deepmd/dpmodel/utils/default_neighbor_list.py b/deepmd/dpmodel/utils/default_neighbor_list.py index 03b4289795..3628664c5a 100644 --- a/deepmd/dpmodel/utils/default_neighbor_list.py +++ b/deepmd/dpmodel/utils/default_neighbor_list.py @@ -6,6 +6,9 @@ from deepmd.dpmodel.array_api import ( Array, ) +from deepmd.dpmodel.utils.neighbor_list import ( + EdgeNeighborList, +) from .neighbor_list import ( NeighborList, @@ -33,7 +36,12 @@ def build( box: Array | None, rcut: float, sel: list[int], - ) -> tuple[Array, Array, Array, Array]: + return_mode: str = "extended", + ) -> tuple[Array, Array, Array, Array] | EdgeNeighborList: + if return_mode != "extended": + raise NotImplementedError( + "DefaultNeighborList only supports the extended-coordinate contract." + ) xp = array_api_compat.array_namespace(coord, atype) nframes, nloc = atype.shape[:2] if box is not None: diff --git a/deepmd/dpmodel/utils/neighbor_list.py b/deepmd/dpmodel/utils/neighbor_list.py index 49504cc7c3..e63f5099dd 100644 --- a/deepmd/dpmodel/utils/neighbor_list.py +++ b/deepmd/dpmodel/utils/neighbor_list.py @@ -10,11 +10,47 @@ neighbor list was built. """ +from dataclasses import ( + dataclass, +) +from typing import ( + Literal, +) + from deepmd.dpmodel.array_api import ( Array, ) +@dataclass +class EdgeNeighborList: + """Edge-vector neighbor-list contract. + + The model consumes geometry only through ``edge_vec``. Builders that own + periodic-image shifts compute those shifts before constructing this object; + callers that receive already-shifted ghost coordinates use zero-shift edge + vectors computed from the provided coordinates. + """ + + coord: Array + """Coordinates of the scatter domain with shape ``(nf, nscatter, 3)``.""" + + atype: Array + """Local owner atom types with shape ``(nf, nloc)``.""" + + edge_index: Array + """Message-passing edge indices with shape ``(2, nedge)`` in owner space.""" + + edge_vec: Array + """Per-edge displacement vectors with shape ``(nedge, 3)``.""" + + edge_scatter_index: Array + """Force/virial scatter indices with shape ``(2, nedge)`` in scatter space.""" + + edge_mask: Array + """Boolean edge-validity mask with shape ``(nedge,)``.""" + + class NeighborList: """Strategy that builds the extended neighbor environment from local atoms. @@ -32,7 +68,8 @@ def build( box: Array | None, rcut: float, sel: list[int], - ) -> tuple[Array, Array, Array, Array]: + return_mode: Literal["extended", "edges"] = "extended", + ) -> tuple[Array, Array, Array, Array] | EdgeNeighborList: """Build the extended system and a candidate neighbor list. Parameters @@ -47,6 +84,10 @@ def build( cutoff radius. sel number of selected neighbors per type. + return_mode + ``"extended"`` returns the historical extended-coordinate quartet. + ``"edges"`` returns :class:`EdgeNeighborList`, where ``edge_vec`` is + the only geometric displacement consumed by the model. Returns ------- diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index d28d7789f7..c97c4a07ad 100644 --- a/deepmd/pt/entrypoints/freeze_pt2.py +++ b/deepmd/pt/entrypoints/freeze_pt2.py @@ -58,6 +58,9 @@ from deepmd.pt.utils.env import ( DEVICE, ) +from deepmd.pt_expt.utils.edge_schema import ( + edge_schema_from_extended, +) from deepmd.utils.model_branch_dict import ( get_model_dict, ) @@ -97,13 +100,13 @@ def _model_has_message_passing(model: torch.nn.Module) -> bool: def _strip_shape_assertions(graph_module: torch.nn.Module) -> None: - """Remove deferred shape assertions from spin export graphs. + """Remove deferred shape assertions from SeZM export graphs. - The spin lower path slices tensors using both ``nall`` and ``nloc`` after - virtual atom expansion. ``torch.export`` may turn valid dynamic cases into - deferred ``Ne(nall, nloc)`` assertions, even though the graph works for both - NoPBC and ghost-atom inputs. The generic pt_expt spin exporter applies the - same cleanup. + SeZM lower inputs intentionally keep extended-atom and local-atom axes + independent: regular exports pass ghost coordinates through ``coord`` while + ``atype`` remains local-only, and spin exports slice both ``nall`` and + ``nloc`` after virtual atom expansion. ``torch.export`` may turn these valid + dynamic cases into deferred ``Ne(nall, nloc)`` assertions. """ graph = graph_module.graph for node in list(graph.nodes): @@ -257,18 +260,20 @@ def _collect_metadata( "intensive": vdef.intensive, } ) + exports_atomic_virial = True if not is_spin else bool(do_atomic_virial) metadata = { "type_map": list(model.get_type_map()), "ntypes": _get_model_ntypes(model), "rcut": float(model.get_rcut()), "sel": [int(s) for s in model.get_sel()], + "lower_input_kind": "nlist" if is_spin else "edge_vec", "dim_fparam": int(model.get_dim_fparam()), "dim_aparam": int(model.get_dim_aparam()), "dim_chg_spin": int(model.get_dim_chg_spin()), "mixed_types": bool(model.mixed_types()), "has_message_passing": _model_has_message_passing(model), "has_comm_artifact": False, - "do_atomic_virial": bool(do_atomic_virial), + "do_atomic_virial": exports_atomic_virial, "nnei": int(sum(model.get_sel())), "has_default_fparam": bool(model.has_default_fparam()), "default_fparam": _to_py_list(model.get_default_fparam()), @@ -377,7 +382,24 @@ def _make_sample_inputs( aparam, charge_spin, ) - return ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin + formatted_nlist: torch.Tensor = model.format_nlist(ext_coord, ext_atype, nlist_t) + edge_schema = edge_schema_from_extended( + ext_coord, + ext_atype[:, :nloc], + formatted_nlist, + mapping_t, + ) + return ( + edge_schema.coord, + edge_schema.atype, + edge_schema.edge_index, + edge_schema.edge_vec, + edge_schema.edge_scatter_index, + edge_schema.edge_mask, + fparam, + aparam, + charge_spin, + ) def _resolve_nframes( @@ -423,29 +445,20 @@ def _resolve_nframes( def _build_dynamic_shapes( sample_inputs: tuple[torch.Tensor | None, ...], ) -> tuple: - """Positional ``dynamic_shapes`` for the traced - ``(ext_coord, ext_atype, nlist, mapping, fparam, aparam)`` signature. - """ + """Build positional dynamic-shape constraints for the traced lower input.""" nframes_dim = torch.export.Dim("nframes", min=1) has_spin = ( len(sample_inputs) >= 7 and sample_inputs[2] is not None and sample_inputs[2].is_floating_point() ) - has_charge_spin = (has_spin and len(sample_inputs) == 8) or ( - not has_spin and len(sample_inputs) == 7 - ) - # Spin export currently generates a valid lower-bound guard from its - # virtual-atom split/concat pattern. Matching the bound keeps export strict, - # while `_strip_shape_assertions` removes the spurious deferred guards later. nall_dim = torch.export.Dim("nall", min=4 if has_spin else 1) nloc_dim = torch.export.Dim("nloc", min=1) - fparam = sample_inputs[5] if has_spin else sample_inputs[4] - aparam = sample_inputs[6] if has_spin else sample_inputs[5] - charge_spin = None - if has_charge_spin: - charge_spin = sample_inputs[7] if has_spin else sample_inputs[6] + nedge_dim = torch.export.Dim("nedge", min=2) if has_spin: + fparam = sample_inputs[5] + aparam = sample_inputs[6] + charge_spin = sample_inputs[7] if len(sample_inputs) == 8 else None shapes = ( {0: nframes_dim, 1: nall_dim}, # extended_coord {0: nframes_dim, 1: nall_dim}, # extended_atype @@ -455,18 +468,23 @@ def _build_dynamic_shapes( {0: nframes_dim} if fparam is not None else None, {0: nframes_dim, 1: nloc_dim} if aparam is not None else None, ) - if has_charge_spin: + if len(sample_inputs) == 8: shapes = (*shapes, {0: nframes_dim} if charge_spin is not None else None) return shapes + fparam = sample_inputs[6] + aparam = sample_inputs[7] + charge_spin = sample_inputs[8] if len(sample_inputs) == 9 else None shapes = ( {0: nframes_dim, 1: nall_dim}, # extended_coord: (nframes, nall, 3) - {0: nframes_dim, 1: nall_dim}, # extended_atype: (nframes, nall) - {0: nframes_dim, 1: nloc_dim}, # nlist: (nframes, nloc, nnei) - {0: nframes_dim, 1: nall_dim}, # mapping: (nframes, nall) + {0: nframes_dim, 1: nloc_dim}, # atype + {1: nedge_dim}, # edge_index + {0: nedge_dim}, # edge_vec + {1: nedge_dim}, # edge_scatter_index + {0: nedge_dim}, # edge_mask {0: nframes_dim} if fparam is not None else None, {0: nframes_dim, 1: nloc_dim} if aparam is not None else None, ) - if has_charge_spin: + if len(sample_inputs) == 9: shapes = (*shapes, {0: nframes_dim} if charge_spin is not None else None) return shapes @@ -561,27 +579,29 @@ def freeze_sezm_to_pt2( fparam=fparam, aparam=aparam, charge_spin=charge_spin, - do_atomic_virial=atomic_virial, ) else: ( - ext_coord, - ext_atype, - nlist_t, - mapping_t, + coord, + atype, + edge_index, + edge_vec, + edge_scatter_index, + edge_mask, fparam, aparam, charge_spin, ) = sample_inputs_cpu traced = model.forward_common_lower_exportable( - ext_coord, - ext_atype, - nlist_t, - mapping_t, + coord, + atype, + edge_index, + edge_vec, + edge_scatter_index, + edge_mask, fparam=fparam, aparam=aparam, charge_spin=charge_spin, - do_atomic_virial=atomic_virial, ) # Output key order is taken from a concrete run; Python dict order @@ -598,8 +618,7 @@ def freeze_sezm_to_pt2( strict=False, prefer_deferred_runtime_asserts_over_guards=True, ) - if is_spin: - _strip_shape_assertions(exported.graph_module) + _strip_shape_assertions(exported.graph_module) # move_to_device_pass handles FakeTensor device propagation cleanly; # a naive .to(device) on the exported program does not. diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 7d54c7ef01..e3f195ac67 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -253,7 +253,11 @@ def __init__( self._has_spin = getattr(self.dp.model["Default"], "has_spin", False) if callable(self._has_spin): self._has_spin = self._has_spin() - self._has_hessian = self.model_def_script.get("hessian_mode", False) + selected_model_params = getattr(self, "input_param", self.model_def_script) + self._has_hessian = selected_model_params.get("hessian_mode", False) + self._uses_edge_schema = ( + _is_sezm_model_params(selected_model_params) and not self._has_spin + ) self._setup_nlist_backend(nlist_backend) def _setup_nlist_backend(self, nlist_backend: str) -> None: @@ -743,34 +747,59 @@ def _eval_lower_strategy( ) -> dict[str, torch.Tensor]: """Evaluate via the selected O(N) ``NeighborList`` strategy. - Builds the extended representation with ``self._nlist_builder`` (vesin or - nv), runs the model's ``forward_common_lower``, and maps the extended - outputs back to local atoms with ``communicate_extended_output``. + Uses the selected O(N) builder (vesin or nv). Models that declare the + edge-vector contract consume it directly; other energy models keep the + historical extended-coordinate contract and fold extended outputs back + to local atoms. Returns a dict keyed by backend names, matching the normal ``model()`` - output so the caller's extraction is unchanged. ``requires_grad`` is set - on the extended coordinates internally, exactly as on the native path, so - forces/virials are produced identically. + output so the caller's extraction is unchanged. """ inner = self.dp.model["Default"] - ext_coord, ext_atype, nlist, mapping = self._nlist_builder.build( - coord, atype, box, self.rcut, list(inner.get_sel()) - ) - model_lower = inner.forward_common_lower( - ext_coord, - ext_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - charge_spin=charge_spin, - ) - predict = communicate_extended_output( - model_lower, - inner.model_output_def(), - mapping, - do_atomic_virial=do_atomic_virial, - ) + if self._uses_edge_schema: + edge_schema = self._nlist_builder.build( + coord, + atype, + box, + self.rcut, + list(inner.get_sel()), + return_mode="edges", + ) + predict = inner.forward_common_lower( + edge_schema.coord, + edge_schema.atype, + edge_schema.edge_index, + edge_schema.edge_vec, + edge_schema.edge_scatter_index, + edge_schema.edge_mask, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + input_prec=coord.dtype, + ) + else: + ext_coord, ext_atype, nlist, mapping = self._nlist_builder.build( + coord, + atype, + box, + self.rcut, + list(inner.get_sel()), + ) + model_lower = inner.forward_common_lower( + ext_coord, + ext_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, + ) + predict = communicate_extended_output( + model_lower, + self.output_def, + mapping, + do_atomic_virial=do_atomic_virial, + ) return { backend: predict[internal] for internal, backend in self._OUTDEF_DP2BACKEND.items() diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index a6cb1f538d..a369fd028a 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -101,13 +101,13 @@ ``core_compute`` so that make_fx sees a pure tensor-in / tensor-out function: -* ``core_compute`` rebuilds a compact, GPU-friendly edge list from the - padded DeePMD neighbor list (``build_edge_list_from_nlist``), with - two masked dummy edges appended so the edge tensor has a non-singular - symbolic lower bound (NOTE 10). Edge vectors are gathered from the - extended coordinate tensor and then **detached into a fresh leaf** - (``edge_vec.detach().requires_grad_(True)``), so the gather lives - entirely outside the autograd region (NOTE 11). +* ``core_compute`` consumes the unified edge-vector schema prepared by its + caller. Two masked dummy edges are appended by the schema builders so the + edge tensor has a non-singular symbolic lower bound (NOTE 10). + Edge vectors are then **detached into a fresh leaf** + (``edge_vec.detach().requires_grad_(True)``), so neighbor construction, + shift application and coordinate gathers live outside the autograd region + (NOTE 11). * The SeZM descriptor and the analytical ZBL term (``InterPotential``) both consume that edge-vector leaf, so the energy depends on coordinates *only* through ``edge_vec``. @@ -236,8 +236,9 @@ One Inductor option set governs both backends: ``torch.compile`` takes it as ``options=`` for training, and the eval path (NOTE 13) applies it via -``torch._inductor.config.patch`` around ``compile_fx_inner``, adding -``triton.max_tiles=1``. The options are: +``torch._inductor.config.patch`` around ``compile_fx_inner``. The set +includes ``triton.max_tiles=1``, so the 1D launch-grid constraint applies to +both graphs. The options are: * ``max_autotune=False`` Autotune regresses on dynamic shapes because each recompile rolls @@ -350,8 +351,8 @@ NOTE 10 -- Tail dummy edges --------------------------- -``build_edge_list_from_nlist`` appends two masked edges at the end of -every batch. Real edge compaction happens via +The edge-schema builders append two masked edges at the end of every batch. +Real edge compaction happens via ``torch.nonzero(valid_mask)``, whose output length is data-dependent and can be zero in sparse or single-atom systems (e.g. isolated-atom reference frames in training data). make_fx cannot trace an @@ -467,6 +468,7 @@ from jaxtyping import Float, Int from torch import Tensor + from deepmd.dpmodel.utils.neighbor_list import EdgeNeighborList, NeighborList from deepmd.pt.model.atomic_model.sezm_atomic_model import ( SeZMAtomicModel, @@ -484,7 +486,6 @@ BaseModel, ) from deepmd.pt.model.model.transform_output import ( - communicate_extended_output, edge_energy_deriv, ) from deepmd.pt.utils import ( @@ -503,13 +504,17 @@ strip_saved_tensor_detach, trace_pad_dim, ) -from deepmd.pt.utils.nlist import ( - extend_input_and_build_neighbor_list, -) from deepmd.pt.utils.nv_nlist import ( NvNeighborList, is_nv_available, ) +from deepmd.pt_expt.utils.edge_schema import ( + edge_schema_from_extended, +) +from deepmd.pt_expt.utils.vesin_neighbor_list import ( + VesinNeighborList, + is_vesin_torch_available, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -518,11 +523,46 @@ SeZMModel_ = make_model(SeZMAtomicModel) -# Local-atom counts above which Toolkit-Ops replaces the dense all-pairs builder. -# Non-periodic systems switch when dense all-pairs transients become memory-heavy, -# even though the dense path remains slightly faster at medium sizes. -SEZM_NV_NLIST_THRESHOLD = 1024 -SEZM_NV_NONPERIODIC_NLIST_THRESHOLD = 2048 + +def _select_neighbor_builder(nf: int, device: torch.device) -> NeighborList: + """Select the O(N) neighbor builder for the given batch shape and device. + + A single CPU frame uses ``vesin``: its lightweight cell list has no + per-frame batching to amortize and avoids the ``nvalchemiops`` (Warp) + launch overhead. Every other case -- any CUDA input or any multi-frame + batch -- uses ``nvalchemiops``, whose batched kernel amortizes the launch + cost across frames. Both builders keep every neighbor within ``rcut``, so + the resulting edge set never depends on ``sel``. ``vesin`` also serves as + the fallback when ``nvalchemiops`` is absent. + + Parameters + ---------- + nf : int + Number of frames in the batch. + device : torch.device + Device of the input coordinates. + + Returns + ------- + NeighborList + A builder instance whose ``build`` method produces the neighbor list. + + Raises + ------ + RuntimeError + If neither ``nvalchemiops`` nor ``vesin`` is importable. + """ + if device.type == "cpu" and nf == 1 and is_vesin_torch_available(): + return VesinNeighborList() + if is_nv_available(): + return NvNeighborList() + if is_vesin_torch_available(): + return VesinNeighborList() + raise RuntimeError( + "SeZM neighbor-list construction requires either 'nvalchemiops' or " + "'vesin', but neither is importable." + ) + # Apply the process-global PyTorch workarounds the compile pipeline relies on # (autotune log suppression, DDP optimiser, and the 2.12 divisibility repair) @@ -715,7 +755,7 @@ def __init__( f"DP_TF32_INFER must be one of 0/1/2, got {tf32_infer_env!r}" ) self._tf32_infer_precision = _TF32_INFER_PRECISION_CHOICES[tf32_infer_env] - if self.use_compile or self._env_use_compile_infer is True: + if self._env_use_compile_infer is True: check_compile_torch_version() # === Bridging (optional short-range zone bridging) === @@ -898,31 +938,192 @@ def forward_common( if cc.ndim == 2: cc = cc.view(nf, nloc, 3) - # === Step 2. Build neighbor list === + # === Step 2. Build geometry schema === with nvtx_range("SeZM/build_neighbor_list"): - # extended_coord: (nf, nall, 3), extended_atype: (nf, nall) - # mapping: (nf, nall), nlist: (nf, nloc, nsel) - extended_coord, extended_atype, mapping, nlist = ( - self.build_neighbor_list(cc, atype, bb) - ) + if self.get_active_mode() == "dens": + # extended_coord: (nf, nall, 3), extended_atype: (nf, nall) + # nlist: (nf, nloc, nsel), mapping: (nf, nall) + extended_coord, extended_atype, nlist, mapping = ( + self.build_extended_neighbor_list(cc, atype, bb) + ) + else: + edge_schema = self.build_neighbor_list(cc, atype, bb) - # === Step 3. Run the shared extended-input path === - return self.forward_common_after_nlist( - extended_coord, - extended_atype, - mapping, - nlist, - atype, - fp, - ap, - input_prec, - do_atomic_virial=do_atomic_virial, - force_input=force_input, - noise_mask=noise_mask, + # === Step 3. Run the model compute path === + if self.get_active_mode() == "dens": + return self.forward_common_lower_dens( + extended_coord, + extended_atype, + mapping, + nlist, + atype, + fp, + ap, + input_prec, + force_input=force_input, + noise_mask=noise_mask, + charge_spin=charge_spin, + ) + return self.forward_common_lower( + edge_schema.coord, + edge_schema.atype, + edge_schema.edge_index, + edge_schema.edge_vec, + edge_schema.edge_scatter_index, + edge_schema.edge_mask, + fparam=fp, + aparam=ap, charge_spin=charge_spin, + input_prec=input_prec, ) - def forward_common_after_nlist( + def forward_common_lower( + self, + coord: torch.Tensor, + atype: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_scatter_index: torch.Tensor, + edge_mask: torch.Tensor, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + extended_coord_corr: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + input_prec: torch.dtype | None = None, + use_compile: bool | None = None, + ) -> dict[str, torch.Tensor]: + """ + Run the conservative SeZM lower interface on explicit edge vectors. + + ``edge_vec`` is the only coordinate-dependent geometry entering the + descriptor. ``edge_scatter_index`` defines the force/virial scatter + domain, which may be local atoms for Python inference or local-plus-ghost + slots for LAMMPS. + """ + del comm_dict + coord, _, fp, ap, inferred_input_prec = self._input_type_cast( + coord, + fparam=fparam, + aparam=aparam, + ) + if input_prec is None: + input_prec = inferred_input_prec + if coord.ndim == 2: + coord = coord.reshape(atype.shape[0], -1, 3) + atype = atype.to(device=coord.device, dtype=torch.long) + edge_index = edge_index.to(device=coord.device, dtype=torch.long) + edge_vec = edge_vec.to(device=coord.device, dtype=coord.dtype) + edge_scatter_index = edge_scatter_index.to( + device=coord.device, dtype=torch.long + ) + edge_mask = edge_mask.to(device=coord.device, dtype=torch.bool) + if extended_coord_corr is not None: + extended_coord_corr = extended_coord_corr.to( + device=coord.device, dtype=coord.dtype + ) + if extended_coord_corr.ndim == 2: + extended_coord_corr = extended_coord_corr.reshape(atype.shape[0], -1, 3) + nf = atype.shape[0] + should_compile = ( + self.should_use_compile() if use_compile is None else use_compile + ) + charge_spin = self.convert_charge_spin( + charge_spin, + nf=nf, + dtype=coord.dtype, + device=coord.device, + ) + with self.tf32_precision_ctx(): + if should_compile: + fp, ap = self.convert_fp_ap( + fp, + ap, + nf=nf, + nloc=atype.shape[1], + dtype=coord.dtype, + device=coord.device, + ) + has_coord_corr = extended_coord_corr is not None + cache_key = (bool(self.training), has_coord_corr) + if cache_key not in self.compiled_core_compute_cache: + self.trace_and_compile( + coord, + atype, + edge_index, + edge_vec, + edge_scatter_index, + edge_mask, + fp, + ap, + charge_spin, + extended_coord_corr=extended_coord_corr, + ) + compiled_core_compute = self.compiled_core_compute_cache[cache_key] + task_buf_vals = get_task_buffer_values( + self, + self._task_buf_order_cache[cache_key], + ) + grad_ctx: Any = nullcontext() if self.training else torch.no_grad() + with nvtx_range("SeZM/core_compute"), grad_ctx: + if extended_coord_corr is None: + model_predict = compiled_core_compute( + coord, + atype, + edge_index, + edge_vec, + edge_scatter_index, + edge_mask, + fp, + ap, + charge_spin, + *task_buf_vals, + ) + else: + model_predict = compiled_core_compute( + coord, + atype, + edge_index, + edge_vec, + edge_scatter_index, + edge_mask, + fp, + ap, + charge_spin, + extended_coord_corr, + *task_buf_vals, + ) + if ( + self._core_compute_pending_compile_t0 is not None + and self._core_compute_pending_compile_key == cache_key + ): + if torch.cuda.is_available(): + torch.cuda.synchronize() + log.info( + "SeZM: finished compiling (mode=%s, coord_corr=%s) in %.2fs", + "train" if self.training else "eval", + has_coord_corr, + time.perf_counter() - self._core_compute_pending_compile_t0, + ) + self._core_compute_pending_compile_t0 = None + self._core_compute_pending_compile_key = None + else: + with nvtx_range("SeZM/core_compute"): + model_predict = self.core_compute( + coord, + atype, + edge_index, + edge_vec, + edge_scatter_index, + edge_mask, + fparam=fp, + aparam=ap, + charge_spin=charge_spin, + extended_coord_corr=extended_coord_corr, + ) + return self._output_type_cast(model_predict, input_prec) + + def forward_common_lower_dens( self, extended_coord: torch.Tensor, extended_atype: torch.Tensor, @@ -933,14 +1134,12 @@ def forward_common_after_nlist( ap: torch.Tensor | None, input_prec: torch.dtype, *, - do_atomic_virial: bool = False, force_input: torch.Tensor | None = None, noise_mask: torch.Tensor | None = None, - extended_coord_corr: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """ - Run SeZM from already-built extended inputs. + Run the denoising/direct-force path from already-built extended inputs. Parameters ---------- @@ -960,16 +1159,12 @@ def forward_common_after_nlist( Cast atomic parameters with shape (nf, nloc, nda), or None. input_prec Original input precision used for output casting. - do_atomic_virial - Whether to compute per-atom virial. force_input Optional atom-wise force input for the ``dens`` path with shape (nf, nloc, 3). noise_mask Optional atom-wise corruption mask for the ``dens`` path with shape (nf, nloc). - extended_coord_corr - Coordinate correction for virial with shape (nf, nall, 3), or None. charge_spin Frame-level charge and spin conditions with shape `(nf, 2)`. @@ -985,171 +1180,68 @@ def forward_common_after_nlist( dtype=extended_coord.dtype, device=extended_coord.device, ) - active_mode = self.get_active_mode() - if active_mode == "dens": - # === Step 1. `dens` path (no coordinate gradients needed) === - extended_coord = extended_coord.detach() - force_input, noise_mask = self.canonicalize_dens_inputs( - force_input, - noise_mask, - nf=nf, - nloc=nloc, - dtype=extended_coord.dtype, - device=extended_coord.device, - ) + # === Step 1. Denoising path (no coordinate gradients needed) === + extended_coord = extended_coord.detach() + force_input, noise_mask = self.canonicalize_dens_inputs( + force_input, + noise_mask, + nf=nf, + nloc=nloc, + dtype=extended_coord.dtype, + device=extended_coord.device, + ) - with self.tf32_precision_ctx(): - if self.should_use_compile(): - fp, ap = self.convert_fp_ap( - fp, - ap, - nf=nf, - nloc=nloc, - dtype=extended_coord.dtype, - device=extended_coord.device, - ) - if self.compiled_dens_compute is None or not self._dens_compiled: - self.compile_dens() - with nvtx_range("SeZM/core_compute_dens"): - compute_ret = self.compiled_dens_compute( - extended_coord, - extended_atype, - nlist, - mapping, - force_input=force_input, - noise_mask=noise_mask, - fparam=fp, - aparam=ap, - charge_spin=charge_spin, - ) - if self._dens_pending_compile_t0 is not None: - if torch.cuda.is_available(): - torch.cuda.synchronize() - log.info( - "SeZM: finished compiling dens path in %.2fs", - time.perf_counter() - self._dens_pending_compile_t0, - ) - self._dens_pending_compile_t0 = None - else: - with nvtx_range("SeZM/core_compute_dens"): - compute_ret = self.core_compute_dens( - extended_coord, - extended_atype, - nlist, - mapping, - force_input=force_input, - noise_mask=noise_mask, - fparam=fp, - aparam=ap, - charge_spin=charge_spin, - ) - with nvtx_range("SeZM/post_process"): - model_predict = self.post_process_output_dens( - compute_ret, - atype, - noise_mask=noise_mask, + with self.tf32_precision_ctx(): + if self.should_use_compile(): + fp, ap = self.convert_fp_ap( + fp, + ap, + nf=nf, + nloc=nloc, + dtype=extended_coord.dtype, + device=extended_coord.device, ) - else: - # === Step 1. `ener` path (edges built inside core_compute) === - with self.tf32_precision_ctx(): - if self.should_use_compile(): - fp, ap = self.convert_fp_ap( - fp, - ap, - nf=nf, - nloc=nloc, - dtype=extended_coord.dtype, - device=extended_coord.device, + if self.compiled_dens_compute is None or not self._dens_compiled: + self.compile_dens() + with nvtx_range("SeZM/core_compute_dens"): + compute_ret = self.compiled_dens_compute( + extended_coord, + extended_atype, + nlist, + mapping, + force_input=force_input, + noise_mask=noise_mask, + fparam=fp, + aparam=ap, + charge_spin=charge_spin, ) - has_coord_corr = extended_coord_corr is not None - cache_key = (bool(self.training), has_coord_corr) - if cache_key not in self.compiled_core_compute_cache: - self.trace_and_compile( - extended_coord, - extended_atype, - nlist, - mapping, - fp, - ap, - charge_spin, - extended_coord_corr=extended_coord_corr, - ) - compiled_core_compute = self.compiled_core_compute_cache[cache_key] - # Read current values of per-task buffers (optimizer steps - # update them in-place; out-of-place replacements from - # model_change_out_bias are captured because we read fresh - # each call rather than caching the values at compile time). - _task_buf_vals = get_task_buffer_values( - self, - self._task_buf_order_cache[cache_key], + if self._dens_pending_compile_t0 is not None: + if torch.cuda.is_available(): + torch.cuda.synchronize() + log.info( + "SeZM: finished compiling dens path in %.2fs", + time.perf_counter() - self._dens_pending_compile_t0, ) - # NOTE: Inference needs no autograd tape -- the force - # (-dE/dx) is already materialised as forward ops in the - # traced graph, so keeping the tape would only make - # AOTAutograd save the full forward activation set for a - # backward eval never runs (see NOTE 13). Training keeps it - # for the force-loss second derivative. - grad_ctx: Any = nullcontext() if self.training else torch.no_grad() - with nvtx_range("SeZM/core_compute"), grad_ctx: - if extended_coord_corr is None: - model_predict_lower = compiled_core_compute( - extended_coord, - extended_atype, - nlist, - mapping, - fp, - ap, - charge_spin, - *_task_buf_vals, - ) - else: - model_predict_lower = compiled_core_compute( - extended_coord, - extended_atype, - nlist, - mapping, - fp, - ap, - charge_spin, - extended_coord_corr, - *_task_buf_vals, - ) - if ( - self._core_compute_pending_compile_t0 is not None - and self._core_compute_pending_compile_key == cache_key - ): - if torch.cuda.is_available(): - torch.cuda.synchronize() - log.info( - "SeZM: finished compiling " - "(mode=%s, coord_corr=%s) in %.2fs", - "train" if self.training else "eval", - has_coord_corr, - time.perf_counter() - self._core_compute_pending_compile_t0, - ) - self._core_compute_pending_compile_t0 = None - self._core_compute_pending_compile_key = None - else: - with nvtx_range("SeZM/core_compute"): - model_predict_lower = self.core_compute( - extended_coord, - extended_atype, - nlist, - mapping=mapping, - fparam=fp, - aparam=ap, - charge_spin=charge_spin, - extra_nlist_sort=self.need_sorted_nlist_for_lower(), - extended_coord_corr=extended_coord_corr, - ) - - with nvtx_range("SeZM/communicate_output"): - model_predict = communicate_extended_output( - model_predict_lower, - self.model_output_def(), - mapping, - do_atomic_virial=do_atomic_virial, - ) + self._dens_pending_compile_t0 = None + else: + with nvtx_range("SeZM/core_compute_dens"): + compute_ret = self.core_compute_dens( + extended_coord, + extended_atype, + nlist, + mapping, + force_input=force_input, + noise_mask=noise_mask, + fparam=fp, + aparam=ap, + charge_spin=charge_spin, + ) + with nvtx_range("SeZM/post_process"): + model_predict = self.post_process_output_dens( + compute_ret, + atype, + noise_mask=noise_mask, + ) # === Step 2. Type cast output === with nvtx_range("SeZM/output_type_cast"): @@ -1158,35 +1250,40 @@ def forward_common_after_nlist( def core_compute( self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, + coord: torch.Tensor, + atype: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_scatter_index: torch.Tensor, + edge_mask: torch.Tensor, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, - extra_nlist_sort: bool = False, extended_coord_corr: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """ - Compute SeZM lower outputs from extended inputs. + Compute SeZM lower outputs from the unified edge-vector schema. - Builds compact sparse edges, runs descriptor and fitting evaluation, - applies output masking and the optional analytical pair potential, - then calls ``edge_energy_deriv`` (edge-force scatter) for force / - virial / per-atom virial. + The caller owns neighbor-list construction. Periodic images enter the + model only through ``edge_vec``; LAMMPS ghost atoms are the zero-shift + instance of the same contract. Parameters ---------- - extended_coord - Coordinates in extended region with shape (nf, nall, 3). - extended_atype - Atom types in extended region with shape (nf, nall). - nlist - Neighbor list with shape (nf, nloc, nsel). - mapping - Extended-to-local mapping with shape (nf, nall), or ``None``. + coord + Coordinates that define the force-scatter domain with shape + ``(nf, nscatter, 3)``. + atype + Local atom types with shape ``(nf, nloc)``. + edge_index + Message-passing indices in flattened local-owner space. + edge_vec + Per-edge displacement vectors in Å. + edge_scatter_index + Force/virial scatter indices in flattened ``coord`` space. + edge_mask + Boolean validity mask aligned with ``edge_vec``. fparam Frame parameters with shape (nf, ndf), or ``None``. aparam @@ -1195,10 +1292,9 @@ def core_compute( Frame-level charge and spin conditions with shape `(nf, 2)`. comm_dict Communication data for parallel inference. Currently unused. - extra_nlist_sort - Whether to forcibly sort the nlist. extended_coord_corr - Coordinates correction for virial with shape (nf, nall, 3) or ``None``. + Coordinates correction for virial with shape ``(nf, nscatter, 3)`` or + ``None``. Returns ------- @@ -1209,32 +1305,23 @@ def core_compute( it. """ del comm_dict - nlist = self.format_nlist( - extended_coord, extended_atype, nlist, extra_nlist_sort=extra_nlist_sort - ) - nf, nloc, _ = nlist.shape - atype = extended_atype[:, :nloc] + nf, nloc = atype.shape[:2] + nscatter = coord.shape[1] descriptor_model = self.atomic_model.descriptor - # === Step 1. Build compact sparse edges === - edge_index, edge_vec, edge_mask, edge_index_ext = ( - self.build_edge_list_from_nlist( - extended_coord=extended_coord, - nlist=nlist, - mapping=mapping, - ) - ) - # Edge displacements are the autograd leaf for the force / virial - # backward. The coordinate gather that produced ``edge_vec`` stays a - # pure forward op, so the differentiated region is the function - # ``(edge_vec, theta) -> E``; this keeps the make_fx symbolic trace and - # second-order lowering clean (see doc/outisli/dpa4.md §12.4). + # === Step 1. Establish the force-autograd endpoint === + # Neighbor-list construction and periodic-image resolution are explicit + # caller responsibilities. Once the edge displacements are supplied, + # SeZM differentiates only the pure map ``(edge_vec, theta) -> E``. + # This keeps coordinate gathering and shift application outside the + # differentiated region while preserving conservative forces through the + # scatter indices below. edge_vec = edge_vec.detach().requires_grad_(True) # === Step 2. Descriptor forward === with nvtx_range("SeZM/descriptor"): descriptor, _ = descriptor_model.forward_with_edges( - extended_coord=extended_coord[:, :nloc, :], + extended_coord=coord[:, :nloc, :], extended_atype=atype, edge_index=edge_index, edge_vec=edge_vec, @@ -1263,8 +1350,7 @@ def core_compute( fit_ret = self.atomic_model.apply_out_stat(fit_ret, atype) # === Step 4. Apply atom mask === - ext_atom_mask = self.atomic_model.make_atom_mask(extended_atype) - atom_mask = ext_atom_mask[:, :nloc].to(torch.int32) + atom_mask = self.atomic_model.make_atom_mask(atype).to(torch.int32) if self.atomic_model.atom_excl is not None: atom_mask *= self.atomic_model.atom_excl(atype) for key in fit_ret.keys(): @@ -1303,15 +1389,14 @@ def core_compute( energy_redu = torch.sum( energy_atom.to(env.GLOBAL_PT_ENER_FLOAT_PRECISION), dim=1 ) - nall = extended_coord.shape[1] energy_derv_r, energy_derv_c, energy_derv_c_redu = edge_energy_deriv( energy_redu, edge_vec, - edge_index_ext[0], - edge_index_ext[1], + edge_scatter_index[0], + edge_scatter_index[1], edge_mask, nf, - nall, + nscatter, create_graph=self.training, extended_coord_corr=extended_coord_corr, ) @@ -1440,29 +1525,36 @@ def core_compute_dens( @torch.jit.export def forward_lower( self, - extended_coord: Float[Tensor, "nf nall_x3"] | Float[Tensor, "nf nall 3"], - extended_atype: Int[Tensor, "nf nall"], - nlist: Int[Tensor, "nf nloc nsel"], - mapping: Int[Tensor, "nf nall"] | None = None, + coord: Float[Tensor, "nf nscatter_x3"] | Float[Tensor, "nf nscatter 3"], + atype: Int[Tensor, "nf nloc"], + edge_index: Int[Tensor, "two nedge"], + edge_vec: Float[Tensor, "nedge 3"], + edge_scatter_index: Int[Tensor, "two nedge"], + edge_mask: torch.Tensor, fparam: Float[Tensor, "nf ndf"] | None = None, - aparam: Float[Tensor, "nf nall nda"] | None = None, + aparam: Float[Tensor, "nf nloc nda"] | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, charge_spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """ - Lower-level public forward using the DeePMD lower-interface contract. + Lower-level public forward using the compact-edge contract. Parameters ---------- - extended_coord - Extended coordinates with shape (nf, nall*3) or (nf, nall, 3) in Å. - extended_atype - Extended atom types with shape (nf, nall). - nlist - Neighbor list with shape (nf, nloc, nsel). - mapping - Mapping indices with shape (nf, nall), or None. + coord + Coordinates defining the force-scatter domain with shape + ``(nf, nscatter*3)`` or ``(nf, nscatter, 3)`` in Å. + atype + Local atom types with shape (nf, nloc). + edge_index + Message-passing edge indices in flattened local-atom space. + edge_vec + Edge displacement vectors in Å. + edge_scatter_index + Force-scatter edge indices in flattened extended-atom space. + edge_mask + Boolean edge-validity mask aligned with `edge_index`. fparam Frame parameters with shape (nf, ndf) or None. aparam @@ -1494,22 +1586,18 @@ def forward_lower( raise NotImplementedError( "SeZM `forward_lower` only supports the conservative `ener` mode." ) - cc_ext, _, fp, ap, input_prec = self._input_type_cast( - extended_coord, fparam=fparam, aparam=aparam - ) model_ret = self.forward_common_lower( - cc_ext, - extended_atype, - nlist, - mapping, - fparam=fp, - aparam=ap, - do_atomic_virial=do_atomic_virial, + coord, + atype, + edge_index, + edge_vec, + edge_scatter_index, + edge_mask, + fparam=fparam, + aparam=aparam, comm_dict=comm_dict, - extra_nlist_sort=self.need_sorted_nlist_for_lower(), charge_spin=charge_spin, ) - model_ret = self._output_type_cast(model_ret, input_prec) if self.get_fitting_net() is not None: model_predict: dict[str, torch.Tensor] = {} @@ -1543,61 +1631,18 @@ def forward_lower( model_predict = model_ret return model_predict - def forward_common_lower( - self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, - fparam: torch.Tensor | None = None, - aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, - comm_dict: dict[str, torch.Tensor] | None = None, - extra_nlist_sort: bool = False, - extended_coord_corr: torch.Tensor | None = None, - charge_spin: torch.Tensor | None = None, - ) -> dict[str, torch.Tensor]: - """Public lower interface with dtype casting around ``core_compute()``.""" - cc_ext, _, fp, ap, input_prec = self._input_type_cast( - extended_coord, fparam=fparam, aparam=aparam - ) - extended_atype = extended_atype.to(device=cc_ext.device, dtype=torch.long) - cc_ext = cc_ext.reshape(extended_atype.shape[0], -1, 3) - if extended_coord_corr is not None and extended_coord_corr.ndim == 2: - extended_coord_corr = extended_coord_corr.reshape( - extended_atype.shape[0], -1, 3 - ) - nf = extended_atype.shape[0] - charge_spin = self.convert_charge_spin( - charge_spin, - nf=nf, - dtype=cc_ext.dtype, - device=cc_ext.device, - ) - model_predict = self.core_compute( - cc_ext, - extended_atype, - nlist, - mapping=mapping, - fparam=fp, - aparam=ap, - charge_spin=charge_spin, - comm_dict=comm_dict, - extra_nlist_sort=extra_nlist_sort, - extended_coord_corr=extended_coord_corr, - ) - return self._output_type_cast(model_predict, input_prec) - # ========================================================================= # Compile Utilities # ========================================================================= def trace_and_compile( self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor, + coord: torch.Tensor, + atype: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_scatter_index: torch.Tensor, + edge_mask: torch.Tensor, fp: torch.Tensor, ap: torch.Tensor, charge_spin: torch.Tensor, @@ -1615,6 +1660,7 @@ def trace_and_compile( compiled callable is stored outside the ``nn.Module`` tree so FSDP/DDP cannot see or shard its duplicated parameters. """ + check_compile_torch_version() from torch._decomp import ( get_decompositions, ) @@ -1726,10 +1772,12 @@ def _prepare_coord_for_trace(coord: torch.Tensor) -> torch.Tensor: if extended_coord_corr is None: def compute_fn( - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor, + coord: torch.Tensor, + atype: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_scatter_index: torch.Tensor, + edge_mask: torch.Tensor, fp: torch.Tensor, ap: torch.Tensor, charge_spin: torch.Tensor, @@ -1738,14 +1786,15 @@ def compute_fn( _saved = _patch_task_bufs(task_buf_vals) try: return self.core_compute( - _prepare_coord_for_trace(extended_coord), - extended_atype, - nlist, - mapping=mapping, + _prepare_coord_for_trace(coord), + atype, + edge_index, + edge_vec, + edge_scatter_index, + edge_mask, fparam=fp, aparam=ap, charge_spin=charge_spin, - extra_nlist_sort=self.need_sorted_nlist_for_lower(), ) finally: _restore_task_bufs(_saved) @@ -1753,10 +1802,12 @@ def compute_fn( else: def compute_fn( # type: ignore[misc] - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor, + coord: torch.Tensor, + atype: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_scatter_index: torch.Tensor, + edge_mask: torch.Tensor, fp: torch.Tensor, ap: torch.Tensor, charge_spin: torch.Tensor, @@ -1769,14 +1820,15 @@ def compute_fn( # type: ignore[misc] _saved = _patch_task_bufs(task_buf_vals) try: return self.core_compute( - _prepare_coord_for_trace(extended_coord), - extended_atype, - nlist, - mapping=mapping, + _prepare_coord_for_trace(coord), + atype, + edge_index, + edge_vec, + edge_scatter_index, + edge_mask, fparam=fp, aparam=ap, charge_spin=charge_spin, - extra_nlist_sort=self.need_sorted_nlist_for_lower(), extended_coord_corr=extended_coord_corr, ) finally: @@ -1785,7 +1837,7 @@ def compute_fn( # type: ignore[misc] # Trace dims are pairwise-distinct primes >= 5 so ``make_fx`` neither # unifies two axes onto one symbol (duck-shape) nor specializes an axis # on a literal; ``next_safe_prime`` documents why. The forbidden set - # adds the model-contracted dims (``nsel``, fparam / aparam widths, + # adds the model-contracted dims (fparam / aparam widths, # charge_spin) and the promoted task-buffer dims so the chosen primes # never collide with them. _forbidden: set[int] = {1, 2, 3, 9} @@ -1793,54 +1845,45 @@ def compute_fn( # type: ignore[misc] for _d in _tbv.shape: if _d > 1: _forbidden.add(int(_d)) - # Model-contracted dims kept at their real values (changing them - # would break the model's own assertions about ``sel``, fparam / - # aparam widths, charge_spin dim). Add to forbidden so primes - # picked for free dims do not collide. - _nsel_real = int(nlist.shape[2]) + # Model-contracted dims kept at their real values. Add them to the + # forbidden set so free symbolic axes do not collide. _dim_fp = int(fp.shape[1]) _dim_ap = int(ap.shape[2]) _dim_cs = int(charge_spin.shape[1]) - for _d in (_nsel_real, _dim_fp, _dim_ap, _dim_cs): + for _d in (_dim_fp, _dim_ap, _dim_cs): if _d > 1: _forbidden.add(_d) - # Pick primes in physical order ``nf < nloc < nall``. The order - # ``trace_nloc < trace_nall`` matters: the model slices - # ``extended_atype[:, :nloc]`` to get local atoms; if - # ``trace_nloc > trace_nall`` the slice silently truncates at - # trace time, breaking the captured symbolic shape relation - # ``atype.shape[1] == nloc``. + # Pick distinct primes for free axes. ``nscatter`` may be larger than + # ``nloc`` for LAMMPS ghost domains, so keep it separately symbolic. trace_nf = next_safe_prime(5, _forbidden) _forbidden.add(trace_nf) trace_nloc = next_safe_prime(trace_nf + 1, _forbidden) _forbidden.add(trace_nloc) - trace_nall = next_safe_prime(trace_nloc + 1, _forbidden) + trace_nscatter = next_safe_prime(trace_nloc + 1, _forbidden) + _forbidden.add(trace_nscatter) + trace_nedge = next_safe_prime(trace_nscatter + 1, _forbidden) # Build trace inputs by padding/trimming real-data tensors into the # chosen prime shapes; ``trace_pad_dim`` documents how index-bearing # tensors keep valid values. - coord_for_trace = trace_pad_dim(extended_coord[:1], 0, trace_nf) - coord_for_trace = trace_pad_dim(coord_for_trace, 1, trace_nall) - atype_for_trace = trace_pad_dim(extended_atype[:1], 0, trace_nf) - atype_for_trace = trace_pad_dim(atype_for_trace, 1, trace_nall) - nlist_for_trace = trace_pad_dim(nlist[:1], 0, trace_nf) - nlist_for_trace = trace_pad_dim(nlist_for_trace, 1, trace_nloc) - # Real nlist values are in ``[-1, real_nall)`` (``-1`` marks - # padded slots, non-negative entries index into extended_coord). - # After trimming ``nall`` down to ``trace_nall`` some of those - # values can exceed ``trace_nall``, which would produce - # out-of-range gather indices in ``coord_flat.index_select(0, - # src_ext)`` during the trace pass. Clamp the upper bound to - # ``trace_nall - 1`` (the ``-1`` padding stays untouched since - # clamp only caps the high side). - nlist_for_trace = torch.clamp(nlist_for_trace, max=trace_nall - 1) - mapping_for_trace = trace_pad_dim(mapping[:1], 0, trace_nf) - mapping_for_trace = trace_pad_dim(mapping_for_trace, 1, trace_nall) - # Real mapping values are in ``[0, real_nloc)``. If - # ``trace_nloc < real_nloc`` they can exceed ``trace_nloc`` and - # silently propagate into ``src_local`` (used as a local-atom - # index downstream). Clamp to ``trace_nloc - 1``. - mapping_for_trace = torch.clamp(mapping_for_trace, min=0, max=trace_nloc - 1) + coord_for_trace = trace_pad_dim(coord[:1], 0, trace_nf) + coord_for_trace = trace_pad_dim(coord_for_trace, 1, trace_nscatter) + atype_for_trace = trace_pad_dim(atype[:1], 0, trace_nf) + atype_for_trace = trace_pad_dim(atype_for_trace, 1, trace_nloc) + edge_index_for_trace = trace_pad_dim(edge_index, 1, trace_nedge) + edge_index_for_trace = torch.clamp( + edge_index_for_trace, + min=0, + max=trace_nf * trace_nloc - 1, + ) + edge_scatter_for_trace = trace_pad_dim(edge_scatter_index, 1, trace_nedge) + edge_scatter_for_trace = torch.clamp( + edge_scatter_for_trace, + min=0, + max=trace_nf * trace_nscatter - 1, + ) + edge_vec_for_trace = trace_pad_dim(edge_vec, 0, trace_nedge) + edge_mask_for_trace = trace_pad_dim(edge_mask, 0, trace_nedge) fp_for_trace = trace_pad_dim(fp[:1], 0, trace_nf) ap_for_trace = trace_pad_dim(ap[:1], 0, trace_nf) ap_for_trace = trace_pad_dim(ap_for_trace, 1, trace_nloc) @@ -1849,15 +1892,17 @@ def compute_fn( # type: ignore[misc] trace_args = [ coord_for_trace, atype_for_trace, - nlist_for_trace, - mapping_for_trace, + edge_index_for_trace, + edge_vec_for_trace, + edge_scatter_for_trace, + edge_mask_for_trace, fp_for_trace, ap_for_trace, charge_spin_for_trace, ] if extended_coord_corr is not None: corr_for_trace = trace_pad_dim(extended_coord_corr[:1], 0, trace_nf) - corr_for_trace = trace_pad_dim(corr_for_trace, 1, trace_nall) + corr_for_trace = trace_pad_dim(corr_for_trace, 1, trace_nscatter) trace_args.append(corr_for_trace) # Append task-buffer values last so they map to the *task_buf_vals # varargs in compute_fn. Their shapes are static (they don't vary @@ -1982,11 +2027,10 @@ def compute_fn( # type: ignore[misc] def _inductor_inference_compiler( fx_gm: torch.fx.GraphModule, fx_inputs: list[Any] ) -> Any: - # max_tiles=1 keeps pointwise grids 1D so the data-dependent - # edge axis stays on Triton's x grid (limit 2**31); the default - # tiling places it on the y/z grid (limit 65535), which - # overflows for large systems. - with _ind_cfg.patch({**compile_options, "triton.max_tiles": 1}): + # ``compile_options`` already includes ``triton.max_tiles=1`` + # from ``build_inductor_compile_options``, so the 1D launch-grid + # constraint applies to both the training and evaluation graphs. + with _ind_cfg.patch(compile_options): return compile_fx_inner(fx_gm, fx_inputs) # select_decomp_table keeps the decomposition set aligned with @@ -1995,7 +2039,7 @@ def _inductor_inference_compiler( # device: AOTAutograd's PhiloxStateTracker allocates an RNG-state # tensor without an explicit device, which otherwise lands on a # stray default device and raises "invalid device ordinal". - with torch.no_grad(), torch.device(extended_coord.device): + with torch.no_grad(), torch.device(coord.device): _compiled_flat = aot_module_simplified( traced, example_inputs, @@ -2046,6 +2090,7 @@ def compiled(*args: Any, _fn: Any = _compiled_flat) -> dict[str, Any]: def compile_dens(self) -> None: """Compile the direct-force `dens` path.""" + check_compile_torch_version() from torch._inductor import config as inductor_config log.info("SeZM: start compiling dens path") @@ -2091,6 +2136,7 @@ def _trace_lower_exportable( *sample_inputs: torch.Tensor | None, ) -> torch.nn.Module: """Trace a lower-interface closure into an exportable FX graph.""" + check_compile_torch_version() from torch._decomp import ( get_decompositions, ) @@ -2106,15 +2152,15 @@ def _trace_lower_exportable( def forward_common_lower_exportable( self, - extended_coord: torch.Tensor, - extended_atype: torch.Tensor, - nlist: torch.Tensor, - mapping: torch.Tensor | None = None, + coord: torch.Tensor, + atype: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_scatter_index: torch.Tensor, + edge_mask: torch.Tensor, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, - *, - do_atomic_virial: bool = False, ) -> torch.nn.Module: """Trace ``forward_common_lower`` into an exportable FX ``GraphModule``. @@ -2134,13 +2180,14 @@ def forward_common_lower_exportable( ) model = self - extra_sort = self.need_sorted_nlist_for_lower() def lower_fn( - ext_coord: torch.Tensor, - ext_atype: torch.Tensor, - nlist_: torch.Tensor, - mapping_: torch.Tensor | None, + coord_: torch.Tensor, + atype_: torch.Tensor, + edge_index_: torch.Tensor, + edge_vec_: torch.Tensor, + edge_scatter_index_: torch.Tensor, + edge_mask_: torch.Tensor, fparam_: torch.Tensor | None, aparam_: torch.Tensor | None, charge_spin_: torch.Tensor | None, @@ -2150,33 +2197,42 @@ def lower_fn( # endpoint is the per-edge ``edge_vec`` leaf created inside # ``core_compute`` (edge-force scatter), so the coordinates carry # no grad endpoint here. - ext_coord = ext_coord.detach() + coord_ = coord_.detach() + edge_vec_ = edge_vec_.detach() return model.forward_common_lower( - ext_coord, - ext_atype, - nlist_, - mapping_, + coord_, + atype_, + edge_index_, + edge_vec_, + edge_scatter_index_, + edge_mask_, fparam=fparam_, aparam=aparam_, - do_atomic_virial=do_atomic_virial, - extra_nlist_sort=extra_sort, charge_spin=charge_spin_, + # Export tracing must capture the eager lower graph itself. The + # runtime compile cache is a deployment optimisation around that + # graph and must not be entered while make_fx is tracing it. + use_compile=False, ) def fn( - ext_coord: torch.Tensor, - ext_atype: torch.Tensor, - nlist_: torch.Tensor, - mapping_: torch.Tensor | None, + coord_: torch.Tensor, + atype_: torch.Tensor, + edge_index_: torch.Tensor, + edge_vec_: torch.Tensor, + edge_scatter_index_: torch.Tensor, + edge_mask_: torch.Tensor, fparam_: torch.Tensor | None, aparam_: torch.Tensor | None, charge_spin_: torch.Tensor | None, ) -> dict[str, torch.Tensor]: return lower_fn( - ext_coord, - ext_atype, - nlist_, - mapping_, + coord_, + atype_, + edge_index_, + edge_vec_, + edge_scatter_index_, + edge_mask_, fparam_, aparam_, charge_spin_, @@ -2185,18 +2241,20 @@ def fn( if self.get_dim_chg_spin() > 0: charge_spin = self.convert_charge_spin( charge_spin, - nf=extended_atype.shape[0], - dtype=extended_coord.dtype, - device=extended_coord.device, + nf=atype.shape[0], + dtype=coord.dtype, + device=coord.device, ) # Always include the charge_spin slot (possibly None) so the traced - # module's forward signature matches the 7-tuple the freeze pipeline + # module's forward signature matches the freeze pipeline # passes at runtime, regardless of whether the model is conditioned. trace_inputs = ( - extended_coord, - extended_atype, - nlist, - mapping, + coord, + atype, + edge_index, + edge_vec, + edge_scatter_index, + edge_mask, fparam, aparam, charge_spin, @@ -2227,64 +2285,43 @@ def build_neighbor_list( coord: Float[Tensor, "nf nloc 3"] | Float[Tensor, "nf nloc_x3"], atype: Int[Tensor, "nf nloc"], box: Float[Tensor, "nf 9"] | None, - ) -> tuple[ - Float[Tensor, "nf nall 3"], - Int[Tensor, "nf nall"], - Int[Tensor, "nf nall"], - Int[Tensor, "nf nloc nsel"], - ]: - """ - Build extended inputs and the neighbor list for the ``forward`` entry. + ) -> EdgeNeighborList: + """Build the unified edge-vector schema for the ``forward`` entry.""" + nf, nloc = atype.shape[:2] + return _select_neighbor_builder(nf, coord.device).build( + coord.view(nf, nloc, 3), + atype, + box, + self.get_rcut(), + self.get_sel(), + return_mode="edges", + ) - Used when the model constructs its own neighbor list from ``coord`` / - ``box``, as opposed to ``forward_lower`` which receives an externally - built nlist (e.g. from LAMMPS or an inference ``NeighborList`` strategy). - Large CUDA systems use the Toolkit-Ops neighbor list - (:class:`NvNeighborList`); all other cases use the dense all-pairs - builder. The non-periodic Toolkit-Ops path uses a larger threshold - because the dense builder is still faster at small sizes. + def build_extended_neighbor_list( + self, + coord: Float[Tensor, "nf nloc 3"] | Float[Tensor, "nf nloc_x3"], + atype: Int[Tensor, "nf nloc"], + box: Float[Tensor, "nf 9"] | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Build the historical extended-coordinate representation. - Parameters - ---------- - coord - Coordinates with shape (nf, nloc, 3) in Å. - atype - Atom types with shape (nf, nloc). - box - Box tensor with shape (nf, 9) in Å, or None. + This helper is retained for modes whose physical preprocessing still + requires explicit extended atoms, such as the denoising path and the spin + virtual-atom transform. The conservative SeZM energy path converts the + result to :class:`EdgeNeighborList` before entering ``core_compute``. - Returns - ------- - tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] - Extended coordinates, extended atom types, mapping, and neighbor list. + Returns the :class:`~deepmd.dpmodel.utils.neighbor_list.NeighborList` + contract order ``(extended_coord, extended_atype, nlist, mapping)``. """ - nloc = atype.shape[1] - nv_threshold = ( - SEZM_NV_NLIST_THRESHOLD - if box is not None - else SEZM_NV_NONPERIODIC_NLIST_THRESHOLD - ) - if coord.is_cuda and nloc >= nv_threshold and is_nv_available(): - # Large systems: the device-resident Toolkit-Ops neighbor list avoids - # the dense all-pairs ghost expansion. It already keeps - # the nearest sum(sel) neighbors (fixed width, like the standard - # builder); only its (nlist, mapping) order is swapped to this - # method's (mapping, nlist) contract. - extended_coord, extended_atype, nlist, mapping = NvNeighborList().build( - coord.view(atype.shape[0], nloc, 3), - atype, - box, - self.get_rcut(), - self.get_sel(), - ) - return extended_coord, extended_atype, mapping, nlist - return extend_input_and_build_neighbor_list( - coord, + nf, nloc = atype.shape[:2] + return _select_neighbor_builder(nf, coord.device).build( + coord.view(nf, nloc, 3), atype, + box, self.get_rcut(), self.get_sel(), - mixed_types=True, - box=box, + return_mode="extended", ) def build_edge_list_from_nlist( @@ -2295,15 +2332,11 @@ def build_edge_list_from_nlist( mapping: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """ - Build a compact edge list from DeePMD padded neighbor list. + Build a compact edge list from a DeePMD padded neighbor list. - Edge vectors are gathered from ``extended_coord`` with ghosts' - real (periodic-image) coordinates, so ``edge_vec`` is the true - minimum-image displacement and the edge-force-scatter virial is - PBC-correct without any extended-coordinate trick. Two masked dummy - edges are always appended to avoid data-dependent empty-edge branches - that ``make_fx`` cannot trace and singular edge-axis guards in - Inductor's batched matmul lowering. + This adapter is retained for the denoising path and legacy tests. The + conservative energy path builds :class:`EdgeNeighborList` directly and + enters ``forward_common_lower`` with explicit edge vectors. Two index spaces are returned. ``edge_index`` uses *local* node indices ``[0, nf * nloc)`` (neighbours mapped to their local image) @@ -2331,127 +2364,28 @@ def build_edge_list_from_nlist( Edge vectors with shape (E+2, 3). edge_mask Boolean mask with shape (E+2,). The two trailing elements are ``False``. - edge_index_ext - Extended (src, dst) indices with shape (2, E+2), values in - ``[0, nf * nall)``, aligned 1:1 with ``edge_index`` / ``edge_vec``. + edge_scatter_index + Scatter-domain (src, dst) indices with shape (2, E+2), aligned + 1:1 with ``edge_index`` and ``edge_vec``. """ - nf, nloc, nsel = nlist.shape - device = extended_coord.device - nall = extended_coord.shape[1] - - # === Step 1. Build per-edge geometry via gather === - # Edge vectors come from ``torch.gather`` rather than advanced indexing - # ``coord_flat[...]``: gather lowers to an explicit, symbolic-shape - # -friendly op under make_fx, while advanced indexing under symbolic - # shapes can silently truncate gradients. ``core_compute`` detaches the - # result into the ``edge_vec`` autograd leaf, so this gather is a pure - # forward op. ``torch.where(valid_flat, neighbor_flat, 0)`` sanitises - # padded ``-1`` entries before indexing so we never hit an - # out-of-range gather; the corresponding edges are filtered out below. - neighbor_flat = nlist.reshape(-1) - # ``dst_actual = arange(N*K) // K`` produces the same value - # sequence as ``arange(N).repeat_interleave(K)`` but its length - # is derived from ``neighbor_flat.shape[0]`` -- a single symbolic - # source shared with the ``torch.where`` below. The previous - # ``arange(nf*nloc).repeat_interleave(nsel)`` chain could - # decouple from ``nlist.numel()`` in the FX graph if any - # upstream code path ever specialized ``nloc`` at trace time; - # deriving from ``neighbor_flat.shape[0]`` makes the equality - # structural and survives any future change in trace-shape - # selection in ``trace_and_compile``. - dst_actual = ( - torch.arange(neighbor_flat.shape[0], device=device, dtype=torch.long) - // nsel - ) - f_idx = dst_actual // nloc - dst_local = dst_actual % nloc - valid_flat = neighbor_flat >= 0 - neighbor_safe = torch.where( - valid_flat, neighbor_flat, torch.zeros_like(neighbor_flat) + nloc = nlist.shape[1] + atype = torch.empty( + (nlist.shape[0], nloc), + dtype=torch.long, + device=extended_coord.device, ) - # Gather coordinates within each frame instead of flattening - # ``(nf, nall)`` into one index space. The flattened form is - # mathematically correct, but Inductor may lower - # ``coord_flat.index_select(0, f_idx * nall + local_idx)`` to a kernel - # that asserts the composite index against ``nall`` rather than - # ``nf * nall`` when ``nf > 1``. Frame-local gather keeps the same - # differentiable path to coordinates while making every indirect index - # visibly bounded by the atom axis. - neighbor_safe_2d = neighbor_safe.to(dtype=torch.long).view(nf, nloc * nsel) - nei_coord = torch.gather( + edge_schema = edge_schema_from_extended( extended_coord, - 1, - neighbor_safe_2d.unsqueeze(-1).expand(-1, -1, 3), - ).reshape(-1, 3) - dst_coord = torch.gather( - extended_coord[:, :nloc, :], - 1, - dst_local.view(nf, -1).unsqueeze(-1).expand(-1, -1, 3), - ).reshape(-1, 3) - diff = nei_coord - dst_coord - edge_len2 = torch.sum(diff * diff, dim=-1) - - # === Step 2. Build compact src/dst (local indices) === - if mapping is None: - src_local = neighbor_safe.to(dtype=torch.long) - else: - src_local = torch.gather(mapping, 1, neighbor_safe_2d).reshape(-1) - src_actual = f_idx * nloc + src_local.to(dtype=torch.long) - - # Extended-index counterparts for the edge-force scatter. The - # neighbour keeps its ghost identity (``neighbor_safe`` indexes - # ``[0, nall)``) while the centre, always a local atom, occupies its - # own slot ``dst_local`` in the extended layout. Scattering the edge - # gradient onto these indices yields per-ghost extended force / - # virial, the exact contract ``communicate_extended_output`` reduces. - src_ext = f_idx * nall + neighbor_safe.to(dtype=torch.long) - dst_ext = f_idx * nall + dst_local - - # Filter: valid nlist entry AND src in [0, nloc) AND non-zero distance. - src_local_valid = (src_local >= 0) & (src_local < nloc) - len_positive = edge_len2 > 1e-10 - edge_mask_actual = valid_flat & src_local_valid & len_positive - - valid_idx = torch.nonzero(edge_mask_actual, as_tuple=False).flatten() - - # === Step 3. Compact edges + append masked dummies === - # NOTE: Always append two masked dummy edges. - # ``torch.nonzero(edge_mask_actual)`` produces a data-dependent - # number of valid edges, which can be zero on sparse or - # single-type systems (e.g. isolated-atom reference frames). - # make_fx cannot trace an ``if n_edges == 0: skip`` branch - # symbolically; without the dummies it would fall back to - # concrete shape specialisation and break - # ``torch.compile(dynamic=True)`` for later batches. Two dummy - # slots also give Inductor's batched matmul lowering a static - # ``E >= 2`` edge-axis bound, avoiding data-dependent layout - # guards on ``E == 1`` that would otherwise trigger an extra - # recompile when the first batch contains only a single edge. - # Each dummy copies entry 0 (any in-range index is fine) and - # carries ``edge_mask=False`` so every downstream sum, gather - # or scatter ignores it. - dummy_count = 2 - padded_idx = torch.cat( - [valid_idx, torch.zeros(dummy_count, dtype=torch.long, device=device)] - ) - src_sel = src_actual.index_select(0, padded_idx) - dst_sel = dst_actual.index_select(0, padded_idx) - edge_vec_sel = diff.index_select(0, padded_idx) - edge_index = torch.stack([src_sel, dst_sel], dim=0) - edge_index_ext = torch.stack( - [ - src_ext.index_select(0, padded_idx), - dst_ext.index_select(0, padded_idx), - ], - dim=0, + atype, + nlist, + mapping, ) - edge_mask = torch.cat( - [ - torch.ones(valid_idx.shape[0], dtype=torch.bool, device=device), - torch.zeros(dummy_count, dtype=torch.bool, device=device), - ] + return ( + edge_schema.edge_index, + edge_schema.edge_vec, + edge_schema.edge_mask, + edge_schema.edge_scatter_index, ) - return edge_index, edge_vec_sel, edge_mask, edge_index_ext # ========================================================================= # Input Canonicalization diff --git a/deepmd/pt/model/model/sezm_spin_model.py b/deepmd/pt/model/model/sezm_spin_model.py index d36ee36fd9..af973ea529 100644 --- a/deepmd/pt/model/model/sezm_spin_model.py +++ b/deepmd/pt/model/model/sezm_spin_model.py @@ -37,6 +37,9 @@ from deepmd.pt.utils.utils import ( to_torch_tensor, ) +from deepmd.pt_expt.utils.edge_schema import ( + edge_schema_from_extended, +) from deepmd.utils.path import ( DPPath, ) @@ -161,8 +164,8 @@ def forward_common( cc = cc.view(nf, nloc, 3) spin = spin.to(dtype=cc.dtype, device=cc.device).reshape(nf, nloc, 3) - extended_coord, extended_atype, mapping, nlist = self.build_neighbor_list( - cc, atype, bb + extended_coord, extended_atype, nlist, mapping = ( + self.build_extended_neighbor_list(cc, atype, bb) ) extended_spin = torch.gather( spin, @@ -184,18 +187,31 @@ def forward_common( ) if ap is not None: ap = self.expand_aparam(ap, nloc * 2) - model_ret = self.forward_common_after_nlist( + nlist_updated = self.format_nlist( extended_coord_updated, extended_atype_updated, - mapping_updated, nlist_updated, + extra_nlist_sort=self.need_sorted_nlist_for_lower(), + ) + edge_schema = edge_schema_from_extended( + extended_coord_updated, extended_atype_updated[:, : nloc * 2], - fp, - ap, - input_prec, - do_atomic_virial=do_atomic_virial, - extended_coord_corr=extended_coord_corr, + nlist_updated, + mapping_updated, + scatter_to_local=True, + ) + model_ret = super().forward_common_lower( + edge_schema.coord, + edge_schema.atype, + edge_schema.edge_index, + edge_schema.edge_vec, + edge_schema.edge_scatter_index, + edge_schema.edge_mask, + fparam=fp, + aparam=ap, + extended_coord_corr=extended_coord_corr[:, : nloc * 2, :], charge_spin=charge_spin, + input_prec=input_prec, ) return self._split_spin_common_output(model_ret, atype, nloc) @@ -221,7 +237,6 @@ def forward_lower( mapping=mapping, fparam=fparam, aparam=aparam, - do_atomic_virial=do_atomic_virial, comm_dict=comm_dict, charge_spin=charge_spin, extra_nlist_sort=self.need_sorted_nlist_for_lower(), @@ -253,7 +268,6 @@ def forward_common_lower( mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, - do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, extra_nlist_sort: bool = False, charge_spin: torch.Tensor | None = None, @@ -278,16 +292,28 @@ def forward_common_lower( ) if aparam is not None: aparam = self.expand_aparam(aparam, nloc * 2) - model_ret = super().forward_common_lower( + nlist_updated = self.format_nlist( extended_coord_updated, extended_atype_updated, nlist_updated, - mapping=mapping_updated, + extra_nlist_sort=extra_nlist_sort, + ) + edge_schema = edge_schema_from_extended( + extended_coord_updated, + extended_atype_updated[:, : nloc * 2], + nlist_updated, + mapping_updated, + ) + model_ret = super().forward_common_lower( + edge_schema.coord, + edge_schema.atype, + edge_schema.edge_index, + edge_schema.edge_vec, + edge_schema.edge_scatter_index, + edge_schema.edge_mask, fparam=fparam, aparam=aparam, - do_atomic_virial=do_atomic_virial, comm_dict=comm_dict, - extra_nlist_sort=extra_nlist_sort, extended_coord_corr=extended_coord_corr, charge_spin=charge_spin, ) @@ -303,8 +329,6 @@ def forward_common_lower_exportable( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, - *, - do_atomic_virial: bool = False, ) -> torch.nn.Module: """Trace the spin lower interface into an exportable FX graph.""" extra_sort = self.need_sorted_nlist_for_lower() @@ -328,7 +352,6 @@ def lower_fn( mapping_, fparam=fparam_, aparam=aparam_, - do_atomic_virial=do_atomic_virial, extra_nlist_sort=extra_sort, charge_spin=charge_spin_, ) diff --git a/deepmd/pt/utils/compile_compat.py b/deepmd/pt/utils/compile_compat.py index 3cdfd54da6..b1b533f740 100644 --- a/deepmd/pt/utils/compile_compat.py +++ b/deepmd/pt/utils/compile_compat.py @@ -42,6 +42,7 @@ "get_task_buffer_values", "is_prime", "next_safe_prime", + "patch_inductor_force_int64_indexing", "patch_inductor_symbolic_divisibility", "rebuild_graph_module", "strip_saved_tensor_detach", @@ -83,6 +84,10 @@ def apply_global_compile_patches() -> None: dynamo_config.optimize_ddp = False + # Force int64 tensor indexing in every compiled kernel. Applies on all + # supported PyTorch versions and is independent of runtime shapes. + patch_inductor_force_int64_indexing() + # The symbolic-divisibility regression exists only on PyTorch 2.12; the # 2.11 backend evaluates the same predicate correctly and must not be # patched. @@ -90,6 +95,36 @@ def apply_global_compile_patches() -> None: patch_inductor_symbolic_divisibility() +def patch_inductor_force_int64_indexing() -> None: + """Force Inductor to emit int64 tensor indexing in every compiled kernel. + + Inductor selects the index dtype from static size hints. The compiled + ``core_compute`` graph is traced with the small placeholder shapes returned + by :func:`next_safe_prime`, from which Inductor infers that the + data-dependent edge and node axes fit in int32. At runtime those axes grow + large enough that the flattened index of a tensor such as ``(E, D, D, C)`` + exceeds ``2**31`` and wraps to an out-of-range address, which surfaces + asynchronously as a CUDA illegal memory access. Forcing int64 indexing + removes this dependence on the trace-time size hints at the cost of a small + amount of additional address arithmetic. The patch is idempotent and + complements ``triton.max_tiles=1`` in :func:`build_inductor_compile_options`. + """ + try: + from torch._inductor.codegen.simd import ( + SIMDScheduling, + ) + except Exception: + return + + if getattr(SIMDScheduling, "_dp_force_int64_patched", False): + return + + # ``can_use_32bit_indexing`` gates int32 selection; returning ``False`` + # forces int64 indexing in every generated kernel. + SIMDScheduling.can_use_32bit_indexing = staticmethod(lambda numel, buffers: False) + SIMDScheduling._dp_force_int64_patched = True + + def check_compile_torch_version() -> None: """Fail fast when ``torch.compile`` is requested on an unsupported PyTorch.""" version = Version(torch.__version__).release @@ -276,6 +311,13 @@ def build_inductor_compile_options() -> dict[str, Any]: # shapes on PyTorch 2.11 and earlier (pytorch/pytorch#174379, #178080, # #179494); the edge count is exactly that kind of shape. "triton.mix_order_reduction": False, + # Constrain every generated kernel to a 1D launch grid. The default + # 2D/3D tiling can place the data-dependent edge or node axis on the y + # or z launch dimension, whose limit is 65535; a larger axis then + # launches an out-of-range grid that surfaces as a CUDA illegal memory + # access. A 1D grid keeps that axis on the x dimension (limit 2**31-1). + # The option is shared by the training and evaluation graphs. + "triton.max_tiles": 1, } try: from torch._inductor import config as inductor_config diff --git a/deepmd/pt/utils/nlist.py b/deepmd/pt/utils/nlist.py index 7f74e65f26..ccea0be79c 100644 --- a/deepmd/pt/utils/nlist.py +++ b/deepmd/pt/utils/nlist.py @@ -18,6 +18,7 @@ def extend_input_and_build_neighbor_list( sel: list[int], mixed_types: bool = False, box: torch.Tensor | None = None, + cap_neighbors: bool = True, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: nframes, nloc = atype.shape[:2] if box is not None: @@ -39,6 +40,7 @@ def extend_input_and_build_neighbor_list( rcut, sel, distinguish_types=(not mixed_types), + cap_neighbors=cap_neighbors, ) extended_coord = extended_coord.view(nframes, -1, 3) return extended_coord, extended_atype, mapping, nlist @@ -51,6 +53,7 @@ def build_neighbor_list( rcut: float, sel: int | list[int], distinguish_types: bool = True, + cap_neighbors: bool = True, ) -> torch.Tensor: """Build neighbor list for a single frame. keeps nsel neighbors. @@ -72,6 +75,11 @@ def build_neighbor_list( types. distinguish_types : bool distinguish different types. + cap_neighbors : bool + If True (default), keep at most ``sum(sel)`` nearest neighbors per atom, + the historical fixed-width behavior. If False, keep every neighbor + within ``rcut`` and size the neighbor axis from the data so no neighbor + is dropped; supported only together with ``distinguish_types=False``. Returns ------- @@ -121,7 +129,14 @@ def build_neighbor_list( idx = torch.arange(diag_len, device=rr.device, dtype=torch.int) rr[:, idx, idx] -= 1.0 - nsel = sum(sel) + if cap_neighbors: + nsel = sum(sel) + else: + # Keep every neighbor within ``rcut``: size the axis from the densest + # local environment. Each local atom counts itself once (its diagonal + # distance was decremented above), so subtract one to exclude self. + nsel = int((rr <= rcut).sum(dim=-1).max().item()) - 1 + nsel = max(nsel, 1) nnei = rr.shape[-1] top_k = min(nsel + 1, nnei) rr, nlist = torch.topk(rr, top_k, largest=False) @@ -131,7 +146,7 @@ def build_neighbor_list( nlist = nlist[:, :, 1:] return _trim_mask_distinguish_nlist( - is_vir, atype, rr, nlist, rcut, sel, distinguish_types + is_vir, atype, rr, nlist, rcut, sel, distinguish_types, nsel=nsel ) @@ -143,9 +158,11 @@ def _trim_mask_distinguish_nlist( rcut: float, sel: list[int], distinguish_types: bool, + nsel: int | None = None, ) -> torch.Tensor: """Trim the size of nlist, mask if any central atom is virtual, distinguish types if necessary.""" - nsel = sum(sel) + if nsel is None: + nsel = sum(sel) # nloc x nsel batch_size, nloc, nnei = rr.shape assert batch_size == is_vir_cntl.shape[0] diff --git a/deepmd/pt/utils/nv_nlist.py b/deepmd/pt/utils/nv_nlist.py index 188d4f557d..b7143a9ac6 100644 --- a/deepmd/pt/utils/nv_nlist.py +++ b/deepmd/pt/utils/nv_nlist.py @@ -20,6 +20,8 @@ import contextlib import logging +import os +import sys from typing import ( TYPE_CHECKING, Any, @@ -28,14 +30,22 @@ import torch from deepmd.dpmodel.utils.neighbor_list import ( + EdgeNeighborList, NeighborList, ) from deepmd.pt.utils.region import ( normalize_coord, ) +from deepmd.pt_expt.utils.edge_schema import ( + edge_schema_from_neighbor_matrix, +) NV_CELL_LIST_THRESHOLD = 1024 NV_NONPERIODIC_CELL_LIST_THRESHOLD = 4096 +# CPU has far less parallelism than CUDA, so the O(N^2) ``batch_naive`` method +# is overtaken by the O(N) ``batch_cell_list`` at a much smaller atom count; +# switch over early regardless of periodicity. +NV_CPU_CELL_LIST_THRESHOLD = 128 log = logging.getLogger(__name__) @@ -45,32 +55,76 @@ ) +@contextlib.contextmanager +def _suppress_native_stderr() -> Iterator[None]: + """Redirect the process ``stderr`` file descriptor to ``os.devnull``. + + ``nvalchemiops`` initializes NVIDIA Warp on first import, which probes for a + CUDA driver and prints a native ``Warp CUDA error 100`` line straight to the + ``stderr`` fd on CPU-only hosts. That line bypasses Python logging, so the + only way to mute it is at the descriptor level around the triggering import. + """ + try: + stderr_fd = sys.stderr.fileno() + except (AttributeError, OSError, ValueError): + # stderr is not a real file descriptor (e.g. captured in tests); the + # native chatter cannot be redirected, so import without suppression. + yield + return + saved_fd = os.dup(stderr_fd) + with open(os.devnull, "w") as devnull: + os.dup2(devnull.fileno(), stderr_fd) + try: + yield + finally: + os.dup2(saved_fd, stderr_fd) + os.close(saved_fd) + + def is_nv_available() -> bool: """Whether the ``nvalchemiops`` Toolkit-Ops neighbor list is importable.""" + # Warp's one-time CUDA probe prints to the native stderr on CPU-only hosts; + # mute it there without hiding diagnostics on machines that have a GPU. + import_ctx = ( + _suppress_native_stderr() + if not torch.cuda.is_available() + else contextlib.nullcontext() + ) try: - import nvalchemiops.torch.neighbors # noqa: F401 + with import_ctx: + import nvalchemiops.torch.neighbors # noqa: F401 except (ImportError, OSError, RuntimeError) as err: log.debug("nvalchemiops Toolkit-Ops neighbor list is unavailable: %s", err) return False return True -def choose_nv_nlist_method(nloc: int, *, periodic: bool = True) -> str: +def choose_nv_nlist_method( + nloc: int, *, periodic: bool = True, device: torch.device | None = None +) -> str: """Choose the Toolkit-Ops neighbor method for a homogeneous batch. Parameters ---------- nloc Number of local atoms per frame. + periodic + Whether the batch is periodic. + device + Target device. CPU uses a lower cell-list threshold than CUDA because + the ``batch_naive`` method does not parallelize well there. Returns ------- str Toolkit-Ops method name. """ - threshold = ( - NV_CELL_LIST_THRESHOLD if periodic else NV_NONPERIODIC_CELL_LIST_THRESHOLD - ) + if device is not None and device.type == "cpu": + threshold = NV_CPU_CELL_LIST_THRESHOLD + elif periodic: + threshold = NV_CELL_LIST_THRESHOLD + else: + threshold = NV_NONPERIODIC_CELL_LIST_THRESHOLD if nloc >= threshold: return "batch_cell_list" return "batch_naive" @@ -103,7 +157,8 @@ def build( box: Any, rcut: float, sel: list[int], - ) -> tuple[Any, Any, Any, Any]: + return_mode: str = "extended", + ) -> tuple[Any, Any, Any, Any] | EdgeNeighborList: """Build the extended system and neighbor list. See :meth:`deepmd.dpmodel.utils.neighbor_list.NeighborList.build`. The @@ -133,7 +188,17 @@ def build( nf, dtype=torch.int32, device=device ).repeat_interleave(nloc) batch_ptr = torch.arange(nf + 1, dtype=torch.int32, device=device) * nloc - method = choose_nv_nlist_method(nloc, periodic=periodic) + method = choose_nv_nlist_method(nloc, periodic=periodic, device=device) + + # ``batch_naive`` otherwise derives ``max_atoms_per_system`` from + # ``batch_ptr`` with a ``.max().item()`` device->host sync on every + # call. Our batches are homogeneous (``nloc`` atoms per frame), so the + # value is known on the host; passing it explicitly removes that + # per-call sync. ``batch_cell_list`` neither accepts the argument nor + # has a ``**kwargs`` catch-all, so the override is guarded on method. + extra_nl_kwargs: dict[str, Any] = {} + if method == "batch_naive": + extra_nl_kwargs["max_atoms_per_system"] = int(nloc) # Grow the search capacity until all neighbors fit so the distance-sort # below selects the true nearest ``sum(sel)``. @@ -149,6 +214,7 @@ def build( max_neighbors=int(search_capacity), return_neighbor_list=False, wrap_positions=False, + **extra_nl_kwargs, ) if len(nlist_result) == 2: neighbor_matrix, num_neighbors = nlist_result @@ -166,6 +232,21 @@ def build( break search_capacity = max(max_found, _grow_search_capacity(search_capacity)) + if return_mode == "edges": + return edge_schema_from_neighbor_matrix( + coord=coord, + atype=atype, + cell=cell, + neighbor_matrix=neighbor_matrix, + num_neighbors=num_neighbors, + shifts=shifts, + rcut=float(rcut), + ) + if return_mode != "extended": + raise ValueError( + f"Unsupported neighbor-list return_mode: {return_mode!r}" + ) + extended_coord, extended_atype, mapping, nlist = _matrix_to_extended_inputs( coord=coord, atype=atype, diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 42e4181d65..97bba3d4a5 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -54,6 +54,9 @@ from deepmd.pt.utils.auto_batch_size import ( AutoBatchSize, ) +from deepmd.pt_expt.utils.edge_schema import ( + edge_schema_from_extended, +) from deepmd.pt_expt.utils.vesin_neighbor_list import ( VesinNeighborList, is_vesin_torch_available, @@ -76,6 +79,18 @@ def _reshape_charge_spin( ) from err +def _is_pt_backend_dpa4_params(model_params: dict[str, Any]) -> bool: + """Return whether a training checkpoint should be loaded by the pt backend.""" + model_type = str(model_params.get("type", "")).lower() + if model_type in {"sezm", "dpa4", "sezm_spin"}: + return True + descriptor = model_params.get("descriptor") + if isinstance(descriptor, dict): + descriptor_type = str(descriptor.get("type", "")).lower() + return descriptor_type in {"sezm", "dpa4"} + return False + + class DeepEval(DeepEvalBackend): """PyTorch Exportable backend implementation of DeepEval. @@ -408,6 +423,13 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: state_dict = head_state model_params = head_params + if _is_pt_backend_dpa4_params(model_params): + raise ValueError( + "DPA4/SeZM `.pt` checkpoints belong to the regular `pt` backend. " + "Use the `pt` backend for eager checkpoint inference, or export " + "the checkpoint to `.pt2` / `.pte` before loading it with `pt_expt`." + ) + model = get_model(deepcopy(model_params)).to(DEVICE) # Strip the `_CompiledModel` wrapper that pt_expt training applies @@ -502,6 +524,7 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: else None ), "is_spin": self._is_spin, + "lower_input_kind": "nlist", } if self._is_spin: self.metadata["ntypes_spin"] = model.spin.get_ntypes_spin() @@ -1095,7 +1118,29 @@ def _build_nlist_ase_single( return extended_coord, extended_atype, nlist, mapping - def _prepare_inputs( + @staticmethod + def _build_edge_inputs_from_nlist( + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Convert a padded neighbor list into the compact-edge schema.""" + nloc = nlist.shape[1] + schema = edge_schema_from_extended( + extended_coord, + extended_atype[:, :nloc], + nlist, + mapping, + ) + return ( + schema.edge_index, + schema.edge_vec, + schema.edge_scatter_index, + schema.edge_mask, + ) + + def _prepare_nlist_inputs( self, coords: np.ndarray, cells: np.ndarray | None, @@ -1104,7 +1149,7 @@ def _prepare_inputs( aparam: np.ndarray | None, charge_spin: np.ndarray | None = None, ) -> tuple: - """Prepare tensor inputs for model evaluation. + """Prepare the extended-coordinate and padded-neighbor-list inputs. Returns ------- @@ -1208,16 +1253,75 @@ def _prepare_inputs( natoms, ) - def _eval_model( + def _prepare_inputs( self, coords: np.ndarray, cells: np.ndarray | None, atom_types: np.ndarray, fparam: np.ndarray | None, aparam: np.ndarray | None, - request_defs: list[OutputVariableDef], charge_spin: np.ndarray | None = None, - ) -> tuple[np.ndarray, ...]: + ) -> tuple[tuple[torch.Tensor | None, ...], torch.Tensor, int, int]: + """Prepare lower-interface inputs and the output fold-back mapping.""" + if ( + self.metadata.get("lower_input_kind") == "edge_vec" + and self._nlist_builder is not None + and self.neighbor_list is None + ): + from deepmd.pt_expt.utils.env import ( + DEVICE, + ) + + nframes = coords.shape[0] + if len(atom_types.shape) == 1: + natoms = len(atom_types) + atom_types = np.tile(atom_types, nframes).reshape(nframes, -1) + else: + natoms = len(atom_types[0]) + coord_t = torch.tensor( + coords.reshape(nframes, natoms, 3), + dtype=torch.float64, + device=DEVICE, + ) + atype_t = torch.tensor(atom_types, dtype=torch.int64, device=DEVICE) + cells_t = ( + torch.tensor(cells, dtype=torch.float64, device=DEVICE) + if cells is not None + else None + ) + edge_schema = self._nlist_builder.build( + coord_t, + atype_t, + cells_t, + self._rcut, + self._sel, + return_mode="edges", + ) + fparam_t, aparam_t = self._prepare_optional_lower_inputs( + fparam, + aparam, + nframes, + natoms, + DEVICE, + ) + charge_spin_t = self._make_charge_spin_input(nframes, charge_spin) + model_inputs = ( + edge_schema.coord, + edge_schema.atype, + edge_schema.edge_index, + edge_schema.edge_vec, + edge_schema.edge_scatter_index, + edge_schema.edge_mask, + fparam_t, + aparam_t, + charge_spin_t, + ) + mapping_t = torch.arange(natoms, dtype=torch.int64, device=DEVICE).reshape( + 1, natoms + ) + mapping_t = mapping_t.expand(nframes, -1).contiguous() + return model_inputs, mapping_t, nframes, natoms + ( ext_coord_t, ext_atype_t, @@ -1228,17 +1332,99 @@ def _eval_model( charge_spin_t, nframes, natoms, - ) = self._prepare_inputs(coords, cells, atom_types, fparam, aparam, charge_spin) + ) = self._prepare_nlist_inputs( + coords, cells, atom_types, fparam, aparam, charge_spin + ) + if self.metadata.get("lower_input_kind") == "edge_vec": + edge_index_t, edge_vec_t, edge_scatter_t, edge_mask_t = ( + self._build_edge_inputs_from_nlist( + ext_coord_t, + ext_atype_t, + nlist_t, + mapping_t, + ) + ) + model_inputs = ( + ext_coord_t, + ext_atype_t[:, :natoms], + edge_index_t, + edge_vec_t, + edge_scatter_t, + edge_mask_t, + fparam_t, + aparam_t, + charge_spin_t, + ) + else: + model_inputs = ( + ext_coord_t, + ext_atype_t, + nlist_t, + mapping_t, + fparam_t, + aparam_t, + charge_spin_t, + ) + return model_inputs, mapping_t, nframes, natoms - # Call the model (forward_common_lower interface, internal keys) - model_inputs = ( - ext_coord_t, - ext_atype_t, - nlist_t, - mapping_t, - fparam_t, - aparam_t, - charge_spin_t, + def _prepare_optional_lower_inputs( + self, + fparam: np.ndarray | None, + aparam: np.ndarray | None, + nframes: int, + natoms: int, + device: torch.device, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Prepare optional frame and atomic parameters for lower interfaces.""" + if fparam is not None: + fparam_t = torch.tensor( + fparam.reshape(nframes, self.get_dim_fparam()), + dtype=torch.float64, + device=device, + ) + elif self.get_dim_fparam() > 0: + default_fp = self.metadata.get("default_fparam") + if default_fp is None: + raise ValueError( + f"fparam is required for this model (dim_fparam={self.get_dim_fparam()}) " + "but was not provided, and no default_fparam is stored in the model." + ) + fparam_t = ( + torch.tensor(default_fp, dtype=torch.float64, device=device) + .unsqueeze(0) + .expand(nframes, -1) + .contiguous() + ) + else: + fparam_t = None + + if aparam is not None: + aparam_t = torch.tensor( + aparam.reshape(nframes, natoms, self.get_dim_aparam()), + dtype=torch.float64, + device=device, + ) + elif self.get_dim_aparam() > 0: + raise ValueError( + f"aparam is required for this model (dim_aparam={self.get_dim_aparam()}) " + "but was not provided." + ) + else: + aparam_t = None + return fparam_t, aparam_t + + def _eval_model( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + fparam: np.ndarray | None, + aparam: np.ndarray | None, + request_defs: list[OutputVariableDef], + charge_spin: np.ndarray | None = None, + ) -> tuple[np.ndarray, ...]: + model_inputs, mapping_t, nframes, natoms = self._prepare_inputs( + coords, cells, atom_types, fparam, aparam, charge_spin ) if self._is_pt2: # AOTInductor's __call__ unflattens output using stored out_spec, @@ -1599,7 +1785,9 @@ def eval_descriptor( charge_spin_t, _nframes, _natoms, - ) = self._prepare_inputs(coords, cells, atom_types, fparam, aparam, charge_spin) + ) = self._prepare_nlist_inputs( + coords, cells, atom_types, fparam, aparam, charge_spin + ) with torch.no_grad(): descriptor, *_ = dp_am.descriptor( ext_coord_t, @@ -1668,7 +1856,9 @@ def eval_fitting_last_layer( charge_spin_t, _nframes, natoms, - ) = self._prepare_inputs(coords, cells, atom_types, fparam, aparam, charge_spin) + ) = self._prepare_nlist_inputs( + coords, cells, atom_types, fparam, aparam, charge_spin + ) with torch.no_grad(): descriptor, rot_mat, g2, h2, _sw = dp_am.descriptor( ext_coord_t, diff --git a/deepmd/pt_expt/utils/edge_schema.py b/deepmd/pt_expt/utils/edge_schema.py new file mode 100644 index 0000000000..5532bc9b6e --- /dev/null +++ b/deepmd/pt_expt/utils/edge_schema.py @@ -0,0 +1,306 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Edge-vector neighbor-list helpers for SeZM-style models.""" + +from __future__ import ( + annotations, +) + +import torch + +from deepmd.dpmodel.utils.neighbor_list import ( + EdgeNeighborList, +) + +_DUMMY_EDGE_COUNT = 2 + + +def _append_dummy_edges( + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_scatter_index: torch.Tensor, +) -> EdgeNeighborList: + """Append masked in-range edges so exported graphs never see empty inputs.""" + device = edge_index.device + dummy_index = torch.zeros( + (2, _DUMMY_EDGE_COUNT), + dtype=edge_index.dtype, + device=device, + ) + dummy_vec = torch.zeros( + (_DUMMY_EDGE_COUNT, 3), + dtype=edge_vec.dtype, + device=device, + ) + edge_index = torch.cat([edge_index, dummy_index], dim=1) + edge_vec = torch.cat([edge_vec, dummy_vec], dim=0) + edge_scatter_index = torch.cat([edge_scatter_index, dummy_index], dim=1) + edge_mask = torch.cat( + [ + torch.ones( + edge_vec.shape[0] - _DUMMY_EDGE_COUNT, dtype=torch.bool, device=device + ), + torch.zeros(_DUMMY_EDGE_COUNT, dtype=torch.bool, device=device), + ] + ) + # ``coord`` and ``atype`` are filled by public constructors. + return EdgeNeighborList( + coord=torch.empty(0, dtype=edge_vec.dtype, device=device), + atype=torch.empty(0, dtype=torch.long, device=device), + edge_index=edge_index, + edge_vec=edge_vec, + edge_scatter_index=edge_scatter_index, + edge_mask=edge_mask, + ) + + +def edge_schema_from_extended( + coord: torch.Tensor, + atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None, + *, + scatter_to_local: bool = False, +) -> EdgeNeighborList: + """Build the unified edge schema from an extended-coordinate neighbor list. + + This is the zero-shift form used by callers that already have periodic-image + coordinates, such as LAMMPS ghost atoms or the native fallback builder. + + Contract: ``nlist`` is assumed to be already truncated to the model cutoff + before this conversion is used. DeePMD's native builders select within + ``rcut``; exported LAMMPS nlist paths keep cutoff filtering inside the + traced model graph. + Unlike the candidate-list builders (:func:`edge_schema_from_neighbor_matrix`, + :func:`edge_schema_from_ij_shifts`), this function therefore applies no + ``edge_len <= rcut`` upper bound -- doing so would be a redundant op on the + native training hot path, and any residual out-of-range edge is already + zeroed by the descriptor's smooth cutoff envelope. + """ + nf, nloc, nsel = nlist.shape + device = coord.device + nall = coord.shape[1] + + neighbor_flat = nlist.reshape(-1) + dst_actual = ( + torch.arange(neighbor_flat.shape[0], device=device, dtype=torch.long) // nsel + ) + frame_idx = dst_actual // nloc + dst_local = dst_actual % nloc + valid_flat = neighbor_flat >= 0 + neighbor_safe = torch.where( + valid_flat, neighbor_flat, torch.zeros_like(neighbor_flat) + ) + neighbor_safe_2d = neighbor_safe.to(dtype=torch.long).view(nf, nloc * nsel) + + neighbor_coord = torch.gather( + coord, + 1, + neighbor_safe_2d.unsqueeze(-1).expand(-1, -1, 3), + ).reshape(-1, 3) + dst_coord = torch.gather( + coord[:, :nloc, :], + 1, + dst_local.view(nf, -1).unsqueeze(-1).expand(-1, -1, 3), + ).reshape(-1, 3) + edge_vec_all = neighbor_coord - dst_coord + edge_len2 = torch.sum(edge_vec_all * edge_vec_all, dim=-1) + + if mapping is None: + src_local = neighbor_safe.to(dtype=torch.long) + else: + src_local = torch.gather(mapping, 1, neighbor_safe_2d).reshape(-1) + src_actual = frame_idx * nloc + src_local.to(dtype=torch.long) + src_scatter = frame_idx * nall + neighbor_safe.to(dtype=torch.long) + dst_scatter = frame_idx * nall + dst_local + + # No ``edge_len2 <= rcut**2`` upper bound here: ``nlist`` is contractually + # cutoff-truncated by the caller (see the docstring). Only padding (-1), + # ghost-only neighbours, and coincident pairs are dropped. + edge_keep = valid_flat & (src_local >= 0) & (src_local < nloc) & (edge_len2 > 1e-10) + valid_idx = torch.nonzero(edge_keep, as_tuple=False).flatten() + edge_index = torch.stack( + [ + src_actual.index_select(0, valid_idx), + dst_actual.index_select(0, valid_idx), + ], + dim=0, + ) + if scatter_to_local: + edge_scatter_index = edge_index + else: + edge_scatter_index = torch.stack( + [ + src_scatter.index_select(0, valid_idx), + dst_scatter.index_select(0, valid_idx), + ], + dim=0, + ) + schema = _append_dummy_edges( + edge_index, + edge_vec_all.index_select(0, valid_idx), + edge_scatter_index, + ) + schema.coord = coord[:, :nloc, :].contiguous() if scatter_to_local else coord + schema.atype = atype[:, :nloc] + return schema + + +def edge_schema_from_neighbor_matrix( + coord: torch.Tensor, + atype: torch.Tensor, + cell: torch.Tensor | None, + neighbor_matrix: torch.Tensor, + num_neighbors: torch.Tensor, + shifts: torch.Tensor, + rcut: float, +) -> EdgeNeighborList: + """Build edge schema from a dense neighbor matrix and integer shifts.""" + nf, nloc = atype.shape[:2] + total_atoms, max_neighbors = neighbor_matrix.shape + device = coord.device + slot = torch.arange(max_neighbors, dtype=torch.long, device=device).expand( + total_atoms, max_neighbors + ) + valid = (slot < num_neighbors.unsqueeze(1)).reshape(-1) + edge_idx = torch.nonzero(valid, as_tuple=False).flatten() + if edge_idx.numel() == 0: + empty = _append_dummy_edges( + torch.zeros((2, 0), dtype=torch.long, device=device), + torch.zeros((0, 3), dtype=coord.dtype, device=device), + torch.zeros((2, 0), dtype=torch.long, device=device), + ) + empty.coord = coord + empty.atype = atype + return empty + + dst = edge_idx // max_neighbors + src = neighbor_matrix.reshape(-1).index_select(0, edge_idx).to(dtype=torch.long) + shift = shifts.reshape(-1, 3).index_select(0, edge_idx) + src_local = src % nloc + frame_idx = dst // nloc + src_actual = frame_idx * nloc + src_local + coord_flat = coord.reshape(nf * nloc, 3) + edge_vec_all = coord_flat.index_select(0, src_actual) - coord_flat.index_select( + 0, dst + ) + + if cell is not None: + shifted_idx = torch.nonzero( + torch.any(shift != 0, dim=1), as_tuple=False + ).flatten() + if shifted_idx.numel() > 0: + shift_cart = torch.bmm( + shift.index_select(0, shifted_idx).to(dtype=coord.dtype).unsqueeze(1), + cell.index_select(0, frame_idx.index_select(0, shifted_idx)), + ).squeeze(1) + edge_vec_all.index_add_(0, shifted_idx, shift_cart) + + edge_len2 = torch.sum(edge_vec_all * edge_vec_all, dim=-1) + edge_keep = (edge_len2 > 1e-10) & (edge_len2 <= float(rcut) * float(rcut)) + valid_idx = torch.nonzero(edge_keep, as_tuple=False).flatten() + schema = _append_dummy_edges( + torch.stack( + [ + src_actual.index_select(0, valid_idx), + dst.index_select(0, valid_idx), + ], + dim=0, + ), + edge_vec_all.index_select(0, valid_idx), + torch.stack( + [ + src_actual.index_select(0, valid_idx), + dst.index_select(0, valid_idx), + ], + dim=0, + ), + ) + schema.coord = coord + schema.atype = atype + return schema + + +def edge_schema_from_ij_shifts( + positions: torch.Tensor, + atype: torch.Tensor, + cell: torch.Tensor | None, + ii: torch.Tensor, + jj: torch.Tensor, + shifts: torch.Tensor, + rcut: float, +) -> EdgeNeighborList: + """Build a single-frame edge schema from ``vesin`` ``(i, j, S)`` output.""" + device = positions.device + nloc = positions.shape[0] + if ii.numel() == 0: + empty = _append_dummy_edges( + torch.zeros((2, 0), dtype=torch.long, device=device), + torch.zeros((0, 3), dtype=positions.dtype, device=device), + torch.zeros((2, 0), dtype=torch.long, device=device), + ) + empty.coord = positions.reshape(1, nloc, 3) + empty.atype = atype.reshape(1, nloc) + return empty + + ii = ii.to(dtype=torch.long) + jj = jj.to(dtype=torch.long) + shifts = shifts.to(dtype=positions.dtype) + edge_vec_all = positions.index_select(0, jj) - positions.index_select(0, ii) + if cell is not None: + shifted_idx = torch.nonzero( + torch.any(shifts != 0, dim=1), as_tuple=False + ).flatten() + if shifted_idx.numel() > 0: + edge_vec_all.index_add_( + 0, + shifted_idx, + shifts.index_select(0, shifted_idx) @ cell, + ) + edge_len2 = torch.sum(edge_vec_all * edge_vec_all, dim=-1) + edge_keep = (edge_len2 > 1e-10) & (edge_len2 <= float(rcut) * float(rcut)) + valid_idx = torch.nonzero(edge_keep, as_tuple=False).flatten() + schema = _append_dummy_edges( + torch.stack( + [ + jj.index_select(0, valid_idx), + ii.index_select(0, valid_idx), + ], + dim=0, + ), + edge_vec_all.index_select(0, valid_idx), + torch.stack( + [ + jj.index_select(0, valid_idx), + ii.index_select(0, valid_idx), + ], + dim=0, + ), + ) + schema.coord = positions.reshape(1, nloc, 3) + schema.atype = atype.reshape(1, nloc) + return schema + + +def merge_frame_edge_schemas(frames: list[EdgeNeighborList]) -> EdgeNeighborList: + """Merge per-frame local edge schemas into a batched flattened schema.""" + if not frames: + raise ValueError("at least one frame schema is required") + coord = torch.cat([frame.coord for frame in frames], dim=0) + atype = torch.cat([frame.atype for frame in frames], dim=0) + nloc = coord.shape[1] + edge_indices: list[torch.Tensor] = [] + edge_vecs: list[torch.Tensor] = [] + scatter_indices: list[torch.Tensor] = [] + for frame_idx, frame in enumerate(frames): + real = frame.edge_mask + offset = frame_idx * nloc + edge_indices.append(frame.edge_index[:, real] + offset) + scatter_indices.append(frame.edge_scatter_index[:, real] + offset) + edge_vecs.append(frame.edge_vec[real]) + edge_index = torch.cat(edge_indices, dim=1) + edge_vec = torch.cat(edge_vecs, dim=0) + edge_scatter_index = torch.cat(scatter_indices, dim=1) + schema = _append_dummy_edges(edge_index, edge_vec, edge_scatter_index) + schema.coord = coord + schema.atype = atype + return schema diff --git a/deepmd/pt_expt/utils/vesin_neighbor_list.py b/deepmd/pt_expt/utils/vesin_neighbor_list.py index 1d55548c56..6b1a165b98 100644 --- a/deepmd/pt_expt/utils/vesin_neighbor_list.py +++ b/deepmd/pt_expt/utils/vesin_neighbor_list.py @@ -25,8 +25,13 @@ import torch from deepmd.dpmodel.utils.neighbor_list import ( + EdgeNeighborList, NeighborList, ) +from deepmd.pt_expt.utils.edge_schema import ( + edge_schema_from_ij_shifts, + merge_frame_edge_schemas, +) def is_vesin_torch_available() -> bool: @@ -54,7 +59,8 @@ def build( box: Any, rcut: float, sel: list[int], - ) -> tuple[Any, Any, Any, Any]: + return_mode: str = "extended", + ) -> tuple[Any, Any, Any, Any] | EdgeNeighborList: """Build the extended system + candidate neighbor list with vesin. See :meth:`deepmd.dpmodel.utils.neighbor_list.NeighborList.build`. The @@ -81,6 +87,31 @@ def build( if box_t is not None: box_t = box_t.reshape(nframes, 3, 3) + if return_mode == "edges": + frame_edges = [ + _build_single_edges( + coord_t[ff], + box_t[ff] if box_t is not None else None, + atype_t[ff], + rcut, + sel, + ) + for ff in range(nframes) + ] + schema = merge_frame_edge_schemas(frame_edges) + if is_numpy: + return EdgeNeighborList( + coord=schema.coord.detach().cpu().numpy(), + atype=schema.atype.cpu().numpy(), + edge_index=schema.edge_index.cpu().numpy(), + edge_vec=schema.edge_vec.detach().cpu().numpy(), + edge_scatter_index=schema.edge_scatter_index.cpu().numpy(), + edge_mask=schema.edge_mask.cpu().numpy(), + ) + return schema + if return_mode != "extended": + raise ValueError(f"Unsupported neighbor-list return_mode: {return_mode!r}") + frame_results = [ _build_single( coord_t[ff], @@ -230,3 +261,50 @@ def _build_single( ) return extended_coord, extended_atype, nlist, mapping + + +def _build_single_edges( + positions: torch.Tensor, + cell: torch.Tensor | None, + atype: torch.Tensor, + rcut: float, + sel: list[int], +) -> EdgeNeighborList: + """Single-frame ``vesin`` output converted directly to edge vectors.""" + import vesin.torch + + device = positions.device + nsel = sum(sel) + nloc = positions.shape[0] + if nloc == 0: + return edge_schema_from_ij_shifts( + positions, + atype, + cell, + torch.zeros(0, dtype=torch.int64, device=device), + torch.zeros(0, dtype=torch.int64, device=device), + torch.zeros(0, 3, dtype=positions.dtype, device=device), + rcut, + ) + + periodic = cell is not None + box = ( + cell if periodic else torch.zeros((3, 3), dtype=positions.dtype, device=device) + ) + nl = vesin.torch.NeighborList(cutoff=rcut, full_list=True) + with torch.device(device): + ii, jj, ss = nl.compute( + points=positions.detach(), + box=box.detach(), + periodic=periodic, + quantities="ijS", + ) + return edge_schema_from_ij_shifts( + positions=positions, + atype=atype, + cell=box if periodic else None, + ii=ii, + jj=jj, + shifts=ss, + rcut=rcut, + ) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 48c711a10b..ba1a2cb347 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -366,10 +366,10 @@ def descrpt_se_a_args() -> list[Argument]: ) def descrpt_se_zm_args() -> list[Argument]: # Follows exact order of docstring in sezm.py DescrptSeZM class - doc_sel = 'The maximum number of neighbors. It can be:\n\n\ - - `int`: the total maximum number of neighbors within `rcut` (all types combined)\n\n\ - - `list[int]`: sel[i] specifies the maximum number of type-i neighbors within `rcut`\n\n\ - - `str`: Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wrapped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".' + doc_sel = 'The neighbor-search capacity, with a default of 256. The conservative energy path keeps every neighbor within `rcut` regardless of this value, so `sel` only sets the initial search capacity of the O(N) `nvalchemiops` builder (which grows on demand) and never truncates the energy-path neighbor list. The denoising (`dens`) and spin paths still cap the neighbor list at `sum(sel)`, so for those modes `sel` must cover the true maximum neighbor count. It can be:\n\n\ + - `int`: the total capacity across all atom types.\n\n\ + - `list[int]`: `sel[i]` is the type-i capacity; only `sum(sel)` is used.\n\n\ + - `str`: "auto" or "auto:factor" sizes `sel` from the training data via neighbor statistics (`factor` larger than 1, rounded up to a multiple of 4; "auto" equals "auto:1.1"). This requires the neighbor-statistics pass and is therefore unavailable under `--skip-neighbor-stat`.' doc_rcut = "The cut-off radius." doc_env_exp = ( "C^3 cutoff envelope exponents `[rbf_env_exp, edge_env_exp]`. " @@ -636,9 +636,7 @@ def descrpt_se_zm_args() -> list[Argument]: doc_trainable = "If the parameters in the descriptor are trainable." doc_seed = "Random seed for parameter initialization." return [ - Argument( - "sel", [int, list[int], str], optional=True, default="auto", doc=doc_sel - ), + Argument("sel", [int, list[int], str], optional=True, default=256, doc=doc_sel), Argument("rcut", float, optional=True, default=6.0, doc=doc_rcut), Argument( "env_exp", diff --git a/examples/water/dpa4/input.json b/examples/water/dpa4/input.json index 415f6a5be0..9eb48e201c 100644 --- a/examples/water/dpa4/input.json +++ b/examples/water/dpa4/input.json @@ -7,7 +7,6 @@ "H" ], "descriptor": { - "sel": 120, "rcut": 6.0, "channels": 32, "n_radial": 16, diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index 461fc6f33c..669fa99e71 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -307,9 +307,12 @@ class DeepPotPTExpt : public DeepPotBackend { std::vector output_keys; // sorted internal output key names bool do_atomic_virial; // whether model was exported with atomic virial corr int nnei; // expected nlist nnei dimension (= sum(sel)) + bool lower_input_is_edge_ = false; NeighborListData nlist_data; - at::Tensor mapping_tensor; // cached mapping tensor (LAMMPS path) - at::Tensor firstneigh_tensor; // cached nlist tensor (LAMMPS path) + at::Tensor mapping_tensor; // cached mapping tensor (LAMMPS path) + at::Tensor firstneigh_tensor; // cached nlist tensor (LAMMPS path) + at::Tensor edge_index_tensor; // cached local edge graph (LAMMPS path) + at::Tensor edge_index_ext_tensor; // cached extended edge graph (LAMMPS path) std::unique_ptr loader; // Optional second AOTInductor artifact for the multi-rank GNN code // path (Phase 4). Loaded only if the .pt2 metadata reports @@ -384,6 +387,17 @@ class DeepPotPTExpt : public DeepPotBackend { const torch::Tensor& aparam, const torch::Tensor& charge_spin); + std::vector run_model_edges( + const torch::Tensor& coord, + const torch::Tensor& atype, + const torch::Tensor& edge_index, + const torch::Tensor& edge_vec, + const torch::Tensor& edge_scatter_index, + const torch::Tensor& edge_mask, + const torch::Tensor& fparam, + const torch::Tensor& aparam, + const torch::Tensor& charge_spin); + /** * @brief Run the with-comm .pt2 artifact with comm tensors appended. * diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index 61b823ae62..637ade70b4 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -4,7 +4,9 @@ #ifdef BUILD_PYTORCH #include +#include #include +#include #include #include "common.h" @@ -140,6 +142,213 @@ inline torch::Tensor createNlistTensor( return flat_tensor.view({1, nloc, nnei}); } +struct EdgeTensorPack { + torch::Tensor edge_index; + torch::Tensor edge_vec; + torch::Tensor edge_index_ext; + torch::Tensor edge_mask; +}; + +/** + * @brief Build compact edge tensors from a neighbor list. + * + * The returned tensors are aligned by edge: + * - edge_index uses flattened local-atom indices and drives descriptor message + * passing. + * - edge_index_ext uses flattened extended-atom indices and drives force and + * virial scatter. + * - edge_mask marks physical edges. When geometry is requested, two masked + * dummy edges are appended so the exported graph never observes a singular + * edge dimension. + * + * @param nlist Neighbor-list rows. By default row i is center atom i; callers + * that compact LAMMPS rows must pass row_centers. + * @param coord Extended coordinates shaped as nall x 3. + * @param mapping Extended-to-local atom map with length nall. + * @param nloc Number of local atoms. + * @param nall Number of extended atoms. + * @param device Target device for the returned tensors. + * @param with_geometry Whether to also materialize edge_vec, edge_mask and + * model-input dummy edges. The cached LAMMPS path keeps only the real skin + * topology and compacts it on-device every step, so it passes ``false``. + * The returned edge_vec and edge_mask are left undefined in that case. + * @param row_centers Optional center atom index for each neighbor-list row. + */ +template +inline EdgeTensorPack createEdgeTensors( + const std::vector>& nlist, + const std::vector& coord, + const std::vector& mapping, + const int nloc, + const int nall, + const torch::Device& device, + const bool with_geometry = true, + const std::vector* row_centers = nullptr) { + std::vector src; + std::vector dst; + std::vector src_ext; + std::vector dst_ext; + std::vector edge_vec; + size_t reserve_size = with_geometry ? 2 : 0; + for (const auto& row : nlist) { + reserve_size += row.size(); + } + src.reserve(reserve_size); + dst.reserve(reserve_size); + src_ext.reserve(reserve_size); + dst_ext.reserve(reserve_size); + if (with_geometry) { + edge_vec.reserve(reserve_size * 3); + } + + // Real edges: use row_centers when LAMMPS has compacted away empty rows. + for (int ii = 0; ii < static_cast(nlist.size()); ++ii) { + if (row_centers != nullptr && + static_cast(ii) >= row_centers->size()) { + continue; + } + const int center = + row_centers == nullptr ? ii : (*row_centers)[static_cast(ii)]; + if (center < 0 || center >= nloc || center >= nall) { + continue; + } + const size_t center_offset = static_cast(center) * 3; + for (const int jj : nlist[ii]) { + if (jj < 0 || jj >= nall) { + continue; + } + const std::int64_t src_local = mapping[static_cast(jj)]; + if (src_local < 0 || src_local >= nloc) { + continue; + } + const size_t neighbor_offset = static_cast(jj) * 3; + const VALUETYPE dx = coord[neighbor_offset] - coord[center_offset]; + const VALUETYPE dy = + coord[neighbor_offset + 1] - coord[center_offset + 1]; + const VALUETYPE dz = + coord[neighbor_offset + 2] - coord[center_offset + 2]; + const VALUETYPE rr = dx * dx + dy * dy + dz * dz; + if (rr <= static_cast(1e-10)) { + continue; + } + src.push_back(src_local); + dst.push_back(center); + src_ext.push_back(jj); + dst_ext.push_back(center); + if (with_geometry) { + edge_vec.push_back(dx); + edge_vec.push_back(dy); + edge_vec.push_back(dz); + } + } + } + + const size_t real_edges = src.size(); + if (with_geometry) { + // Dummy edges keep exported edge tensors non-empty without affecting + // output. + for (int ii = 0; ii < 2; ++ii) { + src.push_back(0); + dst.push_back(0); + src_ext.push_back(0); + dst_ext.push_back(0); + edge_vec.push_back(0); + edge_vec.push_back(0); + edge_vec.push_back(0); + } + } + const size_t nedge = src.size(); + std::vector edge_index(2 * nedge); + std::vector edge_index_ext(2 * nedge); + // Materialize local-owner and extended scatter index spaces side by side. + for (size_t ii = 0; ii < nedge; ++ii) { + edge_index[ii] = src[ii]; + edge_index[nedge + ii] = dst[ii]; + edge_index_ext[ii] = src_ext[ii]; + edge_index_ext[nedge + ii] = dst_ext[ii]; + } + + auto int_options = torch::TensorOptions().dtype(torch::kInt64); + EdgeTensorPack pack; + if (nedge == 0) { + pack.edge_index = torch::empty({2, 0}, int_options).to(device); + pack.edge_index_ext = torch::empty({2, 0}, int_options).to(device); + } else { + pack.edge_index = + torch::from_blob(edge_index.data(), + {2, static_cast(nedge)}, int_options) + .clone() + .to(device); + pack.edge_index_ext = + torch::from_blob(edge_index_ext.data(), + {2, static_cast(nedge)}, int_options) + .clone() + .to(device); + } + if (with_geometry) { + pack.edge_vec = + torch::from_blob( + edge_vec.data(), {static_cast(nedge), 3}, + torch::TensorOptions().dtype(std::is_same::value + ? torch::kFloat32 + : torch::kFloat64)) + .clone() + .to(device); + std::vector edge_mask(nedge, 0); + std::fill(edge_mask.begin(), edge_mask.begin() + real_edges, + static_cast(1)); + pack.edge_mask = + torch::from_blob(edge_mask.data(), {static_cast(nedge)}, + torch::TensorOptions().dtype(torch::kUInt8)) + .clone() + .to(torch::kBool) + .to(device); + } + return pack; +} + +/** + * @brief Compact a cached LAMMPS skin topology to the current cutoff edge set. + * + * LAMMPS rebuilds neighbor topology only when its skin list is refreshed. The + * SeZM lower graph, however, should see only the current model-cutoff edges. + * This helper keeps the cached skin topology immutable, recomputes edge + * vectors from the current coordinates on the target device, filters by + * ``rr <= rcut**2``, then appends the two masked dummy edges required by the + * exported graph contract. + */ +inline EdgeTensorPack compactEdgeTensors(const torch::Tensor& edge_index, + const torch::Tensor& edge_index_ext, + const torch::Tensor& coord, + const double rcut) { + const auto coord_flat = coord.reshape({-1, 3}); + const auto src_ext = edge_index_ext.select(0, 0); + const auto dst_ext = edge_index_ext.select(0, 1); + const auto edge_vec_all = + coord_flat.index_select(0, src_ext) - coord_flat.index_select(0, dst_ext); + const auto rr = (edge_vec_all * edge_vec_all).sum(1); + const auto keep = (rr > 1e-10) & (rr <= rcut * rcut); + const auto real_idx = torch::nonzero(keep).reshape({-1}); + + EdgeTensorPack pack; + const auto real_edge_index = edge_index.index_select(1, real_idx); + const auto real_edge_index_ext = edge_index_ext.index_select(1, real_idx); + const auto real_edge_vec = edge_vec_all.index_select(0, real_idx); + + const auto dummy_index = torch::zeros({2, 2}, edge_index.options()); + const auto dummy_vec = torch::zeros({2, 3}, edge_vec_all.options()); + pack.edge_index = torch::cat({real_edge_index, dummy_index}, 1); + pack.edge_index_ext = torch::cat({real_edge_index_ext, dummy_index}, 1); + pack.edge_vec = torch::cat({real_edge_vec, dummy_vec}, 0); + + const auto real_mask = torch::ones( + {real_idx.size(0)}, + torch::TensorOptions().dtype(torch::kBool).device(coord.device())); + const auto dummy_mask = torch::zeros({2}, real_mask.options()); + pack.edge_mask = torch::cat({real_mask, dummy_mask}, 0); + return pack; +} + } // namespace deepmd #endif // BUILD_PYTORCH diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 394ad89301..673468a172 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -151,6 +151,13 @@ void DeepPotPTExpt::init(const std::string& model, nnei += v.as_int(); } } + if (metadata.obj_val.count("lower_input_kind")) { + const std::string lower_input_kind = + metadata["lower_input_kind"].as_string(); + lower_input_is_edge_ = lower_input_kind == "edge_vec"; + } else { + lower_input_is_edge_ = false; + } type_map.clear(); for (const auto& v : metadata["type_map"].as_array()) { @@ -258,6 +265,30 @@ std::vector DeepPotPTExpt::run_model( return loader->run(inputs); } +std::vector DeepPotPTExpt::run_model_edges( + const torch::Tensor& coord, + const torch::Tensor& atype, + const torch::Tensor& edge_index, + const torch::Tensor& edge_vec, + const torch::Tensor& edge_scatter_index, + const torch::Tensor& edge_mask, + const torch::Tensor& fparam, + const torch::Tensor& aparam, + const torch::Tensor& charge_spin) { + std::vector inputs = { + coord, atype, edge_index, edge_vec, edge_scatter_index, edge_mask}; + if (dfparam > 0) { + inputs.push_back(fparam); + } + if (daparam > 0) { + inputs.push_back(aparam); + } + if (dchgspin > 0) { + inputs.push_back(charge_spin); + } + return loader->run(inputs); +} + std::vector DeepPotPTExpt::run_model_with_comm( const torch::Tensor& coord, const torch::Tensor& atype, @@ -428,14 +459,14 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, // LAMMPS sets ago=0 on every nlist rebuild (neighbor rebuild, re-partition, // atom exchange between subdomains), so `ago > 0` implies the cached // mapping and nlist tensors are still valid. Rebuild only on ago==0. + std::vector mapping; if (ago == 0) { nlist_data.copy_from_nlist(lmp_list, nall - nghost); nlist_data.shuffle_exclude_empty(fwd_map); - nlist_data.padding(); // Rebuild mapping tensor if (lmp_list.mapping) { - std::vector mapping(nall_real); + mapping.resize(nall_real); for (int ii = 0; ii < nall_real; ii++) { mapping[ii] = fwd_map[lmp_list.mapping[bkw_map[ii]]]; } @@ -453,7 +484,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, // features via border_op and ignores this tensor for ghost // gather — see deepmd/pt_expt/descriptor/ // repflows.py::_exchange_ghosts). - std::vector mapping(nall_real); + mapping.resize(nall_real); for (int ii = 0; ii < nall_real; ii++) { mapping[ii] = ii; } @@ -463,9 +494,22 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, .to(device); } - // Flatten raw nlist — the .pt2 model sorts by distance on-device. - firstneigh_tensor = - createNlistTensor(nlist_data.jlist, nnei).to(torch::kInt64).to(device); + if (lower_input_is_edge_) { + // Cache only the real skin topology. The model-cutoff topology and + // model-input dummy edges are rebuilt on-device from current coordinates + // every step, so rcut-crossing skin atoms are handled without carrying + // out-of-cutoff edges into the exported graph. + const auto edge_tensors = + createEdgeTensors(nlist_data.jlist, dcoord, mapping, nloc, nall_real, + device, /*with_geometry=*/false, &nlist_data.ilist); + edge_index_tensor = edge_tensors.edge_index; + edge_index_ext_tensor = edge_tensors.edge_index_ext; + } else { + nlist_data.padding(); + firstneigh_tensor = createNlistTensor(nlist_data.jlist, nnei) + .to(torch::kInt64) + .to(device); + } } // Build fparam/aparam tensors (cast to float64 for the model) @@ -558,6 +602,12 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, std::vector remapped_sendlist_ptrs; std::vector remapped_sendnum, remapped_recvnum; if (use_with_comm) { + if (lower_input_is_edge_) { + throw deepmd::deepmd_exception( + "SeZM edge-schema .pt2 inference requires the regular single-rank " + "AOTInductor artifact. Multi-rank inference must use an artifact " + "whose lower input schema includes explicit communication tensors."); + } bool has_null_atoms = (nall_real < nall); std::vector comm_tensors; if (has_null_atoms) { @@ -574,9 +624,20 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, coord_Tensor, atype_Tensor, firstneigh_tensor, mapping_tensor, fparam_tensor, aparam_tensor, charge_spin_tensor, comm_tensors); } else { - flat_outputs = - run_model(coord_Tensor, atype_Tensor, firstneigh_tensor, mapping_tensor, - fparam_tensor, aparam_tensor, charge_spin_tensor); + if (lower_input_is_edge_) { + const auto edge_tensors = + compactEdgeTensors(edge_index_tensor, edge_index_ext_tensor, + coord_Tensor, static_cast(rcut)); + flat_outputs = + run_model_edges(coord_Tensor, atype_Tensor.slice(1, 0, nloc), + edge_tensors.edge_index, edge_tensors.edge_vec, + edge_tensors.edge_index_ext, edge_tensors.edge_mask, + fparam_tensor, aparam_tensor, charge_spin_tensor); + } else { + flat_outputs = run_model(coord_Tensor, atype_Tensor, firstneigh_tensor, + mapping_tensor, fparam_tensor, aparam_tensor, + charge_spin_tensor); + } } // Map flat outputs to internal keys @@ -790,14 +851,20 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, torch::from_blob(atype_64.data(), {1, nall}, int_options) .clone() .to(device); - // Flatten raw nlist — the .pt2 model sorts by distance on-device. - at::Tensor nlist_tensor = - createNlistTensor(nlist_raw, nnei).to(torch::kInt64).to(device); std::vector mapping_64(mapping_vec.begin(), mapping_vec.end()); at::Tensor mapping_tensor = torch::from_blob(mapping_64.data(), {1, nall}, int_options) .clone() .to(device); + at::Tensor nlist_tensor; + EdgeTensorPack edge_tensors; + if (lower_input_is_edge_) { + edge_tensors = createEdgeTensors(nlist_raw, coord_cpy_d, mapping_64, nloc, + nall, device); + } else { + nlist_tensor = + createNlistTensor(nlist_raw, nnei).to(torch::kInt64).to(device); + } // Build fparam/aparam tensors (cast to float64 for the model) auto valuetype_options = std::is_same::value @@ -873,9 +940,18 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, } // 5. Run the .pt2 model - auto flat_outputs = - run_model(coord_Tensor, atype_Tensor, nlist_tensor, mapping_tensor, - fparam_tensor, aparam_tensor, charge_spin_tensor); + std::vector flat_outputs; + if (lower_input_is_edge_) { + flat_outputs = + run_model_edges(coord_Tensor, atype_Tensor.slice(1, 0, nloc), + edge_tensors.edge_index, edge_tensors.edge_vec, + edge_tensors.edge_index_ext, edge_tensors.edge_mask, + fparam_tensor, aparam_tensor, charge_spin_tensor); + } else { + flat_outputs = + run_model(coord_Tensor, atype_Tensor, nlist_tensor, mapping_tensor, + fparam_tensor, aparam_tensor, charge_spin_tensor); + } // 6. Map flat outputs to internal keys std::map output_map; diff --git a/source/api_cc/tests/test_deeppot_pt.cc b/source/api_cc/tests/test_deeppot_pt.cc index 7f527296b1..0560e976dd 100644 --- a/source/api_cc/tests/test_deeppot_pt.cc +++ b/source/api_cc/tests/test_deeppot_pt.cc @@ -411,6 +411,88 @@ TYPED_TEST(TestInferDeepPotAPt, cpu_lmp_nlist_2rc) { } } +TYPED_TEST(TestInferDeepPotAPt, cpu_lmp_nlist_skin_below_model_width) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_f = this->expected_f; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + float rc = dp.cutoff(); + int nloc = coord.size() / 3; + std::vector coord_cpy; + std::vector atype_cpy, mapping; + std::vector > nlist_wide; + _build_nlist(nlist_wide, coord_cpy, atype_cpy, mapping, coord, + atype, box, rc * 2); + + const double rc2 = static_cast(rc) * static_cast(rc); + bool has_skin_neighbor = false; + std::vector > nlist_skin(nlist_wide.size()); + for (size_t ii = 0; ii < nlist_wide.size(); ++ii) { + bool added_skin_neighbor = false; + for (const int jj : nlist_wide[ii]) { + const double dx = static_cast(coord_cpy[jj * 3]) - + static_cast(coord_cpy[ii * 3]); + const double dy = static_cast(coord_cpy[jj * 3 + 1]) - + static_cast(coord_cpy[ii * 3 + 1]); + const double dz = static_cast(coord_cpy[jj * 3 + 2]) - + static_cast(coord_cpy[ii * 3 + 2]); + const double rr = dx * dx + dy * dy + dz * dz; + if (rr <= rc2) { + nlist_skin[ii].push_back(jj); + } else if (!added_skin_neighbor) { + nlist_skin[ii].push_back(jj); + added_skin_neighbor = true; + has_skin_neighbor = true; + } + } + } + ASSERT_TRUE(has_skin_neighbor); + + int nall = coord_cpy.size() / 3; + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_skin); + + double ener; + std::vector force_(nall * 3, 0.0), virial(9, 0.0); + dp.compute(ener, force_, virial, coord_cpy, atype_cpy, box, nall - nloc, + inlist, 0); + std::vector force; + _fold_back(force, force_, mapping, nloc, nall, 3); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } + + ener = 0.; + std::fill(force_.begin(), force_.end(), 0.0); + std::fill(virial.begin(), virial.end(), 0.0); + dp.compute(ener, force_, virial, coord_cpy, atype_cpy, box, nall - nloc, + inlist, 1); + _fold_back(force, force_, mapping, nloc, nall, 3); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + TYPED_TEST(TestInferDeepPotAPt, cpu_lmp_nlist_type_sel) { using VALUETYPE = TypeParam; std::vector& coord = this->coord; diff --git a/source/api_cc/tests/test_deeppot_ptexpt.cc b/source/api_cc/tests/test_deeppot_ptexpt.cc index cbf1303633..a4b925849c 100644 --- a/source/api_cc/tests/test_deeppot_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_ptexpt.cc @@ -388,6 +388,73 @@ TYPED_TEST(TestInferDeepPotAPtExpt, cpu_lmp_nlist_2rc) { } } +TYPED_TEST(TestInferDeepPotAPtExpt, cpu_lmp_nlist_skin_below_model_width) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_f = this->expected_f; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + float rc = dp.cutoff(); + int nloc = coord.size() / 3; + std::vector coord_cpy; + std::vector atype_cpy, mapping; + std::vector > nlist_wide; + _build_nlist(nlist_wide, coord_cpy, atype_cpy, mapping, coord, + atype, box, rc * 2); + + const double rc2 = static_cast(rc) * static_cast(rc); + bool has_skin_neighbor = false; + std::vector > nlist_skin(nlist_wide.size()); + for (size_t ii = 0; ii < nlist_wide.size(); ++ii) { + bool added_skin_neighbor = false; + for (const int jj : nlist_wide[ii]) { + const double dx = static_cast(coord_cpy[jj * 3]) - + static_cast(coord_cpy[ii * 3]); + const double dy = static_cast(coord_cpy[jj * 3 + 1]) - + static_cast(coord_cpy[ii * 3 + 1]); + const double dz = static_cast(coord_cpy[jj * 3 + 2]) - + static_cast(coord_cpy[ii * 3 + 2]); + const double rr = dx * dx + dy * dy + dz * dz; + if (rr <= rc2) { + nlist_skin[ii].push_back(jj); + } else if (!added_skin_neighbor) { + nlist_skin[ii].push_back(jj); + added_skin_neighbor = true; + has_skin_neighbor = true; + } + } + } + ASSERT_TRUE(has_skin_neighbor); + + int nall = coord_cpy.size() / 3; + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_skin); + + double ener; + std::vector force_(nall * 3, 0.0), virial(9, 0.0); + dp.compute(ener, force_, virial, coord_cpy, atype_cpy, box, nall - nloc, + inlist, 0); + std::vector force; + _fold_back(force, force_, mapping, nloc, nall, 3); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + TYPED_TEST(TestInferDeepPotAPtExpt, cpu_lmp_nlist_oversized) { using VALUETYPE = TypeParam; std::vector& coord = this->coord; diff --git a/source/api_cc/tests/test_neighbor_list_data.cc b/source/api_cc/tests/test_neighbor_list_data.cc index 3e5198d30d..3ddf8dcc26 100644 --- a/source/api_cc/tests/test_neighbor_list_data.cc +++ b/source/api_cc/tests/test_neighbor_list_data.cc @@ -8,6 +8,9 @@ #include #include "common.h" +#ifdef BUILD_PYTORCH +#include "commonPT.h" +#endif #include "neighbor_list.h" namespace deepmd { @@ -136,4 +139,65 @@ TEST(TestNeighborListData, RoundTripWithEmptyRows) { EXPECT_EQ(out.numneigh[3], 1); } +#ifdef BUILD_PYTORCH +TEST(TestEdgeTensorPack, CreateEdgeTensorsUsesRowCenters) { + const torch::Device device(torch::kCPU); + const std::vector> nlist = {{0}, {1}}; + const std::vector centers = {2, 0}; + const std::vector coord = { + 0.0, 0.0, 0.0, // atom 0 + 1.0, 0.0, 0.0, // atom 1 + 2.0, 0.0, 0.0, // atom 2 + }; + const std::vector mapping = {0, 1, 2}; + + const auto pack = createEdgeTensors(nlist, coord, mapping, 3, 3, device, + /*with_geometry=*/true, ¢ers); + + ASSERT_EQ(pack.edge_index.size(1), 4); + EXPECT_EQ(pack.edge_index.select(0, 0).select(0, 0).item(), 0); + EXPECT_EQ(pack.edge_index.select(0, 1).select(0, 0).item(), 2); + EXPECT_EQ(pack.edge_index_ext.select(0, 1).select(0, 0).item(), 2); + EXPECT_DOUBLE_EQ(pack.edge_vec.select(0, 0).select(0, 0).item(), + -2.0); + EXPECT_EQ(pack.edge_index.select(0, 0).select(0, 1).item(), 1); + EXPECT_EQ(pack.edge_index.select(0, 1).select(0, 1).item(), 0); + EXPECT_DOUBLE_EQ(pack.edge_vec.select(0, 1).select(0, 0).item(), 1.0); +} + +TEST(TestEdgeTensorPack, CompactFiltersSkinTopologyAndAppendsDummies) { + const torch::Device device(torch::kCPU); + const std::vector> nlist = {{1, 2}, {0}}; + const std::vector coord = { + 0.0, 0.0, 0.0, // atom 0 + 0.5, 0.0, 0.0, // atom 1, inside cutoff + 2.0, 0.0, 0.0, // atom 2, skin-only ghost + }; + const std::vector mapping = {0, 1, 0}; + + const auto skin_topology = + createEdgeTensors(nlist, coord, mapping, 2, 3, device, + /*with_geometry=*/false); + ASSERT_EQ(skin_topology.edge_index.size(1), 3); + ASSERT_EQ(skin_topology.edge_index_ext.size(1), 3); + EXPECT_FALSE(skin_topology.edge_vec.defined()); + EXPECT_FALSE(skin_topology.edge_mask.defined()); + + const auto coord_tensor = + torch::tensor(coord, torch::TensorOptions().dtype(torch::kFloat64)) + .view({1, 3, 3}); + const auto compact = + compactEdgeTensors(skin_topology.edge_index, skin_topology.edge_index_ext, + coord_tensor, /*rcut=*/1.0); + + ASSERT_EQ(compact.edge_index.size(1), 4); + ASSERT_EQ(compact.edge_index_ext.size(1), 4); + ASSERT_EQ(compact.edge_vec.size(0), 4); + ASSERT_EQ(compact.edge_mask.size(0), 4); + EXPECT_EQ(compact.edge_mask.sum().item(), 2); + EXPECT_FALSE(compact.edge_mask.select(0, 2).item()); + EXPECT_FALSE(compact.edge_mask.select(0, 3).item()); +} +#endif + } // namespace deepmd diff --git a/source/tests/pt/model/test_descriptor_sezm.py b/source/tests/pt/model/test_descriptor_sezm.py index ec5a549903..1f155f1b40 100644 --- a/source/tests/pt/model/test_descriptor_sezm.py +++ b/source/tests/pt/model/test_descriptor_sezm.py @@ -1520,45 +1520,6 @@ def _assert_cutoff_near_energy_curve_is_smooth( ), ) - def _assert_bridged_boundary_energy_curve_is_smooth( - self, - n_atten_head: int, - *, - use_amp: bool, - n_focus: int, - nearest_distance: float, - boundary_label: str, - ) -> None: - """Check that one bridged boundary probe keeps one smooth minimum.""" - model = self._build_random_weight_model( - n_atten_head, - use_amp=use_amp, - n_focus=n_focus, - bridging_method="ZBL", - bridging_r_inner=self.BRIDGING_R_INNER, - bridging_r_outer=self.BRIDGING_R_OUTER, - ) - displacements, energies = self._scan_total_energy_curve( - model, - nearest_distance=nearest_distance, - ) - self.assertTrue(torch.isfinite(energies).all().item()) - - stats = self._collect_curve_statistics(energies, displacements) - self._assert_curve_has_usable_signal( - stats, - label=f"Bridged {boundary_label} (use_amp={use_amp}, n_focus={n_focus})", - n_atten_head=n_atten_head, - ) - self.assertEqual( - stats["curve_kind"], - "minimum", - ( - f"Bridged {boundary_label} probe should form one symmetric repulsive bowl " - f"for n_atten_head={n_atten_head}, use_amp={use_amp}, n_focus={n_focus}: {stats}" - ), - ) - def test_scaled_cutoff_near_energy_curve_is_smooth_across_attention_modes( self, ) -> None: @@ -1575,41 +1536,78 @@ def test_scaled_cutoff_near_energy_curve_is_smooth_across_attention_modes( n_focus=n_focus, ) - def test_scaled_bridging_inner_energy_curve_is_smooth_across_attention_modes( + def _assert_bridging_force_consistent_across_switch( self, + model: torch.nn.Module, + *, + eps: float = 1.0e-5, ) -> None: - """Check the bridged near-r_inner PES shape across attention and AMP modes.""" - for use_amp in (False, True): - for n_atten_head in (0, 1, 2): - for n_focus in (1, 2): - with self.subTest( - n_atten_head=n_atten_head, use_amp=use_amp, n_focus=n_focus - ): - self._assert_bridged_boundary_energy_curve_is_smooth( - n_atten_head, - use_amp=use_amp, - n_focus=n_focus, - nearest_distance=self.BRIDGING_R_INNER, - boundary_label="r_inner", - ) + """Assert the bridged force matches a finite difference of the energy. + + A finite, isolated cluster keeps the neighbor count bounded and equal to + the physical neighbors within ``rcut``, so the check is independent of + the periodic-image count and of ``sel``. Atom 1 slides along ``x`` from + below ``r_inner`` to above ``r_outer`` while the spectator atoms stay + beyond ``r_outer``; if the ``BridgingSwitch`` blend kinked at either + boundary, the analytical force (``-dE/dx``) would diverge from the + central finite difference there. + """ + distances = torch.linspace(0.70, 1.30, 25, dtype=self.dtype, device=self.device) + # Isolated cluster: atom 0 anchor, atom 1 probe (slides on x), spectators + # fixed beyond r_outer but inside rcut so only the probe crosses a switch. + template = torch.tensor( + [ + [0.0, 0.0, 0.0], + [0.0, 0.0, 0.0], + [0.0, 2.6, 0.0], + [0.0, 0.0, 2.8], + [-3.0, 0.2, 0.0], + ], + dtype=self.dtype, + device=self.device, + ) + atype = torch.tensor([[0, 1, 0, 1, 0]], dtype=torch.int32, device=self.device) + # Three coordinate frames per scan point: the probe distance and its + # +/- eps neighbors for a central difference of the total energy. + offsets = torch.tensor([0.0, eps, -eps], dtype=self.dtype, device=self.device) + probe_x = (distances.unsqueeze(1) + offsets.unsqueeze(0)).reshape(-1) + n_frame = probe_x.shape[0] + coord = template.unsqueeze(0).repeat(n_frame, 1, 1) + coord[:, 1, 0] = probe_x + result = model(coord, atype.expand(n_frame, -1), box=None) + energy = result["energy"][:, 0].reshape(-1, 3) + self.assertTrue(torch.isfinite(energy).all().item()) + fd_force = -(energy[:, 1] - energy[:, 2]) / (2.0 * eps) + analytical_force = result["force"][::3, 1, 0] + torch.testing.assert_close(analytical_force, fd_force, atol=1.0e-6, rtol=1.0e-4) + # Below r_inner the gate freezes the network term and the ZBL repulsion + # dominates, so the closest probe must be pushed outward (+x). + self.assertGreater( + analytical_force[0].item(), + 0.0, + "bridged short-range force should be repulsive", + ) - def test_scaled_bridging_outer_energy_curve_is_smooth_across_attention_modes( - self, - ) -> None: - """Check the bridged near-r_outer PES shape across attention and AMP modes.""" + def test_bridging_force_consistent_across_switch_boundaries(self) -> None: + """The bridged total energy stays conservative and C1-smooth across both + switch boundaries on a bounded, sel-independent isolated cluster, and is + repulsive at short range. + """ for use_amp in (False, True): for n_atten_head in (0, 1, 2): for n_focus in (1, 2): with self.subTest( n_atten_head=n_atten_head, use_amp=use_amp, n_focus=n_focus ): - self._assert_bridged_boundary_energy_curve_is_smooth( + model = self._build_random_weight_model( n_atten_head, use_amp=use_amp, n_focus=n_focus, - nearest_distance=self.BRIDGING_R_OUTER, - boundary_label="r_outer", + bridging_method="ZBL", + bridging_r_inner=self.BRIDGING_R_INNER, + bridging_r_outer=self.BRIDGING_R_OUTER, ) + self._assert_bridging_force_consistent_across_switch(model) class TestSourceFreezePropagationGate(TestDescriptorEnergyCurveSmoothness): diff --git a/source/tests/pt/model/test_nv_nlist.py b/source/tests/pt/model/test_nv_nlist.py index 5a198cc06d..b09ff22cb2 100644 --- a/source/tests/pt/model/test_nv_nlist.py +++ b/source/tests/pt/model/test_nv_nlist.py @@ -13,6 +13,7 @@ patch, ) +import numpy as np import torch from deepmd.pt.utils import ( @@ -25,6 +26,13 @@ NvNeighborList, _input_device_context, ) +from deepmd.pt_expt.utils.edge_schema import ( + edge_schema_from_extended, +) +from deepmd.pt_expt.utils.vesin_neighbor_list import ( + VesinNeighborList, + is_vesin_torch_available, +) _NV_AVAILABLE = nv_nlist.is_nv_available() _TEST_DEVICES = [torch.device("cpu")] @@ -182,6 +190,7 @@ def _assert_nv_matches_native( with ( patch.object(nv_nlist, "NV_CELL_LIST_THRESHOLD", 1), patch.object(nv_nlist, "NV_NONPERIODIC_CELL_LIST_THRESHOLD", 1), + patch.object(nv_nlist, "NV_CPU_CELL_LIST_THRESHOLD", 1), ): nv = builder.build(coord, atype, box, rcut, sel) else: @@ -267,3 +276,115 @@ def test_nonperiodic_cell_list_matches_native(self) -> None: sel=[8], force_cell_list=True, ) + + +def _to_numpy(x) -> np.ndarray: + """Detach a tensor or pass a numpy array through to a numpy array.""" + return x.detach().cpu().numpy() if isinstance(x, torch.Tensor) else np.asarray(x) + + +def _canonical_edges(schema) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Return the real edges of an ``EdgeNeighborList`` in a builder-independent + order as ``(src, dst, edge_vec)``. + + Edges are keyed by ``(src, dst, round(edge_vec))`` so that the multiple + periodic images of one neighbor pair (distinct displacements that share a + ``(src, dst)``) are disambiguated and ordered deterministically. + """ + mask = _to_numpy(schema.edge_mask).astype(bool) + edge_index = _to_numpy(schema.edge_index) + src = edge_index[0][mask].astype(np.int64) + dst = edge_index[1][mask].astype(np.int64) + edge_vec = _to_numpy(schema.edge_vec)[mask].astype(np.float64) + keys = np.round(edge_vec * 1.0e6).astype(np.int64) + order = np.lexsort((keys[:, 2], keys[:, 1], keys[:, 0], dst, src)) + return src[order], dst[order], edge_vec[order] + + +def _dense_keepall_edges( + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None, + rcut: float, +): + """Reference edge schema from the dense builder with no ``sel`` cap.""" + ext_coord, ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( + coord, atype, rcut, [1], mixed_types=True, box=box, cap_neighbors=False + ) + return edge_schema_from_extended( + ext_coord, atype, nlist, mapping, scatter_to_local=True + ) + + +class TestEdgeSchemaConsistency(unittest.TestCase): + """``nv``, ``vesin``, and the sel-free dense builder must agree on the edge + set, so that dropping ``sel`` and switching builders changes neither the + neighbor topology nor the per-edge geometry. + """ + + def _dense_case( + self, nframes: int, device: torch.device + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # rcut > L / 2 makes each atom see several periodic images, exercising + # the shifted-edge bookkeeping that differs most between builders. + side = 5.0 + generator = torch.Generator(device="cpu").manual_seed(7) + coord = ( + torch.rand( + nframes, 8, 3, generator=generator, dtype=torch.float64, device="cpu" + ) + * side + ) + atype = torch.randint( + 0, 2, (nframes, 8), generator=generator, dtype=torch.int64, device="cpu" + ) + box = ( + (torch.eye(3, dtype=torch.float64, device="cpu") * side) + .reshape(1, 9) + .repeat(nframes, 1) + ) + return coord.to(device), atype.to(device), box.to(device) + + def _assert_builders_match( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None, + rcut: float, + sel: list[int], + ) -> None: + ref_src, ref_dst, ref_vec = _canonical_edges( + _dense_keepall_edges(coord, atype, box, rcut) + ) + builders = [] + if _NV_AVAILABLE: + builders.append(("nv", NvNeighborList())) + if is_vesin_torch_available(): + builders.append(("vesin", VesinNeighborList())) + self.assertTrue(builders, "no accelerated neighbor builder is available") + for name, builder in builders: + with self.subTest(builder=name): + src, dst, vec = _canonical_edges( + builder.build(coord, atype, box, rcut, sel, return_mode="edges") + ) + np.testing.assert_array_equal(src, ref_src) + np.testing.assert_array_equal(dst, ref_dst) + np.testing.assert_allclose(vec, ref_vec, atol=1.0e-9, rtol=1.0e-9) + + def test_periodic_single_frame(self) -> None: + for device in _TEST_DEVICES: + with self.subTest(device=str(device)): + coord, atype, box = self._dense_case(1, device) + self._assert_builders_match(coord, atype, box, 4.0, [64]) + + def test_periodic_multi_frame(self) -> None: + for device in _TEST_DEVICES: + with self.subTest(device=str(device)): + coord, atype, box = self._dense_case(3, device) + self._assert_builders_match(coord, atype, box, 4.0, [64]) + + def test_nonperiodic_single_frame(self) -> None: + for device in _TEST_DEVICES: + with self.subTest(device=str(device)): + coord, atype, _ = self._dense_case(1, device) + self._assert_builders_match(coord, atype, None, 4.0, [64]) diff --git a/source/tests/pt/model/test_sezm_export.py b/source/tests/pt/model/test_sezm_export.py index 7398cf7cee..a233ef8406 100644 --- a/source/tests/pt/model/test_sezm_export.py +++ b/source/tests/pt/model/test_sezm_export.py @@ -29,6 +29,7 @@ import numpy as np import torch +from packaging.version import parse as parse_version from deepmd.pt.entrypoints.freeze_pt2 import ( _build_dynamic_shapes, @@ -61,6 +62,15 @@ "energy_derv_c", "energy_derv_c_redu", } +_TORCH_VERSION = parse_version(torch.__version__) +_SKIP_OFF_COMPILE_TORCH = (_TORCH_VERSION.major, _TORCH_VERSION.minor) not in { + (2, 11), + (2, 12), +} +_SKIP_OFF_COMPILE_TORCH_REASON = ( + "SeZM's torch.compile/export path is only supported on torch 2.11.x and " + f"2.12.x; current torch is {torch.__version__}." +) def _tiny_sezm_model_params() -> dict: @@ -195,21 +205,32 @@ def _eager_forward( sample_inputs: tuple, ) -> dict[str, torch.Tensor]: """Mirror the trace closure: fresh leaf coord + ``requires_grad=True``.""" - ext_coord, ext_atype, nlist, mapping, fparam, aparam, charge_spin = sample_inputs + ( + ext_coord, + atype, + edge_index, + edge_vec, + edge_scatter_index, + edge_mask, + fparam, + aparam, + charge_spin, + ) = sample_inputs eager_coord = ext_coord.detach().clone().requires_grad_(True) return model.forward_common_lower( eager_coord, - ext_atype, - nlist, - mapping=mapping, + atype, + edge_index, + edge_vec, + edge_scatter_index, + edge_mask, fparam=fparam, aparam=aparam, charge_spin=charge_spin, - do_atomic_virial=True, - extra_nlist_sort=model.need_sorted_nlist_for_lower(), ) +@unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) class TestSeZMExportPipeline(_ClearDefaultDeviceTestCase): """Bitwise trace / export / ``.pte`` round-trip parity (``rtol=1e-10``). @@ -262,7 +283,6 @@ def _build_pipeline( ]: traced = model.forward_common_lower_exportable( *sample_inputs, - do_atomic_virial=True, ) exported = torch.export.export( traced, @@ -456,6 +476,7 @@ def test_archive_metadata(self) -> None: "ntypes", "rcut", "sel", + "lower_input_kind", "dim_fparam", "dim_aparam", "dim_chg_spin", @@ -473,6 +494,7 @@ def test_archive_metadata(self) -> None: self.assertEqual(metadata["ntypes"], len(self.params["type_map"])) self.assertEqual(metadata["rcut"], self.params["descriptor"]["rcut"]) self.assertEqual(list(metadata["sel"]), list(self.params["descriptor"]["sel"])) + self.assertEqual(metadata["lower_input_kind"], "edge_vec") self.assertTrue(metadata["mixed_types"]) self.assertFalse(metadata["is_spin"]) self.assertEqual(metadata["dim_fparam"], 0) @@ -710,7 +732,7 @@ def test_charge_spin_export_sample_has_runtime_input_slot(self) -> None: metadata = _collect_metadata(model, ["energy"]) dynamic_shapes = _build_dynamic_shapes(sample_inputs) - self.assertEqual(len(sample_inputs), 7) + self.assertEqual(len(sample_inputs), 9) self.assertEqual(sample_inputs[-1].shape, (5, 2)) self.assertEqual(len(dynamic_shapes), len(sample_inputs)) self.assertEqual(metadata["dim_chg_spin"], 2) @@ -756,6 +778,7 @@ def test_freeze_requires_head_for_multi_task(self) -> None: with self.assertRaises(ValueError): freeze_sezm_to_pt2(str(ckpt_path), str(out)) + @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) def test_freeze_accepts_multi_task_dpa4_head(self) -> None: """Multitask DPA4 checkpoints should export the selected branch.""" @@ -793,6 +816,7 @@ def fake_compile(_exported: torch.export.ExportedProgram, package_path: str): self.assertEqual(model_def["type"], "dpa4") + @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) def test_freeze_accepts_spin_checkpoint_metadata(self) -> None: """SeZM spin checkpoints should export a spin-compatible pt2 contract.""" @@ -818,6 +842,7 @@ def fake_compile(_exported: torch.export.ExportedProgram, package_path: str): ) self.assertTrue(metadata["is_spin"]) + self.assertEqual(metadata["lower_input_kind"], "nlist") self.assertEqual(metadata["type_map"], params["type_map"]) self.assertEqual(metadata["ntypes"], len(params["type_map"])) self.assertEqual(metadata["dim_chg_spin"], 0) diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index afdd1bb72b..87b0a1cc31 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -565,8 +565,8 @@ def test_fixed_edge_geometry_matches_standard_cache(self) -> None: del fp, ap if cc.ndim == 2: cc = cc.view(coord.shape[0], atype.shape[1], 3) - extended_coord, extended_atype, mapping, nlist = model.build_neighbor_list( - cc, atype, bb + extended_coord, extended_atype, nlist, mapping = ( + model.build_extended_neighbor_list(cc, atype, bb) ) atype_loc = extended_atype[:, : nlist.shape[1]] type_ebed = descriptor.type_embedding(atype_loc).reshape(