diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index a4726c2cc8..280da63686 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -17,6 +17,9 @@ jobs: test_cuda: name: Test Python and C++ on CUDA runs-on: gpu + # The full CUDA suite (serial pytest ~3h43m + C++ + LAMMPS) exceeds the + # default 360-min job limit; the self-hosted GPU runner has no hard cap. + timeout-minutes: 480 # https://github.com/deepmodeling/deepmd-kit/pull/2884#issuecomment-1744216845 # container: # image: nvidia/cuda:12.9.1-cudnn-devel-ubuntu22.04 diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 8d68ab4869..36c9fa65e5 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -381,6 +381,7 @@ def forward_common_atomic_graph( fparam: Array | None = None, aparam: Array | None = None, charge_spin: Array | None = None, + comm_dict: dict | None = None, ) -> dict: """Graph analogue of :meth:`forward_common_atomic` on the flat node axis. @@ -406,6 +407,11 @@ def forward_common_atomic_graph( charge_spin charge/spin conditioning. Unused by the dpa1 graph path; accepted so the interface stays stable for charge/spin-conditioned descriptors. + comm_dict + MPI communication metadata forwarded to :meth:`forward_atomic_graph` + (and, from there, the descriptor's ``call_graph``). ``None`` for + non-parallel inference (default). Mirrors :meth:`forward_common_atomic`'s + ``comm_dict`` on the dense route. Returns ------- @@ -420,6 +426,7 @@ def forward_common_atomic_graph( fparam=fparam, aparam=aparam, charge_spin=charge_spin, + comm_dict=comm_dict, ) return self._finalize_atomic_ret(ret_dict, output_mask, atype) diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index ba8b224d19..9f8f06f82d 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -301,6 +301,7 @@ def forward_atomic_graph( fparam: Array | None = None, aparam: Array | None = None, charge_spin: Array | None = None, + comm_dict: dict | None = None, ) -> dict[str, Array]: """Graph analogue of :meth:`forward_atomic` on the flat node axis. @@ -322,6 +323,11 @@ def forward_atomic_graph( charge_spin charge/spin conditioning. Unused by the dpa1 graph path; accepted so the interface stays stable for charge/spin-conditioned descriptors. + comm_dict + MPI communication metadata forwarded to the descriptor's + ``call_graph`` (the message-passing part). ``None`` for + non-parallel inference (default). Mirrors :meth:`forward_atomic`'s + ``comm_dict`` on the dense route. Returns ------- @@ -338,14 +344,16 @@ def forward_atomic_graph( xp = array_api_compat.array_namespace(graph.edge_vec) type_embedding = self.descriptor.type_embedding.call() gg, rot_mat = self.descriptor.call_graph( - graph, atype, type_embedding=type_embedding + graph, atype, type_embedding=type_embedding, comm_dict=comm_dict ) fparam_node = None if fparam is not None: - frame_id = frame_id_from_n_node( - graph.n_node, - n_total=atype.shape[0], - ) + # Pass the STATIC flat node count (``atype.shape[0] == N``) so the + # helper does not fall back to ``int(sum(n_node))``: that int() on a + # traced tensor breaks make_fx / torch.export + # (``GuardOnDataDependentSymNode``) for the graph .pt2 export and + # compiled-training paths when ``numb_fparam > 0``. + frame_id = frame_id_from_n_node(graph.n_node, n_total=atype.shape[0]) fparam_node = xp.take(fparam, frame_id, axis=0) # (N, ndf) aparam_node = aparam if aparam is not None and graph.n_local is not None and aparam.ndim == 3: diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index c1cb5a4823..24e85ae678 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -459,6 +459,23 @@ def uses_graph_lower(self) -> bool: ) return self.se_atten.tebd_input_mode in ("concat", "strip") + def uses_compact_edge_pairs(self) -> bool: + """Returns whether the graph lower traces compact edge pairs. + + The transformer attention lower (``attn_layer > 0``) enumerates + neighbor pairs via the compact ``center_edge_pairs`` realization + (unbacked-SymInt ``nonzero``/``repeat`` sizes, ``pairs.py``); + the factorizable lower (``attn_layer == 0``) traces with backed + symbols only. ``check_graph_trace_torch_version`` keys its + torch >= 2.6 requirement on this capability. + + Returns + ------- + bool + Whether tracing :meth:`call_graph` runs ``center_edge_pairs``. + """ + return self.se_atten.attn_layer > 0 + def disable_graph_lower(self) -> None: """Force the legacy dense lower for this descriptor. @@ -781,6 +798,7 @@ def call_graph( atype: Array, type_embedding: Array | None = None, static_nnei: int | None = None, + comm_dict: dict | None = None, ) -> tuple[Array, Array]: """Descriptor-level graph-native forward. @@ -817,6 +835,12 @@ def call_graph( (N,) flat LOCAL atom types where ``N = sum(n_node)``. type_embedding (ntypes_with_padding, tebd_dim) type-embedding table. + comm_dict + MPI communication metadata. Accepted for ABI parity with + :meth:`DescrptDPA2.call_graph` (uniform ``forward_atomic_graph`` + call site), but UNUSED: a single se_atten descriptor has no + cross-rank message passing (``has_message_passing_across_ranks()`` + is ``False``), so this is always ``None`` in practice. Returns ------- diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index 696c912c63..06a86956ee 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -19,6 +19,7 @@ ) from deepmd.dpmodel.common import ( cast_precision, + get_xp_precision, to_numpy_array, ) from deepmd.dpmodel.utils import ( @@ -602,6 +603,9 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: self.trainable = trainable self.add_tebd_to_repinit_out = add_tebd_to_repinit_out self.compress = False + # graph-native lower opt-out flag (mirrors DescrptDPA1); not + # serialized, re-derived structurally at construction/deserialization. + self._graph_lower_disabled = False self.repinit_out_dim = self.repinit.dim_out if self.repinit_args.use_three_body: @@ -709,6 +713,84 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.env_protection + def uses_graph_lower(self) -> bool: + """Returns whether this descriptor supports the graph-native lower. + + Graph-eligible: repinit (``tebd_input_mode`` in ``{"concat", + "strip"}``) + repformers, with ALL repformer update toggles + included (attention rides ``center_edge_pairs``, see PR-D). + Ineligible (fall back to the legacy dense path): three-body repinit + (needs the angle machinery -- PR-G-dpa3), compressed descriptors + (geo/tebd tabulation is dense-only), and the explicit disable flag + (used by e.g. the spin model wrapper). + + Returns + ------- + uses_graph_lower : bool + Whether the graph-native lower (:meth:`call_graph`) is supported + for the current descriptor configuration. + """ + if self._graph_lower_disabled: + return False + if self.compress: + return False + if self.use_three_body: + return False + # Nonzero-mean checkpoints stay on the legacy route: the dense + # repinit accumulates the (sel - n_real) padding slots' -davg/dstd + # environment rows, which the sel-free carry-all lower has no + # phantom-row counterpart for -- a checkpoint trained against the + # dense expression (set_davg_zero=False + compute_input_stats) would + # otherwise be silently evaluated with a different function + # (observed ~4e-2 energy shift at NON-binding sel, attention off). + # Reproducing the dense padding term on the graph route is a + # versioned training/migration change, not a default flip. + if not self.repinit_args.set_davg_zero: + return False + if not self.repformer_args.set_davg_zero: + return False + return self.repinit.tebd_input_mode in ("concat", "strip") + + def uses_compact_edge_pairs(self) -> bool: + """Returns whether the graph lower traces compact edge pairs. + + The compact ``center_edge_pairs`` realization (unbacked-SymInt + ``nonzero``/``repeat`` sizes, ``pairs.py``) backs the repformer's + nnei-x-nnei attention (``update_g2_has_attn``: Atten2Map / + MultiHeadApply) and the equivariant pair update (``update_h2``: + EquiVarApply). ``check_graph_trace_torch_version`` keys its + torch >= 2.6 requirement on this capability. + + Derived from the EFFECTIVE per-layer flags, not the configured + arguments: the LAST repformer layer is built with + ``update_chnnl_2=False``, which forces its ``update_g2_has_attn`` + and ``update_h2`` off -- so e.g. ``nlayers=1`` never runs + ``center_edge_pairs`` even when the arguments enable it, and old + torch stays usable. This mirrors the gate + ``DescrptBlockRepformers.call_graph`` itself uses to decide whether + to build the pairs. + + Returns + ------- + bool + Whether tracing :meth:`call_graph` runs ``center_edge_pairs``. + """ + return any( + ll.update_g2_has_attn or ll.update_h2 for ll in self.repformers.layers + ) + + def disable_graph_lower(self) -> None: + """Force the legacy dense lower for this descriptor. + + This is an explicit opt-out knob used by contexts where the + graph-native lower is unsupported or undesirable (e.g. spin + models). After calling this, :meth:`uses_graph_lower` returns + ``False`` regardless of the descriptor configuration. The flag is + not serialized; it is re-derived structurally at spin-model + construction/deserialization. + """ + self._graph_lower_disabled = True + def share_params( self, base_class: Any, shared_level: int, resume: bool = False ) -> NoReturn: @@ -883,6 +965,366 @@ def call( The smooth switch function. shape: nf x nloc x nnei """ + # The dense ``call`` always runs the legacy dense body -- it is the + # cross-backend consistency reference and must match the tf/pt/pd/jax + # dense descriptors bit-for-bit. Unlike ``DescrptDPA1.call`` (whose + # graph adapter is bit-exact because its default repinit uses + # ``attn_layer > 0``, where the real attention masks padding), DPA2's + # repinit is hardwired to ``attn_layer == 0``. There the dense + # se_atten body leaks a phantom padding-neighbor ``-davg/dstd`` + # residual that the graph path deliberately does NOT reproduce (see + # :meth:`call_graph`'s Notes; the graph output is the + # physically-correct one). That makes ``_call_graph_adapter`` diverge + # from ``_call_dense`` for any model with non-trivial statistics + # (nonzero ``davg`` or non-unit ``dstd``) -- an accepted graph/dense + # divergence (see ``KNOWN_GRAPH_DENSE_DIVERGENT``). The graph-native + # route is therefore reached exclusively through :meth:`call_graph` + # (pt_expt ``forward_atomic_graph`` and the C++ graph ``.pt2``), never + # through ``call``. ``_call_graph_adapter`` is retained as the + # bit-exact-regime reference exercised by ``TestDPA2AdapterBitExact``. + return self._call_dense( + coord_ext, + atype_ext, + nlist, + mapping=mapping, + fparam=fparam, + comm_dict=comm_dict, + charge_spin=charge_spin, + ) + + def _call_graph_adapter( + self, + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None, + ) -> tuple[Array, Array, None, None, Array]: + """Dense-quartet -> shape-static graph -> call_graph -> dense-ABI 5-tuple. + + Builds a NeighborGraph from the dense quartet with the SHAPE-STATIC + converter (``compact=False``, jit/export-traceable -- no + ``nonzero``), runs :meth:`call_graph`, and reconstructs the dense + 5-tuple ABI. The per-block edge masks in :meth:`call_graph` (dist + filter + slot filter against ``static_nnei``) replicate + ``build_multiple_neighbor_list``'s ``nlist[:, :, :ns]`` slicing. + + Bit-exact vs :meth:`_call_dense` **only in the trivial-statistics + regime** (``davg == 0`` and ``dstd == 1``, i.e. ``set_davg_zero`` with + actually-zero stats). For any model with nonzero ``davg`` or non-unit + ``dstd`` this adapter DIFFERS from ``_call_dense``: DPA2's repinit is + always ``attn_layer == 0``, where the dense body leaks a phantom + padding-neighbor ``-davg/dstd`` residual that the graph path + deliberately omits (see :meth:`call_graph`'s Notes). This is why the + default dense :meth:`call` does NOT route here -- the graph-native + route lives in :meth:`call_graph`. This method is retained as the + bit-exact-regime reference exercised by ``TestDPA2AdapterBitExact``. + + Parameters + ---------- + coord_ext + The extended coordinates of atoms. shape: nf x (nall x 3) + atype_ext + The extended atom types. shape: nf x nall + nlist + The neighbor list. shape: nf x nloc x nnei + mapping + The index mapping from extended to local region. shape: nf x nall. + ``None`` is allowed only when nall == nloc (identity mapping). + + Returns + ------- + descriptor + The descriptor. shape: nf x nloc x (ng x axis_neuron) + gr + The rotationally equivariant single-particle representation. + shape: nf x nloc x ng x 3 + g2 + ``None`` for this descriptor (graph-native repformers carries g2 + internally; the dense 5-tuple ABI never surfaces it for dpa2). + h2 + ``None`` for this descriptor. + sw + The smooth switch function, at the REPFORMER block's own + ``nsel`` width (matching :meth:`_call_dense`, whose ``sw`` comes + from the repformer sub-nlist, narrower than the outer ``nlist``). + shape: nf x nloc x repformers.get_nsel() + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + graph_from_dense_quartet, + ) + + xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist) + nframes, nloc, nnei = nlist.shape + graph, atype_local = graph_from_dense_quartet( + coord_ext, atype_ext, nlist, mapping + ) + g1, rot_mat, sw = self.call_graph( + graph, atype_local, static_nnei=nnei, return_sw=True + ) + g1 = xp.reshape(g1, (nframes, nloc, *g1.shape[1:])) + rot_mat = xp.reshape(rot_mat, (nframes, nloc, *rot_mat.shape[1:])) + # call_graph's per-block truncation already narrows the repformer + # edge axis to the repformer block's own nsel width (matching the + # dense body's sw, which comes from a nlist already truncated to + # repformers.get_nsel() columns) -- a plain reshape suffices. + sw = xp.reshape(sw, (nframes, nloc, -1)) + return g1, rot_mat, None, None, sw + + def call_graph( + self, + graph: Any, + atype: Array, + type_embedding: Array | None = None, + static_nnei: int | None = None, + comm_dict: dict | None = None, + return_sw: bool = False, + ) -> tuple[Array, Array] | tuple[Array, Array, Array]: + """Descriptor-level graph-native forward: one carry-all graph at the + model rcut (== repinit rcut), per-block edge masks in place of + ``build_multiple_neighbor_list``. + + This is what :meth:`~deepmd.dpmodel.atomic_model.dp_atomic_model. + DPAtomicModel.forward_atomic_graph` calls. Geometry enters only + through ``graph.edge_vec``; the descriptor is graph-native from + repinit through repformers (three-body repinit is graph-ineligible + and gated out by :meth:`uses_graph_lower`). + + Notes + ----- + **Dense's ``attn_layer=0`` padding-slot leak is intentionally NOT + reproduced here.** When ``set_davg_zero=False`` (nonzero ``mean``) + AND ``exclude_types == []``, the dense + :meth:`DescrptBlockSeAtten.call` that backs ``repinit`` (always + ``attn_layer=0`` in DPA2) leaks a data-dependent residual from its + padding neighbor slots into the output: ``PairExcludeMask. + build_type_exclude_mask`` short-circuits to an all-ones mask when no + types are excluded, which is the *only* real/padding mask that code + path applies at ``attn_layer=0`` (the identity ``NeighborGatedAttention`` + with zero layers masks nothing). The padding rows' geometry is + already weight-zeroed inside ``_make_env_mat``, but ``EnvMat.call`` + subtracts ``davg`` AFTER that zeroing, so every padding slot carries + a deterministic ``-davg/dstd`` residual (padding-count/sel-dependent) + that nothing re-masks before it reaches ``gr``. This graph path + masks/zeros every + padding and excluded edge before each ``segment_sum``, so it does NOT + reproduce that leak. In this regime :meth:`call_graph` (and the dense + adapter, :meth:`_call_graph_adapter`) deliberately DIFFER from + :meth:`_call_dense` -- the graph output is the physically correct + one, and the divergence is a pre-existing dense-body bug (present + since ``attn_layer=0`` was introduced, not caused by this task) that + is documented, not bit-matched, here. Bit-exactness vs the dense body + (as exercised by ``TestDPA2AdapterBitExact``) holds in every OTHER + regime: ``set_davg_zero=True``, or non-empty ``exclude_types`` (which + supplies a real mask independent of ``davg``), or configurations with + no padding slots within ``rcut``. + + Parameters + ---------- + graph + A :class:`~deepmd.dpmodel.utils.neighbor_graph.NeighborGraph`. + atype + (N,) flat LOCAL atom types where ``N = sum(n_node)``. + type_embedding + (ntypes_with_padding, tebd_dim) type-embedding table. Defaults + to ``self.type_embedding.call()`` when not provided. + static_nnei + When the graph uses the shape-static center-major layout + (``graph_from_dense_quartet`` / ``from_dense_quartet(compact= + False)``, ``E = n_center * nnei``), pass ``nnei`` so the + per-block edge masks and the repformer attention edge-pair + enumeration stay jit/export-traceable (no ``nonzero``). ``None`` + (carry-all / compact graphs) selects the dynamic eager form + (sel is normalization-only there, decision #9). + comm_dict + MPI communication metadata forwarded to the repformer block + (the message-passing part). ``None`` for non-parallel inference + (default). + return_sw + When True, also return the repformer block's smooth switch + function on the flat edge axis. + + Returns + ------- + g1 : Array + (N, dim_out) descriptor, flat node axis. + rot_mat : Array + (N, dim_emb, 3) equivariant single-particle representation, + flat node axis. + sw : Array + (E,) smooth switch, zeroed on padding/ineligible edges. Only + returned when ``return_sw`` is True. + """ + import dataclasses + + from deepmd.dpmodel.utils.safe_gradient import ( + safe_for_vector_norm, + ) + + xp = array_api_compat.array_namespace(graph.edge_vec, atype) + dev = array_api_compat.device(graph.edge_vec) + # manual @cast_precision: the decorator casts array ARGUMENTS, but + # the graph's only float input (edge_vec) is inside the + # NeighborGraph dataclass, invisible to it. Cast edge_vec down to + # the descriptor precision on entry and the outputs back to the + # caller's dtype on exit (dpa1 precedent). + in_dtype = graph.edge_vec.dtype + prec = get_xp_precision(xp, self.precision) + if in_dtype != prec: + graph = dataclasses.replace(graph, edge_vec=xp.astype(graph.edge_vec, prec)) + dist = safe_for_vector_norm(graph.edge_vec, axis=-1) # (E,) + e_ax = graph.edge_mask.shape[0] + + def _block_graph(rc: float, ns: int) -> tuple[Any, int | None]: + """Per-block view of the model graph (local closure). + + Graph analogue of ``build_multiple_neighbor_list`` + (nlist.py:408): dist mask always; slot TRUNCATION only in the + shape-static dense-adapter layout, replicating the dense + ``nlist[:, :, :ns]`` slicing. + + Parameters + ---------- + rc + The block cutoff radius. + ns + The block sel (dense sub-nlist width). + + Returns + ------- + block_graph : NeighborGraph + The block-restricted graph view. + block_static_nnei : int | None + The block's static width (``ns``) in the shape-static + layout, ``None`` for carry-all graphs. + + Notes + ----- + The genuine array-width SLICE (rather than an edge_mask AND) + keeps the shape-static pair enumeration + (``center_edge_pairs``/``static_nnei``) at exactly the dense + sub-nlist width ``ns``: the pair count stays minimal, and the + attention softmax sees the same present-slot layout as the + dense body. (Before the fixed-phantom-count compensation in + ``segment_softmax`` this width-match was also required for + numeric parity -- extra always-masked pairs each contributed + ``exp(-attnw_shift)`` to the denominator; the compensation has + since made the denominator width-independent.) Carry-all graphs + (``static_nnei is None``) have no static width at all: sel is + normalization-only there (spec decision #9), and the compact + attention pairing groups per-center dynamically, so no + truncation is needed or possible. + """ + if static_nnei is None or ns >= static_nnei: + m = graph.edge_mask & (dist <= rc) + return dataclasses.replace(graph, edge_mask=m), static_nnei + n_center = e_ax // static_nnei + ei = xp.reshape(graph.edge_index, (2, n_center, static_nnei))[:, :, :ns] + ei = xp.reshape(ei, (2, n_center * ns)) + ev = xp.reshape(graph.edge_vec, (n_center, static_nnei, 3))[:, :ns, :] + ev = xp.reshape(ev, (n_center * ns, 3)) + em = xp.reshape(graph.edge_mask, (n_center, static_nnei))[:, :ns] + em = xp.reshape(em, (n_center * ns,)) + dd = xp.reshape(dist, (n_center, static_nnei))[:, :ns] + dd = xp.reshape(dd, (n_center * ns,)) + em = em & (dd <= rc) + sliced = dataclasses.replace( + graph, edge_index=ei, edge_vec=ev, edge_mask=em + ) + return sliced, ns + + tebd_table = ( + type_embedding if type_embedding is not None else self.type_embedding.call() + ) + # NB: no xp.asarray(..., device=) wrap -- it detaches the + # type-embedding gradient under torch (the dpa1 lesson). + atype_local = xp.asarray(atype, device=dev) + g1_inp = xp.take(tebd_table, atype_local, axis=0) # (N, tebd_dim) + # repinit (attn_layer == 0 se_atten block; call_graph exists since PR-A) + repinit_graph, repinit_static_nnei = _block_graph( + self.repinit.get_rcut(), self.repinit.get_nsel() + ) + g1, _ = self.repinit.call_graph( + repinit_graph, + atype, + type_embedding=tebd_table, + static_nnei=repinit_static_nnei, + ) + # three-body is graph-ineligible (uses_graph_lower gates it out) + g1 = self.g1_shape_tranform(g1) + if self.add_tebd_to_repinit_out: + assert self.tebd_transform is not None + g1 = g1 + self.tebd_transform(g1_inp) + # dense does the mapping gather to g1_ext here; the graph has ONE + # flat node space -- nothing to do (extended multi-rank ghost refresh + # happens inside the repformer layer loop via + # `_exchange_ghosts_graph`). + repformer_graph, repformer_static_nnei = _block_graph( + self.repformers.get_rcut(), self.repformers.get_nsel() + ) + g1, g2, h2, rot_mat, sw = self.repformers.call_graph( + repformer_graph, + atype, + g1, + comm_dict=comm_dict, + static_nnei=repformer_static_nnei, + ) + if self.concat_output_tebd: + g1 = xp.concat([g1, g1_inp], axis=-1) + if in_dtype != prec: + g1 = xp.astype(g1, in_dtype) + rot_mat = xp.astype(rot_mat, in_dtype) + sw = xp.astype(sw, in_dtype) + if return_sw: + return g1, rot_mat, sw + return g1, rot_mat + + def _call_dense( + self, + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None = None, + fparam: Array | None = None, + comm_dict: dict | None = None, + charge_spin: Array | None = None, + ) -> tuple[Array, Array, Array, Array, Array]: + """Legacy dense descriptor body (the ineligible/no-mapping-ghost + ``call`` path). + + Parameters + ---------- + coord_ext + The extended coordinates of atoms. shape: nf x (nallx3) + atype_ext + The extended aotm types. shape: nf x nall + nlist + The neighbor list. shape: nf x nloc x nnei + mapping + The index mapping, maps extended region index to local region. + comm_dict + MPI communication metadata for parallel inference. Forwarded to + the repformer block (the message-passing part). The repinit + sub-block does no message passing and does not receive it. + ``None`` for non-parallel inference (default). + + Returns + ------- + descriptor + The descriptor. shape: nf x nloc x (ng x axis_neuron) + gr + The rotationally equivariant and permutationally invariant single particle + representation. shape: nf x nloc x ng x 3 + g2 + The rotationally invariant pair-partical representation. + shape: nf x nloc x nnei x ng + h2 + The rotationally equivariant pair-partical representation. + shape: nf x nloc x nnei x 3 + sw + The smooth switch function. shape: nf x nloc x nnei + + """ + del fparam, charge_spin xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist) use_three_body = self.use_three_body nframes, nloc, nnei = nlist.shape diff --git a/deepmd/dpmodel/descriptor/repformers.py b/deepmd/dpmodel/descriptor/repformers.py index d20f7703f7..09d4309bd5 100644 --- a/deepmd/dpmodel/descriptor/repformers.py +++ b/deepmd/dpmodel/descriptor/repformers.py @@ -527,6 +527,173 @@ def _exchange_ghosts( g1_ext = xp_take_along_axis(g1, mapping_tiled, axis=1) return xp.where(mapping_mask, g1_ext, xp.zeros_like(g1_ext)) + def _exchange_ghosts_graph( + self, + g1: Array, + comm_dict: dict | None, + n_total: int, + ) -> Array: + """Per-layer node-channel refresh on the graph path. + + Ghost-free graphs (Python single-rank; src IS the local owner) and + extended single-process graphs need NO exchange -- identity. The + pt_expt subclass overrides this to overwrite ghost rows via + ``deepmd_export::border_op`` when ``comm_dict`` is provided (C++ + multi-rank extended-region graphs). + + Parameters + ---------- + g1 + Flat node-wise atomic invariant rep, with shape [n_total, ng1]. + comm_dict + MPI communication metadata; must be ``None`` on this (dpmodel) + path. + n_total + Total number of nodes (unused here; the pt_expt override needs + it). + + Returns + ------- + g1 : Array + The (unchanged) node channel, with shape [n_total, ng1]. + + Raises + ------ + NotImplementedError + If ``comm_dict`` is not ``None``. + """ + del n_total + if comm_dict is not None: + raise NotImplementedError( + "comm_dict on the dpmodel graph path requires the pt_expt " + "_exchange_ghosts_graph override (MPI border_op)" + ) + return g1 + + def call_graph( + self, + graph: Any, + atype: Array, + g1_input: Array, + comm_dict: dict | None = None, + static_nnei: int | None = None, + ) -> tuple[Array, Array, Array, Array, Array]: + """Graph twin of :meth:`call` on the flat node/edge axes. + + Parameters + ---------- + graph + A :class:`~deepmd.dpmodel.utils.neighbor_graph.NeighborGraph` whose + ``edge_index = [src, dst]`` (src = neighbor local owner, dst = center), + ``edge_vec = r_src - r_dst`` and ``edge_mask`` marks real edges. + atype + (N,) flat node atom types (``N = sum(graph.n_node)``). + g1_input + (N, g1_dim) node channel BEFORE the block activation -- the flat + twin of the dense ``atype_embd`` (``atype_embd_ext`` sliced to the + first ``nloc`` atoms); the block applies ``self.act`` to it, exactly + like the dense body. + comm_dict + MPI communication metadata forwarded to :meth:`_exchange_ghosts_graph`. + ``None`` for non-parallel inference (dpmodel default: identity). + static_nnei + When the graph uses the shape-static center-major layout + (``graph_from_dense_quartet`` / ``from_dense_quartet(compact=False)``, + ``E = n_center * nnei``), pass ``nnei`` so the attention edge-pair + enumeration stays jit/export-traceable (no ``nonzero``). ``None`` + (carry-all / compact graphs) selects the dynamic eager form. + + Returns + ------- + g1 : Array + (N, dim_out) updated node channel. + g2 : Array + (E, ng2) updated edge channel, invariant. + h2 : Array + (E, 3) updated edge channel, equivariant. + rot_mat : Array + (N, dim_emb, 3) rotationally equivariant single-particle + representation. + sw : Array + (E,) smooth switch, zeroed on padding edges. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + apply_pair_exclusion, + center_edge_pairs, + edge_env_mat, + ) + + xp = array_api_compat.array_namespace(graph.edge_vec, atype) + n_total = atype.shape[0] + # descriptor-block exclude_types: same canonical transform as dense + # nlist erasure (repformers.call:541-543) + graph = apply_pair_exclusion(graph, atype, self.emask) + src = graph.edge_index[0, :] + dst = graph.edge_index[1, :] + center_type = xp.take(atype, dst, axis=0) + rr, sw_e = edge_env_mat( + graph.edge_vec, + center_type, + self.mean[:, 0, :], + self.stddev[:, 0, :], + self.rcut, + self.rcut_smth, + protection=self.env_protection, + edge_mask=graph.edge_mask, + return_sw=True, + ) # (E, 4), (E, 1) + sw = sw_e[:, 0] # (E,), zeroed on padding + g1 = self.act(g1_input) # (N, g1_dim) + if not self.direct_dist: + g2, h2 = rr[:, :1], rr[:, 1:] + else: + g2 = safe_for_vector_norm(graph.edge_vec, axis=-1, keepdims=True) + g2 = g2 / self.rcut + h2 = graph.edge_vec / self.rcut + g2 = self.act(self.g2_embd(g2)) # (E, ng2) + pairs = None + # DescrptBlockRepformers has no block-wide `update_chnnl_2` (that flag + # lives per-layer: the last layer always sets it False, see __init__'s + # `update_chnnl_2=(ii != nlayers - 1)`). Precompute the shared + # (topology-only) edge pairs iff ANY layer actually consumes them. + if any(ll.update_g2_has_attn or ll.update_h2 for ll in self.layers): + pairs = center_edge_pairs( + dst, + graph.edge_mask, + n_total, + include_self=True, + ordered=True, + static_nnei=static_nnei, + ) + for ll in self.layers: + g1 = self._exchange_ghosts_graph(g1, comm_dict, n_total) + g1, g2, h2 = ll.call_graph( + g1, + g2, + h2, + src, + dst, + graph.edge_mask, + sw, + n_total, + self.nnei, + pairs, + ) + h2g2 = _cal_hg_graph( + g2, + h2, + graph.edge_mask, + sw, + dst, + n_total, + self.nnei, + smooth=self.smooth, + epsilon=self.epsilon, + use_sqrt_nnei=self.use_sqrt_nnei, + ) # (N, 3, ng2) + rot_mat = xp.matrix_transpose(h2g2) # (N, ng2, 3) + return g1, g2, h2, rot_mat, sw + def call( self, nlist: Array, @@ -964,6 +1131,151 @@ def symmetrization_op( return grrg +def _cal_hg_graph( + g: Array, + h: Array, + edge_mask: Array, + sw: Array, + dst: Array, + n_total: int, + nnei: int, + smooth: bool = True, + epsilon: float = 1e-4, + use_sqrt_nnei: bool = True, +) -> Array: + """Graph twin of :func:`_cal_hg`: hg[n, a, b] = sum_{e: dst(e)=n} h_e[a] g_e[b]. + + Parameters + ---------- + g + Flat edge-wise invariant rep, with shape [n_edge, ng]. + h + Flat edge-wise equivariant rep, with shape [n_edge, 3]. + edge_mask + Edge mask, where zero means no edge, with shape [n_edge]. + sw + The switch function per edge, with shape [n_edge]. + dst + Destination (center) node index of each edge, with shape [n_edge]. + n_total + Total number of nodes. + nnei + The block sel (the smooth-branch normalization constant, spec + decision #9: sel = normalization only). + smooth + Whether to use the smooth (sw-weighted, sel-normalized) branch. + epsilon + Degree-normalization protection for the non-smooth branch. + use_sqrt_nnei + Whether to normalize by sqrt(nnei) instead of nnei. + + Returns + ------- + hg : Array + Transposed rotation matrix per node, with shape [n_total, 3, ng]. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + segment_sum, + ) + + xp = array_api_compat.array_namespace(g, h) + g = g * xp.astype(edge_mask[:, None], g.dtype) + if not smooth: + deg = segment_sum(xp.astype(edge_mask, g.dtype), dst, n_total) # (N,) + if not use_sqrt_nnei: + invnnei = 1.0 / (epsilon + deg) + else: + invnnei = 1.0 / (epsilon + xp.sqrt(deg)) + invnnei = invnnei[:, None, None] + else: + g = g * sw[:, None] + invnnei = 1.0 / float(nnei) if not use_sqrt_nnei else 1.0 / (float(nnei) ** 0.5) + hg = segment_sum(h[:, :, None] * g[:, None, :], dst, n_total) # (N, 3, ng) + return hg * invnnei + + +def _cal_grrg_graph(hg: Array, axis_neuron: int) -> Array: + """Graph twin of :func:`_cal_grrg` (node-local, no neighbor axis). + + Parameters + ---------- + hg + Transposed rotation matrix per node, with shape [n_total, 3, ng]. + axis_neuron + Size of the submatrix of hg (embedding matrix). + + Returns + ------- + grrg : Array + Atomic invariant rep, with shape [n_total, axis_neuron * ng]. + """ + xp = array_api_compat.array_namespace(hg) + n, _, ng = hg.shape + hgm = hg[..., :axis_neuron] # (N, 3, axis) + grrg = xp.matmul(xp.matrix_transpose(hgm), hg) / (3.0**1) # (N, axis, ng) + return xp.reshape(grrg, (n, axis_neuron * ng)) + + +def symmetrization_op_graph( + g: Array, + h: Array, + edge_mask: Array, + sw: Array, + dst: Array, + n_total: int, + nnei: int, + axis_neuron: int, + smooth: bool = True, + epsilon: float = 1e-4, + use_sqrt_nnei: bool = True, +) -> Array: + """Graph twin of :func:`symmetrization_op`. + + Parameters + ---------- + g + Flat edge-wise invariant rep, with shape [n_edge, ng]. + h + Flat edge-wise equivariant rep, with shape [n_edge, 3]. + edge_mask + Edge mask, where zero means no edge, with shape [n_edge]. + sw + The switch function per edge, with shape [n_edge]. + dst + Destination (center) node index of each edge, with shape [n_edge]. + n_total + Total number of nodes. + nnei + The block sel (the smooth-branch normalization constant). + axis_neuron + Size of the submatrix of hg (embedding matrix). + smooth + Whether to use the smooth (sw-weighted, sel-normalized) branch. + epsilon + Degree-normalization protection for the non-smooth branch. + use_sqrt_nnei + Whether to normalize by sqrt(nnei) instead of nnei. + + Returns + ------- + grrg : Array + Atomic invariant rep, with shape [n_total, axis_neuron * ng]. + """ + hg = _cal_hg_graph( + g, + h, + edge_mask, + sw, + dst, + n_total, + nnei, + smooth=smooth, + epsilon=epsilon, + use_sqrt_nnei=use_sqrt_nnei, + ) + return _cal_grrg_graph(hg, axis_neuron) + + class Atten2Map(NativeOP): r"""Masked and smoothed angular attention map. @@ -1082,6 +1394,123 @@ def call( ret = xp_transpose_01342(ret) return ret + def call_graph( + self, + g2: Array, + h2: Array, + sw: Array, + q_e: Array, + k_e: Array, + pair_mask: Array, + e_tot: int, + sel: int, + ) -> Array: + """Graph twin of :meth:`call`: per-pair attention map over edges sharing a center. + + The dense (nnei, nnei) square becomes ordered+self edge pairs; the + softmax over the key axis becomes a segment_softmax grouped by the + query edge. The dense post-softmax query-row AND key-column zeroing + both collapse to ``pair_mask`` (a pair is real iff BOTH edges are + real). + + Parameters + ---------- + g2 + Flat edge-wise pair invariant rep, with shape [n_edge, ng2]. + h2 + Flat edge-wise pair equivariant rep, with shape [n_edge, 3]. + sw + The switch function per edge, with shape [n_edge]. + q_e + Query edge index of each pair, with shape [n_pair]. + k_e + Key edge index of each pair, with shape [n_pair]. + pair_mask + Pair validity (both edges real), with shape [n_pair]. + e_tot + Total number of edges (the number of softmax segments). + sel + The block sel: the fixed dense softmax width used for the + phantom-count compensation of the smooth branch. + + Returns + ------- + attn : Array + Attention map on the flat pair axis, with shape [n_pair, nh]. + + Notes + ----- + The smooth branch reproduces dense's FIXED-WIDTH softmax denominator + exactly: dense keeps every one of its ``sel`` key slots in the + denominator, the non-real ones each at ``exp(-attnw_shift)``. Here the + masked pairs are excluded from the softmax and replaced by + ``sel - n_real`` SIGNED phantom denominator terms per query edge, so + the phantom COUNT is geometry-independent. Without this, one phantom + per PRESENT pair would make the denominator term count change when an + edge enters/leaves the graph at the model cutoff -- an + ``O(exp(-attnw_shift))`` energy/force discontinuity -- and differ from + dense by ``O(sel * exp(-attnw_shift))``. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + segment_softmax, + segment_sum, + ) + + xp = array_api_compat.array_namespace(g2, h2) + e_ax = g2.shape[0] + nd, nh = self.hidden_dim, self.head_num + g2qk = self.mapqk(g2) # (E, nd * nh * 2) + g2qk = xp.reshape(g2qk, (e_ax, nd, nh * 2)) # heads LAST, dense order + g2q = g2qk[:, :, :nh] # (E, nd, nh) + g2k = g2qk[:, :, nh:] # (E, nd, nh) + # per-pair, per-head logits + attnw = ( + xp.sum(xp.take(g2q, q_e, axis=0) * xp.take(g2k, k_e, axis=0), axis=1) + / nd**0.5 + ) # (P, nh) + if self.has_gate: + gate = xp.sum( + xp.take(h2, q_e, axis=0) * xp.take(h2, k_e, axis=0), axis=-1 + ) # (P,) + attnw = attnw * gate[:, None] + sw_q = xp.take(sw, q_e, axis=0) # (P,) + sw_k = xp.take(sw, k_e, axis=0) + if self.smooth: + attnw = (attnw + self.attnw_shift) * sw_q[:, None] * sw_k[ + :, None + ] - self.attnw_shift + # dense-width phantom compensation (see docstring): real pairs per + # query edge, then sel - n_real phantoms at exp(-shift) + pm = xp.astype(pair_mask, attnw.dtype) + n_real = segment_sum(pm, q_e, e_tot) # (e_tot,) + # SIGNED count: beyond sel it subtracts the excess phantom-level + # terms, so the boundary-pair denominator increment vanishes for + # ARBITRARY degree (a clamped count left a finite + # exp(-attnw_shift) step at the cutoff whenever n_real > sel). + phantom = sel - n_real + attnw = segment_softmax( + attnw, + q_e, + e_tot, + mask=pair_mask, + phantom_count=phantom, + phantom_logit=-self.attnw_shift, + # per-pair slot weight: the KEY edge's envelope (the pair's + # logit reaches -attnw_shift when the key edge exits; sw_q + # is constant within the segment and applied post-softmax) + slot_weight=sw_k, + ) + else: + attnw = segment_softmax(attnw, q_e, e_tot, mask=pair_mask) + attnw = attnw * xp.astype(pair_mask[:, None], attnw.dtype) + if self.smooth: + attnw = attnw * sw_q[:, None] * sw_k[:, None] + h2h2t = ( + xp.sum(xp.take(h2, q_e, axis=0) * xp.take(h2, k_e, axis=0), axis=-1) + / 3.0**0.5 + ) # (P,) + return attnw * h2h2t[:, None] # (P, nh) + def serialize(self) -> dict: """Serialize the networks to a dict. @@ -1185,6 +1614,42 @@ def call( # nf x nloc x nnei x ng2 return self.head_map(ret) + def call_graph( + self, AA: Array, g2: Array, q_e: Array, k_e: Array, e_tot: int + ) -> Array: + """Graph twin of :meth:`call`: out[q] = sum_k AA[q,k] g2v[k] per head. + + Parameters + ---------- + AA + Attention map on the flat pair axis, with shape [n_pair, nh]. + g2 + Flat edge-wise pair invariant rep, with shape [n_edge, ng2]. + q_e + Query edge index of each pair, with shape [n_pair]. + k_e + Key edge index of each pair, with shape [n_pair]. + e_tot + Total number of edges. + + Returns + ------- + g2_new : Array + Attention-updated edge channel, with shape [n_edge, ng2]. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + segment_sum, + ) + + xp = array_api_compat.array_namespace(AA, g2) + e_ax, ng2 = g2.shape + nh = self.head_num + g2v = xp.reshape(self.mapv(g2), (e_ax, ng2, nh)) + contrib = AA[:, None, :] * xp.take(g2v, k_e, axis=0) # (P, ng2, nh) + ret = segment_sum(contrib, q_e, e_tot) # (E, ng2, nh) + ret = xp.reshape(ret, (e_ax, ng2 * nh)) # nh-fastest == dense reshape + return self.head_map(ret) # (E, ng2) + def serialize(self) -> dict: """Serialize the networks to a dict. @@ -1280,6 +1745,39 @@ def call( # nf x nloc x nnei x 3 return xp.squeeze(self.head_map(ret), axis=-1) + def call_graph( + self, AA: Array, h2: Array, q_e: Array, k_e: Array, e_tot: int + ) -> Array: + """Graph twin of :meth:`call` (heads applied to the equivariant channel). + + Parameters + ---------- + AA + Attention map on the flat pair axis, with shape [n_pair, nh]. + h2 + Flat edge-wise pair equivariant rep, with shape [n_edge, 3]. + q_e + Query edge index of each pair, with shape [n_pair]. + k_e + Key edge index of each pair, with shape [n_pair]. + e_tot + Total number of edges. + + Returns + ------- + h2_new : Array + Attention-updated equivariant edge channel, with shape + [n_edge, 3]. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + segment_sum, + ) + + xp = array_api_compat.array_namespace(AA, h2) + contrib = AA[:, None, :] * xp.take(h2, k_e, axis=0)[:, :, None] # (P, 3, nh) + ret = segment_sum(contrib, q_e, e_tot) # (E, 3, nh) + return xp.squeeze(self.head_map(ret), axis=-1) # (E, 3) + def serialize(self) -> dict: """Serialize the networks to a dict. @@ -1426,6 +1924,95 @@ def call( ret = self.head_map(ret) return ret + def call_graph( + self, + g1: Array, + gg1: Array, + edge_mask: Array, + sw: Array, + dst: Array, + n_total: int, + sel: int, + ) -> Array: + """Graph twin of :meth:`call`: per-center attention over its edges. + + The dense softmax over the neighbor axis becomes a segment_softmax + over ``dst``. + + Parameters + ---------- + g1 + Flat node-wise atomic invariant rep, with shape [n_total, ng1]. + gg1 + Neighbor (source-node) g1 gathered per edge, with shape + [n_edge, ng1]. + edge_mask + Edge mask, where zero means no edge, with shape [n_edge]. + sw + The switch function per edge, with shape [n_edge]. + dst + Destination (center) node index of each edge, with shape + [n_edge]. + n_total + Total number of nodes (the number of softmax segments). + sel + The block sel: the fixed dense softmax width used for the + phantom-count compensation of the smooth branch. + + Returns + ------- + g1_attn : Array + Attention-updated node channel, with shape [n_total, ng1]. + + Notes + ----- + The smooth branch uses the dense-width phantom compensation (see + :meth:`Atten2Map.call_graph`): masked edges are excluded from the + softmax and ``sel - n_real`` SIGNED phantom denominator terms at + ``exp(-attnw_shift)`` are added per center, keeping the denominator + term count geometry-independent (continuous at the cutoff, and equal + to the dense softmax at non-binding ``sel``). + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + segment_softmax, + segment_sum, + ) + + xp = array_api_compat.array_namespace(g1, gg1) + ni, nd, nh = self.input_dim, self.hidden_dim, self.head_num + g1q = xp.reshape(self.mapq(g1), (n_total, nd, nh)) + gg1kv = xp.reshape(self.mapkv(gg1), (gg1.shape[0], nd + ni, nh)) + gg1k = gg1kv[:, :nd, :] # (E, nd, nh) + gg1v = gg1kv[:, nd:, :] # (E, ni, nh) + attnw = xp.sum(xp.take(g1q, dst, axis=0) * gg1k, axis=1) / nd**0.5 # (E, nh) + if self.smooth: + attnw = (attnw + self.attnw_shift) * sw[:, None] - self.attnw_shift + em = xp.astype(edge_mask, attnw.dtype) + n_real = segment_sum(em, dst, n_total) # (n_total,) + # SIGNED count: see Atten2Map.call_graph -- beyond sel the excess + # is subtracted so the boundary-edge denominator increment + # vanishes for arbitrary degree. + phantom = sel - n_real + attnw = segment_softmax( + attnw, + dst, + n_total, + mask=edge_mask, + phantom_count=phantom, + phantom_logit=-self.attnw_shift, + slot_weight=sw, + ) + else: + attnw = segment_softmax(attnw, dst, n_total, mask=edge_mask) + attnw = attnw * xp.astype(edge_mask[:, None], attnw.dtype) + if self.smooth: + attnw = attnw * sw[:, None] + # out[n, h, :] = sum_{e in n} w[e, h] v[e, :, h] (head-major, dense order) + contrib = attnw[:, :, None] * xp.matrix_transpose(gg1v) # (E, nh, ni) + ret = segment_sum(contrib, dst, n_total) # (N, nh, ni) + ret = xp.reshape(ret, (n_total, nh * ni)) + return self.head_map(ret) # (N, ng1) + def serialize(self) -> dict: """Serialize the networks to a dict. @@ -1851,6 +2438,68 @@ def _update_g1_conv( g1_11 = xp.sum(g2 * gg1, axis=2) * invnnei return g1_11 + def _update_g1_conv_graph( + self, + gg1: Array, + g2: Array, + edge_mask: Array, + sw: Array, + dst: Array, + n_total: int, + nnei: int, + ) -> Array: + """ + Graph twin of :meth:`_update_g1_conv` (segment_sum over dst). + + Parameters + ---------- + gg1 + Edge-wise atomic invariant rep, with shape [n_edge, ng1]. + g2 + Edge-wise pair invariant rep, with shape [n_edge, ng2]. + edge_mask + Edge mask, where zero means no edge, with shape [n_edge]. + sw + The switch function, with shape [n_edge]. + dst + Destination (center) node index of each edge, with shape [n_edge]. + n_total + Total number of nodes. + nnei + The block sel: the smooth-branch normalization constant AND the + fixed dense softmax width for the attention phantom-count + compensation (see :meth:`Atten2Map.call_graph`). + + Returns + ------- + g1_conv : Array + Convolution-updated node channel, with shape [n_total, ng1]. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + segment_sum, + ) + + xp = array_api_compat.array_namespace(gg1, g2, edge_mask, sw) + assert self.proj_g1g2 is not None + if not self.g1_out_conv: + # gg1 : n_edge x ng2 + gg1 = self.proj_g1g2(gg1) + # else gg1 : n_edge x ng1 + gg1 = gg1 * xp.astype(edge_mask[:, None], gg1.dtype) + if not self.smooth: + # normalized by number of neighbors, not smooth + deg = segment_sum(xp.astype(edge_mask, gg1.dtype), dst, n_total) # (N,) + invnnei = (1.0 / (self.epsilon + deg))[:, None] # (N, 1) + else: + gg1 = gg1 * sw[:, None] + invnnei = 1.0 / float(nnei) + if not self.g1_out_conv: + # (N, ng2) + return segment_sum(g2 * gg1, dst, n_total) * invnnei + # (N, ng1) + g2p = self.proj_g1g2(g2) + return segment_sum(g2p * gg1, dst, n_total) * invnnei + def _update_g2_g1g1( self, g1: Array, # nf x nloc x ng1 @@ -1881,6 +2530,43 @@ def _update_g2_g1g1( ret = _apply_switch(ret, sw) return ret + def _update_g2_g1g1_graph( + self, + g1: Array, + src: Array, + dst: Array, + edge_mask: Array, + sw: Array, + ) -> Array: + """ + Graph twin of :meth:`_update_g2_g1g1`: per-edge g1_center * g1_neighbor. + + Parameters + ---------- + g1 + Atomic invariant rep, with shape [n_total, ng1]. + src + Source (neighbor) node index of each edge, with shape [n_edge]. + dst + Destination (center) node index of each edge, with shape [n_edge]. + edge_mask + Edge mask, where zero means no edge, with shape [n_edge]. + sw + The switch function, with shape [n_edge]. + + Returns + ------- + g1g1 : Array + Per-edge center-neighbor g1 product, with shape [n_edge, ng1]. + """ + xp = array_api_compat.array_namespace(g1, edge_mask, sw) + # (n_edge, ng1) + ret = xp.take(g1, dst, axis=0) * xp.take(g1, src, axis=0) + ret = ret * xp.astype(edge_mask[:, None], ret.dtype) + if self.smooth: + ret = ret * sw[:, None] + return ret + def call( self, g1_ext: Array, # nf x nall x ng1 @@ -2023,6 +2709,174 @@ def call( g1_new = self.list_update(g1_update, "g1") return g1_new, g2_new, h2_new + def call_graph( + self, + g1: Array, + g2: Array, + h2: Array, + src: Array, + dst: Array, + edge_mask: Array, + sw: Array, + n_total: int, + nnei: int, + pairs: tuple[Array, Array, Array] | None = None, + ) -> tuple[Array, Array, Array]: + """Graph twin of :meth:`call`. One residual update of (g1, g2, h2) on + the flat node/edge axes; the neighbor axis of every dense reduction is + replaced by segment ops over ``dst``. + + Parameters + ---------- + g1 + Flat node-wise atomic invariant rep, with shape [n_total, ng1]. + Unlike :meth:`call`, this is the FULL node channel with no + local/ghost slice: ghost-free graphs have no ghosts, and extended + (multi-rank) graphs deliberately compute ghost rows here and let + the later communication step overwrite them. + g2 + Flat edge-wise pair invariant rep, with shape [n_edge, ng2]. + h2 + Flat edge-wise pair equivariant rep, with shape [n_edge, 3]. + src + Source (neighbor) node index of each edge, with shape [n_edge]. + dst + Destination (center) node index of each edge, with shape [n_edge]. + edge_mask + Edge mask, where zero means no edge, with shape [n_edge]. + sw + The switch function, with shape [n_edge]. + n_total + Total number of nodes. + nnei + The block sel: the smooth-branch normalization constant AND the + fixed dense softmax width for the attention phantom-count + compensation (see :meth:`Atten2Map.call_graph`). + pairs + ``(q_e, k_e, pair_mask)`` from + :func:`~deepmd.dpmodel.utils.neighbor_graph.center_edge_pairs`. + Required iff ``self.update_g2_has_attn or self.update_h2``. + + Returns + ------- + g1_new : [n_total, ng1] updated node channel + g2_new : [n_edge, ng2] updated edge channel, invariant + h2_new : [n_edge, 3] updated edge channel, equivariant + """ + xp = array_api_compat.array_namespace(g1, g2, h2) + cal_gg1 = ( + self.update_g1_has_drrd + or self.update_g1_has_conv + or self.update_g1_has_attn + or self.update_g2_has_g1g1 + ) + e_tot = g2.shape[0] + + g2_update: list[Array] = [g2] + h2_update: list[Array] = [h2] + g1_update: list[Array] = [g1] + g1_mlp: list[Array] = [g1] if not self.g1_out_mlp else [] + if self.g1_out_mlp: + assert self.g1_self_mlp is not None + g1_update.append(self.act(self.g1_self_mlp(g1))) + + gg1 = xp.take(g1, src, axis=0) if cal_gg1 else None # (E, ng1) + + if self.update_chnnl_2: + assert self.linear2 is not None + g2_update.append(self.act(self.linear2(g2))) + + if self.update_g2_has_g1g1: + assert self.proj_g1g1g2 is not None + g2_update.append( + self.proj_g1g1g2( + self._update_g2_g1g1_graph(g1, src, dst, edge_mask, sw) + ) + ) + + if self.update_g2_has_attn or self.update_h2: + assert self.attn2g_map is not None + assert pairs is not None, ( + "center_edge_pairs required for g2 attention on the graph path" + ) + q_e, k_e, pair_mask = pairs + AAg = self.attn2g_map.call_graph( + g2, h2, sw, q_e, k_e, pair_mask, e_tot, nnei + ) # (P, nh) + if self.update_g2_has_attn: + assert self.attn2_mh_apply is not None + assert self.attn2_lm is not None + g2_2 = self.attn2_mh_apply.call_graph(AAg, g2, q_e, k_e, e_tot) + g2_update.append(self.attn2_lm(g2_2)) + if self.update_h2: + assert self.attn2_ev_apply is not None + h2_update.append( + self.attn2_ev_apply.call_graph(AAg, h2, q_e, k_e, e_tot) + ) + + if self.update_g1_has_conv: + assert gg1 is not None + g1_conv = self._update_g1_conv_graph( + gg1, g2, edge_mask, sw, dst, n_total, nnei + ) + if not self.g1_out_conv: + g1_mlp.append(g1_conv) + else: + g1_update.append(g1_conv) + + if self.update_g1_has_grrg: + g1_mlp.append( + symmetrization_op_graph( + g2, + h2, + edge_mask, + sw, + dst, + n_total, + nnei, + self.axis_neuron, + smooth=self.smooth, + epsilon=self.epsilon, + use_sqrt_nnei=self.use_sqrt_nnei, + ) + ) + + if self.update_g1_has_drrd: + assert gg1 is not None + g1_mlp.append( + symmetrization_op_graph( + gg1, + h2, + edge_mask, + sw, + dst, + n_total, + nnei, + self.axis_neuron, + smooth=self.smooth, + epsilon=self.epsilon, + use_sqrt_nnei=self.use_sqrt_nnei, + ) + ) + + g1_1 = self.act(self.linear1(xp.concat(g1_mlp, axis=-1))) + g1_update.append(g1_1) + + if self.update_g1_has_attn: + assert gg1 is not None + assert self.loc_attn is not None + g1_update.append( + self.loc_attn.call_graph(g1, gg1, edge_mask, sw, dst, n_total, nnei) + ) + + if self.update_chnnl_2: + g2_new = self.list_update(g2_update, "g2") + h2_new = self.list_update(h2_update, "h2") + else: + g2_new, h2_new = g2, h2 + g1_new = self.list_update(g1_update, "g1") + return g1_new, g2_new, h2_new + def list_update_res_avg( self, update_list: list[Array], diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 4734a6be66..1216347f08 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -839,6 +839,15 @@ def call_graph( d1 = xp.reshape(descriptor, (n, 1, nd)) a1 = xp.reshape(atype, (n, 1)) g1 = None if gr is None else xp.reshape(gr, (n, 1, gr.shape[-2], 3)) + if aparam is not None and len(aparam.shape) != 2: + # enforce the flat contract loudly: a rectangular (nf, nloc, nda) + # aparam with nf*nloc == N would silently reshape into the right + # element order here but hand torch.export an unprovable + # N == nf*nloc relation (and misalign rows for any other layout). + raise ValueError( + "graph-route aparam must be flat (N, nda) on the node axis; " + f"got a rank-{len(aparam.shape)} array of shape {aparam.shape}" + ) ap1 = None if aparam is None else xp.reshape(aparam, (n, 1, aparam.shape[-1])) # fparam: dense API expects (nf, nfp); here nf'=N single-atom frames, so the # node-level (N, nfp) IS the per-(pseudo)frame param -- tiled over nloc'=1. diff --git a/deepmd/dpmodel/model/edge_transform_output.py b/deepmd/dpmodel/model/edge_transform_output.py index f9cc49f874..ce50eba87b 100644 --- a/deepmd/dpmodel/model/edge_transform_output.py +++ b/deepmd/dpmodel/model/edge_transform_output.py @@ -39,11 +39,49 @@ ) +def node_ownership_mask(n_node: Array, n_local: Array, n_total: int) -> Array: + """Derive the ``(n_total,)`` owned-node mask from per-frame local counts. + + Owned-prefix layout: frame ``f`` owns the first ``n_local[f]`` of its + contiguous ``n_node[f]``-node block (the remainder, if any, are ghost/ + ghost nodes owned by another rank). Local helper matching the (not yet + merged) ``#5758`` name/semantics so a later rebase converges. + + Parameters + ---------- + n_node + (nf,) per-frame REAL (local + ghost) node counts. + n_local + (nf,) per-frame OWNED node counts, ``n_local[f] <= n_node[f]``. + n_total + Size of the flat node axis ``N`` (``== sum(n_node)`` when unpadded). + + Returns + ------- + mask + (n_total,) boolean mask, ``True`` for nodes owned by this rank. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + frame_id_from_n_node, + ) + + xp = array_api_compat.array_namespace(n_node, n_local) + device = array_api_compat.device(n_node) + node_index = xp.arange(n_total, dtype=n_node.dtype, device=device) + frame_id = frame_id_from_n_node(n_node, n_total=n_total) + frame_end = xp.cumulative_sum(n_node) + frame_start = frame_end - n_node + index_in_frame = node_index - xp.take(frame_start, frame_id, axis=0) + local_count = xp.take(n_local, frame_id, axis=0) + return index_in_frame < local_count + + def fit_output_to_model_output_graph( fit_ret: dict[str, Array], fit_output_def: FittingOutputDef, graph: NeighborGraph, mask: Array | None = None, + n_local: Array | None = None, ) -> dict[str, Array]: """Flat-N analogue of :func:`~deepmd.dpmodel.model.transform_output.fit_output_to_model_output`. @@ -58,6 +96,16 @@ def fit_output_to_model_output_graph( ``graph.n_node`` is used (the node->frame map for the reduction). mask the ``(N,)`` real-node mask for the intensive-output denominator. + n_local + ``(nf,)`` per-frame OWNED node counts for multi-rank ghost exclusion + (owned-prefix layout, :func:`node_ownership_mask`). When given, every + reducible per-node value is masked to zero on ghost rows (index + ``>= n_local[frame]``) BEFORE the per-frame ``segment_sum`` -- each + ghost atom is owned (and counted) on another rank, so its contribution + must not double-count here. The per-node output (````) itself is + left FULL/unmasked (the C++ caller slices owned rows itself; ghost + partial forces are reverse-commed by LAMMPS). ``None`` (default): + unchanged single-rank behavior. Returns ------- @@ -75,6 +123,10 @@ def fit_output_to_model_output_graph( xp = array_api_compat.get_namespace(n_node) nf = n_node.shape[0] frame_id = frame_id_from_n_node(n_node) + n_total = next(iter(fit_ret.values())).shape[0] + owned = ( + node_ownership_mask(n_node, n_local, n_total) if n_local is not None else None + ) model_ret = dict(fit_ret.items()) for kk, vv in fit_ret.items(): vdef = fit_output_def[kk] @@ -82,12 +134,18 @@ def fit_output_to_model_output_graph( continue kk_redu = get_reduce_name(kk) vv_e = xp.astype(vv, GLOBAL_ENER_FLOAT_PRECISION) + if owned is not None: + owned_e = xp.astype(owned, GLOBAL_ENER_FLOAT_PRECISION) + vv_e = vv_e * xp.reshape(owned_e, (n_total, *([1] * (vv_e.ndim - 1)))) redu = segment_sum(vv_e, frame_id, nf) # (nf, *shape) if vdef.intensive: if mask is not None: - cnt = segment_sum( - xp.astype(mask, GLOBAL_ENER_FLOAT_PRECISION), frame_id, nf - ) + cnt_mask = xp.astype(mask, GLOBAL_ENER_FLOAT_PRECISION) + if owned is not None: + cnt_mask = cnt_mask * owned_e + cnt = segment_sum(cnt_mask, frame_id, nf) + elif owned is not None: + cnt = segment_sum(owned_e, frame_id, nf) else: cnt = xp.astype(n_node, GLOBAL_ENER_FLOAT_PRECISION) redu = redu / xp.reshape(cnt, (nf, *([1] * (redu.ndim - 1)))) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 3f4465a2de..b7ff1a0293 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -516,7 +516,12 @@ def _call_common_graph( edge_vec=ng.edge_vec, edge_mask=ng.edge_mask, fparam=fp, - aparam=ap, + # graph-lower ABI: aparam is FLAT on the node axis, (N, nda). + aparam=( + xp.reshape(ap, (nf * nloc, ap.shape[-1])) + if ap is not None + else None + ), ) # Public ABI is rectangular (nf, nloc, *); the lower is flat # (N=nf*nloc, *). Unravel per-atom keys here at the boundary. @@ -676,7 +681,11 @@ def forward_common_atomic_graph( edge_mask (E,) boolean/0-1 valid-edge mask. n_local - Per-frame owned node counts. Halo fitting outputs are masked. + Per-rank local (owned) atom counts for multi-rank inference, + ``(nf,)``. When given, ghost rows (index ``>= n_local[frame]``) + are excluded from ``_redu`` (see + :func:`fit_output_to_model_output_graph`); ``None`` (default) + is the single-rank/all-owned behavior. fparam Frame parameter, ``(nf, ndf)``. aparam @@ -708,6 +717,7 @@ def forward_common_atomic_graph( self.atomic_output_def(), graph, mask=atomic_ret["mask"] if "mask" in atomic_ret else None, + n_local=n_local, ) def call_common_lower_graph( @@ -747,7 +757,10 @@ def call_common_lower_graph( edge_mask (E,) boolean/0-1 valid-edge mask. n_local - Per-frame owned node counts. Halo fitting outputs are masked. + Per-rank local (owned) atom counts for multi-rank inference, + ``(nf,)``. When given, ghost rows (index ``>= n_local[frame]``) + are excluded from ``_redu``; ``None`` (default) is the + single-rank/all-owned behavior. fparam Frame parameter, ``(nf, ndf)``. aparam diff --git a/deepmd/dpmodel/utils/neighbor_graph/segment.py b/deepmd/dpmodel/utils/neighbor_graph/segment.py index f95671d05c..fcaa22ce57 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/segment.py +++ b/deepmd/dpmodel/utils/neighbor_graph/segment.py @@ -56,27 +56,280 @@ def segment_max(data: Array, segment_ids: Array, num_segments: int) -> Array: return xp_maximum_at(out, segment_ids, data) +def _slot_occupancy( + slot_weight: Array, + segment_ids: Array, + num_segments: int, + capacity: Array, +) -> Array: + """Per-entry fractional slot occupancy by C1 capped water-filling. + + Solves, independently per segment, ``theta_j = h(w_j * u)`` with the + C1 plateau saturator ``h(t) = t * (2 - t)`` for ``t < 1`` and + ``h(t) = 1`` for ``t >= 1``, and the largest inverse water level + ``u >= 0`` such that ``sum_j theta_j <= capacity``. Because + ``h'(1) = 0``, both the per-entry saturation boundary and the + water-filling active-set transitions are C1 in the weights -- the + min-based projection's kinks are exactly what this form removes. + + When the capacity is not binding (``count(w_j > 0) <= capacity``) + every positive-weight entry gets ``theta_j == 1`` BITWISE (a literal + 1.0 through the plateau branch); zero-weight entries get 0. + + Parameters + ---------- + slot_weight + Non-negative per-entry weights (the smooth cutoff envelope ``sw``), + with shape ``(n,)``. + segment_ids + Segment id of each entry (int64), with shape ``(n,)``. + num_segments + Number of segments. + capacity + Per-segment slot capacity (``sel``), with shape ``(num_segments,)``. + + Returns + ------- + theta : Array + Per-entry occupancy in ``[0, 1]``, with shape ``(n,)``. C1 in + ``slot_weight``, with ``theta_j -> 0`` as ``w_j -> 0`` whenever the + segment's capacity is binding. + + Notes + ----- + For a cut with ``k`` saturated (plateau) entries the water equation + ``k + sum_unsat (2 w u - w^2 u^2) = capacity`` is quadratic in ``u``; + the physical (smaller) root in cancellation-free form is + ``u = (capacity - k) / (S + sqrt(S^2 - Q (capacity - k)))`` with + ``S = sum_unsat w`` and ``Q = sum_unsat w^2``. The per-rank suffix + sums that SELECT the cut feed only comparisons (no gradient -- a + gradient-carrying cumsum on the data-dependent entry axis breaks + torch.export, see the inline comment), while the differentiable + ``S``/``Q`` at the chosen cut are rebuilt as masked ``segment_sum``. + """ + xp = array_api_compat.array_namespace(slot_weight) + dev = array_api_compat.device(slot_weight) + n = slot_weight.shape[0] + dt = slot_weight.dtype + # NOTE: no ``n == 0`` early-out -- the entry count is a data-dependent + # (unbacked) symbol under torch.export and a Python branch on it raises + # GuardOnDataDependentSymNode; every op below is empty-safe as-is. + # sort by (segment asc, weight desc): stable argsort twice + order_w = xp.argsort(slot_weight, descending=True, stable=True) + order = xp.take( + order_w, + xp.argsort(xp.take(segment_ids, order_w, axis=0), stable=True), + axis=0, + ) + sw_s = xp.take(slot_weight, order, axis=0) + seg_s = xp.take(segment_ids, order, axis=0) + ones = xp.ones((n,), dtype=dt, device=dev) + n_pos = segment_sum( + xp.astype(slot_weight > 0.0, dt), segment_ids, num_segments + ) # (S,) + # The ENTIRE cut selection below runs in float64: the suffix moments are + # formed as ``total - prefix`` and the discriminant as ``S^2 - Q*room``, + # both ill-conditioned subtractions -- in float32, cutoff weights + # spanning ordinary decades (e.g. [1, 0.1, 0.01, 5e-7]) lose the small + # suffixes entirely, the valid saturated cut is rejected, and the + # fallback root violates the capacity equation by O(1). The selection + # feeds ONLY comparisons (kstar), so the upcast costs no gradient and no + # export surface; the differentiable quantities at the chosen cut are + # rebuilt in the compute dtype from positive-only masked sums. + f64 = xp.float64 + sw64 = xp.astype(slot_weight, f64) + ones64 = xp.ones((n,), dtype=f64, device=dev) + sw64_s = xp.take(sw64, order, axis=0) + counts64 = segment_sum(ones64, segment_ids, num_segments) # (S,) + total_s = segment_sum(sw64, segment_ids, num_segments) # (S,) + total_q = segment_sum(sw64 * sw64, segment_ids, num_segments) + # within-segment 1-based rank of each sorted entry + offsets = xp.cumulative_sum(counts64) - counts64 # exclusive prefix + iota = xp.cumulative_sum(ones64) - 1.0 # arange(n) as float + rank = iota - xp.take(offsets, seg_s, axis=0) + 1.0 # (n,) 1-based + # Per-rank suffix sums S_k / Q_k (weights strictly below rank k). These + # cumsums feed ONLY the cut-selection comparisons below, so no gradient + # ever flows through them -- deliberately: the entry axis is a + # data-dependent (unbacked) size under torch.export, and a + # gradient-carrying cumsum there guards ``numel <= 1`` in its backward + # (and padding workarounds trip inductor's unbacked_bindings on the + # slice). The differentiable sums at the chosen cut are rebuilt as + # masked segment_sum below. + prefix_s = xp.cumulative_sum(sw64_s) + prefix_q = xp.cumulative_sum(sw64_s * sw64_s) + off_s = xp.take(xp.cumulative_sum(total_s) - total_s, seg_s, axis=0) + off_q = xp.take(xp.cumulative_sum(total_q) - total_q, seg_s, axis=0) + s_k = xp.take(total_s, seg_s, axis=0) - (prefix_s - off_s) # (n,) + q_k = xp.take(total_q, seg_s, axis=0) - (prefix_q - off_q) # (n,) + cap_e = xp.astype(xp.take(capacity, seg_s, axis=0), f64) # (n,) + room = cap_e - rank + disc_k = s_k * s_k - q_k * room + disc_pos = disc_k > 0.0 + sq_k = xp.where(disc_pos, xp.sqrt(xp.where(disc_pos, disc_k, ones64)), 0.0 * ones64) + den_k = s_k + sq_k + den_k_pos = den_k > 0.0 + u_k = xp.where( + den_k_pos, room / xp.where(den_k_pos, den_k, ones64), xp.zeros_like(room) + ) + # cut k (top-k entries saturated) is feasible iff a real water level + # exists (disc >= 0) and the k-th largest weight is still saturated + # (w_k * u_k >= 1); taking the LARGEST feasible k also puts every + # unsaturated weight below the plateau (verified against a bisection + # reference in the unit tests) + valid = xp.logical_and( + xp.logical_and(room > 0.0, disc_k >= 0.0), sw64_s * u_k >= 1.0 + ) + kstar64 = segment_max( + xp.where(valid, rank, xp.zeros_like(rank)), seg_s, num_segments + ) + kstar64 = xp.maximum(kstar64, xp.zeros_like(kstar64)) # empty: -inf -> 0 + kstar = xp.astype(kstar64, dt) # small integer, exact in any dtype + # differentiable S/Q of the unsaturated set (rank > kstar; the whole + # segment when kstar == 0): masked index_add of POSITIVE terms -- no + # cancellation, safe in the compute dtype + unsat = xp.astype(rank > xp.take(kstar64, seg_s, axis=0), dt) + s_star = segment_sum(sw_s * unsat, seg_s, num_segments) + q_star = segment_sum(sw_s * sw_s * unsat, seg_s, num_segments) + cap_rem = capacity - kstar + disc = s_star * s_star - q_star * cap_rem + # safe-where every branch BEFORE the sqrt/division (an unselected + # sqrt(0) backward is inf, and 0 * inf = NaN in float32 autodiff) + dpos = disc > 0.0 + one_g = xp.ones_like(disc) + sq = xp.where(dpos, xp.sqrt(xp.where(dpos, disc, one_g)), xp.zeros_like(disc)) + den_u = s_star + sq + upos = xp.logical_and(den_u > 0.0, cap_rem > 0.0) + u = xp.where(upos, cap_rem / xp.where(upos, den_u, one_g), xp.zeros_like(disc)) + # binding <=> even full saturation of every positive weight exceeds cap + binding = n_pos > capacity + t = slot_weight * xp.take(u, segment_ids, axis=0) + theta_soft = xp.where(t >= 1.0, ones, t * (2.0 - t)) + return xp.where( + xp.take(binding, segment_ids, axis=0), + theta_soft, + xp.astype(slot_weight > 0.0, dt), + ) + + def segment_softmax( data: Array, segment_ids: Array, num_segments: int, mask: Array | None = None, + phantom_count: Array | None = None, + phantom_logit: float = 0.0, + slot_weight: Array | None = None, ) -> Array: """Softmax over entries sharing a segment id, numerically stable. Mirrors the dense ``np_softmax`` max-subtraction trick with a PER-SEGMENT - max. ``mask`` (bool, per entry) removes masked entries from the softmax - entirely (zero weight AND excluded from the denominator). Empty or - fully-masked segments produce all-zero weights (no NaN). + max. Empty or fully-masked segments produce all-zero weights (no NaN). + + Parameters + ---------- + data + Per-entry logits, with shape ``(n, *f)``. + segment_ids + Segment id of each entry (int64), with shape ``(n,)``. + num_segments + Number of segments. + mask + Optional per-entry bool mask, with shape ``(n,)``. Masked entries are + removed from the softmax entirely (zero weight AND excluded from the + denominator). + phantom_count + Optional per-segment SIGNED count of virtual entries, with shape + ``(num_segments,)``. Each unit adds (or, when negative, removes) one + ``exp(phantom_logit)`` term to the segment's DENOMINATOR only (no + output rows are produced for them). Negative counts express the + sel-free carry-all convention beyond ``sel`` -- see Notes. + phantom_logit + The raw logit of every phantom entry. + slot_weight + Per-entry smooth slot weight (the cutoff envelope ``sw``), with + shape ``(n,)``. REQUIRED whenever ``phantom_count`` is given: it + drives the soft slot occupancy that keeps the phantom denominator + strictly positive and cutoff-continuous (see Notes). Must vanish + exactly when the entry's logit reaches ``phantom_logit`` at the + cutoff (the smooth-envelope invariant of every caller). + + Returns + ------- + weights : Array + Per-entry softmax weights, with shape ``(n, *f)``. + + Raises + ------ + ValueError + If ``phantom_count`` is given without ``slot_weight``. + + Notes + ----- + The phantom machinery reproduces the dense smooth-attention convention + -- a fixed ``sel``-width softmax whose padding slots each hold + ``exp(-attnw_shift)`` -- on a ragged edge set whose entry count varies + with geometry, via SOFT SLOT OCCUPANCY: each entry occupies + ``theta_j`` of the segment's ``sel`` slots (C1 capped water-filling of + the envelope weights, see :func:`_slot_occupancy`) and contributes the + C1 bracket ``theta_j e_j + (1 - theta_j)(e_j - P) chi_j`` to the + denominator (``chi = (e/P)**(1/theta)`` below ``P``, else 1), with the + unoccupied capacity contributing ``relu(sel - sum theta) * P`` + (``e_j = exp(l_j)``, ``P = exp(phantom_logit)``). Properties: + + - ``n_real <= sel``: every ``theta == 1`` bitwise, giving the EXACT + dense fixed-width denominator ``sum e_j + (sel - n) P`` term for + term, for arbitrary logits -- including logits below + ``phantom_logit``. + - in-design (all logits >= ``phantom_logit``): ``chi == 1`` and the + total telescopes to the signed + ``sum e_j - (n - sel) P``, independent of ``theta`` -- the sel-free + carry-all compensation whose per-entry contribution vanishes at the + cutoff (the boundary logit equals ``phantom_logit``), keeping the + weights continuous when an entry enters or leaves for ARBITRARY + degree. + - strictly positive for any finite logits: the bracket is bounded below + by ``theta (1 - 1/e0) e`` -- the raw signed denominator's zero + crossing (reachable for negative counts because the envelope maps + ``raw < -attnw_shift`` to ``l < phantom_logit`` at ``sw > 0``) cannot + occur; no floor term is needed. + - C1 (first-derivative-continuous) in the logits and envelope: the + damped tail meets the in-design branch at ``e == P`` with matching + slope, and the occupancy's plateau saturator makes water-filling + active-set changes kink-free -- weight GRADIENTS, hence forces, are + continuous across the below-phantom surface (second derivatives + still jump there: exact in-design linearity and smoothness beyond C1 + are mutually exclusive with bitwise dense parity). + - cutoff-continuous across the count 0 -> -1 transition for arbitrary + logits: the entering entry has BOTH ``e -> P`` and ``theta -> 0``, so + its bracket vanishes; per-entry weights are bounded by + ``(e0/(e0 - 1)) / theta_j`` (``n_real / sel``-scale), and the + caller's post-softmax ``sw`` factor keeps the boundary entry's own + contribution vanishing. """ xp = array_api_compat.array_namespace(data) + dev = array_api_compat.device(data) if mask is not None: + # broadcast mask (n,) over any trailing feature dims of data (n, *f) + mask_b = xp.reshape(mask, mask.shape + (1,) * (data.ndim - 1)) # keep masked entries out of the per-segment max: send them to -inf neg = xp.full_like(data, -xp.inf) - data_for_max = xp.where(mask, data, neg) + data_for_max = xp.where(mask_b, data, neg) else: data_for_max = data seg_max = segment_max(data_for_max, segment_ids, num_segments) + ph_b = None + if phantom_count is not None: + # phantom entries participate in the max exactly like dense's padding + # slots do in np_softmax's row max (dense: m = max(real, -shift)); + # this also guards exp-overflow when every real logit < phantom_logit. + ph_b = xp.reshape( + phantom_count, phantom_count.shape + (1,) * (seg_max.ndim - 1) + ) + seg_max = xp.where( + ph_b > 0, + xp.maximum(seg_max, xp.full_like(seg_max, phantom_logit)), + seg_max, + ) # guard -inf (empty / fully-masked segments) so gather doesn't yield inf-inf seg_max = xp.where(xp.isinf(seg_max), xp.zeros_like(seg_max), seg_max) # shift data_for_max (masked entries already -inf), NOT the raw data: @@ -89,8 +342,110 @@ def segment_softmax( if mask is not None: # defensive no-op after the -inf shift (exp(-inf) == 0); kept so the # zero-weight guarantee never depends on the shift implementation - ex = ex * xp.astype(mask, ex.dtype) - denom = segment_sum(ex, segment_ids, num_segments) + ex = ex * xp.astype(mask_b, ex.dtype) + if ph_b is not None: + if slot_weight is None: + raise ValueError( + "segment_softmax: phantom_count requires slot_weight (the " + "smooth per-entry envelope) -- the sel-slot occupancy that " + "keeps the phantom denominator positive and cutoff-continuous " + "cannot be derived from the logits alone" + ) + # exponent clamped at 80 (exp(80) ~ 5.5e34 is finite in BOTH float32 + # and float64): pl - seg_max > 80 means every real logit sits > 80 + # below the phantom logit; an unclamped exp would overflow (inf in + # float32 already at ~88) and poison the sum with inf - inf. + ph_arg = xp.minimum( + xp.full_like(seg_max, phantom_logit) - seg_max, + xp.full_like(seg_max, 80.0), + ) + ph_term = xp.exp(ph_arg) + # SOFT SLOT OCCUPANCY: real logits are not bounded below by + # ``phantom_logit`` (the smooth envelope maps ``raw < -shift`` to + # ``l < -shift`` at ``sw > 0``), so the plain SIGNED denominator + # ``sum_j e_j + (sel - n) * P`` crosses zero for negative counts -- + # a pole reachable with finite, valid model parameters. Instead, + # each entry occupies ``theta_j`` of the segment's ``sel`` dense + # slots (theta = C1 capped water-filling of the envelope weights, + # see :func:`_slot_occupancy`) and contributes the C1 bracket + # built below. Properties (details in the class docstring Notes): + # bitwise dense for n <= sel; theta-independent signed formula + # in-design; strictly positive; C1 in logits and envelope (forces + # carry no kink at the below-phantom surface or at water-filling + # active-set changes); cutoff-continuous at every crossing. + if mask is not None: + m_f = xp.astype(mask, ex.dtype) + n_real_seg = segment_sum(m_f, segment_ids, num_segments) + sw_eff = xp.astype(slot_weight, ex.dtype) * m_f + else: + n_real_seg = segment_sum( + xp.ones(data.shape[:1], dtype=ex.dtype, device=dev), + segment_ids, + num_segments, + ) + sw_eff = xp.astype(slot_weight, ex.dtype) + cap = n_real_seg + xp.astype( + xp.reshape(phantom_count, (-1,)), ex.dtype + ) # (S,) == sel + theta = _slot_occupancy(sw_eff, segment_ids, num_segments, cap) + theta_b = xp.reshape(theta, theta.shape + (1,) * (ex.ndim - 1)) + ph_e = xp.take(ph_term, segment_ids, axis=0) # (n, *f) + # Per-entry bracket, C1 in the logit: + # + # B = theta*e + (1 - theta) * (e - P) * chi, + # chi = (e/P)**(1/theta) for e < P, else 1 + # + # For e >= P this is the exact linear in-design branch (telescopes + # to the signed formula, theta-independent). For e < P the + # (e/P)**(1/theta) damping meets the linear branch at e == P with + # value theta*P AND slope 1 (chi -> 1, (e-P)*chi' -> 0), so the + # below-phantom surface carries no force kink (the relu form jumped + # the slope theta*P -> P there). The damped correction is bounded: + # with a = e/P, factoring out e gives + # B = e * [theta + (1-theta)(a-1)a**(1/theta - 1)], and on (0, 1) + # the factor (a-1)a**(1/theta - 1) attains its minimum + # -theta*(1-theta)**(1/theta - 1) at a = 1 - theta, so + # B/e >= theta*(1 - (1-theta)**(1/theta)) >= theta*(1 - 1/e0) + # (e0 = Euler's number; (1-theta)**(1/theta) increases to 1/e0 as + # theta -> 0). Hence B > 0.63*theta*e -- strictly positive with + # per-entry weights bounded by ~1.6/theta. theta == 1 is BITWISE + # ``e`` with no special case: the (1 - theta) = 0 factor kills the + # finite tail term exactly. theta -> 0 or e -> 0 send the tail to + # zero. + # + # chi is computed ENTIRELY IN LOG SPACE: log(e/P) is the already + # available ``shifted - ph_arg``, so no log() and no division by + # ``ph_e`` ever enter the autodiff graph. This matters for + # in-design HIGH logits (count >= 0): ``ph_e = exp(pl - m)`` goes + # subnormal when the segment max is large, and a log(ex_t / ph_e) + # formulation -- even with ex_t where-substituted to ph_e -- pushes + # ``1 / ph_e`` (torch) or ``ph_e**2`` (jax) past the float32 range + # in the UNSELECTED branch's backward, whose 0 * inf poisons the + # selected gradient with NaN. The log-ratio is clamped at -1e4 + # (exp(-1e4) == 0 in both precisions) so masked entries' -inf + # shifted logits cannot reach the division either; the only division + # is by the where-substituted theta. + one_t = xp.ones_like(ex) + log_ratio = xp.maximum( + shifted - xp.take(ph_arg, segment_ids, axis=0), + xp.full_like(ex, -1.0e4), + ) + below = log_ratio < 0.0 + tail_on = xp.logical_and(below, theta_b > 0.0) + lr_t = xp.where(tail_on, log_ratio, xp.zeros_like(ex)) + th_t = xp.where(tail_on, theta_b * one_t, one_t) + chi = xp.where(tail_on, xp.exp(lr_t / th_t), one_t) + # zero-occupancy entries contribute nothing below P + dead = xp.logical_and(below, xp.logical_not(tail_on)) + chi = xp.where(dead, xp.zeros_like(chi), chi) + theta_dead = xp.where(dead, xp.zeros_like(theta_b), theta_b * one_t) + bracket = theta_dead * ex + (1.0 - theta_dead) * (ex - ph_e) * chi + denom = segment_sum(bracket, segment_ids, num_segments) + occ = segment_sum(theta, segment_ids, num_segments) # (S,) + free = xp.maximum(cap - occ, xp.zeros_like(cap)) + denom = denom + xp.reshape(free, free.shape + (1,) * (denom.ndim - 1)) * ph_term + else: + denom = segment_sum(ex, segment_ids, num_segments) denom_e = xp.take(denom, segment_ids, axis=0) safe = xp.where(denom_e > 0, denom_e, xp.ones_like(denom_e)) return ex / safe diff --git a/deepmd/pt/utils/compile_compat.py b/deepmd/pt/utils/compile_compat.py index 0e64f6b040..a3cc33982f 100644 --- a/deepmd/pt/utils/compile_compat.py +++ b/deepmd/pt/utils/compile_compat.py @@ -152,6 +152,57 @@ def is_prime(n: int) -> bool: return True +def forbidden_dims_from_model( + model: torch.nn.Module, + task_buf_vals: tuple[torch.Tensor, ...] = (), +) -> set[int]: + """Prime-collision set for trace-dim selection. + + Collects every ``> 1`` dim of the model's parameters/buffers (so + :func:`next_safe_prime` never aliases an internal dim like ``g2_dim`` / + ``axis_neuron`` / ``attn_head`` without a hardcoded list), plus + ``dim_fparam``/``dim_aparam`` and the task-buffer dims. Shared by the + compiled-training traces (``_trace_and_compile`` / + ``_trace_and_compile_graph``) and the graph ``.pt2`` export trace + (``_trace_and_export``); each caller adds its path-specific dims + (nall/nloc/nsel for dense, charge_spin for both) on top of this base set. + + Parameters + ---------- + model : torch.nn.Module + The model whose parameter/buffer/conditioning dims to collect. + task_buf_vals : tuple of torch.Tensor + Per-task buffers promoted to FX placeholders (multi-task compiled + training); their dims join the forbidden set. + + Returns + ------- + set of int + Every ``> 1`` dimension a trace-time size must not collide with. + """ + forbidden: set[int] = { + int(_d) + for _src in (model.parameters(), model.buffers()) + for _p in _src + for _d in _p.shape + if _d > 1 + } + for _getter_name in ("get_dim_fparam", "get_dim_aparam"): + try: + # resolve inside the try: a model without the accessor must fall + # through the best-effort path, not raise during tuple building + _dim = getattr(model, _getter_name)() + if _dim > 1: + forbidden.add(int(_dim)) + except Exception: + pass # best-effort: dim unavailable -> nothing to forbid + for _tbv in task_buf_vals: + for _d in _tbv.shape: + if _d > 1: + forbidden.add(int(_d)) + return forbidden + + def next_safe_prime(start: int, forbidden: set[int]) -> int: """Return the smallest prime ``>= max(start, 5)`` not in ``forbidden``. diff --git a/deepmd/pt_expt/descriptor/dpa1.py b/deepmd/pt_expt/descriptor/dpa1.py index c1520402fb..b6a7f9f84f 100644 --- a/deepmd/pt_expt/descriptor/dpa1.py +++ b/deepmd/pt_expt/descriptor/dpa1.py @@ -273,6 +273,58 @@ def _type_pair_table(desc: Any, type_embedding: torch.Tensor) -> torch.Tensor: class DescrptDPA1(DescrptDPA1DP): _update_sel_cls = UpdateSel + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + # Persisted graph-routing knob (first-class training configuration): + # ``disable_graph_lower()`` used to flip only the plain dpmodel bool, + # which a Trainer checkpoint restart silently reset (the fresh model + # is rebuilt from config before ``load_state_dict``, and neither the + # state-dict keys nor ``_extra_state.model_params`` carried the + # choice) -- on a binding-sel system that switched the training + # equation and gradients without warning. A persistent buffer rides + # every pt_expt state_dict, so save/restart round-trips it. + torch.nn.Module.register_buffer( + self, + "graph_lower_disabled", + torch.zeros((), dtype=torch.bool, device="cpu"), + ) + + def disable_graph_lower(self) -> None: + """Persisted variant of the dpmodel escape hatch (see base class). + + The buffer (and the routing bool) are PER-TASK state: multi-task + ``share_params`` shares network submodules, not this buffer, so + disabling the graph lower on one task branch does not propagate to + branches sharing the same descriptor weights -- each branch owns + its routing decision. + """ + super().disable_graph_lower() + self.graph_lower_disabled.fill_(True) + + def _load_from_state_dict( + self, + state_dict: dict[str, Any], + prefix: str, + *args: Any, + **kwargs: Any, + ) -> None: + # Back-compat: checkpoints written before the knob was persisted lack + # the buffer; default to the fresh module's value (graph enabled) + # instead of failing the strict load. + key = prefix + "graph_lower_disabled" + if key not in state_dict: + state_dict[key] = self.graph_lower_disabled.detach().clone() + else: + # Re-sync the dpmodel-side routing bool from the RESTORED value + # here, at load time, where the incoming tensor is real. The + # routing predicate itself must stay a plain python bool: + # ``uses_graph_lower()`` runs inside traced forwards (the dense + # adapter gate), and reading the buffer there would emit a + # data-dependent ``bool(FakeTensor)`` guard that breaks + # torch.export (GuardOnDataDependentSymNode Eq(u0, 1)). + self._graph_lower_disabled = bool(state_dict[key]) + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + def share_params( self, base_class: Any, @@ -451,6 +503,7 @@ def call_graph( atype: torch.Tensor, type_embedding: torch.Tensor | None = None, static_nnei: int | None = None, + comm_dict: dict | None = None, ) -> Any: """Graph-native forward, routed through the fused edge kernel when eligible. @@ -499,10 +552,15 @@ def call_graph( graph = dataclasses.replace(graph, edge_vec=graph.edge_vec.to(prec)) # Triton edge convolution: serves the uncompressed concat / strip lower; # a tabulated (geo_compress) descriptor stays on the table path below. + # Descriptor-level exclusions stay on the reference path: the fused + # edge kernel bypasses the reference ``apply_pair_exclusion(graph, + # atype, emask)`` (descriptor exclude_types is independent of the + # model-level pair_exclude_types already applied at graph build). if ( fused and triton_infer_level() >= 1 and not self.geo_compress + and not self.se_atten.exclude_types and self._fused_eligible("triton") ): return self._call_graph_triton(graph, atype, type_embedding) @@ -511,7 +569,12 @@ def call_graph( if self.geo_compress: return self._call_graph_compress_reference(graph, atype, type_embedding) return DescrptDPA1DP.call_graph( - self, graph, atype, type_embedding=type_embedding, static_nnei=static_nnei + self, + graph, + atype, + type_embedding=type_embedding, + static_nnei=static_nnei, + comm_dict=comm_dict, ) def _fused_eligible(self, backend: str) -> bool: @@ -589,6 +652,12 @@ def _fused_eligible(self, backend: str) -> bool: return ( TRITON_AVAILABLE and se.tebd_input_mode in ("strip", "concat") + # No exclude_types condition here: exclusion handling differs per + # dispatch site. The dense ``_call_triton`` applies the descriptor + # emask itself (``_env_mat`` masks nlist / sw / rr), so it serves + # excluded models; the graph ``_call_graph_triton`` bypasses the + # reference ``apply_pair_exclusion`` and is gated at its dispatch + # site in :meth:`call_graph` instead. and str(layers[-1].activation_function).lower() in ACT_CODES ) diff --git a/deepmd/pt_expt/descriptor/dpa2.py b/deepmd/pt_expt/descriptor/dpa2.py index 01a7dfde32..a9943635e7 100644 --- a/deepmd/pt_expt/descriptor/dpa2.py +++ b/deepmd/pt_expt/descriptor/dpa2.py @@ -33,6 +33,58 @@ class DescrptDPA2(DescrptDPA2DP): _update_sel_cls = UpdateSel + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + # Persisted graph-routing knob (first-class training configuration): + # ``disable_graph_lower()`` used to flip only the plain dpmodel bool, + # which a Trainer checkpoint restart silently reset (the fresh model + # is rebuilt from config before ``load_state_dict``, and neither the + # state-dict keys nor ``_extra_state.model_params`` carried the + # choice) -- on a binding-sel system that switched the training + # equation and gradients without warning. A persistent buffer rides + # every pt_expt state_dict, so save/restart round-trips it. + torch.nn.Module.register_buffer( + self, + "graph_lower_disabled", + torch.zeros((), dtype=torch.bool, device="cpu"), + ) + + def disable_graph_lower(self) -> None: + """Persisted variant of the dpmodel escape hatch (see base class). + + The buffer (and the routing bool) are PER-TASK state: multi-task + ``share_params`` shares network submodules, not this buffer, so + disabling the graph lower on one task branch does not propagate to + branches sharing the same descriptor weights -- each branch owns + its routing decision. + """ + super().disable_graph_lower() + self.graph_lower_disabled.fill_(True) + + def _load_from_state_dict( + self, + state_dict: dict[str, Any], + prefix: str, + *args: Any, + **kwargs: Any, + ) -> None: + # Back-compat: checkpoints written before the knob was persisted lack + # the buffer; default to the fresh module's value (graph enabled) + # instead of failing the strict load. + key = prefix + "graph_lower_disabled" + if key not in state_dict: + state_dict[key] = self.graph_lower_disabled.detach().clone() + else: + # Re-sync the dpmodel-side routing bool from the RESTORED value + # here, at load time, where the incoming tensor is real. The + # routing predicate itself must stay a plain python bool: + # ``uses_graph_lower()`` runs inside traced forwards (the dense + # adapter gate), and reading the buffer there would emit a + # data-dependent ``bool(FakeTensor)`` guard that breaks + # torch.export (GuardOnDataDependentSymNode Eq(u0, 1)). + self._graph_lower_disabled = bool(state_dict[key]) + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + def share_params( self, base_class: "DescrptDPA2", diff --git a/deepmd/pt_expt/descriptor/repformers.py b/deepmd/pt_expt/descriptor/repformers.py index 9b8ddb4a85..91a65dc5b2 100644 --- a/deepmd/pt_expt/descriptor/repformers.py +++ b/deepmd/pt_expt/descriptor/repformers.py @@ -92,6 +92,63 @@ def _exchange_ghosts( return concat_switch_virtual(real_ext, virt_ext, real_nloc) return exchanged + def _exchange_ghosts_graph( + self, + g1: torch.Tensor, + comm_dict: dict | None, + n_total: int, + ) -> torch.Tensor: + """Graph-path per-layer ghost refresh via ``border_op``. + + Flat single-frame counterpart of :meth:`_exchange_ghosts`: ``g1`` is + already ``(N, ng1)`` with ``N == nlocal + nghost`` (owned prefix + first), so no squeeze/pad/unsqueeze dance is needed -- ``border_op`` + overwrites the ghost rows ``[nlocal, N)`` in place with the owner + rows and returns the same tensor. Identity without ``comm_dict`` + (ghost-free Python graphs / extended single-process graphs). Spin + models never route the graph lower (``disable_graph_lower``), so a + ``has_spin`` comm_dict reaching here is a programming error. + + Parameters + ---------- + g1 + Flat node-wise atomic invariant rep, with shape [n_total, ng1]. + comm_dict + MPI communication metadata (``send_list``, ``send_proc``, + ``recv_proc``, ``send_num``, ``recv_num``, ``communicator``, + ``nlocal``, ``nghost``); ``None`` for single-process graphs. + n_total + Total number of nodes (``nlocal + nghost``). + + Returns + ------- + g1 : torch.Tensor + The node channel with ghost rows refreshed, with shape + [n_total, ng1]. + + Raises + ------ + NotImplementedError + If ``comm_dict`` carries ``has_spin``. + """ + if comm_dict is None: + return super()._exchange_ghosts_graph(g1, comm_dict, n_total) + if "has_spin" in comm_dict: + raise NotImplementedError( + "spin models do not route the graph lower (disable_graph_lower)" + ) + return torch.ops.deepmd_export.border_op( + comm_dict["send_list"], + comm_dict["send_proc"], + comm_dict["recv_proc"], + comm_dict["send_num"], + comm_dict["recv_num"], + g1, + comm_dict["communicator"], + comm_dict["nlocal"], + comm_dict["nghost"], + ) + register_dpmodel_mapping( DescrptBlockRepformersDP, diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 66efeddb89..1b97ef7f50 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -1832,6 +1832,22 @@ def _eval_model_graph( ) if self.metadata.get("lower_input_kind") == "dpa1_canonical": + # The canonical ABI has NO fparam/aparam/charge_spin slots; the + # export gate (fitting_eligible) rejects such models today, so + # this is unreachable -- assert it loudly so a future loosening + # of the eligibility fails here instead of silently dropping + # conditioning inputs at inference. + if ( + self.get_dim_fparam() > 0 + or self.get_dim_aparam() > 0 + or int(self.metadata.get("dim_chg_spin", 0) or 0) > 0 + ): + raise NotImplementedError( + "dpa1_canonical artifacts carry no fparam/aparam/" + "charge_spin inputs; a model requiring them must not be " + "frozen with lower_kind='dpa1_canonical' (the export " + "eligibility gate should have rejected it)." + ) from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, ) @@ -1866,6 +1882,11 @@ def _eval_model_graph( fparam_t, aparam_t = self._prepare_optional_lower_inputs( fparam, aparam, nframes, natoms, DEVICE ) + if aparam_t is not None: + # graph-lower ABI: aparam is FLAT on the node axis, (N, nda) + # -- the same axis as ``atype`` (the shared helper above + # returns the dense rectangular (nf, natoms, nda) layout). + aparam_t = aparam_t.reshape(nframes * natoms, -1) charge_spin_t = self._make_charge_spin_input(nframes, charge_spin) model_inputs = ( atype_t, diff --git a/deepmd/pt_expt/model/edge_transform_output.py b/deepmd/pt_expt/model/edge_transform_output.py index 24c469c0a5..52761db645 100644 --- a/deepmd/pt_expt/model/edge_transform_output.py +++ b/deepmd/pt_expt/model/edge_transform_output.py @@ -12,6 +12,9 @@ get_deriv_name, get_reduce_name, ) +from deepmd.dpmodel.model.edge_transform_output import ( + node_ownership_mask, +) from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, edge_force_virial, @@ -151,6 +154,7 @@ def fit_output_to_model_output_graph( create_graph: bool = True, mask: torch.Tensor | None = None, node_capacity: int | None = None, + n_local: torch.Tensor | None = None, force_precision: torch.dtype | None = None, ) -> dict[str, torch.Tensor]: """Graph analogue of the dense pt_expt ``fit_output_to_model_output``. @@ -194,6 +198,20 @@ def fit_output_to_model_output_graph( input node axis rather than a re-derived shape -- hardening; the actual CUDA out-of-bounds device-assert is prevented by the index clamp in :func:`~deepmd.dpmodel.utils.neighbor_graph.derivatives.edge_force_virial`. + n_local + ``(nf,)`` per-frame OWNED node counts for multi-rank ghost exclusion + (owned-prefix layout, :func:`~deepmd.dpmodel.model.edge_transform_output.node_ownership_mask`). + When given, every reducible per-node value is masked to zero on ghost + rows (index ``>= n_local[frame]``) BEFORE the per-frame + ``segment_sum`` -- each ghost atom is owned (and counted) on another + rank, so it must not double-count into THIS rank's differentiated + energy. Critically, the mask is applied BEFORE ``edge_energy_deriv`` + differentiates the reduced value, so ``grad(energy, edge_vec)`` (and + therefore force/virial/atom-virial) only carries owned-energy terms. + The per-node output (````) itself stays FULL/unmasked (the C++ + caller slices owned rows itself; ghost partial forces are + reverse-commed by LAMMPS -- dpa1-MP precedent). ``None`` (default): + unchanged single-rank behavior. force_precision Compute precision (model dtype) in which to assemble the force / virial during inference, decoupled from the fp64 ``edge_vec`` leaf; see @@ -230,6 +248,13 @@ def fit_output_to_model_output_graph( frame_id = frame_id_from_n_node( n_node, n_total=N ) # (N,) int64 frame index per atom + # owned-node (multi-rank ghost) mask: (N,) bool, True for owned rows. + # Computed once (array-API pure, works directly on torch tensors) and + # applied to every reducible per-node value BEFORE its segment_sum, so + # the downstream force/virial autograd (which differentiates the + # ALREADY-masked ``_redu``) only carries owned-energy terms. + owned = node_ownership_mask(n_node, n_local, N) if n_local is not None else None + owned_e = owned.to(redu_prec) if owned is not None else None model_ret: dict[str, torch.Tensor] = dict(fit_ret.items()) for kk, vv in fit_ret.items(): vdef = fit_output_def[kk] @@ -239,13 +264,22 @@ def fit_output_to_model_output_graph( kk_redu = get_reduce_name(kk) # segment_sum reduces axis 0 (the flat atom axis) per frame vv_e = vv.to(redu_prec) # (N, *shape) + if owned_e is not None: + vv_e = vv_e * owned_e.reshape(N, *([1] * (vv_e.ndim - 1))) redu = segment_sum(vv_e, frame_id, nf) # (nf, *shape) if vdef.intensive: if mask is not None: # real-atom count per frame: segment_sum of the mask - cnt = segment_sum(mask.to(redu_prec), frame_id, nf) # (nf,) + cnt_mask = mask.to(redu_prec) + if owned_e is not None: + cnt_mask = cnt_mask * owned_e + cnt = segment_sum(cnt_mask, frame_id, nf) # (nf,) # broadcast cnt to (nf, 1, ..., 1) to match redu shape cnt = cnt.reshape(nf, *([1] * (redu.ndim - 1))) + elif owned_e is not None: + cnt = segment_sum(owned_e, frame_id, nf).reshape( + nf, *([1] * (redu.ndim - 1)) + ) else: cnt = n_node.to(redu_prec).reshape(nf, *([1] * (redu.ndim - 1))) redu = redu / cnt diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index d0990ffe17..390afba79b 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -583,3 +583,188 @@ def fn( aparam, charge_spin, ) + + def forward_lower_graph_exportable_with_comm( + self, + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + send_list: torch.Tensor, + send_proc: torch.Tensor, + recv_proc: torch.Tensor, + send_num: torch.Tensor, + recv_num: torch.Tensor, + communicator: torch.Tensor, + nlocal: torch.Tensor, + nghost: torch.Tensor, + do_atomic_virial: bool = False, + **make_fx_kwargs: Any, + ) -> torch.nn.Module: + """Trace ``forward_common_lower_graph`` with comm_dict tensors as + additional positional inputs -- the with-comm counterpart of + :meth:`forward_lower_graph_exportable` for message-passing graph + descriptors (dpa2's repformer block drives cross-rank ghost refresh + via ``deepmd_export::border_op``, see + :meth:`~deepmd.pt_expt.descriptor.repformers. + DescrptBlockRepformers._exchange_ghosts_graph`). + + Mirrors the dense ``forward_common_lower_exportable_with_comm`` + (``pt_expt/model/make_model.py``): packs the 8 trailing positional + comm tensors into a ``comm_dict`` inside the traced function. Also + derives ``n_local`` (the per-frame OWNED node count, reshaped to + ``(1,)``; single-frame -- LAMMPS always drives inference with + ``nf=1``) from the scalar ``nlocal`` tensor, so the differentiated + reduction excludes ghost (not-owned) nodes (see + :meth:`forward_common_lower_graph`'s ``n_local`` parameter). Unlike + the plain-graph export path (which traces + ``forward_common_lower_graph_exportable`` and then wraps a SECOND + make_fx trace around the key-translation closure), this method + traces ONCE: the comm-dict packing, ``n_local`` derivation, the + ``forward_common_lower_graph`` call and the key translation all live + in a single traced ``fn`` -- following the dense with-comm + precedent, which is also a single trace. + + Parameters + ---------- + atype, n_node, edge_index, edge_vec, edge_mask, fparam, aparam, charge_spin, do_atomic_virial + As in :meth:`forward_lower_graph_exportable`. + send_list, send_proc, recv_proc, send_num, recv_num, communicator, nlocal, nghost + The 8 comm tensors (see ``_make_comm_sample_inputs`` in + ``serialization.py``), packed into ``comm_dict`` inside the + traced function. + + Runtime device contract: ALL 8 stay on CPU, symmetric with + the dense with-comm artifact -- they are consumed only by the + opaque ``border_op`` whose HOST code dereferences their + ``data_ptr`` (``send_list`` carries raw host pointers) and + reads ``nlocal``/``nghost`` via cheap host ``.item()`` calls. + Deriving the in-graph owned count from a device-placed + ``nlocal`` instead (the previous design) made every per-layer + ``border_op`` forward AND custom backward pull the scalars + back with synchronizing D2H reads (``4 * nlayers`` per MD + step). The C++ ``run_model_graph_with_comm`` implements this + placement. + n_local + (1,) int64 ON THE MODEL DEVICE: the per-frame OWNED node + count consumed IN-GRAPH by the owned-node energy mask (it + becomes a device kernel operand after + ``move_to_device_pass``, like ``n_node``; a CPU tensor fed + there is read as a device pointer -- CUDA illegal memory + access). Carries the same value as the ``nlocal`` comm + tensor; the two inputs exist precisely to separate the + device-compute role from the host-MPI-control role. + **make_fx_kwargs + Extra keyword arguments forwarded to ``make_fx`` + (e.g. ``tracing_mode="symbolic"``). + + Returns + ------- + torch.nn.Module + A traced module whose ``forward`` accepts ``(atype, n_node, + n_local, edge_index, edge_vec, edge_mask, destination_order, + destination_row_ptr, source_order, source_row_ptr, fparam, + aparam, charge_spin, send_list, send_proc, recv_proc, send_num, + recv_num, communicator, nlocal, nghost)`` and returns a dict with the + SAME public keys as :meth:`forward_lower_graph_exportable` + (``atom_energy``, ``energy``, ``force``, ``virial``, + ``atom_virial`` when ``do_atomic_virial``). + """ + model = self + do_grad_r = self.do_grad_r("energy") + do_grad_c = self.do_grad_c("energy") + + def fn( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + send_list: torch.Tensor, + send_proc: torch.Tensor, + recv_proc: torch.Tensor, + send_num: torch.Tensor, + recv_num: torch.Tensor, + communicator: torch.Tensor, + nlocal: torch.Tensor, + nghost: torch.Tensor, + ) -> dict[str, torch.Tensor]: + comm_dict = { + "send_list": send_list, + "send_proc": send_proc, + "recv_proc": recv_proc, + "send_num": send_num, + "recv_num": recv_num, + "communicator": communicator, + "nlocal": nlocal, + "nghost": nghost, + } + # ``n_local`` (slot 2, DEVICE) is the owned-count input consumed + # by the in-graph owned-node mask; the CPU ``nlocal`` comm + # tensor is host control metadata for border_op only. + model_ret = model.forward_common_lower_graph( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + destination_sorted=True, + do_atomic_virial=do_atomic_virial, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + comm_dict=comm_dict, + ) + return _translate_energy_keys( + model_ret, + do_grad_r=do_grad_r, + do_grad_c=do_grad_c, + do_atomic_virial=do_atomic_virial, + local=True, + ) + + return make_fx(fn, **make_fx_kwargs)( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + send_list, + send_proc, + recv_proc, + send_num, + recv_num, + communicator, + nlocal, + nghost, + ) diff --git a/deepmd/pt_expt/model/graph_lower.py b/deepmd/pt_expt/model/graph_lower.py index fce1a63814..788d9191a1 100644 --- a/deepmd/pt_expt/model/graph_lower.py +++ b/deepmd/pt_expt/model/graph_lower.py @@ -29,6 +29,17 @@ def model_uses_graph_lower(model: Any) -> bool: except (AttributeError, NotImplementedError): return False + # ENERGY-output models only, mirroring ``_resolve_graph_method``'s + # default-flip gate: the compiled-training graph trace is + # energy-specific (``do_grad_r("energy")``, ``_translate_energy_keys``), + # so a non-energy model (property/dos/dipole/polar) on this path raises + # ``KeyError('energy')`` at its first compiled batch. + try: + if "energy" not in model.atomic_output_def().keys(): + return False + except (AttributeError, NotImplementedError): + return False + descriptor = getattr(getattr(model, "atomic_model", None), "descriptor", None) uses_graph_lower = getattr(descriptor, "uses_graph_lower", None) if uses_graph_lower is None: diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 9885d599c6..3c77a62221 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -252,6 +252,167 @@ def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: return energy_redu +def _build_graph_for_method( + method: str, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None, + rcut: float, + pair_excl: Any, + with_csr: bool = False, +) -> Any: + """Build a carry-all ``NeighborGraph`` for the named pt_expt builder. + + Single owning site for the graph-builder dispatch shared by + :meth:`_call_common_graph` and the graph Hessian wrapper + (:class:`_WrapperForwardEnergyGraph`), so both build the graph identically. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + build_neighbor_graph_ase, + ) + + if method == "dense": + return build_neighbor_graph( + coord, atype, box, rcut, with_csr=with_csr, pair_excl=pair_excl + ) + if method == "ase": + return build_neighbor_graph_ase( + coord, atype, box, rcut, with_csr=with_csr, pair_excl=pair_excl + ) + if method == "vesin": + from deepmd.pt_expt.utils.vesin_graph_builder import ( + build_neighbor_graph_vesin, + ) + + return build_neighbor_graph_vesin( + coord, atype, box, rcut, with_csr=with_csr, pair_excl=pair_excl + ) + if method == "nv": + from deepmd.pt_expt.utils.nv_graph_builder import ( + build_neighbor_graph_nv, + ) + + return build_neighbor_graph_nv( + coord, atype, box, rcut, with_csr=with_csr, pair_excl=pair_excl + ) + raise ValueError( + f"unknown neighbor_graph_method {method!r}; use 'dense', 'ase', " + "'vesin', or 'nv'" + ) + + +class _WrapperForwardEnergyGraph: + """Graph twin of :class:`_WrapperForwardEnergy` for the Hessian. + + Given flattened LOCAL coordinates for one frame, rebuilds the carry-all + ``NeighborGraph`` (so ``edge_vec`` tracks the coordinates through + ``autograd.functional.hessian``'s double-backward) and returns the scalar + reduced-output component. Unlike the dense wrapper this differentiates + w.r.t. the ``(nloc, 3)`` LOCAL coordinates directly -- the carry-all graph + has no ghost nodes (PBC images enter only as edges, rebuilt from the same + coords each call), so there is no extended-region ``nall -> nloc`` fold. + """ + + def __init__( + self, + model: Any, + kk: str, + ci: int, + nloc: int, + atype: torch.Tensor, # (1, nloc) + box: torch.Tensor | None, # (1, ...) or None + method: str, + pair_excl: Any, + rcut: float, + fparam: torch.Tensor | None, # (1, ndf) or None + aparam: torch.Tensor | None, # (1, nloc, nda) or None + ) -> None: + self.model = model + self.kk = kk + self.ci = ci + self.nloc = nloc + self.atype = atype + self.box = box + self.method = method + self.pair_excl = pair_excl + self.rcut = rcut + self.fparam = fparam + self.aparam = aparam + + def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: + cc = coord_flat.reshape(1, self.nloc, 3) + ng = _build_graph_for_method( + self.method, cc, self.atype, self.box, self.rcut, self.pair_excl + ) + atomic_ret = self.model.atomic_model.forward_common_atomic_graph( + ng, + self.atype.reshape(-1), + fparam=self.fparam, + # graph-lower ABI: aparam is FLAT on the node axis, (N, nda) + # (N == nloc for the single-frame carry-all graph). + aparam=( + self.aparam.reshape(self.nloc, -1) if self.aparam is not None else None + ), + ) + # atomic_ret[kk]: flat (N, *def), N == nloc for a single-frame carry-all + # graph (all nodes owned); reduced output = sum over the node axis. + atom_out = atomic_ret[self.kk] + return atom_out.sum(dim=0).reshape(-1)[self.ci] + + +def _cal_hessian_ext_graph( + model: Any, + kk: str, + vdef: OutputVariableDef, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + method: str, + pair_excl: Any, + rcut: float, + create_graph: bool = False, +) -> torch.Tensor: + """Graph twin of :func:`_cal_hessian_ext`. + + Computes the Hessian of the reduced output w.r.t. the LOCAL coordinates on + the carry-all graph route. Returns shape ``[nf, *vdef.shape, nloc*3, + nloc*3]`` -- the local-only counterpart of the dense extended Hessian, + already in the same final layout the dense route reaches after + ``communicate_extended_output`` folds ``nall -> nloc`` (the graph route + reduces over owned nodes, so no fold is needed). Node axis is + atom-major/xyz-minor, matching the dense final reshape. + """ + nf, nloc, _ = coord.shape + vsize = math.prod(vdef.shape) + coord_flat = coord.reshape(nf, nloc * 3) + hessians = [] + for ii in range(nf): + for ci in range(vsize): + wrapper = _WrapperForwardEnergyGraph( + model, + kk, + ci, + nloc, + atype[ii : ii + 1], + box[ii : ii + 1] if box is not None else None, + method, + pair_excl, + rcut, + fparam[ii : ii + 1] if fparam is not None else None, + aparam[ii : ii + 1] if aparam is not None else None, + ) + hess = torch.autograd.functional.hessian( + wrapper, + coord_flat[ii], + create_graph=create_graph, + ) # (nloc*3, nloc*3) + hessians.append(hess) + return torch.stack(hessians).reshape(nf, *vdef.shape, nloc * 3, nloc * 3) + + def make_model( T_AtomicModel: type[BaseAtomicModel], T_Bases: tuple[type, ...] = (), @@ -358,6 +519,7 @@ def forward_common_lower_graph( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, + comm_dict: dict | None = None, ) -> dict[str, torch.Tensor]: """Graph-native lower with autograd force/virial (dpa1/se_atten concat-tebd, attention included). @@ -385,7 +547,15 @@ def forward_common_lower_graph( n_node (nf,) per-frame total node counts, including halo nodes. n_local - (nf,) per-frame owned node counts. + (nf,) per-frame owned node counts for multi-rank inference. + When given, ghost rows (index ``>= n_local[frame]``) are + excluded from the DIFFERENTIATED ``_redu`` (and thus + from force/virial/atom-virial, which are ``grad`` of that + reduction) -- each ghost atom is owned, and counted, on + another rank. The per-node output (````) itself stays + FULL/unmasked. ``None`` (default) is the single-rank/ + all-owned behavior. See + :func:`~deepmd.pt_expt.model.edge_transform_output.fit_output_to_model_output_graph`. edge_index (2, E) ``[src, dst]`` edge endpoints (flat local indices). edge_vec @@ -408,10 +578,24 @@ def forward_common_lower_graph( fparam Frame parameter, ``(nf, ndf)``. aparam - Atomic parameter, ``(nf, nloc, nda)``. + Atomic parameter, ``(N, nda)`` -- FLAT on the node axis like + every per-node tensor of the graph ABI (on extended-region + multi-rank graphs the ghost rows are included; their values + are inert under the owned-node mask). charge_spin - Charge/spin conditioning. Accepted for descriptors that expose - this input. + charge/spin conditioning. Ignored in PR-A; accepted for ABI + stability with charge/spin-conditioned descriptors. + comm_dict + MPI communication metadata for parallel inference. ``None`` + (default) for non-parallel inference/training. Forwarded to + the atomic model's ``forward_common_atomic_graph``, which + drives the descriptor's per-layer cross-rank ghost-feature + exchange (``deepmd_export::border_op``, e.g. dpa2's + repformer block via + :meth:`~deepmd.pt_expt.descriptor.repformers. + DescrptBlockRepformers._exchange_ghosts_graph`). Only + meaningful for descriptors whose + ``has_message_passing_across_ranks()`` is ``True``. Returns ------- @@ -452,6 +636,7 @@ def forward_common_lower_graph( fparam=fparam, aparam=aparam, charge_spin=charge_spin, + comm_dict=comm_dict, ) # ``forward_common_atomic_graph`` returns flat ``(N, *)`` output. # Pass directly to the flat-N transform; no rectangular reshape needed. @@ -474,6 +659,7 @@ def forward_common_lower_graph( # shape -- avoids a CUDA out-of-bounds device-assert under # dynamic-edge torch.export. See fit_output_to_model_output_graph. node_capacity=atype.shape[0], + n_local=n_local, ) def _resolve_graph_method( @@ -501,6 +687,18 @@ def _resolve_graph_method( return None if neighbor_graph_method is not None: return neighbor_graph_method + # The DEFAULT-flip is gated to ENERGY-output models: the graph + # lower itself is output-agnostic, but the compiled-training + # trace (``training._trace_and_compile_graph``) and the trainer's + # public-key translation are energy-specific, so a non-energy + # model (property/dos/dipole/polar) default-flipped here would + # route eager through the graph while its compiled twin raises + # ``KeyError('energy')`` -- and gating ONLY the compiled side + # would diverge eager (graph) from compiled (dense) instead. + # An EXPLICIT ``neighbor_graph_method=`` above stays available + # for non-energy models (eager-only, output-agnostic). + if "energy" not in self.atomic_output_def().keys(): + return None # Linear/ZBL atomic models have no single ``descriptor`` -> dense. descriptor = getattr(self.atomic_model, "descriptor", None) uses_graph_lower = getattr(descriptor, "uses_graph_lower", lambda: False) @@ -550,11 +748,6 @@ def _call_common_graph( ``energy_redu``, ``energy_derv_r``, ``energy_derv_c_redu``, and ``energy_derv_c`` when ``do_atomic_virial``). """ - from deepmd.dpmodel.utils.neighbor_graph import ( - build_neighbor_graph, - build_neighbor_graph_ase, - ) - # mirror the dpmodel guard: _resolve_graph_method's eligibility # check only protects the default (None) path; an EXPLICIT # neighbor_graph_method would otherwise reach the builders for @@ -573,57 +766,13 @@ def _call_common_graph( and bool(getattr(descriptor, "geo_compress", False)) ) pair_excl = getattr(self.atomic_model, "pair_excl", None) - if method == "dense": - ng = build_neighbor_graph( - cc, - atype, - bb, - rcut, - with_csr=with_csr, - pair_excl=pair_excl, - ) - elif method == "ase": - ng = build_neighbor_graph_ase( - cc, - atype, - bb, - rcut, - with_csr=with_csr, - pair_excl=pair_excl, - ) - elif method == "vesin": - from deepmd.pt_expt.utils.vesin_graph_builder import ( - build_neighbor_graph_vesin, - ) - - ng = build_neighbor_graph_vesin( - cc, - atype, - bb, - rcut, - with_csr=with_csr, - pair_excl=pair_excl, - ) - elif method == "nv": - from deepmd.pt_expt.utils.nv_graph_builder import ( - build_neighbor_graph_nv, - ) - - ng = build_neighbor_graph_nv( - cc, - atype, - bb, - rcut, - with_csr=with_csr, - pair_excl=pair_excl, - ) - else: - raise ValueError( - f"unknown neighbor_graph_method {method!r}; " - "use 'dense', 'ase', 'vesin', or 'nv'" - ) + ng = _build_graph_for_method( + method, cc, atype, bb, rcut, pair_excl, with_csr=with_csr + ) nf, nloc = atype.shape[:2] atype_flat = atype.reshape(nf * nloc) + # graph-lower ABI: aparam is FLAT on the node axis, (N, nda). + ap_flat = ap.reshape(nf * nloc, ap.shape[-1]) if ap is not None else None model_predict = self.forward_common_lower_graph( atype_flat, ng.n_node, @@ -638,7 +787,7 @@ def _call_common_graph( destination_sorted=ng.destination_sorted, do_atomic_virial=do_atomic_virial, fparam=fp, - aparam=ap, + aparam=ap_flat, ) # ``forward_common_lower_graph`` returns flat ``(N, *)`` per-atom # outputs (N = nf * nloc for a carry-all rectangular graph). @@ -655,6 +804,30 @@ def _call_common_graph( and v.shape[:1] == torch.Size([N]) ): model_predict[k] = v.reshape(nf, nloc, *v.shape[1:]) + # Graph-native Hessian (parallel to the dense ``forward_common_atomic`` + # loop): differentiate the reduced output w.r.t. the LOCAL coords by + # rebuilding the graph inside the wrapper. Added AFTER the unravel so + # its ``(nf, *def, nloc, 3, nloc, 3)`` shape is returned as-is. + # Eager-only, like the dense Hessian (autograd.functional.hessian + # does not export/compile). + aod = self.atomic_output_def() + for kk in aod.keys(): + vdef = aod[kk] + if vdef.reducible and vdef.r_hessian: + model_predict[get_hessian_name(kk)] = _cal_hessian_ext_graph( + self, + kk, + vdef, + cc, + atype, + bb, + fp, + ap, + method, + pair_excl, + rcut, + create_graph=self.training, + ) return model_predict def forward_common_atomic( diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 08218835d7..c837fce6a0 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -615,8 +615,6 @@ def _trace_and_compile_graph( Per-task buffers promoted to FX placeholders (see :func:`_detect_task_buffers`). """ - import math - from torch._decomp import ( get_decompositions, ) @@ -676,10 +674,6 @@ def _trace_and_compile_graph( while (trace_nf * nloc_trace) in (_forbidden | {trace_nf}): nloc_trace += 1 trace_N = trace_nf * nloc_trace - # Static edge capacity, prime-padded to stay distinct from nf and N. - nnei = sum(model.get_sel()) - e_max_base = max(math.ceil(1.25 * nloc_trace * nnei), 7) - e_max = _next_safe_prime(e_max_base, _forbidden | {trace_nf, trace_N}) # Shared with the .pt2 export trace (serialization.py) so the two graph # traces can never desync on the input schema. Training uses the run-time @@ -687,16 +681,39 @@ def _trace_and_compile_graph( from deepmd.pt_expt.utils.serialization import ( build_synthetic_graph_inputs, check_graph_trace_torch_version, + count_synthetic_graph_edges, ) check_graph_trace_torch_version(model) + + # Static edge capacity: derived from the ACTUAL edge count of the + # synthetic trace system (the carry-all builder is sel-free; a + # sel-derived estimate overflows whenever the real degree exceeds sel), + # then prime-padded to stay distinct from nf and N. ``+ 2`` keeps at + # least two masked padding rows so the padded-tail branch is traced. + # Trace on the MODEL's device, not the global ``DEVICE``: make_fx keeps the + # real model parameters (``_allow_non_fake_inputs``), so the synthetic trace + # inputs must live where the model does. A CUDA training run keeps the model + # on ``DEVICE`` (these match), but callers that trace a CPU-placed model + # (e.g. the graph .pt2/export path, which moves the model to CPU to dodge a + # CUDA autograd-stream limitation) would otherwise mix a CPU model with + # CUDA inputs and fail only on a GPU host. + _trace_device = next(model.parameters()).device + e_real = count_synthetic_graph_edges( + model, + nframes=trace_nf, + nloc=nloc_trace, + dtype=GLOBAL_PT_FLOAT_PRECISION, + device=_trace_device, + ) + e_max = _next_safe_prime(e_real + 2, _forbidden | {trace_nf, trace_N}) sample = build_synthetic_graph_inputs( model, e_max=e_max, nframes=trace_nf, nloc=nloc_trace, dtype=GLOBAL_PT_FLOAT_PRECISION, - device=DEVICE, + device=_trace_device, want_fparam=fparam is not None, want_aparam=aparam is not None, want_charge_spin=charge_spin is not None, @@ -1053,14 +1070,20 @@ def forward( *task_buf_vals, ) - # Translate forward_lower keys -> forward keys. - # ``extended_force`` lives on all extended atoms (nf, nall, 3). - # Ghost-atom forces must be scatter-summed back to local atoms - # via ``mapping`` — the same operation ``communicate_extended_output`` - # performs in the uncompiled path. + # Translate forward_lower keys -> forward keys. OUTPUT-AGNOSTIC: + # every key passes through unchanged (energy models emit + # atom_energy/energy/virial/..., property/dos/... models their own + # keys -- the hardcoded energy-key copy here used to KeyError on any + # non-energy fitting), EXCEPT the extended-region keys + # ``extended_force`` (nf, nall, 3) and ``extended_virial`` + # (nf, nall, 9): their ghost rows are scatter-summed back onto + # local owners via ``mapping`` -- the same fold + # ``communicate_extended_output`` performs in the uncompiled path + # (which exposes them as ``force`` / ``atom_virial``). Folding + # both keeps the compiled and uncompiled outputs key-for-key + # consistent, including for a future atom-virial training + # objective. out: dict[str, torch.Tensor] = {} - out["atom_energy"] = result["atom_energy"] - out["energy"] = result["energy"] if "extended_force" in result: ext_force = result["extended_force"] # (nf, nall, 3) idx = mapping.unsqueeze(-1).expand_as(ext_force) # (nf, nall, 3) @@ -1069,14 +1092,21 @@ def forward( ) force.scatter_add_(1, idx, ext_force) out["force"] = force - if "virial" in result: - out["virial"] = result["virial"] if "extended_virial" in result: - out["extended_virial"] = result["extended_virial"] - if "atom_virial" in result: - out["atom_virial"] = result["atom_virial"] - if "mask" in result: - out["mask"] = result["mask"] + ext_virial = result["extended_virial"] # (nf, nall, 9) + idx = mapping.unsqueeze(-1).expand_as(ext_virial) # (nf, nall, 9) + atom_virial = torch.zeros( + nframes, + nloc, + ext_virial.shape[-1], + dtype=ext_virial.dtype, + device=ext_virial.device, + ) + atom_virial.scatter_add_(1, idx, ext_virial) + out["atom_virial"] = atom_virial + for key, val in result.items(): + if key not in ("extended_force", "extended_virial"): + out[key] = val return out def _forward_graph( @@ -1108,6 +1138,12 @@ def _forward_graph( coord_3d = coord.detach().reshape(nframes, nloc, 3) box_flat = box.detach().reshape(nframes, 9) if box is not None else None + # graph-lower ABI: aparam is FLAT on the node axis, (N, nda) -- like + # every per-node tensor of the graph schema (the trace sample from + # build_synthetic_graph_inputs is flat too, so the compiled lower's + # input spec expects it). + if aparam is not None: + aparam = aparam.reshape(nframes * nloc, -1) # Mirror the optional-input defaulting of the dense path / eager # call_common: a model configured with fparam / charge_spin substitutes diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index fe8357b858..84d9c53831 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -164,30 +164,36 @@ def check_graph_trace_torch_version(model: torch.nn.Module) -> None: Parameters ---------- model - The graph-eligible model about to be traced. The attention depth is - read from ``model.atomic_model.descriptor.get_numb_attn_layer()``; - models without a single descriptor (linear/zbl/frozen) pass the - check (they take the dense route anyway). + The graph-eligible model about to be traced. The compact-pair usage + is read from the descriptor capability + ``uses_compact_edge_pairs()`` -- ``True`` for dpa1 with + ``attn_layer > 0`` and for dpa2 with ``update_g2_has_attn`` or + ``update_h2`` (keying on ``get_numb_attn_layer()`` alone missed + dpa2, whose repformer attention rides ``center_edge_pairs`` without + implementing that dpa1 accessor). Models without a single + descriptor (linear/zbl/frozen) pass the check (they take the dense + route anyway), as do descriptors without the capability method. Raises ------ RuntimeError - If the descriptor has ``attn_layer > 0`` and the running torch is - older than 2.6. + If the descriptor's graph lower traces compact edge pairs and the + running torch is older than 2.6. """ desc = getattr(getattr(model, "atomic_model", None), "descriptor", None) - get_n_attn = getattr(desc, "get_numb_attn_layer", None) - n_attn = get_n_attn() if get_n_attn is not None else 0 - if n_attn <= 0: + uses_pairs = getattr(desc, "uses_compact_edge_pairs", None) + if uses_pairs is None or not uses_pairs(): return version = torch.__version__.split("+")[0] major_minor = tuple(int(p) for p in version.split(".")[:2] if p.isdigit()) if len(major_minor) == 2 and major_minor < (2, 6): raise RuntimeError( - f"graph-form tracing of attention layers (attn_layer={n_attn}) " - f"requires torch >= 2.6 (unbacked-SymInt support for the compact " + "graph-form tracing of this descriptor's attention/pair updates " + "requires torch >= 2.6 (unbacked-SymInt support for the compact " f"center_edge_pairs realization); found torch {torch.__version__}. " - "Upgrade torch, set 'attn_layer: 0', or use the dense (nlist) path." + "Upgrade torch; or disable the compact-pair consumers " + "(dpa1: set 'attn_layer: 0'; dpa2: set 'update_g2_has_attn: " + "false' and 'update_h2: false'); or use the dense (nlist) path." ) @@ -367,7 +373,7 @@ def _make_sample_inputs( def build_synthetic_graph_inputs( model: torch.nn.Module, - e_max: int, + e_max: int | None, nframes: int = 2, nloc: int = 7, *, @@ -403,8 +409,13 @@ def build_synthetic_graph_inputs( ---------- model : torch.nn.Module The pt_expt energy model (must expose ``get_rcut``/``get_type_map``/...). - e_max : int - Concrete edge-axis size used by the trace sample. + e_max : int or None + Concrete edge-axis size used by the trace sample. Must be at least + the system's real edge count (the carry-all builder raises ``edge + overflow`` otherwise) -- derive it from + :func:`count_synthetic_graph_edges`, never from ``sel`` (the builder + is sel-free). ``None`` selects the dynamic layout (real edges plus + ``min_edges`` guard rows); used by the edge-count probe itself. nframes : int Number of frames in the sample system. nloc : int @@ -467,8 +478,12 @@ def build_synthetic_graph_inputs( if (want_fparam and dim_fparam > 0) else None ) + # aparam is FLAT on the node axis -- (N, nda), the same axis as ``atype`` + # -- like every per-node tensor of the graph ABI (a rectangular + # ``(nf, nloc, nda)`` sample would hand torch.export three independent + # symbols related by ``N == nf * nloc``, which it rejects). aparam = ( - torch.zeros(nframes, nloc, dim_aparam, dtype=dtype, device=device) + torch.zeros(nframes * nloc, dim_aparam, dtype=dtype, device=device) if (want_aparam and dim_aparam > 0) else None ) @@ -582,6 +597,56 @@ def _build_canonical_graph_dynamic_shapes( ) +def count_synthetic_graph_edges( + model: torch.nn.Module, + nframes: int, + nloc: int, + *, + dtype: torch.dtype, + device: torch.device | None = None, +) -> int: + """Count the real (unpadded) edges of the synthetic trace system. + + Probes :func:`build_synthetic_graph_inputs` with ``e_max=None`` (the + dynamic carry-all layout, whose edge axis is the real edge count plus + the ``min_edges`` guard rows) and counts the ``edge_mask`` real prefix + (slot 5 of the 13-tuple). The carry-all builder is sel-free -- edges + are cutoff-determined, ``sel`` is only a normalization constant -- so + the static trace capacity must derive from this geometry-determined + count; a sel-based estimate overflows whenever the synthetic system's + real degree exceeds ``sel`` (small-``sel`` models). + + Parameters + ---------- + model : torch.nn.Module + The pt_expt energy model (must expose ``get_rcut``/``get_type_map``). + nframes : int + Number of frames of the synthetic system; must match the subsequent + :func:`build_synthetic_graph_inputs` call. + nloc : int + Local atoms per frame; must match the subsequent call. + dtype : torch.dtype + Float precision of the probe coordinates; must match the subsequent + call (the edge count is cutoff-thresholded). + device : torch.device, optional + Probe device; must match the subsequent call. + + Returns + ------- + int + Number of real edges of the synthetic system at the model cutoff. + """ + edge_mask = build_synthetic_graph_inputs( + model, + e_max=None, + nframes=nframes, + nloc=nloc, + dtype=dtype, + device=device, + )[5] + return int(edge_mask.sum()) + + def _build_graph_dynamic_shapes( *sample_inputs: torch.Tensor | None, ) -> tuple: @@ -609,7 +674,6 @@ def _build_graph_dynamic_shapes( nframes_dim = torch.export.Dim("nframes", min=1) n_node_total_dim = torch.export.Dim("n_node_total", min=1) nedge_dim = torch.export.Dim("nedge", min=2) - nloc_dim = torch.export.Dim("nloc", min=1) return ( {0: n_node_total_dim}, # atype: (N,) {0: nframes_dim}, # n_node: (nf,) @@ -622,14 +686,77 @@ def _build_graph_dynamic_shapes( {0: nedge_dim}, # source_order: (E,) {0: n_node_total_dim + 1}, # source_row_ptr: (N + 1,) {0: nframes_dim} if fparam is not None else None, # fparam: (nf, ndf) - # aparam: (nf, nloc, nda) — both the frame AND atom axes are dynamic, - # matching the dense ``_build_dynamic_shapes`` (otherwise a dim_aparam>0 - # graph export specializes nloc to the sample size and breaks at runtime). - {0: nframes_dim, 1: nloc_dim} if aparam is not None else None, # aparam + # aparam: (N, nda) — flat on the node axis, SHARING atype's ``N`` + # symbol (the graph fitting consumes aparam per node; an independent + # dim would make torch.export prove/reject the equality). + {0: n_node_total_dim} if aparam is not None else None, # aparam {0: nframes_dim} if charge_spin is not None else None, # charge_spin ) +def _build_graph_dynamic_shapes_with_comm( + *sample_inputs: torch.Tensor | None, +) -> tuple: + """Build dynamic-shape specs for the with-comm graph-form export. + + Same as :func:`_build_graph_dynamic_shapes` (the flat node axis ``N`` + and the edge axis ``E`` stay dynamic -- a with-comm graph carries owned + PLUS ghost nodes in the "owned-prefix" layout) EXCEPT ``nframes`` is + STATIC at ``1``: the pt_expt Repformer with-comm override only supports + ``nf=1`` (LAMMPS always drives multi-rank inference with one frame). + The 8 trailing comm tensors are all STATIC (``None``): ``nswap`` is + baked in at the trace value, same rationale as the dense with-comm + tail of ``_build_dynamic_shapes``. + + Parameters + ---------- + *sample_inputs : torch.Tensor | None + ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, + destination_order, destination_row_ptr, source_order, + source_row_ptr, fparam, aparam, charge_spin, send_list, send_proc, + recv_proc, send_num, recv_num, communicator, nlocal, nghost)`` -- + 21 entries matching ``forward_lower_graph_exportable_with_comm``. + + Returns + ------- + tuple + Per-input dynamic-shape specs (dicts of ``torch.export.Dim`` or + ``None``) in the same order as ``sample_inputs``. + """ + fparam = sample_inputs[10] + aparam = sample_inputs[11] + charge_spin = sample_inputs[12] + nframes_val = 1 + n_node_total_dim = torch.export.Dim("n_node_total", min=1) + nedge_dim = torch.export.Dim("nedge", min=2) + return ( + {0: n_node_total_dim}, # atype: (N,) + {0: nframes_val}, # n_node: (nf,) — nf STATIC at 1 + {0: nframes_val}, # n_local: (nf,) + {1: nedge_dim}, # edge_index: (2, E) — E dynamic + {0: nedge_dim}, # edge_vec: (E, 3) — E dynamic + {0: nedge_dim}, # edge_mask: (E,) — E dynamic + {0: nedge_dim}, # destination_order: (E,) + {0: n_node_total_dim + 1}, # destination_row_ptr: (N + 1,) + {0: nedge_dim}, # source_order: (E,) + {0: n_node_total_dim + 1}, # source_row_ptr: (N + 1,) + {0: nframes_val} if fparam is not None else None, # fparam + # aparam: (N, nda) — flat on the SAME extended node axis as atype + # (owned prefix + ghost rows). + {0: n_node_total_dim} if aparam is not None else None, # aparam + {0: nframes_val} if charge_spin is not None else None, # charge_spin + # 8 comm tensors: static, nswap baked in at the trace value. + None, + None, + None, + None, + None, + None, + None, + None, + ) + + def _build_dynamic_shapes( *sample_inputs: torch.Tensor | None, has_spin: bool = False, @@ -1213,10 +1340,33 @@ def _trace_and_export( raise NotImplementedError( "graph-form .pt2 export is not supported for spin models" ) - if with_comm_dict: + # Defense-in-depth: every production caller (freeze entrypoint, + # compress, _resolve_lower_kind auto) gates on model_uses_graph_lower + # upstream, but a direct programmatic call with lower_kind="graph" + # on a graph-INELIGIBLE model (e.g. set_davg_zero=False dpa2, + # use_three_body, or disable_graph_lower()) would otherwise trace a + # silently divergent artifact. Assert the gate at the innermost + # layer too. + from deepmd.pt_expt.model.graph_lower import ( + model_uses_graph_lower, + ) + + if not model_uses_graph_lower(model): + raise ValueError( + f"lower_kind={lower_kind!r} requested but the model is not " + "graph-lower eligible (model_uses_graph_lower() is False: " + "check uses_graph_lower() gates such as set_davg_zero, " + "compression, use_three_body, disable_graph_lower(), and " + "the energy-output requirement); freeze with the default " + "lower_kind instead." + ) + if with_comm_dict and not hasattr( + model, "forward_lower_graph_exportable_with_comm" + ): raise NotImplementedError( - "graph-form .pt2 export does not support the with-comm artifact " - "required for multi-rank message passing" + f"model {type(model).__name__} has no " + "forward_lower_graph_exportable_with_comm; graph-form " + "with-comm .pt2 export requires an energy model" ) canonical = lower_kind == "dpa1_canonical" required_method = ( @@ -1239,9 +1389,46 @@ def _trace_and_export( "DPA1 energy model" ) + # Trace-time sizes must be pairwise-distinct AND avoid every static + # model dim (dim_fparam / dim_aparam / parameter dims): make_fx's + # duck-shaping merges same-valued dims into ONE symbol, so e.g. + # ``numb_aparam == 2`` traced at ``nframes == 2`` aliases ``nda`` + # with ``nf`` and bakes outputs whose frame axis follows the STATIC + # aparam width -- silently wrong shapes at any other runtime nf. + # Mirrors the compiled-training trace (forbidden set + primes). + from deepmd.pt.utils.compile_compat import ( + forbidden_dims_from_model, + next_safe_prime, + ) + + _forbidden = forbidden_dims_from_model(model) + _dim_cs = model.get_dim_chg_spin() if hasattr(model, "get_dim_chg_spin") else 0 + if _dim_cs > 1: + _forbidden.add(int(_dim_cs)) + nframes_sample = next_safe_prime(5, _forbidden) nloc_sample = 7 - nnei = sum(model.get_sel()) - e_sample = math.ceil(1.25 * nloc_sample * nnei) + while (nframes_sample * nloc_sample) in (_forbidden | {nframes_sample}): + nloc_sample += 1 + n_sample = nframes_sample * nloc_sample + + # The exported edge axis is DYNAMIC: the trace sample only supplies + # representative tensors. The concrete capacity derives from the + # ACTUAL edge count of the synthetic system (the carry-all builder is + # sel-free; a sel-derived estimate overflows for small-sel models): + # 25% headroom keeps the masked padded tail genuinely traced, the + # ``+ 2`` floor guarantees it even for tiny edge counts, and the + # final bump keeps the concrete edge length collision-free under + # duck-sizing. + e_real = count_synthetic_graph_edges( + model, + nframes=nframes_sample, + nloc=nloc_sample, + dtype=torch.float64, + device=torch.device("cpu"), + ) + e_sample = max(math.ceil(1.25 * e_real), e_real + 2) + while e_sample in (_forbidden | {nframes_sample, n_sample}): + e_sample += 1 if canonical: sample_inputs = build_synthetic_canonical_graph_inputs( model, @@ -1255,6 +1442,60 @@ def _trace_and_export( _allow_non_fake_inputs=True, ) dynamic_shapes = _build_canonical_graph_dynamic_shapes(*sample_inputs) + elif with_comm_dict: + # Load libdeepmd_op_pt.so and register border_op fake/autograd + # metadata now, mirroring the dense with-comm precedent below. + from deepmd.pt_expt.utils.comm import ( + ensure_comm_registered, + ) + + ensure_comm_registered() + if not _needs_with_comm_artifact(model): + raise ValueError( + "with_comm_dict=True requested but the model's " + "descriptor does not need cross-rank message passing " + "(has_message_passing_across_ranks() is False) — " + "there's nothing to compile." + ) + # The pt_expt Repformer with-comm override only supports nf=1 + # (LAMMPS always drives multi-rank inference with one frame). + # ``nloc_sample`` is the TOTAL flat node count carried by the + # graph (owned + ghost, "owned-prefix" layout); the comm sample + # splits it into an owned prefix and a ghost suffix so the + # border_op self-send is genuinely exercised at trace time. + # The builder's slot-2 ``n_local`` (clamp(n_node - 1, min=1)) + # already keeps owned != total for the in-graph mask symbols. + nghost_sample = 2 + nlocal_sample = nloc_sample - nghost_sample + edge_dtype = ( + torch.float32 + if metadata["graph_edge_dtype"] == "float32" + else torch.float64 + ) + sample_inputs = build_synthetic_graph_inputs( + model, + e_max=e_sample, + nframes=1, + nloc=nloc_sample, + dtype=torch.float64, + edge_dtype=edge_dtype, + device=torch.device("cpu"), + ) + comm_inputs = _make_comm_sample_inputs( + nloc=nlocal_sample, + nghost=nghost_sample, + device=torch.device("cpu"), + ) + sample_inputs = sample_inputs + comm_inputs + # Trace via make_fx on CPU (decomposes autograd.grad into aten + # ops); single trace, comm tensors packed to comm_dict inside. + traced = model.forward_lower_graph_exportable_with_comm( + *sample_inputs, + do_atomic_virial=do_atomic_virial, + tracing_mode="symbolic", + _allow_non_fake_inputs=True, + ) + dynamic_shapes = _build_graph_dynamic_shapes_with_comm(*sample_inputs) else: edge_dtype = ( torch.float32 @@ -1264,7 +1505,7 @@ def _trace_and_export( sample_inputs = build_synthetic_graph_inputs( model, e_max=e_sample, - nframes=2, + nframes=nframes_sample, nloc=nloc_sample, dtype=torch.float64, edge_dtype=edge_dtype, @@ -1671,6 +1912,9 @@ def _deserialize_to_file_pt2( model_json_override, with_comm_dict=True, do_atomic_virial=do_atomic_virial, + # the nested artifact must consume the SAME lower schema as the + # regular one (graph freeze -> graph with-comm trace). + lower_kind=lower_kind, ) with tempfile.TemporaryDirectory() as td: wc_path = os.path.join(td, "forward_lower_with_comm.pt2") diff --git a/doc/model/dpa2.md b/doc/model/dpa2.md index fbff3aea3d..4cbbc06e2a 100644 --- a/doc/model/dpa2.md +++ b/doc/model/dpa2.md @@ -91,3 +91,32 @@ Model compression is supported when {ref}`repinit/tebd_input_mode 0``) + // force calls added a global synchronization per MD step without any + // added protection. + bool graph_comm_preflight_done_ = false; // Device-resident (ntypes+1)^2 model-level pair-type keep table, uploaded // ONCE in ``init`` from the ``pair_exclude_types`` metadata field (see // ``deepmd::buildPairExcludeTable``). An UNDEFINED tensor => no model-level @@ -618,6 +628,58 @@ class DeepPotPTExpt : public DeepPotBackend { const torch::Tensor& charge_spin, const std::vector& comm_tensors); + /** + * @brief Run the with-comm NeighborGraph (message-passing, e.g. DPA2/DPA3) + * ``.pt2`` artifact with comm tensors appended. + * + * Positional AOTI input order mirrors ``run_model_graph``: ``(atype, + * n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, + * destination_row_ptr, source_order, source_row_ptr, [fparam], [aparam], + * [charge_spin], comm_tensors...)``. The graph is built on the + * EXTENDED region (``fold_to_local=false``): ghost nodes are distinct and + * their embeddings are filled in-place by ``border_op`` inside the + * exported artifact, exactly as the with-comm dense/edge artifacts do. + * + * @param[in] atype Per-node extended-region types, shape ``(N,)`` int64, + * N == nall_real (owned prefix first, then ghost). + * @param[in] n_node Per-frame node count, shape ``(nf,)`` int64. + * @param[in] edge_index Extended-region edge graph ``(2, E)`` int64 + * [src, dst]; ghost-node indices retained (NOT folded to + * owners -- ``fold_to_local=false``). + * @param[in] edge_vec Edge vectors ``(E, 3)`` (neighbour - center). + * @param[in] edge_mask Physical-edge mask ``(E,)`` bool. + * @param[in] aparam Atomic parameters on the flat EXTENDED node axis, + * shape ``(N, dim_aparam)`` (owned prefix filled, ghost rows + * zero-padded -- ghost fitting outputs are masked inside the + * artifact), or an empty tensor when ``dim_aparam == 0``. + * @param[in] comm_tensors 8 comm tensors in canonical positional order: + * send_list, send_proc, recv_proc, send_num, recv_num, + * communicator, nlocal, nghost. All 8 stay on CPU (host + * control metadata for the opaque ``border_op``, symmetric + * with the dense with-comm artifact). + * @param[in] n_local (1,) int64 owned-node count consumed IN-GRAPH by the + * owned-node energy mask; moved to the model device here (its + * consumer is a device kernel after ``move_to_device_pass``). + * Same value as the ``nlocal`` comm tensor -- the two inputs + * separate the device-compute role from the host-MPI-control + * role, so ``border_op`` never pulls device scalars. + */ + std::vector run_model_graph_with_comm( + const torch::Tensor& atype, + const torch::Tensor& n_node, + const torch::Tensor& n_local, + const torch::Tensor& edge_index, + const torch::Tensor& edge_vec, + const torch::Tensor& edge_mask, + const torch::Tensor& destination_order, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& source_order, + const torch::Tensor& source_row_ptr, + const torch::Tensor& fparam, + const torch::Tensor& aparam, + const torch::Tensor& charge_spin, + const std::vector& comm_tensors); + /** * @brief Extract outputs from flat tensor list using output_keys. */ diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index 12bd9754b5..806caab494 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -756,15 +756,22 @@ inline void remap_graph_outputs_to_dense_keys( /** * @brief Remap NeighborGraph public outputs onto the dense internal-key layout - * for the MULTI-RANK (extended-region) non-message-passing path. + * for the MULTI-RANK (extended-region) path. * * Built with ``fold_to_local=false``, the graph has ``N == nall`` nodes: ghost * (halo) atoms are distinct nodes, so the per-node ``force`` is already the * EXTENDED force (one row per extended atom). Ghost reaction forces stay on * their ghost rows and are folded back to their owning rank by LAMMPS * reverse-comm — exactly as the dense path returns its extended force. No - * zero-padding (unlike the single-rank helper) and no with-comm artifact (dpa1 - * is non-MP). + * zero-padding (unlike the single-rank helper). + * + * Shared by both multi-rank graph routes: non-message-passing models (dpa1) + * run this on the plain extended-region artifact (no comm), while + * message-passing models (DPA2/DPA3) run it on the with-comm artifact's + * output -- ``border_op`` inside that artifact fills ghost-node embeddings + * in-place before the model reduces energy/force/virial, so by the time + * outputs reach this function the two cases are indistinguishable: it is + * the ``atom_energy[0:nloc]``-only reduction below that matters either way. * * Key differences from the single-rank helper: * - ``energy_redu`` = sum of the LOCAL atom energies diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 8a81edd284..c1049cdfea 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -29,6 +29,90 @@ using deepmd::ptexpt::read_zip_entry; using namespace deepmd; namespace { +/** + * @brief Normalize a graph-route aparam tensor to the flat node axis. + * + * The NeighborGraph ABI carries atomic parameters FLAT on the node axis -- + * shape (N, daparam) with N == n_node_count, the same axis as ``atype`` + * (mirrors ``build_synthetic_graph_inputs`` / ``_build_graph_dynamic_shapes`` + * on the Python export side). The runtime aparam carries the owned (local) + * rows only (``aparam_nall`` is structurally false for pt_expt models, see + * ``init``); on extended-region graphs (multi-rank routes, N == nall_real) + * the ghost rows are zero-padded here. Ghost fitting outputs are never + * retained -- the with-comm artifact masks non-owned energies before + * reduction, and the plain multi-rank remap sums energy over the owned + * prefix only -- so the padded values are inert. + * + * A ghost-only rank (``nlocal == 0``, ``N > 0``) synthesizes a full zero + * tensor: the graph still carries N nodes and the artifact requires the + * (N, daparam) input, while the owned-node mask keeps its contribution + * exactly zero. A missing aparam on a rank that OWNS atoms, or a width + * mismatch, is an explicit error -- the former silently returned the empty + * tensor (artifact reshape failure mid-collective), and a width mismatch + * used to be absorbed by broadcasting ``copy_`` (silent result corruption + * for daparam > 1). + */ +at::Tensor extend_graph_aparam(const at::Tensor& aparam_tensor, + std::int64_t n_node_count, + std::int64_t nlocal, + std::int64_t daparam) { + if (daparam <= 0) { + return aparam_tensor; // model has no aparam input; passed through empty + } + if (aparam_tensor.numel() == 0) { + if (nlocal > 0) { + throw deepmd::deepmd_exception( + "aparam is required (dim_aparam=" + std::to_string(daparam) + + ") but no values were provided on a rank owning " + + std::to_string(nlocal) + " atoms."); + } + // ghost-only rank: there are no owned rows to supply; zeros are inert + // under the owned-node mask but the artifact needs the full node axis. + return torch::zeros({n_node_count, daparam}, aparam_tensor.options()); + } + if (aparam_tensor.numel() != nlocal * daparam) { + throw deepmd::deepmd_exception( + "aparam holds " + std::to_string(aparam_tensor.numel()) + + " values but the graph route expects nlocal * dim_aparam = " + + std::to_string(nlocal) + " * " + std::to_string(daparam) + "."); + } + at::Tensor owned = aparam_tensor.reshape({nlocal, daparam}); + if (nlocal == n_node_count) { + return owned; // single-rank / folded graph: nothing to pad + } + at::Tensor padded = + torch::zeros({n_node_count, daparam}, aparam_tensor.options()); + padded.slice(0, 0, nlocal).copy_(owned); + return padded; +} + +/** + * @brief Assert the flat graph-route aparam contract at the C++ boundary. + * + * The graph artifacts consume aparam FLAT on the node axis, shape + * (N, daparam) -- the layout ``extend_graph_aparam`` produces. A caller + * hand-rolling a rectangular (1, N, daparam) tensor (the pre-flat + * convention) would otherwise fail DEEP inside the artifact -- or, on a + * GPU-only route, only at deployment where no CPU test can catch it (the + * device-edge branch shipped exactly that bug). Failing loudly here turns + * any future such site into an immediate, self-explanatory error. + */ +void check_graph_aparam_flat(const at::Tensor& aparam, + std::int64_t daparam, + const char* where) { + if (daparam <= 0) { + return; + } + if (aparam.dim() != 2 || aparam.size(1) != daparam) { + std::ostringstream oss; + oss << where + << ": graph-route aparam must be flat (N, daparam) on the node axis " + "(produce it with extend_graph_aparam); got a rank-" + << aparam.dim() << " tensor of shape " << aparam.sizes() + << " for daparam = " << daparam << "."; + throw deepmd::deepmd_exception(oss.str()); + } +} void synchronize_current_accelerator_stream() { #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) @@ -428,6 +512,7 @@ std::vector DeepPotPTExpt::run_model_graph( const torch::Tensor& aparam, const torch::Tensor& charge_spin) { // Graph-input ABI: original edge payload plus destination/source CSR views. + check_graph_aparam_flat(aparam, daparam, "run_model_graph"); std::vector inputs = { atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, @@ -538,6 +623,76 @@ std::vector DeepPotPTExpt::run_model_edges_with_comm( return with_comm_loader->run(inputs); } +std::vector DeepPotPTExpt::run_model_graph_with_comm( + const torch::Tensor& atype, + const torch::Tensor& n_node, + const torch::Tensor& n_local, + const torch::Tensor& edge_index, + const torch::Tensor& edge_vec, + const torch::Tensor& edge_mask, + const torch::Tensor& destination_order, + const torch::Tensor& destination_row_ptr, + const torch::Tensor& source_order, + const torch::Tensor& source_row_ptr, + const torch::Tensor& fparam, + const torch::Tensor& aparam, + const torch::Tensor& charge_spin, + const std::vector& comm_tensors) { + if (!with_comm_loader) { + throw deepmd::deepmd_exception( + "run_model_graph_with_comm called but the with-comm artifact is not " + "available. Either the .pt2 file has no with-comm artifact compiled " + "(programming error: the caller should check has_comm_artifact_ " + "before invoking this path), or the artifact was present in the " + ".pt2 metadata but failed to load at init time (see earlier stderr " + "log). Multi-rank LAMMPS requires a working with-comm artifact."); + } + if (comm_tensors.size() != 8) { + throw deepmd::deepmd_exception( + "run_model_graph_with_comm: comm_tensors must contain exactly 8 " + "tensors (send_list, send_proc, recv_proc, send_num, recv_num, " + "communicator, nlocal, nghost). Got " + + std::to_string(comm_tensors.size()) + "."); + } + check_graph_aparam_flat(aparam, daparam, "run_model_graph_with_comm"); + // NeighborGraph ABI: the 13-input base of run_model_graph (atype, n_node, + // n_local, edge_index, edge_vec, edge_mask, destination_order, + // destination_row_ptr, source_order, source_row_ptr, [fparam], [aparam], + // [charge_spin]) with the 8 comm tensors appended, like + // run_model_with_comm / run_model_edges_with_comm. + // + // Device placement: the base tensors (n_local included) live on the model + // device -- n_local feeds the in-graph owned-node mask, a device kernel + // after ``move_to_device_pass`` (a CPU tensor fed there is read as a + // device pointer -- CUDA illegal memory access, first observed live in + // the dpa2 graph LAMMPS MP test). The 8 comm tensors stay ON CPU, + // symmetric with the dense with-comm route: they are consumed only by the + // opaque ``deepmd_export::border_op`` whose HOST code dereferences their + // ``data_ptr`` directly (``send_list`` even carries raw host pointers) + // and reads ``nlocal``/``nghost`` via cheap host ``.item()`` calls -- + // splitting the compute role (device n_local) from the host-MPI-control + // role removes the per-layer synchronizing D2H scalar reads a + // device-placed nlocal forced inside border_op (2 per layer forward + 2 + // per layer backward). + std::vector inputs = { + atype, n_node, n_local, edge_index, + edge_vec, edge_mask, destination_order, destination_row_ptr, + source_order, source_row_ptr}; + if (dfparam > 0) { + inputs.push_back(fparam); + } + if (daparam > 0) { + inputs.push_back(aparam); + } + if (dchgspin > 0) { + inputs.push_back(charge_spin); + } + for (const auto& ct : comm_tensors) { + inputs.push_back(ct); + } + return with_comm_loader->run(inputs); +} + void DeepPotPTExpt::extract_outputs( std::map& output_map, const std::vector& flat_outputs) { @@ -636,10 +791,24 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, bool multi_rank = (lmp_list.nprocs > 1); bool atom_map_present = (lmp_list.mapping != nullptr); bool use_with_comm = has_comm_artifact_ && multi_rank; + // NeighborGraph multi-rank dispatch: + // - NON-message-passing (dpa1, se_e2_a, ...): the SAME single-rank graph + // .pt2 runs on the EXTENDED region (fold_to_local=false; ghosts are + // distinct nodes whose features come from their real ghost types). No + // with-comm artifact / no border_op is needed; ghost reaction forces are + // folded to their owners by LAMMPS reverse-comm. Handled below. + // - message-passing graph (DPA2/DPA3): multi-rank requires a with-comm + // graph artifact for cross-rank ghost-feature exchange + // (``run_model_graph_with_comm``, below). Fail fast before building any + // tensors when that artifact is missing, so callers get a clear message + // instead of a wrong answer or the loader's own runtime error. if (lower_input_is_graph_ && multi_rank && has_message_passing_) { - throw deepmd::deepmd_exception( - "Multi-rank message-passing graph-input inference requires an " - "edge-input model with its with-comm artifact."); + if (!has_comm_artifact_ || !with_comm_loader) { + throw deepmd::deepmd_exception( + "Multi-rank message-passing graph .pt2 inference requires a " + "with-comm artifact; re-freeze the model with a deepmd-kit that " + "exports it (this .pt2 predates multi-rank graph support)."); + } } // Dispatch guards for message-passing models: // non-message-passing, or nghost == 0: the regular path is always safe. @@ -918,6 +1087,96 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, edge_tensors.edge_index_ext, edge_tensors.edge_mask, fparam_tensor, aparam_tensor, charge_spin_tensor, comm_tensors); } + } else if (lower_input_is_graph_) { + // Message-passing NeighborGraph route (DPA2/DPA3), multi-rank, with-comm + // artifact. Reachable only when the artifact-presence guard earlier in + // ``compute_impl`` passed (``has_comm_artifact_ && with_comm_loader``), + // which the export pipeline only sets for message-passing models, so + // ``use_with_comm`` implies ``has_message_passing_`` here -- non-MP + // graph models (dpa1) never carry a comm artifact and take the plain + // extended-region branch below instead. + // + // Topology construction mirrors the non-comm multi-rank graph branch + // below exactly (extended region, N == nall_real, node types from the + // on-device extended ``atype_Tensor`` slice): ghost nodes are distinct + // and their embeddings are filled in-place by ``border_op`` inside the + // with-comm artifact instead of being read directly from local ghost + // types, so no geometry/mask changes are needed -- only the model + // entry point (and the appended comm tensors) differ. + // Collective preflight: a rank with zero owned+ghost atoms cannot run + // the graph artifact, and a rank-LOCAL failure would leave the + // non-empty peers blocked forever in the per-layer border_op + // collectives. All-reduce the minimum node count over the LAMMPS + // communicator (comm_tensors[5]) so EVERY rank agrees to run -- or + // every rank throws promptly with the same error. The op is an + // identity when the communicator handle is null or MPI is not + // compiled in. + // + // Cached across ``ago > 0`` force calls: the owned+ghost node count + // shares the lifetime of the cached nlist/mapping/edge topology + // (both only change on an ``ago == 0`` rebuild, which is globally + // synchronized by LAMMPS), so re-running the collective on every + // cache-hit MD step added a global synchronization to the hot path + // with no added protection. + if (ago == 0 || !graph_comm_preflight_done_) { + graph_comm_preflight_done_ = false; + const auto allreduce_min = + c10::Dispatcher::singleton() + .findSchemaOrThrow("deepmd_export::allreduce_min_int", "") + .typed(); + at::Tensor local_n_node = + torch::full({1}, static_cast(nall_real), int_option); + const std::int64_t global_min_n_node = + allreduce_min.call(local_n_node, comm_tensors[5].to(torch::kCPU)) + .item(); + if (global_min_n_node == 0) { + throw deepmd::deepmd_exception( + "Multi-rank message-passing graph inference does not yet " + "support a rank with zero owned+ghost atoms (the exported " + "graph artifact requires at least one node, and skipping the " + "run would desync the per-layer MPI ghost exchange; this rank " + "has " + + std::to_string(nall_real) + + " owned+ghost atoms). Use a domain decomposition that keeps " + "every rank non-empty, or a dense .pt2."); + } + graph_comm_preflight_done_ = true; + } + const auto edge_tensors = + compactEdgeTensors(edge_index_tensor, edge_index_ext_tensor, + coord_Tensor, static_cast(rcut)); + const std::int64_t n_node_count = nall_real; + at::Tensor n_node_tensor = + torch::full({1}, n_node_count, int_option).to(device); + // Device owned-count input for the in-graph owned-node mask; the CPU + // nlocal comm tensor stays host-side for border_op. + at::Tensor n_local_tensor = + torch::full({1}, static_cast(nloc), int_option) + .to(device); + at::Tensor node_atype = + atype_Tensor.slice(1, 0, n_node_count).reshape({n_node_count}); + GraphTensorPack graph_pack; + graph_pack.atype = node_atype; + graph_pack.n_node = n_node_tensor; + graph_pack.n_local = n_local_tensor; + graph_pack.edge_index = edge_tensors.edge_index; + graph_pack.edge_vec = graph_edge_fp32_ + ? edge_tensors.edge_vec.to(torch::kFloat32) + : edge_tensors.edge_vec; + // Model-level pair exclusion is a BUILD-time transform (decision + // #18/A4): the exported graph lower consumes a pre-excluded edge_mask + // and never re-applies it -- same seam as the non-comm graph route. + graph_pack.edge_mask = deepmd::applyPairExclusion( + edge_tensors.edge_index, edge_tensors.edge_mask, node_atype, + pair_exclude_table_, ntypes); + canonicalizeGraphPayload(graph_pack, n_node_count); + flat_outputs = run_model_graph_with_comm( + node_atype, n_node_tensor, n_local_tensor, graph_pack.edge_index, + graph_pack.edge_vec, graph_pack.edge_mask, + graph_pack.destination_order, graph_pack.destination_row_ptr, + graph_pack.source_order, graph_pack.source_row_ptr, fparam_tensor, + extend_graph_aparam(aparam_tensor, n_node_count, nloc, daparam), + charge_spin_tensor, comm_tensors); } else { // Model-level pair exclusion is a BUILD-time transform (decision // #18/A4): the exported dense lower consumes a pre-excluded nlist and @@ -978,13 +1237,11 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, at::Tensor n_local_tensor = torch::full({1}, nloc, int_option).to(device); at::Tensor node_atype = atype_Tensor.slice(1, 0, n_node_count).reshape({n_node_count}); - at::Tensor graph_aparam = aparam_tensor; - if (daparam > 0 && n_node_count > nloc) { - graph_aparam = torch::cat( - {aparam_tensor, torch::zeros({1, n_node_count - nloc, daparam}, - aparam_tensor.options())}, - 1); - } + // Flat (N, daparam) graph ABI: validate the width, zero-pad ghost + // rows, and synthesize a zero tensor on a ghost-only rank (see + // extend_graph_aparam). + at::Tensor graph_aparam = + extend_graph_aparam(aparam_tensor, n_node_count, nloc, daparam); GraphTensorPack graph_pack; graph_pack.atype = node_atype; graph_pack.n_node = n_node_tensor; @@ -1035,7 +1292,8 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, if (multi_rank) { // Extended region (N == nall_real): force is already per-extended-atom, // owned energy = sum over local atom energies, no zero-padding. Ghost - // forces fold back via LAMMPS reverse-comm (no with-comm artifact). + // forces fold back via LAMMPS reverse-comm (both the plain and with-comm + // graph routes). deepmd::remap_graph_outputs_to_dense_keys_extended(output_map, nloc, nall_real, atomic); } else { @@ -1394,7 +1652,11 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, graph_tensors.edge_index, graph_tensors.edge_vec, graph_tensors.edge_mask, graph_tensors.destination_order, graph_tensors.destination_row_ptr, graph_tensors.source_order, - graph_tensors.source_row_ptr, fparam_tensor, aparam_tensor, + graph_tensors.source_row_ptr, fparam_tensor, + // standalone path is single-rank (every node owned): this only + // normalizes the rectangular runtime aparam to the flat + // (N, daparam) node axis of the graph ABI. + extend_graph_aparam(aparam_tensor, natoms, natoms, daparam), charge_spin_tensor); } } else { @@ -2154,12 +2416,14 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy, torch::full({1}, static_cast(nnode), opt_i64); const at::Tensor n_local = torch::full({1}, static_cast(nloc), opt_i64); - at::Tensor graph_aparam = aparam_tensor; - if (daparam > 0 && nnode > nloc) { - graph_aparam = torch::cat( - {aparam_tensor, torch::zeros({1, nnode - nloc, daparam}, opt_f64)}, - 1); - } + // Flat (N, daparam) graph ABI: validate the width, zero-pad ghost + // rows, and synthesize a zero tensor on a ghost-only subdomain -- + // the rank-3 (1, nnode, daparam) pad this branch used before is + // rejected by the artifact (GeneralFitting.call_graph requires + // rank-2 aparam), so device-edge graph inference with + // numb_aparam > 0 failed at the artifact boundary. + at::Tensor graph_aparam = + extend_graph_aparam(aparam_tensor, nnode, nloc, daparam); GraphTensorPack graph_pack; graph_pack.atype = atype_t.reshape({nnode}); graph_pack.n_node = n_node; diff --git a/source/api_cc/tests/test_deeppot_dpa1_graph_aparam_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa1_graph_aparam_ptexpt.cc new file mode 100644 index 0000000000..3e56fd0fe2 --- /dev/null +++ b/source/api_cc/tests/test_deeppot_dpa1_graph_aparam_ptexpt.cc @@ -0,0 +1,180 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// C++ graph-route regression for numb_aparam == 2 with DISTINCT per-component +// values (gen_dpa1.py section C). +// +// Two historical silent-corruption bugs on this path: +// 1. select_real_atoms_coord under-sized the selected aparam buffer (missing +// the daparam factor): heap OOB write + a (1, nloc, 1)-shaped tensor whose +// single column then BROADCAST over all daparam columns in the padding +// copy_ -- finite but wrong energy/force. A numb_aparam=1 fixture can +// never expose this; the distinct column values here do. +// 2. The graph ABI carries aparam FLAT on the node axis (N, daparam); +// extend_graph_aparam normalizes the rectangular runtime tensor and +// validates the width instead of broadcasting. +// +// Reference values (deeppot_dpa1_graph_aparam.expected) come from an +// INDEPENDENT nlist (dense-quartet) evaluation of the same weights with the +// same aparam, so a match validates the graph aparam marshalling, not just +// self-consistency. gen_dpa1.py asserts the fixture's aparam columns are +// non-degenerate (column swap shifts the energy), so broadcast corruption +// cannot pass vacuously. +#include + +#include +#include + +#include "DeepPot.h" +#include "DeepPotPTExpt.h" +#include "expected_ref.h" +#include "neighbor_list.h" +#include "test_utils.h" + +namespace { +constexpr const char* kGraphModel = + "../../tests/infer/deeppot_dpa1_graph_aparam.pt2"; +constexpr const char* kRefPath = + "../../tests/infer/deeppot_dpa1_graph_aparam.expected"; +} // namespace + +template +class TestInferDpa1GraphAparamPtExpt : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 00.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector box = {13., 0., 0., 0., 13., 0., 0., 0., 13.}; + // numb_aparam == 2, DISTINCT per-atom AND per-component values; must match + // ``aparam_vals`` in gen_dpa1.py section C (row-major per atom). + std::vector aparam = {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, + 0.7, 0.8, 0.9, 1.0, 1.1, 1.2}; + std::vector fparam = {}; + // Per-atom reference (energy/force/virial) from the .expected sidecar. + std::vector expected_e; + std::vector expected_f; + std::vector expected_v; + int natoms; + double expected_tot_e; + std::vector expected_tot_v; + + static deepmd::DeepPot dp; + + static void SetUpTestSuite() { +#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT + dp.init(kGraphModel); +#endif + } + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + deepmd_test::ExpectedRef ref; + ref.load(kRefPath); + expected_e = ref.get("pbc", "expected_e"); + expected_f = ref.get("pbc", "expected_f"); + expected_v = ref.get("pbc", "expected_v"); + + natoms = expected_e.size(); + EXPECT_EQ(natoms * 3, static_cast(expected_f.size())); + EXPECT_EQ(natoms * 9, static_cast(expected_v.size())); + expected_tot_e = 0.; + expected_tot_v.assign(9, 0.); + for (int ii = 0; ii < natoms; ++ii) { + expected_tot_e += expected_e[ii]; + } + for (int ii = 0; ii < natoms; ++ii) { + for (int dd = 0; dd < 9; ++dd) { + expected_tot_v[dd] += expected_v[ii * 9 + dd]; + } + } + }; + + void TearDown() override {}; + + static void TearDownTestSuite() { dp = deepmd::DeepPot(); } +}; + +template +deepmd::DeepPot TestInferDpa1GraphAparamPtExpt::dp; + +TYPED_TEST_SUITE(TestInferDpa1GraphAparamPtExpt, ValueTypes); + +// Case 1: standalone graph branch (DeepPot builds its own neighbor list). +// Exercises the single-rank extend_graph_aparam normalization +// (nlocal == N: rectangular (1, natoms, 2) -> flat (N, 2)). +TYPED_TEST(TestInferDpa1GraphAparamPtExpt, cpu_build_nlist_aparam2) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& fparam = this->fparam; + std::vector& aparam = this->aparam; + 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; + double ener; + std::vector force, virial; + dp.compute(ener, force, virial, coord, atype, box, fparam, aparam); + + EXPECT_EQ(force.size(), static_cast(natoms * 3)); + EXPECT_EQ(virial.size(), 9u); + + 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 < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +// Case 2: the LAMMPS provided-nlist graph branch (compute_inner). The +// single-rank folded graph has N == nloc, so aparam covers every node; the +// distinct column values guard the width handling on this route too. +TYPED_TEST(TestInferDpa1GraphAparamPtExpt, lammps_nlist_aparam2) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& fparam = this->fparam; + std::vector& aparam = this->aparam; + 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_data; + _build_nlist(nlist_data, coord_cpy, atype_cpy, mapping, coord, + atype, box, rc); + 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_data); + // The graph branch folds ghost neighbours onto their local owners via the + // LAMMPS atom-map; without it periodic (ghost) edges would be dropped. + inlist.mapping = mapping.data(); + + double ener; + std::vector force_, virial; + dp.compute(ener, force_, virial, coord_cpy, atype_cpy, box, nall - nloc, + inlist, 0, fparam, aparam); + std::vector force; + _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 < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} diff --git a/source/api_cc/tests/test_deeppot_universal.cc b/source/api_cc/tests/test_deeppot_universal.cc index 59aeb647c2..b4482868f1 100644 --- a/source/api_cc/tests/test_deeppot_universal.cc +++ b/source/api_cc/tests/test_deeppot_universal.cc @@ -42,6 +42,12 @@ struct VariantDeepPotCase { bool supports_no_pbc_atomic; bool supports_no_pbc_lmp_nlist; bool supports_no_pbc_lmp_nlist_atomic; + // When the model artifact is legitimately absent (e.g. gen_dpa2.py skips the + // AOTInductor graph .pt2 under LeakSanitizer, whose runtime is incompatible + // with the compiled backward kernel), GTEST_SKIP this case instead of + // failing on the missing file. Defaults false: a missing artifact is a hard + // error for every other case. + bool skip_if_artifact_missing = false; }; struct DefaultFParamCase { @@ -187,6 +193,29 @@ std::vector variant_deeppot_cases() { /*supports_no_pbc_atomic=*/false, /*supports_no_pbc_lmp_nlist=*/true, /*supports_no_pbc_lmp_nlist_atomic=*/false}, + {"dpa2_graph_ptexpt", + Backend::PTExpt, + "../../tests/infer/deeppot_dpa2_graph.pt2", + /*convert_pbtxt=*/false, + nullptr, + nullptr, + "../../tests/infer/deeppot_dpa2_graph.expected", + "pbc", + "nopbc", + 1e-10, + 1e-4, + /*supports_float=*/true, + /*supports_finite_difference=*/true, + /*supports_lmp_nlist=*/true, + /*supports_lmp_nlist_atomic=*/true, + /*supports_lmp_nlist_cutoff_twice=*/true, + /*supports_lmp_nlist_type_sel=*/true, + /*supports_print_summary=*/true, + /*supports_no_pbc_simple=*/true, + /*supports_no_pbc_atomic=*/false, + /*supports_no_pbc_lmp_nlist=*/true, + /*supports_no_pbc_lmp_nlist_atomic=*/false, + /*skip_if_artifact_missing=*/true}, {"dpa2_pytorch_pth", Backend::PyTorch, "../../tests/infer/deeppot_dpa2.pth", @@ -384,8 +413,15 @@ class VariantDeepPotTest : public ::testing::TestWithParam { if (!backend_enabled(param.backend)) { GTEST_SKIP() << backend_name(param.backend) << " support is not enabled."; } - ASSERT_TRUE(path_exists(param.model_path)) - << "Model artifact is not available: " << param.model_path; + if (!path_exists(param.model_path)) { + if (param.skip_if_artifact_missing) { + GTEST_SKIP() + << "Optional model artifact not generated (e.g. the " + "AOTInductor graph .pt2 is skipped under LeakSanitizer): " + << param.model_path; + } + FAIL() << "Model artifact is not available: " << param.model_path; + } if (param.builtin_ref != nullptr) { ref = param.builtin_ref; diff --git a/source/install/test_cc_local.sh b/source/install/test_cc_local.sh index 2b1c7fddb2..d3ae4d2886 100755 --- a/source/install/test_cc_local.sh +++ b/source/install/test_cc_local.sh @@ -66,7 +66,13 @@ else: if echo "${CXXFLAGS:-}" | grep -q fsanitize=leak; then _LSAN_LIB=$(gcc -print-file-name=liblsan.so 2>/dev/null || true) if [ -n "${_LSAN_LIB}" ] && [ -f "${_LSAN_LIB}" ]; then - _GEN_ENV="LD_PRELOAD=${_LSAN_LIB} LSAN_OPTIONS=detect_leaks=0" + # DP_GEN_UNDER_SANITIZER: explicit signal for gen scripts that need + # to skip sanitizer-incompatible sections (e.g. gen_dpa2.py's + # AOTInductor graph .pt2 eval, which can SEGV under the LSAN + # runtime). Sniffing LD_PRELOAD inside the gen script is NOT + # reliable: the sanitizer runtime removes its own entry from the + # process environment during startup. + _GEN_ENV="LD_PRELOAD=${_LSAN_LIB} LSAN_OPTIONS=detect_leaks=0 DP_GEN_UNDER_SANITIZER=lsan" fi fi # Run gen scripts in parallel for faster model generation. diff --git a/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py b/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py index 36c62df64a..a574cef821 100644 --- a/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py +++ b/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py @@ -108,6 +108,15 @@ "changes in BOTH directions per rebuild (rank 0 loses one NULL, " "gains another simultaneously).", ) +parser.add_argument( + "--pair-extra", + type=str, + default="", + help="Extra tokens appended to the pair_style line after the model " + "file, e.g. 'aparam 0.25' for a numb_aparam=1 model whose uniform " + "per-atom parameter comes from the pair_style arguments. Used by the " + "empty-subdomain aparam regression.", +) parser.add_argument( "--real-temp", type=float, @@ -178,7 +187,10 @@ else: lammps.velocity(f"nullgroup set {args.null_vx:.6f} 0.0 0.0 units box") -lammps.pair_style(f"deepmd {args.PB_FILE}") +pair_style_line = f"deepmd {args.PB_FILE}" +if args.pair_extra: + pair_style_line += f" {args.pair_extra}" +lammps.pair_style(pair_style_line) lammps.pair_coeff(args.pair_coeff) # Per-atom virial from the pair contribution. ``centroid/stress/atom`` # is parallel-safe (rank-local data, gathered below). LAMMPS computes diff --git a/source/lmp/tests/test_lammps_dpa2_graph_pt2.py b/source/lmp/tests/test_lammps_dpa2_graph_pt2.py new file mode 100644 index 0000000000..dec858d7ea --- /dev/null +++ b/source/lmp/tests/test_lammps_dpa2_graph_pt2.py @@ -0,0 +1,609 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Test LAMMPS with the NeighborGraph (graph-schema) .pt2 DPA2 model. + +The model ``deeppot_dpa2_graph.pt2`` is a dpa2 descriptor (repformer +message-passing, ``use_three_body=False`` so the graph-ineligible +three-body repinit sub-block is off) exported with ``lower_kind="graph"`` +(gen_dpa2.py section B). Unlike dpa1 (non-message-passing), dpa2's +repformer block exchanges ghost-atom features every layer, so the graph +export automatically embeds a nested with-comm AOTI artifact +(``forward_lower_with_comm.pt2``) alongside the plain graph forward. + +Single-rank LAMMPS folds ghosts onto local owners (``fold_to_local=True``, +``N == nloc``) and uses the plain graph artifact. Multi-rank LAMMPS keeps +the extended region (``N == nall_real``) and routes to the with-comm +artifact: the C++ ``DeepPotPTExpt`` drives ``deepmd_export::border_op`` +once per repformer layer to exchange ghost node/edge features across +ranks, masks the fitting reduction to owned nodes only (an owned-atom +energy mask, since the extended region includes ghost nodes that must +not double-count energy), and LAMMPS reverse-comm folds the returned +per-extended-atom forces back onto their owners. This is the first live +exercise of the whole with-comm graph chain (per-layer border_op + +owned-energy mask + reverse force fold) end-to-end through real MPI. + +Reference values come from ``source/tests/infer/gen_dpa2.py`` (the same +``deeppot_dpa2_graph.expected`` the C++ gtest uses). A second, independent +oracle -- ``deeppot_dpa2_graph_nlist_ref.pt2`` (same weights, dense-nlist +lower, NOT graph) -- is also exercised directly through LAMMPS so a +regression that only breaks the *C++* graph ingestion (not the Python +export path already cross-checked at gen-time) still gets caught. +""" + +import importlib.util +import os +import shutil +import signal +import subprocess as sp +import sys +import tempfile +from pathlib import ( + Path, +) + +import constants +import numpy as np +import pytest +from expected_ref import ( + read_expected_ref, +) +from lammps import ( + PyLammps, +) +from write_lmp_data import ( + write_lmp_data, +) + +pb_file = ( + Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot_dpa2_graph.pt2" +) +# Independent dense-nlist oracle exported from the SAME weights +# (gen_dpa2.py B.2, lower_kind="nlist"); at non-binding sel graph and dense +# math are equivalent (gen-time cross-check already enforces atomic energy / +# force / total-virial agreement within 1e-8 at the Python level -- see +# gen_dpa2.py B.4). Comparing through LAMMPS as well exercises the C++ graph +# ingestion path (edge tensors, node atype slicing, pair-exclusion seam) +# independently of that Python-level check. +nlist_ref_file = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa2_graph_nlist_ref.pt2" +) +ref_file = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa2_graph.expected" +) +# numb_aparam=1 sibling of the graph archive (gen_dpa2.py section C); the +# uniform per-atom parameter is fed via the pair_style ``aparam`` argument. +pb_aparam_file = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa2_graph_aparam.pt2" +) +# The MPI runner is backend-agnostic (DATAFILE PB_FILE OUTPUT + flags); reuse +# the DPA3 driver verbatim rather than duplicate it (same pattern as +# test_lammps_dpa1_graph_pt2.py). +mpi_runner = Path(__file__).parent / "run_mpi_pair_deepmd_dpa3_pt2.py" + +# Ceiling for EVERY mpirun invocation (parse mode included): a with-comm +# desync hangs the collective forever, so an unbounded should-succeed +# regression would hang the whole suite. Generous vs the observed ~30 s +# healthy runtime of the heaviest case here. +_MPI_DEFAULT_TIMEOUT = 600.0 + +data_file = Path(__file__).parent / "data_dpa2_graph_pt2.lmp" +# Wide-box, 3-way x-split variant for the genuinely-empty-rank MPI corner +# (``processors 3 1 1``): atoms stay in x in [0.25, 12.83] near the left +# edge of a [0, 90] box. With 3 even x-slabs of width 30, rank 0 owns +# [0, 30) (all atoms), rank 2 owns [60, 90) (empty of local atoms but +# picks up a periodic ghost of the x~0.25 atoms wrapped around the box's +# x=90/x=0 seam, since that distance ~0.25 is well within the ghost cutoff +# rcut(6.0)+skin(2.0)=8.0), and rank 1 (the MIDDLE slab, [30, 60)) borders +# neither the real atoms directly (nearest real atom at distance 30-12.83 +# ~= 17.17 > 8) nor the periodic seam -- so rank 1 is the genuinely empty +# rank (zero owned AND zero ghost atoms) this fixture is built to produce. +# A naive dpa1-style 2-way split (single periodic seam shared by both +# ranks) cannot achieve this: the "empty" rank always picks up a +# wrapped-around ghost of the near-edge atoms (see +# test_lammps_dpa1_graph_pt2.py's empty-subdomain comment) since with only +# two ranks in a periodic dimension the same pair of ranks borders on +# both sides. +data_file_empty_rank = Path(__file__).parent / "data_dpa2_graph_pt2_empty_rank.lmp" + +# Empty-SUBDOMAIN fixture (distinct from the empty-RANK fixture above): a rank +# that owns ZERO local atoms but DOES hold ghost atoms (nlocal == 0, +# nghost > 0). This is the SUPPORTED case -- it bypasses the C++ +# ``nall_real == 0`` guard and enters ``run_model_graph_with_comm`` with a +# real (ghost-only) extended region, exercising the owned-node energy mask +# (n_local == 0) and the per-layer border_op ghost exchange together. A +# 30 x 13 x 13 box with all six atoms clustered in x in [0.25, 12.83] under +# ``processors 2 1 1`` splits at x = 15: rank 1 owns [15, 30) (no atoms) but +# the near-edge atom at x = 12.83 is within rcut+skin (6.0 + 2.0 = 8.0) of the +# split, so rank 1 holds it as a ghost. Mirrors the dense DPA3 empty-subdomain +# fixture (``test_lammps_dpa3_pt2.py``). +data_file_empty_subdomain = ( + Path(__file__).parent / "data_dpa2_graph_pt2_empty_subdomain.lmp" +) + +# Reference values written by source/tests/infer/gen_dpa2.py (PBC case). +# Guarded with try/except because gen_dpa2.py only runs when PyTorch is built. +try: + _ref = read_expected_ref(ref_file)["pbc"] + expected_e = float(np.sum(_ref["expected_e"])) + expected_f = _ref["expected_f"].reshape(6, 3) + # LAMMPS uses the opposite sign convention for virial vs DeepPot. + expected_v = -_ref["expected_v"].reshape(6, 9) +except FileNotFoundError: + expected_e = expected_f = expected_v = None + +# Gate the reference-comparison tests on the generated ``.expected`` fixture so +# they skip cleanly (rather than failing with a ``TypeError`` on ``None``) when +# gen_dpa2.py has not run (e.g. PyTorch not built). The MPI multi-rank tests +# compare against a single-rank run of the same archive and do not need it. +_HAS_REF = expected_e is not None + +box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) +coord = np.array( + [ + [12.83, 2.56, 2.18], + [12.09, 2.87, 2.74], + [0.25, 3.32, 1.68], + [3.36, 3.00, 1.81], + [3.51, 2.51, 2.60], + [4.27, 3.22, 1.56], + ] +) +# Model type_map is ["O", "H"]; gtest atype = [0, 1, 1, 0, 1, 1] -> LAMMPS +# types [1, 2, 2, 1, 2, 2] under identity ``pair_coeff * *``. +type_OH = np.array([1, 2, 2, 1, 2, 2]) + + +def setup_module() -> None: + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip( + "Skip test because PyTorch support is not enabled.", + ) + write_lmp_data(box, coord, type_OH, data_file) + box_empty_rank = np.array([0, 90, 0, 13, 0, 13, 0, 0, 0]) + write_lmp_data(box_empty_rank, coord, type_OH, data_file_empty_rank) + box_empty_subdomain = np.array([0, 30, 0, 13, 0, 13, 0, 0, 0]) + write_lmp_data(box_empty_subdomain, coord, type_OH, data_file_empty_subdomain) + + +def teardown_module() -> None: + for f in [data_file, data_file_empty_rank, data_file_empty_subdomain]: + if f.exists(): + os.remove(f) + + +def _lammps(data_file, units="metal", atom_map: str = "yes") -> PyLammps: + lammps = PyLammps() + lammps.units(units) + lammps.boundary("p p p") + lammps.atom_style("atomic") + if atom_map != "no": + lammps.atom_modify(f"map {atom_map}") + lammps.neighbor("2.0 bin") + lammps.neigh_modify("every 10 delay 0 check no") + lammps.read_data(data_file.resolve()) + lammps.mass("1 16") + lammps.mass("2 2") + lammps.timestep(0.0005) + lammps.fix("1 all nve") + return lammps + + +@pytest.fixture +def lammps(): + lmp = _lammps(data_file=data_file) + yield lmp + lmp.close() + + +@pytest.mark.skipif(not _HAS_REF, reason="gen_dpa2.py .expected fixture not generated") +def test_pair_deepmd(lammps) -> None: + """Single-rank serial run (``atom_modify map yes``): the graph .pt2 + folds ghosts onto local owners (``fold_to_local=True``) and must match + the gen_dpa2.py reference for energy and per-atom force. + """ + lammps.pair_style(f"deepmd {pb_file.resolve()}") + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f[lammps.atoms[ii].id - 1] + ) + lammps.run(1) + + +@pytest.mark.skipif(not _HAS_REF, reason="gen_dpa2.py .expected fixture not generated") +def test_pair_deepmd_virial(lammps) -> None: + """Single-rank per-atom virial via ``centroid/stress/atom``.""" + lammps.pair_style(f"deepmd {pb_file.resolve()}") + lammps.pair_coeff("* *") + lammps.compute("virial all centroid/stress/atom NULL pair") + for ii in range(9): + jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") + lammps.dump( + "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) + ) + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f[lammps.atoms[ii].id - 1] + ) + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 + for ii in range(9): + assert np.array( + lammps.variables[f"virial{ii}"].value + ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) + + +@pytest.mark.skipif( + not nlist_ref_file.exists(), + reason="gen_dpa2.py deeppot_dpa2_graph_nlist_ref.pt2 fixture not generated", +) +def test_pair_deepmd_graph_matches_nlist_ref() -> None: + """Single-rank graph .pt2 vs the independent dense-nlist oracle + (``deeppot_dpa2_graph_nlist_ref.pt2``, same weights, ``lower_kind="nlist"``) + through LAMMPS, energy/force within 1e-6. + + gen_dpa2.py already cross-checks graph vs nlist at the Python + ``DeepPot.eval`` level (B.4, 1e-8) at generation time; this test instead + drives BOTH artifacts through the C++ ``DeepPotPTExpt`` / LAMMPS pair + style, so a regression confined to the C++ graph-ingestion seam (edge + tensor construction, node atype slicing, pair-exclusion application) + that the Python-level gen-time check cannot see is still caught. The + per-atom virial is deliberately NOT compared here: the graph path + assigns each edge's virial contribution fully to the source atom, a + different (equally valid) decomposition than the dense path's for a + message-passing model -- only energy and force (and, at the Python + level already, the *total* virial) are convention-independent. + """ + lmp_graph = _lammps(data_file=data_file) + lmp_graph.pair_style(f"deepmd {pb_file.resolve()}") + lmp_graph.pair_coeff("* *") + lmp_graph.run(0) + e_graph = lmp_graph.eval("pe") + f_graph = np.array([lmp_graph.atoms[ii].force for ii in range(6)]) + id_graph = [lmp_graph.atoms[ii].id for ii in range(6)] + lmp_graph.close() + + lmp_nlist = _lammps(data_file=data_file) + lmp_nlist.pair_style(f"deepmd {nlist_ref_file.resolve()}") + lmp_nlist.pair_coeff("* *") + lmp_nlist.run(0) + e_nlist = lmp_nlist.eval("pe") + f_nlist = np.array([lmp_nlist.atoms[ii].force for ii in range(6)]) + id_nlist = [lmp_nlist.atoms[ii].id for ii in range(6)] + lmp_nlist.close() + + # Same data file, single rank -> identical atom-id ordering; assert this + # rather than silently re-sorting so an ordering change is loud. + assert id_graph == id_nlist + assert e_graph == pytest.approx(e_nlist, rel=0, abs=1e-6) + np.testing.assert_allclose(f_graph, f_nlist, atol=1e-6, rtol=0) + + +# --------------------------------------------------------------------------- +# Multi-rank tests (message-passing with-comm graph route). +# +# dpa2's repformer participates in per-layer ghost exchange, so multi-rank +# LAMMPS routes to the nested with-comm artifact instead of the plain graph +# artifact used above and by dpa1's (non-MP) multi-rank path. These tests +# are the correctness gate for that new machinery: per-layer border_op, +# owned-energy mask, and reverse force fold. +# --------------------------------------------------------------------------- + + +def _run_mpi_subprocess( + extra_args: list[str] | None = None, + nprocs: int = 2, + data_path: Path | None = None, + processors: str | None = None, + runner_args: list[str] | None = None, + pb: Path | None = None, + capture: bool = False, + timeout: float | None = None, +) -> dict: + """Invoke the (backend-agnostic) DPA3 MPI runner under + ``mpirun -n `` against the dpa2 graph .pt2 and return + ``{"pe": float, "forces": (n, 3), "virials": (n, 9)}``. + + ``nprocs == 1`` forces ``--processors 1 1 1`` so the C++ side sees + ``nprocs == 1`` and routes to the plain (single-rank) graph artifact -- + a same-archive reference for the multi-rank comparison. ``pb`` + overrides the model archive (defaults to ``deeppot_dpa2_graph.pt2``). + + EVERY run is bounded: ``timeout`` (seconds, default + ``_MPI_DEFAULT_TIMEOUT``) covers the parse path too, so a deadlocked + collective in a should-succeed regression cannot hang the suite -- on + expiry the WHOLE mpirun process group is SIGKILLed (killing only mpirun + can leave orphaned ranks blocking in the collective and holding the + GPU). With ``capture=True``, return raw subprocess info + (``returncode``, ``stdout``, ``stderr``, ``timed_out``) instead of + parsed output -- used by the fail-fast tests; a timeout there returns + ``timed_out=True`` with ``returncode=None`` for the caller to assert + on. In parse mode a timeout raises ``RuntimeError`` and a nonzero exit + raises ``subprocess.CalledProcessError`` (matching the old + ``check_call`` behavior). + """ + if data_path is None: + data_path = data_file + if pb is None: + pb = pb_file + if timeout is None: + timeout = _MPI_DEFAULT_TIMEOUT + with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: + out_path = f.name + try: + argv = [ + "mpirun", + "-n", + str(nprocs), + sys.executable, + str(mpi_runner), + str(data_path.resolve()), + str(pb.resolve()), + out_path, + ] + if processors is not None: + argv.extend(["--processors", processors]) + elif nprocs == 1: + argv.extend(["--processors", "1 1 1"]) + if extra_args: + argv.extend(extra_args) + if runner_args: + argv.extend(runner_args) + proc = sp.Popen( + argv, + stdout=sp.PIPE if capture else None, + stderr=sp.PIPE if capture else None, + text=True, + start_new_session=True, + ) + try: + stdout, stderr = proc.communicate(timeout=timeout) + except sp.TimeoutExpired: + # Kill the whole process group: killing only mpirun can + # leave the deadlocked ranks orphaned (still blocking in + # the collective and holding the GPU). + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + stdout, stderr = proc.communicate() + if capture: + return { + "returncode": None, + "stdout": stdout or "", + "stderr": stderr or "", + "timed_out": True, + } + raise RuntimeError( + f"mpirun timed out after {timeout}s (process group killed); " + "a should-succeed MPI regression is deadlocked." + ) from None + if capture: + return { + "returncode": proc.returncode, + "stdout": stdout, + "stderr": stderr, + "timed_out": False, + } + if proc.returncode != 0: + raise sp.CalledProcessError(proc.returncode, argv) + with open(out_path) as fh: + lines = fh.read().strip().splitlines() + pe = float(lines[0]) + rows = np.array( + [list(map(float, line.split())) for line in lines[1:]], + dtype=np.float64, + ) + forces = rows[:, :3] + virials = rows[:, 3:] + return {"pe": pe, "forces": forces, "virials": virials} + finally: + if os.path.exists(out_path): + os.remove(out_path) + + +@pytest.mark.skipif( + shutil.which("mpirun") is None, reason="MPI is not installed on this system" +) +@pytest.mark.skipif( + importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" +) +def test_pair_deepmd_mpi_dpa2_graph_matches_single_rank() -> None: + """Multi-rank (``-n 2``) with-comm graph route must equal single-rank + (``-n 1``, plain graph artifact) on the SAME archive and trajectory. + + THE gate on the new with-comm graph machinery: per-layer + ``deepmd_export::border_op`` ghost exchange, the owned-node energy + mask (extended region includes ghost nodes that must not contribute + to the reduced energy), and the reverse-comm force fold back onto + owners. A wrong-but-finite divergence in any of the three would show + up here even though there is no hardcoded reference value. + """ + out_mpi = _run_mpi_subprocess(nprocs=2) + out_ref = _run_mpi_subprocess(nprocs=1) + assert out_mpi["pe"] == pytest.approx(out_ref["pe"], rel=1e-8, abs=1e-10) + np.testing.assert_allclose(out_mpi["forces"], out_ref["forces"], atol=1e-8, rtol=0) + # dpa2 per-atom virial components reach magnitudes ~3.5e3 (unlike dpa1, + # whose small magnitudes let the copied atol=1e-8, rtol=0 criterion pass); + # on CUDA the with-comm route's per-layer ghost exchange plus atomic + # index_add reorders the edge-virial summation, giving an observed max + # abs diff ~1e-6 = 6.9e-10 RELATIVE (10 significant digits agree; forces + # and energy match, and the energy check above already uses rel=1e-8). + # An absolute-only tolerance cannot absorb that at these magnitudes, so + # allow the same 1e-8 relative slack as the energy check. + np.testing.assert_allclose( + out_mpi["virials"], out_ref["virials"], atol=1e-8, rtol=1e-8 + ) + + +@pytest.mark.skipif( + shutil.which("mpirun") is None, reason="MPI is not installed on this system" +) +@pytest.mark.skipif( + importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" +) +def test_pair_deepmd_mpi_dpa2_graph_empty_subdomain_matches_single_rank() -> None: + """Multi-rank with-comm graph route with one rank owning ZERO local + atoms but holding ghosts (nlocal == 0, nghost > 0) must equal the + single-rank reference. + + Distinct from the empty-RANK test below: this is the SUPPORTED + empty-subdomain case that passes the C++ ``nall_real == 0`` guard + (nall_real > 0 because ghosts exist) and enters + ``run_model_graph_with_comm``. It is the DPA2-graph analogue of + ``test_lammps_dpa3_pt2.py::test_pair_deepmd_mpi_dpa3_empty_subdomain`` + and exercises the two pieces of the with-comm route that only fire when + a rank's owned count is zero: the ``n_local == 0`` owned-node energy + mask (the rank must contribute no reduced energy while its ghost nodes + still feed neighbors' features) and the per-layer ``border_op`` ghost + exchange on a ghost-only extended region. Wrong handling would either + desync the collective exchange or leak the empty rank's (masked) + energy, both of which break MP == SP here. + + 30 x 13 x 13 box under ``processors 2 1 1`` -> split at x = 15, rank 1 + owns [15, 30) (empty) but ghosts the near-edge atom at x = 12.83 + (within rcut + skin = 8.0). ``neigh_modify every 100`` also exercises + the cached (ago > 0) dispatch path on the empty-subdomain rank. + """ + runner_args = ["--neigh-every", "100"] + out_mpi = _run_mpi_subprocess( + nprocs=2, + data_path=data_file_empty_subdomain, + extra_args=["--nsteps", "5"], + runner_args=runner_args, + ) + out_ref = _run_mpi_subprocess( + nprocs=1, + data_path=data_file_empty_subdomain, + extra_args=["--nsteps", "5"], + runner_args=runner_args, + ) + assert out_mpi["pe"] == pytest.approx(out_ref["pe"], rel=1e-8, abs=1e-10) + np.testing.assert_allclose(out_mpi["forces"], out_ref["forces"], atol=1e-8, rtol=0) + np.testing.assert_allclose( + out_mpi["virials"], out_ref["virials"], atol=1e-8, rtol=1e-8 + ) + + +@pytest.mark.skipif( + shutil.which("mpirun") is None, reason="MPI is not installed on this system" +) +@pytest.mark.skipif( + importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" +) +@pytest.mark.skipif( + not pb_aparam_file.exists(), + reason="deeppot_dpa2_graph_aparam.pt2 not generated (gen_dpa2.py section C)", +) +def test_pair_deepmd_mpi_dpa2_graph_empty_subdomain_aparam_matches_single_rank() -> ( + None +): + """Empty subdomain (nlocal == 0, nghost > 0) with ``numb_aparam > 0``: + MP must equal SP. + + The ghost-only rank owns no atoms, so its runtime aparam is EMPTY while + the with-comm graph still carries ``N == nghost`` nodes; the C++ + marshalling must SYNTHESIZE a zero ``(N, daparam)`` aparam for it + (``extend_graph_aparam``). Feeding the empty tensor instead made the + artifact's aparam reshape fail rank-LOCALLY -- after the global + empty-rank preflight had already passed (``N > 0``) -- which + desynchronizes the subsequent per-layer ``border_op`` collective + sequence and leaves the peer rank hanging; every ``_run_mpi_subprocess`` + call is therefore bounded by ``_MPI_DEFAULT_TIMEOUT`` so a regression + fails instead of hanging the suite. The uniform ``aparam 0.25`` is + inert on the ghost-only rank (owned-node mask) but feeds every owned + atom on the other rank, so a wrong synthesis (or a dropped aparam) + breaks the MP == SP parity below. + """ + runner_args = ["--neigh-every", "100", "--pair-extra", "aparam 0.25"] + out_mpi = _run_mpi_subprocess( + nprocs=2, + data_path=data_file_empty_subdomain, + extra_args=["--nsteps", "5"], + runner_args=runner_args, + pb=pb_aparam_file, + ) + out_ref = _run_mpi_subprocess( + nprocs=1, + data_path=data_file_empty_subdomain, + extra_args=["--nsteps", "5"], + runner_args=runner_args, + pb=pb_aparam_file, + ) + assert out_mpi["pe"] == pytest.approx(out_ref["pe"], rel=1e-8, abs=1e-10) + np.testing.assert_allclose(out_mpi["forces"], out_ref["forces"], atol=1e-8, rtol=0) + np.testing.assert_allclose( + out_mpi["virials"], out_ref["virials"], atol=1e-8, rtol=1e-8 + ) + + +@pytest.mark.skipif( + shutil.which("mpirun") is None, reason="MPI is not installed on this system" +) +@pytest.mark.skipif( + importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" +) +def test_pair_deepmd_mpi_dpa2_graph_empty_rank_does_not_silently_succeed() -> None: + """A genuinely empty rank (zero owned AND zero ghost atoms) under the + message-passing with-comm graph route must NOT silently produce + wrong-but-plausible numbers. + + Design note (departs from ``test_lammps_dpa1_graph_pt2.py``'s + ``..._empty_subdomain`` test, which asserts the empty-rank run MATCHES + the single-rank reference): dpa1 is non-message-passing, so its + "empty" rank still runs the plain graph artifact over its (non-empty) + ghost region. dpa2's with-comm route instead needs every rank to + participate in the per-layer MPI ghost exchange (``border_op``); a rank + with zero nodes has nothing to export in the traced graph (violates + the exported ``Dim("n_node_total", min=1)`` and would desync the + collective ghost exchange across ranks). The C++ side + (``DeepPotPTExpt.cc``, guard added alongside the with-comm graph route) + throws a clear, actionable error on the empty rank instead of running. + + The failure is COLLECTIVE and PROMPT: before entering the per-layer + ``border_op`` collectives, every rank participates in a communicator- + wide min-reduction of its node count (``deepmd_export:: + allreduce_min_int``), so the non-empty peers detect the empty rank and + throw the same error instead of blocking forever waiting for it. A + timeout is therefore a FAILURE of this test (it would mean the + preflight regressed back into the historical deadlock), and the + documented error message must appear on a nonzero exit. + + ``data_file_empty_rank`` (3-way x-split, ``processors 3 1 1``) was + verified (see the module-level comment above the fixture) to put the + MIDDLE rank in a genuinely empty state -- unlike a naive 2-way + dpa1-style split, whose "empty" rank always picks up a ghost wrapped + around the periodic seam. + """ + out = _run_mpi_subprocess( + nprocs=3, + data_path=data_file_empty_rank, + processors="3 1 1", + capture=True, + timeout=120, + ) + assert not out["timed_out"], ( + "Multi-rank graph run with an empty rank timed out instead of " + "failing promptly: the collective empty-rank preflight " + "(allreduce_min_int) must make every rank throw BEFORE the " + "per-layer border_op collectives." + ) + assert out["returncode"] != 0, ( + "Expected the multi-rank message-passing graph run to fail loudly " + "on a genuinely empty rank, but it exited 0.\n" + f"stdout:\n{out['stdout'][-2000:]}\nstderr:\n{out['stderr'][-2000:]}" + ) + combined = out["stdout"] + out["stderr"] + assert "zero owned+ghost atoms" in combined, ( + "Expected the documented fail-loud message ('zero owned+ghost " + f"atoms'), got:\n{combined[-2000:]}" + ) diff --git a/source/op/pt/comm.cc b/source/op/pt/comm.cc index e00938883a..5a348a5243 100644 --- a/source/op/pt/comm.cc +++ b/source/op/pt/comm.cc @@ -105,6 +105,13 @@ class Border : public torch::autograd::Function { if (cuda_aware == 0) { recv_g1_tensor = torch::empty_like(g1).to(torch::kCPU); recv_g1_tensor.copy_(g1); + } else if (g1.device().is_cuda()) { + // nlocal/nghost are HOST tensors (their .item() reads above no + // longer synchronize the device): explicitly order the prior CUDA + // work that produced g1 before MPI reads the device buffers on the + // CUDA-aware path. The non-CUDA-aware path is ordered by the + // synchronizing copy_ to CPU above. + DPErrcheck(gpuDeviceSynchronize()); } } #endif @@ -130,6 +137,15 @@ class Border : public torch::autograd::Function { world, &request); } if (nsend) { +#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) + // the index_select pack above is an asynchronous kernel launch; + // the pre-loop fence only ordered work launched BEFORE it, so + // CUDA-aware MPI must be ordered behind the pack explicitly or + // MPI_Send races with stale device data + if (cuda_aware != 0 && send_g1_tensor.is_cuda()) { + DPErrcheck(gpuDeviceSynchronize()); + } +#endif MPI_Send(send_g1, nsend * tensor_size, mpi_type, sendproc[iswap], 0, world); } @@ -263,6 +279,11 @@ class Border : public torch::autograd::Function { if (cuda_aware == 0) { d_local_g1_tensor = torch::empty_like(grad_g1).to(torch::kCPU); d_local_g1_tensor.copy_(grad_g1); + } else if (grad_g1.device().is_cuda()) { + // Same explicit CUDA-stream-to-MPI ordering as the forward: the + // host-side nlocal/nghost reads no longer act as an accidental + // barrier before MPI touches the device gradient buffers. + DPErrcheck(gpuDeviceSynchronize()); } } #endif @@ -317,6 +338,15 @@ class Border : public torch::autograd::Function { world, &request); } if (nsend) { +#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) + // chained reverse communication: a PRIOR iteration's asynchronous + // index_add_ may write rows this iteration sends, and the pre-loop + // fence cannot order those later launches -- fence before letting + // CUDA-aware MPI read the device buffer + if (cuda_aware != 0 && d_local_g1_tensor.is_cuda()) { + DPErrcheck(gpuDeviceSynchronize()); + } +#endif MPI_Send(send_g1, nsend * tensor_size, mpi_type, sendproc[iswap], 0, world); } @@ -574,6 +604,40 @@ DEEPMD_MAYBE_UNUSED torch::Tensor border_op_backward_export( communicator_tensor, nlocal_tensor, nghost_tensor) .clone(); } + +/** + * @brief Communicator-wide int64 min-reduction (collective preflight). + * + * Used by the C++ inference runtime to agree ACROSS RANKS whether a + * multi-rank graph run may proceed (e.g. "does any rank have zero + * owned+ghost atoms?") BEFORE entering the per-layer ``border_op`` + * collectives -- a rank-local failure there would leave the peers blocked + * forever. Identity when MPI is not compiled in or the communicator handle + * is null (single-process runs). + */ +DEEPMD_MAYBE_UNUSED torch::Tensor allreduce_min_int_export( + const torch::Tensor& value_tensor, + const torch::Tensor& communicator_tensor) { + torch::Tensor value_cpu = value_tensor.to(torch::kCPU).contiguous(); +#ifdef USE_MPI + torch::Tensor comm_cpu = communicator_tensor.to(torch::kCPU).contiguous(); + if (*comm_cpu.data_ptr() != 0) { + MPI_Comm world; +#ifdef OMPI_MPI_H + std::int64_t* communicator = comm_cpu.data_ptr(); +#else + std::int64_t* ptr = comm_cpu.data_ptr(); + int* communicator = reinterpret_cast(ptr); +#endif + world = reinterpret_cast(*communicator); + std::int64_t local = *value_cpu.data_ptr(); + std::int64_t global_min = local; + MPI_Allreduce(&local, &global_min, 1, MPI_INT64_T, MPI_MIN, world); + return torch::full_like(value_cpu, global_min); + } +#endif + return value_cpu.clone(); +} } // namespace #undef DEEPMD_MAYBE_UNUSED @@ -586,6 +650,7 @@ TORCH_LIBRARY_FRAGMENT(deepmd_export, m) { "border_op_backward(Tensor sendlist, Tensor sendproc, Tensor recvproc, " "Tensor sendnum, Tensor recvnum, Tensor grad_g1, Tensor communicator, " "Tensor nlocal, Tensor nghost) -> Tensor"); + m.def("allreduce_min_int(Tensor value, Tensor communicator) -> Tensor"); } // Register CPU + CUDA implementations under explicit dispatch keys so @@ -594,10 +659,12 @@ TORCH_LIBRARY_FRAGMENT(deepmd_export, m) { TORCH_LIBRARY_IMPL(deepmd_export, CPU, m) { m.impl("border_op", border_op_export); m.impl("border_op_backward", border_op_backward_export); + m.impl("allreduce_min_int", allreduce_min_int_export); } #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) TORCH_LIBRARY_IMPL(deepmd_export, CUDA, m) { m.impl("border_op", border_op_export); m.impl("border_op_backward", border_op_backward_export); + m.impl("allreduce_min_int", allreduce_min_int_export); } #endif diff --git a/source/tests/common/dpmodel/test_descriptor_dpa2.py b/source/tests/common/dpmodel/test_descriptor_dpa2.py index af58d12790..8e331dfe1d 100644 --- a/source/tests/common/dpmodel/test_descriptor_dpa2.py +++ b/source/tests/common/dpmodel/test_descriptor_dpa2.py @@ -59,7 +59,15 @@ def test_self_consistency( em0.repinit.stddev = dstd em0.repformers.mean = davg_2 em0.repformers.stddev = dstd_2 + # this test exercises the legacy dense body's full 5-tuple (real + # g2/h2, not the thinner dense-ABI adapter's `None`/`None`); force + # the dense route explicitly rather than relying on the + # graph-eligibility of this particular config (`uses_graph_lower` + # would otherwise route `.call()` through `_call_graph_adapter`, + # see DescrptDPA2.uses_graph_lower). + em0.disable_graph_lower() em1 = DescrptDPA2.deserialize(em0.serialize()) + em1.disable_graph_lower() mm0 = em0.call(self.coord_ext, self.atype_ext, self.nlist, self.mapping) mm1 = em1.call(self.coord_ext, self.atype_ext, self.nlist, self.mapping) desired_shape = [ diff --git a/source/tests/common/dpmodel/test_dpa2_call_graph.py b/source/tests/common/dpmodel/test_dpa2_call_graph.py new file mode 100644 index 0000000000..57d329579d --- /dev/null +++ b/source/tests/common/dpmodel/test_dpa2_call_graph.py @@ -0,0 +1,1032 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Block-level parity between the graph-native +``DescrptBlockRepformers.call_graph`` (the repformer block kernel: per-edge +env-mat, g2 embedding, layer loop with the ghost-exchange seam, final +``rot_mat``) and the legacy dense ``DescrptBlockRepformers.call`` on the SAME +neighbor list. + +The graph path is ghost-free (``src`` is always a LOCAL owner), so the +per-layer ``_exchange_ghosts_graph`` seam is identity in dpmodel (the pt_expt +MPI override lives elsewhere); the dense path builds a genuine periodic +system with ghosts to also exercise its own (mapping-based) ``_exchange_ghosts``. +""" + +import itertools + +import numpy as np +import pytest + +from deepmd.dpmodel.descriptor.dpa2 import ( + DescrptDPA2, + RepformerArgs, + RepinitArgs, +) +from deepmd.dpmodel.descriptor.repformers import ( + DescrptBlockRepformers, +) +from deepmd.dpmodel.fitting import ( + InvarFitting, +) +from deepmd.dpmodel.model.edge_transform_output import ( + node_ownership_mask, +) +from deepmd.dpmodel.model.ener_model import ( + EnergyModel, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + from_dense_quartet, + graph_from_dense_quartet, +) +from deepmd.dpmodel.utils.nlist import ( + build_multiple_neighbor_list, + extend_input_and_build_neighbor_list, + get_multiple_nlist_key, +) + + +def _make_block(**kwargs) -> DescrptBlockRepformers: + cfg = { + "rcut": 4.0, + "rcut_smth": 0.5, + "sel": 24, + "ntypes": 2, + "nlayers": 2, + "g1_dim": 6, + "g2_dim": 4, + "axis_neuron": 2, + "attn1_hidden": 4, + "attn1_nhead": 2, + "attn2_hidden": 4, + "attn2_nhead": 2, + "precision": "float64", + "seed": 42, + } + cfg.update(kwargs) + return DescrptBlockRepformers(**cfg) + + +def _block_system(seed: int = 0, nf: int = 1, nloc: int = 6, box_size: float = 8.0): + """Small 2-type periodic system.""" + rng = np.random.default_rng(seed) + coord = rng.random((nf, nloc, 3)) * box_size + atype = rng.integers(0, 2, size=(nf, nloc)).astype(np.int64) + box = np.tile((np.eye(3) * box_size).reshape(1, 9), (nf, 1)) + return coord, atype, box + + +def _build_nlist(block: DescrptBlockRepformers, coord, atype, box): + return extend_input_and_build_neighbor_list( + coord, + atype, + block.get_rcut(), + block.get_sel(), + mixed_types=block.mixed_types(), + box=box, + ) + + +def _dense_and_graph_call(block: DescrptBlockRepformers, seed: int = 7): + """Run both the dense and graph block kernels on the SAME nlist. + + Returns dense outputs, graph outputs (reshaped to the dense (nf, nloc, ...) + layout where applicable), and the "real edge" mask (nf, nloc, nnei) that + accounts for BOTH nlist padding AND descriptor-level exclude_types -- + padding/excluded edges are not bit-comparable across the two fill + conventions (dense fills padding slots with atom-0 geometry; the graph + zero-fills), so g2/h2/sw parity is only asserted on real edges. + """ + coord, atype, box = _block_system(seed) + ext_coord, ext_atype, mapping, nlist = _build_nlist(block, coord, atype, box) + nf, nloc = atype.shape + nall = ext_atype.shape[1] + nnei = nlist.shape[-1] + + rng = np.random.default_rng(seed + 1) + g1_local = rng.normal(size=(nf, nloc, block.g1_dim)) + # only the first `nloc` slots of atype_embd_ext are read by the dense + # `call` (see repformers.py:559, `xp_take_first_n(atype_embd_ext, 1, nloc)`); + # the ghost portion is never touched (ghosts are re-derived per layer via + # `_exchange_ghosts`+`mapping`), so it is safe to leave it zero. + atype_embd_ext = np.zeros((nf, nall, block.g1_dim), dtype=np.float64) + atype_embd_ext[:, :nloc, :] = g1_local + + dense_out = block.call( + nlist, + ext_coord, + ext_atype, + atype_embd_ext=atype_embd_ext, + mapping=mapping, + ) + + graph, atype_local = graph_from_dense_quartet(ext_coord, ext_atype, nlist, mapping) + g1_input = np.reshape(g1_local, (-1, block.g1_dim)) + graph_out = block.call_graph(graph, atype_local, g1_input, static_nnei=nnei) + + # the "real" edge mask, INCLUDING descriptor-level exclude_types, mirrors + # the dense nlist-erasure at repformers.call:541-543. + excl = block.emask.build_type_exclude_mask(nlist, ext_atype) + real_mask = (nlist != -1) & np.asarray(excl, dtype=bool) + + return dense_out, graph_out, nf, nloc, nnei, real_mask + + +def _assert_block_parity(block: DescrptBlockRepformers, seed: int = 7) -> None: + ( + (dense_g1, dense_g2, dense_h2, dense_rot_mat, dense_sw), + ( + graph_g1, + graph_g2, + graph_h2, + graph_rot_mat, + graph_sw, + ), + nf, + nloc, + nnei, + real_mask, + ) = _dense_and_graph_call(block, seed) + + # g1: no neighbor axis -> exact reshape parity. + np.testing.assert_allclose( + np.reshape(graph_g1, dense_g1.shape), dense_g1, rtol=1e-12, atol=1e-12 + ) + # rot_mat: no neighbor axis -> exact reshape parity. + np.testing.assert_allclose( + np.reshape(graph_rot_mat, dense_rot_mat.shape), + dense_rot_mat, + rtol=1e-12, + atol=1e-12, + ) + assert not np.any(np.isnan(graph_g1)) + assert not np.any(np.isnan(graph_rot_mat)) + + # g2 / h2 / sw carry a neighbor axis; padding (and excluded) slots use + # DIFFERENT fill conventions between dense (atom-0 geometry via + # `nlist = where(nlist==-1, 0, nlist)`) and the graph (zero-filled + # `edge_vec` on masked edges) -- see edge_env_mat's docstring: "Padding + # edges produce nonzero values but are masked ... downstream." Compare + # only on the real (unmasked, unexcluded) neighbor slots. + ng2 = dense_g2.shape[-1] + graph_g2_dense_shape = np.reshape(graph_g2, (nf, nloc, nnei, ng2)) + graph_h2_dense_shape = np.reshape(graph_h2, (nf, nloc, nnei, 3)) + graph_sw_dense_shape = np.reshape(graph_sw, (nf, nloc, nnei)) + + np.testing.assert_allclose( + graph_g2_dense_shape[real_mask], dense_g2[real_mask], rtol=1e-12, atol=1e-12 + ) + np.testing.assert_allclose( + graph_h2_dense_shape[real_mask], dense_h2[real_mask], rtol=1e-12, atol=1e-12 + ) + np.testing.assert_allclose( + graph_sw_dense_shape[real_mask], dense_sw[real_mask], rtol=1e-12, atol=1e-12 + ) + # sw must be EXACTLY zero (not just unchecked) on padding slots for both + # paths -- this is a hard contract (edge_env_mat / dense nlist_mask). + np.testing.assert_allclose(graph_sw_dense_shape[~real_mask], 0.0, atol=0.0) + np.testing.assert_allclose(dense_sw[~real_mask], 0.0, atol=0.0) + + +class TestRepformersBlockCallGraphParity: + @pytest.mark.parametrize( + "smooth,use_sqrt_nnei,direct_dist", + list(itertools.product([True, False], [True, False], [True, False])), + ) + def test_block_graph_equals_dense(self, smooth, use_sqrt_nnei, direct_dist): + block = _make_block( + smooth=smooth, use_sqrt_nnei=use_sqrt_nnei, direct_dist=direct_dist + ) + _assert_block_parity(block) + + def test_block_graph_equals_dense_exclude_types(self): + block = _make_block(exclude_types=[[0, 1]]) + _assert_block_parity(block) + + +class TestRepformersBlockCallGraphTorch: + def test_call_graph_torch_smoke(self): + import torch + + block = _make_block() + coord, atype, box = _block_system(11) + ext_coord, ext_atype, mapping, nlist = _build_nlist(block, coord, atype, box) + nf, nloc = atype.shape + nnei = nlist.shape[-1] + rng = np.random.default_rng(12) + g1_local = rng.normal(size=(nf, nloc, block.g1_dim)) + + graph, atype_local = graph_from_dense_quartet( + ext_coord, ext_atype, nlist, mapping + ) + g1_input = np.reshape(g1_local, (-1, block.g1_dim)) + + ref = block.call_graph(graph, atype_local, g1_input, static_nnei=nnei) + + t_graph = type(graph)( + n_node=torch.from_numpy(np.asarray(graph.n_node)), + edge_index=torch.from_numpy(np.asarray(graph.edge_index)), + edge_vec=torch.from_numpy(np.asarray(graph.edge_vec)), + edge_mask=torch.from_numpy(np.asarray(graph.edge_mask)), + ) + got = block.call_graph( + t_graph, + torch.from_numpy(np.asarray(atype_local)), + torch.from_numpy(g1_input), + static_nnei=nnei, + ) + for r, g in zip(ref, got, strict=True): + assert isinstance(g, torch.Tensor) + np.testing.assert_allclose(g.numpy(), np.asarray(r), rtol=1e-12, atol=1e-12) + assert not torch.any(torch.isnan(g)) + + +class TestRepformersBlockSlotIndependence: + def test_mean_stddev_slot_independent(self): + """PR-A-proven invariant: EnvMatStatSe stats are slot-independent + (broadcast over the ``nnei`` axis), so reading ``self.mean[:, 0, :]`` / + ``self.stddev[:, 0, :]`` in ``call_graph`` is valid for the FULL + (ntypes, nnei, 4) stat tensor, not just the default zeros/ones. + """ + block = _make_block(set_davg_zero=False) + coord, atype, box = _block_system(3, nloc=8) + sample = {"coord": coord, "atype": atype, "box": box} + block.compute_input_stats([sample]) + + mean = np.asarray(block.mean) + stddev = np.asarray(block.stddev) + assert mean.shape == (block.ntypes, block.nnei, 4) + assert stddev.shape == (block.ntypes, block.nnei, 4) + for i in range(block.nnei): + np.testing.assert_array_equal(mean[:, i, :], mean[:, 0, :]) + np.testing.assert_array_equal(stddev[:, i, :], stddev[:, 0, :]) + + +class TestExchangeGhostsGraphSeam: + def test_identity_when_no_comm_dict(self): + block = _make_block() + g1 = np.random.default_rng(0).normal(size=(5, block.g1_dim)) + out = block._exchange_ghosts_graph(g1, None, n_total=5) + assert out is g1 + + def test_raises_notimplemented_with_comm_dict(self): + block = _make_block() + g1 = np.random.default_rng(0).normal(size=(5, block.g1_dim)) + with pytest.raises(NotImplementedError): + block._exchange_ghosts_graph(g1, {"some": "dict"}, n_total=5) + + +# --------------------------------------------------------------------------- +# Descriptor-level (Task 7): DescrptDPA2.uses_graph_lower / call_graph / +# _call_graph_adapter / the call() routing gate. +# --------------------------------------------------------------------------- + + +def _make_dpa2( + repinit_rcut: float = 4.0, + repinit_nsel: int = 200, + repformer_rcut: float = 2.0, + repformer_nsel: int = 150, + use_three_body: bool = False, + ntypes: int = 2, + repformer_attn: bool = True, + **kwargs, +) -> DescrptDPA2: + repinit_kwargs = { + "rcut": repinit_rcut, + "rcut_smth": 0.5, + "nsel": repinit_nsel, + "neuron": [6, 12], + "axis_neuron": 2, + "tebd_dim": 4, + "use_three_body": use_three_body, + } + if use_three_body: + # keep the three-body block's (rcut, nsel) strictly BELOW the + # repformer block's in the rcut-sorted nsel-ordering check + # (DescrptDPA2.__init__ asserts nsel is non-decreasing with rcut + # across repformer/three_body/repinit). + repinit_kwargs.update( + three_body_rcut=1.0, + three_body_rcut_smth=0.3, + three_body_sel=5, + three_body_neuron=[2, 4], + ) + repinit = RepinitArgs(**repinit_kwargs) + repformer = RepformerArgs( + rcut=repformer_rcut, + rcut_smth=0.5, + nsel=repformer_nsel, + nlayers=2, + g1_dim=6, + g2_dim=4, + axis_neuron=2, + update_g1_has_attn=repformer_attn, + update_g2_has_attn=repformer_attn, + attn1_hidden=4, + attn1_nhead=2, + attn2_hidden=4, + attn2_nhead=2, + ) + cfg = { + "ntypes": ntypes, + "repinit": repinit, + "repformer": repformer, + "precision": "float64", + "seed": 42, + } + cfg.update(kwargs) + return DescrptDPA2(**cfg) + + +def _system(seed: int = 0, nf: int = 1, nloc: int = 8, box_size: float = 6.0): + """Small 2-type periodic system.""" + rng = np.random.default_rng(seed) + coord = rng.random((nf, nloc, 3)) * box_size + atype = rng.integers(0, 2, size=(nf, nloc)).astype(np.int64) + box = np.tile((np.eye(3) * box_size).reshape(1, 9), (nf, 1)) + return coord, atype, box + + +def _dense_quartet(descr: DescrptDPA2, coord, atype, box): + return extend_input_and_build_neighbor_list( + coord, + atype, + descr.get_rcut(), + descr.get_sel(), + mixed_types=descr.mixed_types(), + box=box, + ) + + +class TestDPA2UsesGraphLowerGates: + def test_default_true(self) -> None: + descr = _make_dpa2() + assert descr.uses_graph_lower() is True + + def test_use_three_body_false(self) -> None: + descr = _make_dpa2(use_three_body=True) + assert descr.uses_graph_lower() is False + + def test_disable_graph_lower(self) -> None: + descr = _make_dpa2() + assert descr.uses_graph_lower() is True + descr.disable_graph_lower() + assert descr.uses_graph_lower() is False + # sticky: cannot be re-enabled + assert descr.uses_graph_lower() is False + + def test_compress_gate(self) -> None: + descr = _make_dpa2() + assert descr.uses_graph_lower() is True + descr.compress = True + assert descr.uses_graph_lower() is False + + +class TestDPA2AdapterBitExact: + """The money test: the dense->graph adapter must be BIT-EXACT vs the + dense body for ANY sel, including a deliberately BINDING repformer sel + (the slot mask in ``call_graph``'s ``_block_graph`` replicates + ``build_multiple_neighbor_list``'s ``nlist[:, :, :ns]`` slicing). + """ + + @pytest.mark.parametrize( + "repformer_nsel,box_size,nloc,seed", + [ + ( + 150, + 6.0, + 8, + 21, + ), # non-binding: repformer sel comfortably covers all neighbors + (3, 3.0, 12, 22), # binding: dense cluster, repformer rcut truncates hard + ], + ) + def test_adapter_bitexact_any_sel( + self, repformer_nsel, box_size, nloc, seed + ) -> None: + descr = _make_dpa2(repinit_nsel=200, repformer_nsel=repformer_nsel) + coord, atype, box = _system(seed=seed, nloc=nloc, box_size=box_size) + ext_coord, ext_atype, mapping, nlist = _dense_quartet(descr, coord, atype, box) + nf, n_loc, nnei = nlist.shape + + # confirm the "binding" case actually truncates (otherwise the test + # would vacuously pass without exercising the slot mask): count real + # neighbors within the repformer rcut among ALL nnei outer slots and + # compare against the configured repformer nsel. + rc = descr.repformers.get_rcut() + ns = descr.repformers.get_nsel() + full_sub = build_multiple_neighbor_list(ext_coord, nlist, [rc], [nnei])[ + get_multiple_nlist_key(rc, nnei) + ] + real_counts = (full_sub >= 0).sum(axis=-1) + binds = bool(np.any(real_counts > ns)) + assert binds == (repformer_nsel == 3) + + dense = descr._call_dense(ext_coord, ext_atype, nlist, mapping=mapping) + adapter = descr._call_graph_adapter(ext_coord, ext_atype, nlist, mapping) + + dense_g1, dense_rot_mat, dense_g2, dense_h2, dense_sw = dense + adapter_g1, adapter_rot_mat, adapter_g2, adapter_h2, adapter_sw = adapter + + np.testing.assert_allclose(adapter_g1, dense_g1, rtol=1e-12, atol=1e-12) + np.testing.assert_allclose( + adapter_rot_mat, dense_rot_mat, rtol=1e-12, atol=1e-12 + ) + assert adapter_g2 is None + assert adapter_h2 is None + assert adapter_sw.shape == dense_sw.shape == (nf, n_loc, ns) + np.testing.assert_allclose(adapter_sw, dense_sw, rtol=1e-12, atol=1e-12) + assert not np.any(np.isnan(adapter_g1)) + assert not np.any(np.isnan(adapter_rot_mat)) + + def test_adapter_divergence_davg_nonzero_documented(self) -> None: + """Pin the ONE known bit-exactness exception (see the ``Notes`` + block on :meth:`DescrptDPA2.call_graph`'s docstring): when + ``set_davg_zero=False`` (nonzero ``repinit.mean``) AND + ``exclude_types == []``, the dense + ``DescrptBlockSeAtten.call`` (``attn_layer == 0``, always the case + for DPA2's ``repinit``) leaks a padding-slot residual into its + output -- ``PairExcludeMask.build_type_exclude_mask`` short-circuits + to all-ones when nothing is excluded, so real nlist padding is never + zeroed out of ``rr`` before it reaches ``gr``, and once + ``mean != 0`` the (masked-to-zero-then-mean-subtracted) padding rows + become nonzero garbage that depends on the padding COUNT (i.e. on + ``sel``, not on real physics). The graph path masks padding/excluded + edges out before every ``segment_sum`` and does not reproduce this + leak -- its result is the physically correct one. This regime is + NOT covered by ``test_adapter_bitexact_any_sel``'s bit-exactness + guarantee, and is not fixed here (pre-existing dense-body bug, see + ``.superpowers/sdd/task-7-report.md``'s "Known limitation"). + + This test does not attempt to reproduce the "atom-at-extended- + index-0" framing verbatim -- the leak is triggered by ANY real nlist + padding once ``mean != 0``, independent of where index 0 physically + sits (verified empirically: the same small, non-clustered system + used elsewhere in this file already has padding, since + ``repinit_nsel=200`` comfortably exceeds each atom's real neighbor + count). What matters, and what this test pins, is: (a) with + ``mean == 0`` (the ``set_davg_zero=True``-equivalent control) the + adapter stays 1e-12 bit-exact vs dense, exactly like + ``test_adapter_bitexact_any_sel``; (b) with ``mean != 0`` and + ``exclude_types == []`` and real padding present, the adapter and + dense DIVERGE by a non-trivial amount -- a deliberate, documented + record of the dense-side bug, not a regression in the adapter. + """ + repinit_nsel = 200 + descr = _make_dpa2(repinit_nsel=repinit_nsel, repformer_nsel=150) + assert descr.repinit.exclude_types == [] + coord, atype, box = _system(seed=41, nloc=8, box_size=6.0) + ext_coord, ext_atype, mapping, nlist = _dense_quartet(descr, coord, atype, box) + + # self-check: real nlist padding is actually exercised (otherwise + # this test would be vacuous -- no padding, nothing to leak). + real_counts = (nlist != -1).sum(axis=-1) + assert bool(np.any(real_counts < repinit_nsel)) + + # control: mean == 0 (fresh descriptor's default, i.e. the + # set_davg_zero=True-equivalent state) -> bit-exact, as elsewhere. + dense0 = descr._call_dense(ext_coord, ext_atype, nlist, mapping=mapping) + adapter0 = descr._call_graph_adapter(ext_coord, ext_atype, nlist, mapping) + np.testing.assert_allclose(adapter0[0], dense0[0], rtol=1e-12, atol=1e-12) + np.testing.assert_allclose(adapter0[1], dense0[1], rtol=1e-12, atol=1e-12) + + # set repinit's mean to a nonzero constant (mirrors how other tests + # in this file assign davg/dstd directly; slot-independent per + # TestRepformersBlockSlotIndependence's proven invariant). + nnei_repinit = descr.repinit.get_nsel() + descr.repinit.mean = np.full( + (descr.ntypes, nnei_repinit, 4), 0.3, dtype=np.float64 + ) + + dense1 = descr._call_dense(ext_coord, ext_atype, nlist, mapping=mapping) + adapter1 = descr._call_graph_adapter(ext_coord, ext_atype, nlist, mapping) + max_abs_diff = float(np.max(np.abs(adapter1[0] - dense1[0]))) + # divergence documented, not a regression: dense leaks a + # padding-count-dependent residual that the graph correctly excludes. + assert max_abs_diff > 1e-6 + assert not np.any(np.isnan(adapter1[0])) + assert not np.any(np.isnan(adapter1[1])) + + +class TestDPA2BlockMaskReplicatesMultipleNlist: + def test_block_mask_formula_replicates_multiple_nlist(self) -> None: + """Standalone re-implementation of the ``_block_graph`` mask formula + (dist filter + slot cap) -- kept as a cheap, fast sanity check of the + FORMULA, but it never calls the shipped code, so it cannot catch a + regression in ``_block_graph``'s actual slicing/masking. See + ``test_block_graph_seam_matches_dense_sub_nlist`` below for the test + that exercises the real shipped code path. + """ + descr = _make_dpa2(repinit_nsel=200, repformer_nsel=5) + coord, atype, box = _system(seed=31, nloc=12, box_size=3.0) + ext_coord, ext_atype, mapping, nlist = _dense_quartet(descr, coord, atype, box) + nf, nloc, nnei = nlist.shape + + graph, _ = graph_from_dense_quartet(ext_coord, ext_atype, nlist, mapping) + dist = np.linalg.norm(np.asarray(graph.edge_vec), axis=-1) + edge_mask = np.asarray(graph.edge_mask) + + rc = descr.repformers.get_rcut() + ns = descr.repformers.get_nsel() + e_ax = edge_mask.shape[0] + slot = np.arange(e_ax) % nnei + # graph analogue of DescrptDPA2.call_graph's ``_block_graph``: dist + # mask always, slot mask when static_nnei (== nnei here) is set. + block_mask = edge_mask & (dist <= rc) & (slot < ns) + block_mask = block_mask.reshape(nf, nloc, nnei) + + # binding-sel sanity: without the slot cap, more than `ns` neighbors + # survive the dist filter for at least one center (otherwise this + # test would not exercise the slot mask at all). + dist_only_mask = (edge_mask & (dist <= rc)).reshape(nf, nloc, nnei) + assert np.any(dist_only_mask.sum(axis=-1) > ns) + + sub_nlist = build_multiple_neighbor_list(ext_coord, nlist, [rc], [ns])[ + get_multiple_nlist_key(rc, ns) + ] + expected = np.zeros((nf, nloc, nnei), dtype=bool) + expected[:, :, :ns] = sub_nlist >= 0 + np.testing.assert_array_equal(block_mask, expected) + + def test_block_graph_seam_matches_dense_sub_nlist(self) -> None: + """Exercise the REAL shipped ``_block_graph`` closure (private, + reachable only through ``DescrptDPA2.call_graph``) via the public + seam, on a fixture where the repformer block's ``(rc, ns)`` genuinely + truncate the outer graph. + + Strategy: run the shipped ``descr.call_graph(...)`` end to end (this + is what ``_call_graph_adapter`` calls in production), then + independently reconstruct a REFERENCE for "what the repformers block + saw" by: + + 1. Replaying the repinit stage with the SAME public + ``descr.repinit.call_graph`` call the shipped code makes + (repinit's own ``nsel`` (200) is >= the outer ``static_nnei`` + here, so ``_block_graph`` takes its mask-only fast path -- no + slicing to replicate, just the dist mask), to get the g1 that + feeds into repformers. + 2. Building a graph PRE-TRUNCATED BY HAND to the dense + ``build_multiple_neighbor_list`` sub-nlist (the same width-``ns`` + sub-nlist ``_call_dense`` itself uses), via + ``graph_from_dense_quartet``. + 3. Calling the repformer BLOCK's own public ``call_graph`` on that + pre-truncated graph. + + If ``_block_graph``'s internal slice+mask selected exactly the same + edges as the dense sub-nlist, the shipped end-to-end output and this + by-hand reference must be bit-identical. + """ + import dataclasses + + descr = _make_dpa2(repinit_nsel=200, repformer_nsel=5) + assert descr.add_tebd_to_repinit_out is False # keeps the repinit replay simple + coord, atype, box = _system(seed=31, nloc=12, box_size=3.0) + ext_coord, ext_atype, mapping, nlist = _dense_quartet(descr, coord, atype, box) + nf, nloc, nnei = nlist.shape + + rc = descr.repformers.get_rcut() + ns = descr.repformers.get_nsel() + full_sub = build_multiple_neighbor_list(ext_coord, nlist, [rc], [nnei])[ + get_multiple_nlist_key(rc, nnei) + ] + real_counts = (full_sub >= 0).sum(axis=-1) + # non-vacuity: the slot cap must actually be exercised, otherwise + # this test would pass even with a broken slice. + assert np.any(real_counts > ns) + + graph, atype_local = graph_from_dense_quartet( + ext_coord, ext_atype, nlist, mapping + ) + + # 1) the real, shipped, end-to-end code path. + g1_full, rot_mat_full, sw_full = descr.call_graph( + graph, atype_local, static_nnei=nnei, return_sw=True + ) + + # 2) replay repinit via the same public call the shipped code makes + # (mask-only fast path: repinit's own nsel >= the outer static_nnei). + tebd_table = descr.type_embedding.call() + dist = np.linalg.norm(np.asarray(graph.edge_vec), axis=-1) + repinit_mask = np.asarray(graph.edge_mask) & (dist <= descr.repinit.get_rcut()) + repinit_graph = dataclasses.replace(graph, edge_mask=repinit_mask) + g1_repinit, _ = descr.repinit.call_graph( + repinit_graph, atype_local, type_embedding=tebd_table, static_nnei=nnei + ) + g1_repinit = descr.g1_shape_tranform(g1_repinit) + + # 3) hand-pre-truncate to the dense sub-nlist's kept entries, and run + # the repformer BLOCK's own public call_graph on it directly. + sub_nlist = build_multiple_neighbor_list(ext_coord, nlist, [rc], [ns])[ + get_multiple_nlist_key(rc, ns) + ] + graph_ref, atype_local_ref = graph_from_dense_quartet( + ext_coord, ext_atype, sub_nlist, mapping + ) + np.testing.assert_array_equal( + np.asarray(atype_local_ref), np.asarray(atype_local) + ) + g1_ref, _g2_ref, _h2_ref, rot_mat_ref, sw_ref = descr.repformers.call_graph( + graph_ref, atype_local, g1_repinit, comm_dict=None, static_nnei=ns + ) + if descr.concat_output_tebd: + g1_inp = np.asarray(tebd_table)[np.asarray(atype_local)] + g1_ref = np.concatenate([np.asarray(g1_ref), g1_inp], axis=-1) + + # the shipped internal slice+mask must have selected EXACTLY the + # edges the dense sub-nlist keeps -- bit-identical, not just close. + np.testing.assert_allclose(g1_full, g1_ref, rtol=0.0, atol=0.0) + np.testing.assert_allclose(rot_mat_full, rot_mat_ref, rtol=0.0, atol=0.0) + np.testing.assert_allclose(sw_full, sw_ref, rtol=0.0, atol=0.0) + + +class TestDPA2CallRouting: + def test_call_routing_graph_eligible(self) -> None: + # Even for a graph-ELIGIBLE config, the dense ``call`` runs the dense + # body -- it is the cross-backend consistency reference and must match + # the tf/pt/pd/jax dense descriptors bit-for-bit. DPA2's repinit is + # always ``attn_layer == 0``, where ``_call_graph_adapter`` diverges + # from ``_call_dense`` for non-trivial statistics (the accepted + # padding-leak divergence, see DescrptDPA2.call / call_graph Notes). + # The graph-native route is reached through ``call_graph``, never + # ``call``. + descr = _make_dpa2(repinit_nsel=200, repformer_nsel=150) + assert descr.uses_graph_lower() is True + coord, atype, box = _system(seed=41, nloc=8, box_size=6.0) + ext_coord, ext_atype, mapping, nlist = _dense_quartet(descr, coord, atype, box) + + called = descr.call(ext_coord, ext_atype, nlist, mapping=mapping) + dense = descr._call_dense(ext_coord, ext_atype, nlist, mapping=mapping) + for c, d in zip(called, dense, strict=True): + if c is None or d is None: + assert c is None and d is None + else: + np.testing.assert_array_equal(c, d) + + def test_call_routing_three_body_dense(self) -> None: + descr = _make_dpa2(repinit_nsel=200, repformer_nsel=150, use_three_body=True) + assert descr.uses_graph_lower() is False + coord, atype, box = _system(seed=42, nloc=8, box_size=6.0) + ext_coord, ext_atype, mapping, nlist = _dense_quartet(descr, coord, atype, box) + + called = descr.call(ext_coord, ext_atype, nlist, mapping=mapping) + dense = descr._call_dense(ext_coord, ext_atype, nlist, mapping=mapping) + for c, d in zip(called, dense, strict=True): + if c is None or d is None: + assert c is None and d is None + else: + np.testing.assert_array_equal(c, d) + + +def _make_energy_model( + repinit_nsel: int = 200, repformer_nsel: int = 150 +) -> EnergyModel: + # repformer_attn=False: mirrors test_dpa1_graph_model_energy.py's own + # choice of attn_layer=0 for its carry-all parity fixture. The + # CARRY-ALL graph builder has no padding slots, so smooth attention's + # softmax there is genuinely sel-independent (real neighbors only) -- + # by design DIFFERENT from the dense body, which keeps sel-padding + # slots in its softmax denominator (see DescrptDPA1.call_graph's + # Notes). That divergence is intentional and orthogonal to what this + # test checks (non-attention carry-all/dense parity at non-binding + # sel); the shape-static adapter path (_call_graph_adapter, covered by + # TestDPA2AdapterBitExact) is the one that reproduces the dense + # attention exactly, including with attention enabled. + ds = _make_dpa2( + repinit_nsel=repinit_nsel, repformer_nsel=repformer_nsel, repformer_attn=False + ) + ft = InvarFitting( + "energy", + ds.get_ntypes(), + ds.get_dim_out(), + 1, + mixed_types=ds.mixed_types(), + ) + return EnergyModel(ds, ft, type_map=["foo", "bar"]) + + +class TestDPA2ModelEnergyCarryAll: + """Model-level: ``EnergyModel(dpa2).call_common(..., + neighbor_graph_method="dense")`` vs the dense route at NON-binding sel + (mirrors ``test_dpa1_graph_model_energy.py``). + """ + + @pytest.mark.parametrize("periodic", [True, False]) + def test_energy_parity_non_binding_sel(self, periodic) -> None: + rng = np.random.default_rng(51) + nloc = 8 + coord = rng.normal(size=(1, nloc, 3)) * 1.5 + atype = np.array([[0, 1, 0, 1, 0, 1, 0, 1]], dtype=np.int64) + box = None + if periodic: + # large box so the cell is essentially non-periodic for rcut=4.0 + box = np.eye(3).reshape(1, 9) * 20.0 + model = _make_energy_model() + + dense = model.call_common(coord, atype, box, neighbor_graph_method="legacy") + graph = model.call_common(coord, atype, box, neighbor_graph_method="dense") + + np.testing.assert_allclose( + graph["energy_redu"], dense["energy_redu"], rtol=1e-12, atol=1e-12 + ) + np.testing.assert_allclose( + graph["energy"], dense["energy"], rtol=1e-12, atol=1e-12 + ) + np.testing.assert_array_equal(graph["mask"], dense["mask"]) + + +def _make_attn_energy_model() -> EnergyModel: + """Attention-enabled twin of :func:`_make_energy_model` (both repformer + attention channels ON) for the phantom-count-compensated smooth-attention + tests below. + """ + ds = _make_dpa2(repinit_nsel=200, repformer_nsel=150, repformer_attn=True) + ft = InvarFitting( + "energy", + ds.get_ntypes(), + ds.get_dim_out(), + 1, + mixed_types=ds.mixed_types(), + ) + return EnergyModel(ds, ft, type_map=["foo", "bar"]) + + +class TestDPA2AttentionCarryAllSmoothParity: + """Phantom-count-compensated smooth attention on the carry-all graph route. + + The dense smooth-attention softmax keeps exactly ``sel - n_real`` phantom + terms (each ``exp(-attnw_shift)``) in every denominator -- a count that + never changes with geometry. The graph route instead used one phantom per + PRESENT edge (a geometry-dependent count): outputs differed from dense by + ``O(sel * exp(-attnw_shift)) ~ 1e-7`` even at non-binding sel, and the + energy jumped by ``O(exp(-attnw_shift))`` whenever an edge entered or left + the graph at the model cutoff. With the compensation (masked pairs + excluded from the softmax, ``max(sel - n_real, 0)`` phantoms added), the + denominator is term-for-term the dense one at non-binding sel: parity at + 1e-12 AND exact continuity at the cutoff. + """ + + def test_attention_energy_parity_non_binding_sel(self) -> None: + rng = np.random.default_rng(53) + nloc = 8 + coord = rng.normal(size=(1, nloc, 3)) * 1.5 + atype = np.array([[0, 1, 0, 1, 0, 1, 0, 1]], dtype=np.int64) + model = _make_attn_energy_model() + + dense = model.call_common(coord, atype, None, neighbor_graph_method="legacy") + graph = model.call_common(coord, atype, None, neighbor_graph_method="dense") + + np.testing.assert_allclose( + graph["energy_redu"], dense["energy_redu"], rtol=1e-12, atol=1e-12 + ) + np.testing.assert_allclose( + graph["energy"], dense["energy"], rtol=1e-12, atol=1e-12 + ) + + def test_attention_energy_smooth_at_model_cutoff(self) -> None: + """Graph-route energy is continuous when an atom crosses the model + cutoff (repinit rcut=4.0). A second atom sits inside the repformer + cutoff so the attention softmaxes have >=2 members (the denominator- + jump scenario). Without the compensation the jump is + ``O(exp(-attnw_shift)) ~ 1e-10``; with it, only float-reassociation + noise remains. + """ + model = _make_attn_energy_model() + atype = np.array([[0, 1, 1]], dtype=np.int64) + rcut = 4.0 # _make_dpa2 default repinit (model) rcut + + def energy(r: float) -> float: + coord = np.array( + [[[0.0, 0.0, 0.0], [1.5, 0.0, 0.0], [0.0, r, 0.0]]], + dtype=np.float64, + ) + ret = model.call_common(coord, atype, None, neighbor_graph_method="dense") + return float(np.sum(ret["energy_redu"])) + + eps = 1e-8 + jump = abs(energy(rcut + eps) - energy(rcut - eps)) + assert jump < 1e-13, f"graph-route energy jump {jump:.3e} at rcut" + + +class TestNodeOwnershipMask: + """Direct unit tests of :func:`node_ownership_mask` -- the per-frame + block arithmetic (``cumulative_sum``/``take``) is where bugs hide, so + this is exercised independently of any model. + """ + + def test_single_frame(self) -> None: + n_node = np.array([5], dtype=np.int64) + n_local = np.array([3], dtype=np.int64) + mask = node_ownership_mask(n_node, n_local, n_total=5) + np.testing.assert_array_equal(mask, np.array([True, True, True, False, False])) + + def test_multi_frame_different_n_local(self) -> None: + # frame 0: 3 nodes, owns 2 -> [T, T, F] + # frame 1: 2 nodes, owns 1 -> [T, F] + n_node = np.array([3, 2], dtype=np.int64) + n_local = np.array([2, 1], dtype=np.int64) + mask = node_ownership_mask(n_node, n_local, n_total=5) + np.testing.assert_array_equal(mask, np.array([True, True, False, True, False])) + + def test_all_owned_is_all_true(self) -> None: + n_node = np.array([4, 3], dtype=np.int64) + n_local = n_node.copy() + mask = node_ownership_mask(n_node, n_local, n_total=7) + np.testing.assert_array_equal(mask, np.ones(7, dtype=bool)) + + def test_zero_owned_frame(self) -> None: + # a frame that owns nothing (all ghost) -- degenerate but must not crash. + n_node = np.array([2, 3], dtype=np.int64) + n_local = np.array([0, 2], dtype=np.int64) + mask = node_ownership_mask(n_node, n_local, n_total=5) + np.testing.assert_array_equal(mask, np.array([False, False, True, True, False])) + + +class TestOwnedNodeMaskEnergyReduction: + """``n_local`` (owned-node mask) in the graph output reduction (Task 9): + ghost rows (index >= n_local[frame]) must be excluded from the + DIFFERENTIATED per-frame energy, while ``atom_energy`` itself stays FULL. + """ + + def _make_graph_and_model(self, nloc: int = 5, seed: int = 3): + rng = np.random.default_rng(seed) + coord = rng.normal(size=(1, nloc, 3)) * 1.5 + atype = np.array([[ii % 2 for ii in range(nloc)]], dtype=np.int64) + model = _make_energy_model() + ext_coord, ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( + coord, + atype, + model.get_rcut(), + model.get_sel(), + mixed_types=model.mixed_types(), + box=None, + ) + ng = from_dense_quartet(ext_coord, nlist, mapping) + atype_local = ext_atype.reshape(-1)[:nloc] + return model, ng, atype_local + + def test_owned_mask_energy_reduction(self) -> None: + """1 frame, 5 nodes, n_local=3: energy_redu == sum(atom_energy[:3]), + and the difference from the unmasked (``n_local=None``) energy_redu + equals sum(atom_energy[3:5]). Under the #5758 convention the ghost region + rows of the per-node output are ALSO masked to zero (owned rows + unchanged) -- each ghost atom's fitting output is owned, and + reported, by another rank. + """ + model, ng, atype_local = self._make_graph_and_model(nloc=5) + + out_full = model.call_lower_graph( + atype=atype_local, + n_node=ng.n_node, + edge_index=ng.edge_index, + edge_vec=ng.edge_vec, + edge_mask=ng.edge_mask, + ) + n_local = np.array([3], dtype=np.int64) + out_masked = model.call_lower_graph( + atype=atype_local, + n_node=ng.n_node, + edge_index=ng.edge_index, + edge_vec=ng.edge_vec, + edge_mask=ng.edge_mask, + n_local=n_local, + ) + + # per-node output: owned rows unchanged, ghost rows masked to zero. + np.testing.assert_allclose( + out_masked["energy"].reshape(-1)[:3], + out_full["energy"].reshape(-1)[:3], + rtol=1e-12, + atol=1e-12, + ) + np.testing.assert_allclose( + out_masked["energy"].reshape(-1)[3:], 0.0, rtol=0, atol=1e-12 + ) + + atom_energy = out_full["energy"].reshape(-1) + owned_sum = atom_energy[:3].sum() + halo_sum = atom_energy[3:].sum() + + np.testing.assert_allclose( + out_masked["energy_redu"].reshape(-1), + [owned_sum], + rtol=1e-12, + atol=1e-12, + ) + diff = (out_full["energy_redu"] - out_masked["energy_redu"]).reshape(-1) + np.testing.assert_allclose(diff, [halo_sum], rtol=1e-12, atol=1e-12) + + def test_n_local_none_is_byte_identical(self) -> None: + """Regression: omitting ``n_local`` (default ``None``) reproduces the + pre-Task-9 unmasked reduction exactly. + """ + model, ng, atype_local = self._make_graph_and_model(nloc=5, seed=11) + + out_default = model.call_lower_graph( + atype=atype_local, + n_node=ng.n_node, + edge_index=ng.edge_index, + edge_vec=ng.edge_vec, + edge_mask=ng.edge_mask, + ) + out_explicit_none = model.call_lower_graph( + atype=atype_local, + n_node=ng.n_node, + edge_index=ng.edge_index, + edge_vec=ng.edge_vec, + edge_mask=ng.edge_mask, + n_local=None, + ) + np.testing.assert_array_equal( + out_default["energy_redu"], out_explicit_none["energy_redu"] + ) + np.testing.assert_array_equal( + out_default["energy"], out_explicit_none["energy"] + ) + + +class TestGraphAttentionCutoffContinuityBindingSel: + """Cutoff continuity when the carry-all degree EXCEEDS ``sel``. + + OutisLi repro: repformer ``rcut=2, nsel=1``, one neighbor fixed at + ``r=1`` and a second crossing ``r=2`` -- just inside the cutoff every + center has 2 neighbors > sel=1, so the CLAMPED phantom count was zero + on both sides of the crossing and the departing pair/edge removed a + finite ``exp(-attnw_shift)`` softmax-denominator term: the full-model + energy stepped by ~-6.0e-10 (LocalAtten-only -3.6e-9, Atten2Map-only + +4.9e-9 -- independent defects at repformers.py's two ``call_graph`` + sites). The SIGNED count ``sel - n_real`` subtracts the excess beyond + sel, so the boundary increment vanishes for arbitrary degree. g1 + (LocalAtten) and g2 (Atten2Map) attention are exercised SEPARATELY. + """ + + def _descriptor_delta(self, *, g1_attn: bool, g2_attn: bool, eps: float): + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + ) + + repinit = RepinitArgs( + rcut=4.0, + rcut_smth=0.5, + nsel=20, + neuron=[6, 12], + axis_neuron=2, + tebd_dim=4, + ) + repformer = RepformerArgs( + rcut=2.0, + rcut_smth=0.5, + nsel=1, # BINDING: degree 2 > sel 1 just inside the cutoff + nlayers=2, + g1_dim=6, + g2_dim=4, + axis_neuron=2, + update_g1_has_attn=g1_attn, + update_g2_has_attn=g2_attn, + attn1_hidden=4, + attn1_nhead=2, + attn2_hidden=4, + attn2_nhead=2, + ) + descr = DescrptDPA2( + ntypes=2, + repinit=repinit, + repformer=repformer, + precision="float64", + seed=42, + ) + te = descr.type_embedding.call() + atype = np.array([[0, 1, 1]], dtype=np.int64) + outs = [] + for d in (2.0 - eps, 2.0 + eps): + coord = np.zeros((1, 3, 3), dtype=np.float64) + coord[0, 1, 0] = 1.0 # fixed neighbor at r = 1 + coord[0, 2, 0] = d # crossing neighbor at r = 2 -+ eps + g = build_neighbor_graph(coord, atype, None, descr.get_rcut()) + out, _ = descr.call_graph(g, atype.reshape(-1), type_embedding=te) + outs.append(np.asarray(out)) + # anti-vacuity: just inside the cutoff, the center has 2 repformer + # neighbors (r=1 and r=2-eps) > nsel=1 -- the binding regime where + # the clamped scheme was discontinuous. + return float(np.max(np.abs(outs[0] - outs[1]))) + + @pytest.mark.parametrize( + ("g1_attn", "g2_attn"), + [(True, False), (False, True)], + ids=["g1_local_atten", "g2_atten2map"], + ) + def test_descriptor_continuous_at_repformer_cutoff(self, g1_attn, g2_attn): + # The defect was a CONSTANT step (~5e-9) as eps -> 0; a smooth + # descriptor changes only proportionally to eps. At eps = 1e-12 + # anything above 1e-11 is a genuine discontinuity. + delta_tiny = self._descriptor_delta(g1_attn=g1_attn, g2_attn=g2_attn, eps=1e-12) + assert delta_tiny < 1e-11, ( + f"descriptor step {delta_tiny:.3e} across the repformer cutoff " + "at binding sel (degree > sel): the attention softmax " + "denominator is not smooth" + ) + # scaling check: shrinking eps by 1e4 must shrink the (smooth) + # difference, not plateau at a finite step. + delta_small = self._descriptor_delta(g1_attn=g1_attn, g2_attn=g2_attn, eps=1e-8) + assert delta_tiny < max(delta_small, 1e-13), ( + f"difference plateaus ({delta_small:.3e} -> {delta_tiny:.3e}): " + "finite step at the cutoff" + ) diff --git a/source/tests/common/dpmodel/test_neighbor_graph.py b/source/tests/common/dpmodel/test_neighbor_graph.py index b8b57c8b31..d147ab3dd3 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph.py +++ b/source/tests/common/dpmodel/test_neighbor_graph.py @@ -225,3 +225,44 @@ def test_importable_from_utils(self) -> None: self.assertTrue(callable(from_dense_quartet)) self.assertIsNotNone(NeighborGraph) self.assertIsNotNone(GraphLayout) + + +def test_segment_max_trailing_dims(): + from deepmd.dpmodel.utils.neighbor_graph import ( + segment_max, + ) + + data = np.array([[1.0, 5.0], [3.0, 2.0], [2.0, 9.0]]) + ids = np.array([0, 0, 1], dtype=np.int64) + out = segment_max(data, ids, 2) + np.testing.assert_allclose(out, [[3.0, 5.0], [2.0, 9.0]]) + + +def test_segment_softmax_trailing_dims_matches_columnwise(): + from deepmd.dpmodel.utils.neighbor_graph import ( + segment_softmax, + ) + + rng = np.random.default_rng(0) + data = rng.normal(size=(6, 3)) + ids = np.array([0, 0, 1, 1, 1, 2], dtype=np.int64) + mask = np.array([True, True, True, False, True, True]) + out = segment_softmax(data, ids, 3, mask=mask) + for c in range(3): + col = segment_softmax(data[:, c], ids, 3, mask=mask) + np.testing.assert_allclose(out[:, c], col, rtol=1e-15) + + +def test_segment_softmax_trailing_dims_torch(): + import torch + + from deepmd.dpmodel.utils.neighbor_graph import ( + segment_softmax, + ) + + rng = np.random.default_rng(1) + data = rng.normal(size=(5, 2)) + ids = np.array([0, 1, 1, 0, 1], dtype=np.int64) + ref = segment_softmax(data, ids, 2) + got = segment_softmax(torch.from_numpy(data), torch.from_numpy(ids), 2) + np.testing.assert_allclose(got.numpy(), ref, rtol=1e-15) diff --git a/source/tests/common/dpmodel/test_repformer_graph_ops.py b/source/tests/common/dpmodel/test_repformer_graph_ops.py new file mode 100644 index 0000000000..1b7a4208ab --- /dev/null +++ b/source/tests/common/dpmodel/test_repformer_graph_ops.py @@ -0,0 +1,756 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Per-op parity: repformer graph twins vs the dense reference ops on the +identical (shape-static, center-major) edge layout. Same math, fp64 => +rtol/atol 1e-12. +""" + +import itertools + +import numpy as np +import pytest + +from deepmd.dpmodel.descriptor.repformers import ( + Atten2EquiVarApply, + Atten2Map, + Atten2MultiHeadApply, + LocalAtten, + RepformerLayer, + _cal_hg, + _cal_hg_graph, + _make_nei_g1, + symmetrization_op, + symmetrization_op_graph, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + center_edge_pairs, +) + +NF, NLOC, NNEI, NG = 2, 5, 7, 6 + + +def _mk(seed=0): + rng = np.random.default_rng(seed) + g = rng.normal(size=(NF, NLOC, NNEI, NG)) + h = rng.normal(size=(NF, NLOC, NNEI, 3)) + mask = rng.random((NF, NLOC, NNEI)) > 0.3 + sw = rng.random((NF, NLOC, NNEI)) * mask + n_total = NF * NLOC + dst = np.repeat(np.arange(n_total, dtype=np.int64), NNEI) + return g, h, mask, sw, n_total, dst + + +@pytest.mark.parametrize( + "smooth,use_sqrt_nnei", list(itertools.product([True, False], [True, False])) +) +def test_cal_hg_graph_parity(smooth, use_sqrt_nnei): + g, h, mask, sw, n_total, dst = _mk() + ref = _cal_hg(g, h, mask, sw, smooth=smooth, use_sqrt_nnei=use_sqrt_nnei) + got = _cal_hg_graph( + g.reshape(-1, NG), + h.reshape(-1, 3), + mask.reshape(-1), + sw.reshape(-1), + dst, + n_total, + NNEI, + smooth=smooth, + use_sqrt_nnei=use_sqrt_nnei, + ) + np.testing.assert_allclose(got, ref.reshape(n_total, 3, NG), rtol=1e-12, atol=1e-12) + + +@pytest.mark.parametrize("axis_neuron", [2, 4]) +def test_symmetrization_op_graph_parity(axis_neuron): + g, h, mask, sw, n_total, dst = _mk(1) + ref = symmetrization_op(g, h, mask, sw, axis_neuron) + got = symmetrization_op_graph( + g.reshape(-1, NG), + h.reshape(-1, 3), + mask.reshape(-1), + sw.reshape(-1), + dst, + n_total, + NNEI, + axis_neuron, + ) + np.testing.assert_allclose( + got, ref.reshape(n_total, axis_neuron * NG), rtol=1e-12, atol=1e-12 + ) + + +def test_cal_hg_graph_torch(): + import torch + + g, h, mask, sw, n_total, dst = _mk(2) + ref = _cal_hg_graph( + g.reshape(-1, NG), + h.reshape(-1, 3), + mask.reshape(-1), + sw.reshape(-1), + dst, + n_total, + NNEI, + ) + got = _cal_hg_graph( + torch.from_numpy(g.reshape(-1, NG)), + torch.from_numpy(h.reshape(-1, 3)), + torch.from_numpy(mask.reshape(-1)), + torch.from_numpy(sw.reshape(-1)), + torch.from_numpy(dst), + n_total, + NNEI, + ) + np.testing.assert_allclose(got.numpy(), ref, rtol=1e-12) + + +def _mk_layer(g1_out_conv: bool, seed: int = 0, **kwargs) -> RepformerLayer: + cfg = { + "rcut": 4.0, + "rcut_smth": 0.5, + "sel": NNEI, + "ntypes": 2, + "g1_dim": 8, + "g2_dim": NG, + "axis_neuron": 2, + "update_chnnl_2": True, + "g1_out_conv": g1_out_conv, + "g1_out_mlp": True, + "precision": "float64", + "seed": seed, + } + cfg.update(kwargs) + return RepformerLayer(**cfg) + + +def _mk_g1_nlist(seed=3): + rng = np.random.default_rng(seed) + n_total = NF * NLOC + g1 = rng.normal(size=(n_total, 8)) + # ghost-free: neighbors index local atoms of the SAME frame + nlist = rng.integers(0, NLOC, size=(NF, NLOC, NNEI)).astype(np.int64) + mask = rng.random((NF, NLOC, NNEI)) > 0.3 + nlist = np.where(mask, nlist, -1) + sw = rng.random((NF, NLOC, NNEI)) * mask + # flat-edge view of the same topology + src = (nlist + np.arange(NF)[:, None, None] * NLOC).reshape(-1) + src = np.where(mask.reshape(-1), src, 0).astype(np.int64) + dst = np.repeat(np.arange(n_total, dtype=np.int64), NNEI) + return g1, nlist, mask, sw, src, dst, n_total + + +@pytest.mark.parametrize("g1_out_conv", [True, False]) +def test_update_g1_conv_graph_parity(g1_out_conv): + layer = _mk_layer(g1_out_conv) + g1, nlist, mask, sw, src, dst, n_total = _mk_g1_nlist() + rng = np.random.default_rng(4) + g2 = rng.normal(size=(NF, NLOC, NNEI, NG)) + g1_ext = g1.reshape(NF, NLOC, 8) + gg1 = _make_nei_g1(g1_ext, np.where(mask, nlist, 0)) + ref = layer._update_g1_conv(gg1, g2, mask, sw) + got = layer._update_g1_conv_graph( + np.take(g1, src, axis=0), + g2.reshape(-1, NG), + mask.reshape(-1), + sw.reshape(-1), + dst, + n_total, + NNEI, + ) + np.testing.assert_allclose(got, ref.reshape(n_total, -1), rtol=1e-12, atol=1e-12) + + +def test_update_g2_g1g1_graph_parity(): + layer = _mk_layer(True) + g1, nlist, mask, sw, src, dst, n_total = _mk_g1_nlist(5) + g1_ext = g1.reshape(NF, NLOC, 8) + gg1 = _make_nei_g1(g1_ext, np.where(mask, nlist, 0)) + ref = layer._update_g2_g1g1(g1_ext, gg1, mask, sw) + got = layer._update_g2_g1g1_graph(g1, src, dst, mask.reshape(-1), sw.reshape(-1)) + np.testing.assert_allclose( + got, ref.reshape(n_total * NNEI, -1), rtol=1e-12, atol=1e-12 + ) + + +def test_update_g1_conv_graph_torch(): + import torch + + layer = _mk_layer(True, seed=6) + g1, nlist, mask, sw, src, dst, n_total = _mk_g1_nlist(7) + rng = np.random.default_rng(8) + g2 = rng.normal(size=(NF, NLOC, NNEI, NG)) + ref = layer._update_g1_conv_graph( + np.take(g1, src, axis=0), + g2.reshape(-1, NG), + mask.reshape(-1), + sw.reshape(-1), + dst, + n_total, + NNEI, + ) + got = layer._update_g1_conv_graph( + torch.from_numpy(np.take(g1, src, axis=0)), + torch.from_numpy(g2.reshape(-1, NG)), + torch.from_numpy(mask.reshape(-1)), + torch.from_numpy(sw.reshape(-1)), + torch.from_numpy(dst), + n_total, + NNEI, + ) + np.testing.assert_allclose(got.numpy(), ref, rtol=1e-12, atol=1e-12) + + +def _pairs(mask, dst, n_total): + q_e, k_e, pm = center_edge_pairs( + dst, + mask.reshape(-1), + n_total, + include_self=True, + ordered=True, + static_nnei=NNEI, + ) + return q_e, k_e, pm + + +@pytest.mark.parametrize( + "has_gate,smooth", [(True, True), (False, True), (True, False)] +) +def test_atten2map_parity(has_gate, smooth): + rng = np.random.default_rng(6) + a2m = Atten2Map( + NG, 4, 2, has_gate=has_gate, smooth=smooth, precision="float64", seed=7 + ) + g2 = rng.normal(size=(NF, NLOC, NNEI, NG)) + h2 = rng.normal(size=(NF, NLOC, NNEI, 3)) + mask = rng.random((NF, NLOC, NNEI)) > 0.3 + sw = rng.random((NF, NLOC, NNEI)) * mask + n_total = NF * NLOC + dst = np.repeat(np.arange(n_total, dtype=np.int64), NNEI) + ref = a2m.call(g2, h2, mask, sw) # (nf, nloc, nnei, nnei, nh) + q_e, k_e, pm = _pairs(mask, dst, n_total) + got = a2m.call_graph( + g2.reshape(-1, NG), + h2.reshape(-1, 3), + sw.reshape(-1), + q_e, + k_e, + pm, + NF * NLOC * NNEI, + NNEI, + ) # (P, nh) + slot = np.arange(NF * NLOC * NNEI) % NNEI + ctr = dst[np.asarray(q_e)] + f, l = ctr // NLOC, ctr % NLOC + ref_pairs = ref[f, l, slot[np.asarray(q_e)], slot[np.asarray(k_e)], :] + # NOTE: even for pairs whose query edge is padding, both sides collapse to + # an exact 0.0 (dense: post-softmax where(mask, 0, ...); graph: pair_mask + # multiply, and in the smooth branch the fully-masked segment_softmax + # group also yields exact 0.0), so no restriction to real-query pairs is + # needed here -- verified empirically before relying on it. + np.testing.assert_allclose(np.asarray(got), ref_pairs, rtol=1e-12, atol=1e-12) + + +@pytest.mark.parametrize( + "has_gate,smooth", [(True, True), (False, True), (True, False)] +) +def test_atten2_mh_apply_parity(has_gate, smooth): + rng = np.random.default_rng(9) + nh = 3 + a2m = Atten2Map( + NG, 4, nh, has_gate=has_gate, smooth=smooth, precision="float64", seed=10 + ) + mha = Atten2MultiHeadApply(NG, nh, precision="float64", seed=11) + g2 = rng.normal(size=(NF, NLOC, NNEI, NG)) + h2 = rng.normal(size=(NF, NLOC, NNEI, 3)) + mask = rng.random((NF, NLOC, NNEI)) > 0.3 + sw = rng.random((NF, NLOC, NNEI)) * mask + n_total = NF * NLOC + dst = np.repeat(np.arange(n_total, dtype=np.int64), NNEI) + e_tot = NF * NLOC * NNEI + ref_AA = a2m.call(g2, h2, mask, sw) # (nf, nloc, nnei, nnei, nh) + ref = mha.call(ref_AA, g2) # (nf, nloc, nnei, ng2) + q_e, k_e, pm = _pairs(mask, dst, n_total) + got_AA = a2m.call_graph( + g2.reshape(-1, NG), h2.reshape(-1, 3), sw.reshape(-1), q_e, k_e, pm, e_tot, NNEI + ) + got = mha.call_graph(got_AA, g2.reshape(-1, NG), q_e, k_e, e_tot) + np.testing.assert_allclose( + np.asarray(got), ref.reshape(e_tot, NG), rtol=1e-12, atol=1e-12 + ) + + +@pytest.mark.parametrize( + "has_gate,smooth", [(True, True), (False, True), (True, False)] +) +def test_atten2_ev_apply_parity(has_gate, smooth): + rng = np.random.default_rng(12) + nh = 3 + a2m = Atten2Map( + NG, 4, nh, has_gate=has_gate, smooth=smooth, precision="float64", seed=13 + ) + ev = Atten2EquiVarApply(NG, nh, precision="float64", seed=14) + g2 = rng.normal(size=(NF, NLOC, NNEI, NG)) + h2 = rng.normal(size=(NF, NLOC, NNEI, 3)) + mask = rng.random((NF, NLOC, NNEI)) > 0.3 + sw = rng.random((NF, NLOC, NNEI)) * mask + n_total = NF * NLOC + dst = np.repeat(np.arange(n_total, dtype=np.int64), NNEI) + e_tot = NF * NLOC * NNEI + ref_AA = a2m.call(g2, h2, mask, sw) # (nf, nloc, nnei, nnei, nh) + ref = ev.call(ref_AA, h2) # (nf, nloc, nnei, 3) + q_e, k_e, pm = _pairs(mask, dst, n_total) + got_AA = a2m.call_graph( + g2.reshape(-1, NG), h2.reshape(-1, 3), sw.reshape(-1), q_e, k_e, pm, e_tot, NNEI + ) + got = ev.call_graph(got_AA, h2.reshape(-1, 3), q_e, k_e, e_tot) + np.testing.assert_allclose( + np.asarray(got), ref.reshape(e_tot, 3), rtol=1e-12, atol=1e-12 + ) + + +@pytest.mark.parametrize("smooth", [True, False]) +def test_local_atten_parity(smooth): + la = LocalAtten(8, 4, 2, smooth=smooth, precision="float64", seed=15) + g1, nlist, mask, sw, src, dst, n_total = _mk_g1_nlist(20) + g1_ext = g1.reshape(NF, NLOC, 8) + gg1 = _make_nei_g1(g1_ext, np.where(mask, nlist, 0)) + ref = la.call(g1_ext, gg1, mask, sw) # (nf, nloc, ng1) + got = la.call_graph( + g1, + np.take(g1, src, axis=0), + mask.reshape(-1), + sw.reshape(-1), + dst, + n_total, + NNEI, + ) + np.testing.assert_allclose( + np.asarray(got), ref.reshape(n_total, 8), rtol=1e-12, atol=1e-12 + ) + + +def test_local_atten_below_phantom_dense_parity(): + """OutisLi review: finite valid projection weights can push every smooth + logit below ``-attnw_shift``. With ``n_real == sel`` the signed phantom + count is 0 -- the denominator is a plain positive softmax sum -- and the + graph route must match dense EXACTLY (the always-on floor used to return + ~0.0018 where dense gives 1.0). + """ + la = LocalAtten(1, 1, 1, smooth=True, precision="float64", seed=1) + la.mapq.w = np.array([[1.0]]) + la.mapkv.w = np.array([[-30.0, 1.0]]) # key -30, value 1 + la.head_map.w = np.array([[1.0]]) # identity head + la.head_map.b = np.array([0.0]) + nf, nloc, nnei = 1, 1, 2 + g1 = np.ones((nf * nloc, 1)) + gg1 = np.ones((nf, nloc, nnei, 1)) + mask = np.ones((nf, nloc, nnei), dtype=bool) + sw = np.ones((nf, nloc, nnei)) + ref = la.call(g1.reshape(nf, nloc, 1), gg1, mask, sw) # == [[[1.0]]] + dst = np.zeros(nnei, dtype=np.int64) + got = la.call_graph( + g1, + gg1.reshape(-1, 1), + mask.reshape(-1), + sw.reshape(-1), + dst, + nf * nloc, + nnei, # sel == n_real: phantom count 0 + ) + np.testing.assert_allclose(np.asarray(got), ref.reshape(1, 1), rtol=1e-12, atol=0.0) + # anti-vacuity: the logits really are below -attnw_shift and the dense + # result is the nontrivial value from the review + np.testing.assert_allclose(np.asarray(got), [[1.0]], rtol=1e-12) + + +def test_local_atten_cutoff_crossing_below_phantom_continuous(): + """OutisLi round 5: an edge crossing the cutoff must not JUMP the output + when the surviving logits sit below ``-attnw_shift``. + + Same LocalAtten as :func:`test_local_atten_below_phantom_dense_parity` + (every smooth logit ~= -30 < -attnw_shift) but with ``sel = 1``: outside + the cutoff (one edge, phantom count 0) the output is exactly 1.0; a + second edge entering with ``sw = s -> 0+`` flips the count to -1. The + count-gated denominator floor produced 0.00181599 just inside (limiting + jump -0.998); the slot-occupancy denominator must converge back to the + outside value. + """ + la = LocalAtten(1, 1, 1, smooth=True, precision="float64", seed=1) + la.mapq.w = np.array([[1.0]]) + la.mapkv.w = np.array([[-30.0, 1.0]]) + la.head_map.w = np.array([[1.0]]) + la.head_map.b = np.array([0.0]) + sel = 1 + n_total = 1 + + def run(sw_edges: np.ndarray) -> float: + nnei = sw_edges.shape[0] + g1 = np.ones((n_total, 1)) + gg1 = np.ones((nnei, 1)) + mask = np.ones(nnei, dtype=bool) + dst = np.zeros(nnei, dtype=np.int64) + out = la.call_graph(g1, gg1, mask, sw_edges, dst, n_total, sel) + return float(np.asarray(out)[0, 0]) + + outside = run(np.array([1.0])) + np.testing.assert_allclose(outside, 1.0, rtol=1e-12) + prev_gap = None + for s in [1e-2, 1e-4, 1e-6, 1e-8, 1e-10, 1e-12]: + inside = run(np.array([1.0, s])) + assert np.isfinite(inside) + gap = abs(inside - outside) + if prev_gap is not None: + assert gap < prev_gap # converging, not plateauing at a jump + prev_gap = gap + assert prev_gap < 1e-7 + + +def test_atten2map_below_phantom_dense_parity(): + """Same below-``-attnw_shift`` regime for the pair attention map: with + all key slots real (phantom count 0) graph must equal dense exactly. + """ + a2m = Atten2Map(1, 1, 1, has_gate=False, smooth=True, precision="float64", seed=2) + a2m.mapqk.w = np.array([[1.0, -30.0]]) # query 1, key -30 -> logits -30 + nf, nloc, nnei = 1, 1, 2 + g2 = np.ones((nf, nloc, nnei, 1)) + h2 = np.zeros((nf, nloc, nnei, 3)) + h2[..., 0] = 1.0 + mask = np.ones((nf, nloc, nnei), dtype=bool) + sw = np.ones((nf, nloc, nnei)) + ref = a2m.call(g2, h2, mask, sw) # (nf, nloc, nnei, nnei, nh) + n_total = nf * nloc + dst = np.repeat(np.arange(n_total, dtype=np.int64), nnei) + q_e, k_e, pm = center_edge_pairs( + dst, + mask.reshape(-1), + n_total, + include_self=True, + ordered=True, + static_nnei=nnei, + ) + got = a2m.call_graph( + g2.reshape(-1, 1), + h2.reshape(-1, 3), + sw.reshape(-1), + q_e, + k_e, + pm, + nf * nloc * nnei, + nnei, + ) + ref_pairs = ref[0, 0, np.asarray(q_e) % nnei, np.asarray(k_e) % nnei, :] + np.testing.assert_allclose(np.asarray(got), ref_pairs, rtol=1e-12, atol=0.0) + # anti-vacuity: weight 0.5 per key slot times h2h2t = 1/sqrt(3) + np.testing.assert_allclose(np.asarray(got), 0.5 / np.sqrt(3.0), rtol=1e-12) + + +def test_atten2map_graph_torch(): + import torch + + rng = np.random.default_rng(16) + a2m = Atten2Map(NG, 4, 2, has_gate=True, smooth=True, precision="float64", seed=17) + g2 = rng.normal(size=(NF, NLOC, NNEI, NG)) + h2 = rng.normal(size=(NF, NLOC, NNEI, 3)) + mask = rng.random((NF, NLOC, NNEI)) > 0.3 + sw = rng.random((NF, NLOC, NNEI)) * mask + n_total = NF * NLOC + dst = np.repeat(np.arange(n_total, dtype=np.int64), NNEI) + q_e, k_e, pm = _pairs(mask, dst, n_total) + e_tot = NF * NLOC * NNEI + ref = a2m.call_graph( + g2.reshape(-1, NG), h2.reshape(-1, 3), sw.reshape(-1), q_e, k_e, pm, e_tot, NNEI + ) + got = a2m.call_graph( + torch.from_numpy(g2.reshape(-1, NG)), + torch.from_numpy(h2.reshape(-1, 3)), + torch.from_numpy(sw.reshape(-1)), + torch.from_numpy(q_e), + torch.from_numpy(k_e), + torch.from_numpy(pm), + e_tot, + NNEI, + ) + np.testing.assert_allclose(got.numpy(), ref, rtol=1e-12, atol=1e-12) + + +# -------------------------------------------------------------------------- +# Whole-layer parity: RepformerLayer.call_graph vs RepformerLayer.call +# -------------------------------------------------------------------------- + +# Toggle matrix; each entry is the full set of RepformerLayer kwargs applied +# on top of _mk_layer's base config (g1_out_conv defaults to True unless +# overridden). "conv_only" and "no_g2_attn" additionally exercise +# pairs=None, since neither sets update_g2_has_attn/update_h2 True. +_LAYER_CASES = { + "defaults": {}, + "no_g2_attn": {"update_g2_has_attn": False}, + "update_h2": {"update_h2": True}, + "conv_only": { + "update_g1_has_conv": True, + "g1_out_conv": False, # feed conv output into g1_mlp (else g1_mlp is empty) + "update_g1_has_drrd": False, + "update_g1_has_grrg": False, + "update_g1_has_attn": False, + "update_g2_has_g1g1": False, + "update_g2_has_attn": False, + "update_h2": False, + }, + "no_g1_out": {"g1_out_conv": False, "g1_out_mlp": False}, + "no_chnnl_2_no_gg1": { + # Exercises update_chnnl_2=False (skips g2/h2 updates) and cal_gg1=False (gg1=None). + # With g1_out_mlp=False, g1_mlp starts with [g1] identity seed for xp.concat. + "update_chnnl_2": False, + "update_g1_has_conv": False, + "update_g1_has_grrg": False, + "update_g1_has_drrd": False, + "update_g1_has_attn": False, + "update_g2_has_g1g1": False, + "update_g2_has_attn": False, + "update_h2": False, + "g1_out_conv": False, + "g1_out_mlp": False, + }, +} +# cases where pairs=None is exercised (neither update_g2_has_attn nor +# update_h2 is set, so center_edge_pairs is not required) +_PAIRS_NONE_CASES = {"conv_only", "no_g2_attn"} + + +@pytest.mark.parametrize("case_name", list(_LAYER_CASES)) +def test_repformer_layer_call_graph_parity(case_name): + kwargs = dict(_LAYER_CASES[case_name]) + g1_out_conv = kwargs.pop("g1_out_conv", True) + layer = _mk_layer(g1_out_conv, seed=21, **kwargs) + + g1, nlist, mask, sw, src, dst, n_total = _mk_g1_nlist(23) + rng = np.random.default_rng(24) + g2 = rng.normal(size=(NF, NLOC, NNEI, NG)) + h2 = rng.normal(size=(NF, NLOC, NNEI, 3)) + g1_ext = g1.reshape(NF, NLOC, 8) + nlist0 = np.where(nlist == -1, 0, nlist) + + ref_g1, ref_g2, ref_h2 = layer.call(g1_ext, g2, h2, nlist0, mask, sw) + + g2_flat = g2.reshape(-1, NG) + h2_flat = h2.reshape(-1, 3) + mask_flat = mask.reshape(-1) + sw_flat = sw.reshape(-1) + + if case_name in _PAIRS_NONE_CASES: + assert not (layer.update_g2_has_attn or layer.update_h2) + pairs = None + else: + pairs = _pairs(mask, dst, n_total) + + got_g1, got_g2, got_h2 = layer.call_graph( + g1, g2_flat, h2_flat, src, dst, mask_flat, sw_flat, n_total, NNEI, pairs + ) + + np.testing.assert_allclose( + got_g1, ref_g1.reshape(n_total, -1), rtol=1e-12, atol=1e-12 + ) + np.testing.assert_allclose(got_g2, ref_g2.reshape(-1, NG), rtol=1e-12, atol=1e-12) + np.testing.assert_allclose(got_h2, ref_h2.reshape(-1, 3), rtol=1e-12, atol=1e-12) + + +def test_repformer_layer_call_graph_torch(): + import torch + + layer = _mk_layer(True, seed=25) + g1, nlist, mask, sw, src, dst, n_total = _mk_g1_nlist(26) + rng = np.random.default_rng(27) + g2 = rng.normal(size=(NF, NLOC, NNEI, NG)) + h2 = rng.normal(size=(NF, NLOC, NNEI, 3)) + g2_flat = g2.reshape(-1, NG) + h2_flat = h2.reshape(-1, 3) + mask_flat = mask.reshape(-1) + sw_flat = sw.reshape(-1) + pairs = _pairs(mask, dst, n_total) + + ref = layer.call_graph( + g1, g2_flat, h2_flat, src, dst, mask_flat, sw_flat, n_total, NNEI, pairs + ) + + t_pairs = tuple(torch.from_numpy(np.asarray(x)) for x in pairs) + got = layer.call_graph( + torch.from_numpy(g1), + torch.from_numpy(g2_flat), + torch.from_numpy(h2_flat), + torch.from_numpy(src), + torch.from_numpy(dst), + torch.from_numpy(mask_flat), + torch.from_numpy(sw_flat), + n_total, + NNEI, + t_pairs, + ) + for r, g in zip(ref, got, strict=True): + np.testing.assert_allclose(g.numpy(), r, rtol=1e-12, atol=1e-12) + + +def test_local_atten_gradient_continuous_below_phantom(): + """OutisLi round 6: C1 across the below-phantom surface through the real + LocalAtten path. sel=1, sw=[1, 1], logits [-20 +/- eps, -18], values + [1, 2] (his construction): the relu bracket kept the OUTPUT continuous + (~2.1353353) but jumped d(output)/d(logit) from -0.009157818652523 to + -0.153650933331561. The C1 tail matches the in-design (right) slope + from both sides. + """ + la = LocalAtten(1, 1, 1, smooth=True, precision="float64", seed=1) + la.mapq.w = np.array([[1.0]]) + # key = x (logit = q*k = x); value = 0.5*x + 11 -> v(-20)=1, v(-18)=2 + la.mapkv.w = np.array([[1.0, 0.5]]) + la.mapkv.b = np.array([0.0, 11.0]) + la.head_map.w = np.array([[1.0]]) + la.head_map.b = np.array([0.0]) + n_total, sel = 1, 1 + + def out(x1: float) -> float: + gg1 = np.array([[x1], [-18.0]]) + o = la.call_graph( + np.ones((n_total, 1)), + gg1, + np.ones(2, dtype=bool), + np.ones(2), + np.zeros(2, dtype=np.int64), + n_total, + sel, + ) + return float(np.asarray(o)[0, 0]) + + # value continuity at the surface (his repro printed ~2.135335) + np.testing.assert_allclose(out(-20.0), 2.1353, rtol=1e-3) + d = 1e-6 + left = (out(-20.0 - d) - out(-20.0 - 3.0 * d)) / (2.0 * d) + right = (out(-20.0 + 3.0 * d) - out(-20.0 + d)) / (2.0 * d) + # the in-design branch is unchanged, so the right slope equals his + # measured -0.153650933331561 plus the value-chain term 0.5 * w1 + # (= exp(-2)/2: our value wiring v = 0.5*x + 11 varies with the swept + # feature); the relu bracket jumped the attention part of the left + # slope to -0.009157818652523 + np.testing.assert_allclose( + right, -0.153650933331561 + 0.5 * np.exp(-2.0), rtol=1e-3 + ) + np.testing.assert_allclose(left, right, rtol=5e-3) + + +def test_atten2map_gradient_continuous_below_phantom(): + """Same C1-at-the-surface guarantee through the pair attention map: + sweep one edge feature so a single pair logit crosses -attnw_shift while + the segment stays binding (sel=1 < 2 keys); the FD slope of the + attention map must match across the crossing. + """ + a2m = Atten2Map(1, 1, 1, has_gate=False, smooth=True, precision="float64", seed=2) + # q = g2, k = -g2 -> logit(q,k) = -g_q*g_k: negative-definite pairing + # keeps every pair logit NEAR the -20 phantom level (a positive pairing + # pins self-logits >= 0, drowning the crossing pair 20+ units below the + # segment max where the softmax renders it invisible to FD) + a2m.mapqk.w = np.array([[1.0, -1.0]]) + n_total, nnei = 1, 2 + dst = np.repeat(np.arange(n_total, dtype=np.int64), nnei) + mask = np.ones(nnei, dtype=bool) + q_e, k_e, pm = center_edge_pairs( + dst, mask, n_total, include_self=True, ordered=True, static_nnei=nnei + ) + h2 = np.zeros((nnei, 3)) + h2[:, 0] = 1.0 + + def out(x: float) -> float: + # pair logits: {-x*x, -5x, -5x, -25}; -5x crosses -20 at x = 4 with + # the self pair at -16 in-design and the (1,1) pair statically below + g2 = np.array([[x], [5.0]]) + att = a2m.call_graph(g2, h2, np.ones(nnei), q_e, k_e, pm, n_total * nnei, 1) + return float(np.sum(np.asarray(att))) + + d = 1e-6 + left = (out(4.0 - 3.0 * d) - out(4.0 - d)) / (-2.0 * d) + right = (out(4.0 + d) - out(4.0 + 3.0 * d)) / (-2.0 * d) + assert abs(left) > 1e-6 and abs(right) > 1e-6 + np.testing.assert_allclose(left, right, rtol=5e-3) + + +def test_local_atten_float32_high_logit_backward_finite(): + """OutisLi round 7 model-level repro: float32 LocalAtten.call_graph with + sel == n_real == 2, sw = [1, 1] and raw logits [68, 69] (in-design, + count 0). Dense and graph forwards agree (~68.7311) but the graph + backward produced NaN g1/gg1 gradients through the inactive tail's + ``1 / ph_e`` overflow; gradients must be finite and match dense. + """ + import torch + + la = LocalAtten(1, 1, 1, smooth=True, precision="float32", seed=1) + la.mapq.w = np.array([[1.0]], dtype=np.float32) + la.mapkv.w = np.array([[1.0, 1.0]], dtype=np.float32) + la.mapkv.b = np.array([0.0, 0.0], dtype=np.float32) + la.head_map.w = np.array([[1.0]], dtype=np.float32) + la.head_map.b = np.array([0.0], dtype=np.float32) + n_total, nnei, sel = 1, 2, 2 + + gg1_np = np.array([[68.0], [69.0]], dtype=np.float32) + mask_np = np.ones(nnei, dtype=bool) + sw_np = np.ones(nnei, dtype=np.float32) + + # graph route, torch float32, gradients w.r.t. g1 and gg1 + g1_t = torch.ones(n_total, 1, dtype=torch.float32, requires_grad=True) + gg1_t = torch.from_numpy(gg1_np).clone().requires_grad_(True) + out_g = la.call_graph( + g1_t, + gg1_t, + torch.from_numpy(mask_np), + torch.from_numpy(sw_np), + torch.zeros(nnei, dtype=torch.int64), + n_total, + sel, + ) + out_g.sum().backward() + assert torch.all(torch.isfinite(g1_t.grad)) + assert torch.all(torch.isfinite(gg1_t.grad)) + + # dense reference gradients on the same weights + g1_d = torch.ones(1, 1, 1, dtype=torch.float32, requires_grad=True) + gg1_d = torch.from_numpy(gg1_np.reshape(1, 1, nnei, 1)).clone().requires_grad_(True) + out_d = la.call( + g1_d, + gg1_d, + torch.from_numpy(mask_np.reshape(1, 1, nnei)), + torch.from_numpy(sw_np.reshape(1, 1, nnei)), + ) + out_d.sum().backward() + np.testing.assert_allclose(float(out_g.sum()), float(out_d.sum()), rtol=1e-6) + np.testing.assert_allclose( + g1_t.grad.numpy().reshape(-1), g1_d.grad.numpy().reshape(-1), rtol=1e-4 + ) + np.testing.assert_allclose( + gg1_t.grad.numpy().reshape(-1), gg1_d.grad.numpy().reshape(-1), rtol=1e-4 + ) + + +def test_local_atten_float32_log_spaced_sw_matches_float64(): + """OutisLi round 8 forward repro: float32 LocalAtten.call_graph with + smooth cutoff weights spanning decades (sw = [1, 0.1, 0.01, 5e-7]) and + raw logits -30. The float32 active-set selection used to pick the wrong + water-filling cut (~23% forward error, 0.0612 vs 0.0792); float32 must + now match the float64 evaluation of the same quantized inputs. + """ + n_total, nnei, sel = 1, 4, 3 + sw32 = np.array([1.0, 0.1, 0.01, 5e-7], dtype=np.float32) + + def run(prec: str) -> float: + la = LocalAtten(1, 1, 1, smooth=True, precision=prec, seed=1) + dt = np.float32 if prec == "float32" else np.float64 + la.mapq.w = np.array([[1.0]], dtype=dt) + la.mapkv.w = np.array([[1.0, 1.0]], dtype=dt) + la.mapkv.b = np.array([0.0, 0.0], dtype=dt) + la.head_map.w = np.array([[1.0]], dtype=dt) + la.head_map.b = np.array([0.0], dtype=dt) + o = la.call_graph( + np.ones((n_total, 1), dtype=dt), + np.full((nnei, 1), -30.0, dtype=dt), + np.ones(nnei, dtype=bool), + sw32.astype(dt), + np.zeros(nnei, dtype=np.int64), + n_total, + sel, + ) + return float(np.asarray(o)[0, 0]) + + np.testing.assert_allclose(run("float32"), run("float64"), rtol=1e-3) diff --git a/source/tests/common/dpmodel/test_segment_softmax.py b/source/tests/common/dpmodel/test_segment_softmax.py index a97bd9c1aa..b044ea1b70 100644 --- a/source/tests/common/dpmodel/test_segment_softmax.py +++ b/source/tests/common/dpmodel/test_segment_softmax.py @@ -2,6 +2,7 @@ """segment_max / segment_softmax (NeighborGraph PR-D segment toolkit).""" import numpy as np +import pytest from deepmd.dpmodel.utils.neighbor_graph import ( segment_max, @@ -110,3 +111,597 @@ def test_masked_entry_larger_than_unmasked_max_no_nan() -> None: ref = np.exp([1.0, 2.0]) / np.exp([1.0, 2.0]).sum() np.testing.assert_allclose(out[:2], ref, rtol=1e-12) assert out[2] == 0.0 + + +class TestSignedPhantomStrictPositivity: + """The phantom denominator must stay strictly positive AND + cutoff-continuous for arbitrary finite logits (OutisLi / njzjz-bot + review, rounds 2-5). + + Real logits are not bounded below by ``phantom_logit``: the smooth + envelope maps pre-shift logits ``raw < -attnw_shift`` to + ``l < phantom_logit`` whenever ``sw > 0``. With a negative count the + raw signed denominator ``sum exp(l) + (sel - n) exp(ph)`` then crosses + zero (e.g. ``sel=1``, logits ``[-21, -21]``, ``ph=-20``) -- a pole. A + count-gated denominator floor fixed the pole but switched on + DISCONTINUOUSLY when a boundary edge flipped the count from 0 to -1 + (round 5). The soft slot-occupancy denominator (capped water-filling + of the ``slot_weight`` envelope, see ``segment_softmax`` Notes) is + strictly positive by construction, bit-exact dense for + ``n_real <= sel``, equal to the signed formula in-design, and + continuous at every crossing because an entering entry has BOTH + ``exp(l) -> exp(phantom_logit)`` and occupancy ``theta -> 0``. + """ + + def test_reviewer_repro_finite_positive(self) -> None: + # sel=1, two real entries below the phantom logit -> raw signed + # denominator < 0 (the old pole / where-guard regime). + data = np.array([-21.0, -21.0]) + seg = np.array([0, 0]) + w = segment_softmax( + data, + seg, + 1, + phantom_count=np.array([-1.0]), # sel - n_real = 1 - 2 + phantom_logit=-20.0, + slot_weight=np.ones(2), + ) + assert np.all(np.isfinite(w)) + assert np.all(w >= 0.0) + # occupancy bound: theta = 1/2 each, bracket >= theta*(1-1/e)*ex + # -> w <= (e/(e-1))/theta ~= 3.2 (the old where-guard's [1, 1] came + # from a ZERO denominator instead) + assert np.all(w <= 3.2) + + def test_njzjz_repro_three_entries(self) -> None: + data = np.array([-100.0, -100.0, -100.0]) + seg = np.array([0, 0, 0]) + w = segment_softmax( + data, + seg, + 1, + phantom_count=np.array([-2.0]), # sel=1, n_real=3 + phantom_logit=-20.0, + slot_weight=np.ones(3), + ) + assert np.all(np.isfinite(w)) + assert np.all(w >= 0.0) + # theta = 1/3 each: every deep-suppressed entry sees the denominator + # it would see alone (dense n=1=sel gives w=1); bounded by n/sel. + np.testing.assert_allclose(w, [1.0, 1.0, 1.0], rtol=1e-12) + + def test_missing_slot_weight_raises(self) -> None: + with pytest.raises(ValueError, match="slot_weight"): + segment_softmax( + np.array([-21.0, -21.0]), + np.array([0, 0]), + 1, + phantom_count=np.array([-1.0]), + phantom_logit=-20.0, + ) + + def test_cutoff_crossing_continuous_below_phantom(self) -> None: + """OutisLi round 5: a boundary edge flipping the count 0 -> -1 must + not jump the OTHER (deep-suppressed) entry's weight. + + sel=1; a fixed edge at logit -30 (below phantom_logit -20) and a + crossing edge with raw pre-shift logit -30 entering through the + smooth envelope: ``l_c = (raw + shift) * s - shift = -20 - 10 s``, + slot weight ``s``. Outside (s=0, edge absent, count 0) the fixed + edge's weight is exactly 1. The count-gated floor gave 0.00181599 + just inside (limiting jump -0.998); the occupancy denominator must + converge to the outside value. + """ + outside = segment_softmax( + np.array([-30.0]), + np.array([0]), + 1, + phantom_count=np.array([0.0]), + phantom_logit=-20.0, + slot_weight=np.ones(1), + ) + np.testing.assert_array_equal(outside, [1.0]) + prev_gap = None + # gap decays linearly in s (slope ~ exp(10), the fixed edge's own + # denominator share), so drive s deep enough to see convergence + for s in [1e-2, 1e-4, 1e-6, 1e-8, 1e-10, 1e-12]: + w = segment_softmax( + np.array([-30.0, -20.0 - 10.0 * s]), + np.array([0, 0]), + 1, + phantom_count=np.array([-1.0]), + phantom_logit=-20.0, + slot_weight=np.array([1.0, s]), + ) + assert np.all(np.isfinite(w)) + gap = abs(float(w[0]) - 1.0) + if prev_gap is not None: + assert gap < prev_gap # converging, not plateauing at a jump + prev_gap = gap + assert prev_gap < 1e-7 # the fixed edge recovers its outside weight + + def test_smooth_across_former_zero_crossing(self) -> None: + """Sweep a logit through the raw denominator's zero crossing: the + weights must stay finite and BOUNDED (<= 1/theta = n_real/sel), + and their increments must SCALE with the sweep resolution -- the + smoothness signature a pole cannot fake (at a pole, refining the + grid does not shrink the largest step). + """ + + def _sweep(lo: float, hi: float, num: int) -> np.ndarray: + outs = [] + for l2 in np.linspace(lo, hi, num): + w = segment_softmax( + np.array([-21.0, float(l2)]), + np.array([0, 0]), + 1, + phantom_count=np.array([-1.0]), + phantom_logit=-20.0, + slot_weight=np.ones(2), + ) + assert np.all(np.isfinite(w)) + assert np.all(w >= 0.0) + # occupancy bound: (e/(e-1))/theta ~= 3.2 here (theta=1/2) + assert np.all(w <= 3.2) + outs.append(w) + return np.stack(outs) + + # D_raw(l) = exp(-21) + exp(l) - exp(-20) crosses zero at l = ln( + # exp(-20) - exp(-21)) ~= -20.4587 (the LocalAtten sw-sweep pole + # from the review, expressed directly in logit space). + crossing = np.log(np.exp(-20.0) - np.exp(-21.0)) + coarse = _sweep(crossing - 0.5, crossing + 0.5, 201) + step_coarse = np.abs(np.diff(coarse, axis=0)).max() + # refine 10x: a smooth function's largest step shrinks ~10x; a pole + # (or the old where-guard plateau jump) would keep an O(1) step. + fine = _sweep(crossing - 0.5, crossing + 0.5, 2001) + step_fine = np.abs(np.diff(fine, axis=0)).max() + assert step_fine < 0.2 * step_coarse, ( + f"steps do not shrink under refinement (coarse {step_coarse:.4f} " + f"-> fine {step_fine:.4f}): discontinuity or pole at the crossing" + ) + + def test_zero_phantom_count_below_phantom_exact_dense(self) -> None: + """OutisLi review: ``phantom_count == 0`` means ``n_real == sel`` -- + no phantom term exists and the denominator is a plain positive + softmax sum, so the output must be the EXACT dense softmax even + when every logit sits far below ``phantom_logit`` (the always-on + floor used to return ~0.0009 instead of 0.5 here). + """ + data = np.array([-30.0, -30.0]) + seg = np.array([0, 0], dtype=np.int64) + w = segment_softmax( + data, + seg, + 1, + phantom_count=np.array([0.0]), + phantom_logit=-20.0, + slot_weight=np.ones(2), + ) + np.testing.assert_array_equal(w, [0.5, 0.5]) + # and BIT-equal to the phantom-free primitive on generic data + rng = np.random.default_rng(11) + data = rng.normal(size=9) * 30.0 # spans far below AND above -20 + seg = np.array([0, 0, 0, 1, 1, 1, 1, 2, 2], dtype=np.int64) + w_zero = segment_softmax( + data, + seg, + 3, + phantom_count=np.zeros(3), + phantom_logit=-20.0, + slot_weight=np.ones(9), + ) + w_none = segment_softmax(data, seg, 3) + np.testing.assert_array_equal(w_zero, w_none) + + def test_positive_count_below_phantom_exact_dense(self) -> None: + """Count > 0 with all real logits below ``phantom_logit``: the + denominator is positive (phantom terms only ADD), so the fixed-width + dense softmax must be reproduced exactly -- the floor must not fire. + """ + data = np.array([-30.0, -30.0]) + seg = np.array([0, 0], dtype=np.int64) + w = segment_softmax( + data, + seg, + 1, + phantom_count=np.array([1.0]), + phantom_logit=-20.0, + slot_weight=np.ones(2), + ) + full = np.array([-30.0, -30.0, -20.0]) + ref = np.exp(full - full.max()) + ref = ref / ref.sum() + np.testing.assert_allclose(w, ref[:2], rtol=1e-15, atol=0.0) + + def test_float32_torch_backward_finite(self) -> None: + """OutisLi review: with a tiny ``delta`` the floor's backward forms + ``D / delta**2`` -> inf, and ``0 * inf = NaN`` poisons the gradients + in float32 even though the forward is exact. The gated safe-where + construction must keep float32 autodiff finite and correct, in both + the inactive (reviewer repro) and active (deep-deficit) regimes. + """ + import torch + + # inactive-floor regime: logits far above phantom_logit + data = torch.tensor([20.0, 21.0], dtype=torch.float32, requires_grad=True) + ids = torch.tensor([0, 0], dtype=torch.int64) + w = segment_softmax( + data, + ids, + 1, + phantom_count=torch.tensor([-1.0], dtype=torch.float32), + phantom_logit=-20.0, + slot_weight=torch.ones(2, dtype=torch.float32), + ) + v = torch.tensor([1.0, 2.0]) + (w * v).sum().backward() + assert torch.all(torch.isfinite(data.grad)) + # analytic softmax-jacobian reference: g_i = w_i * (v_i - sum_j w_j v_j) + # (the phantom term exp(-41) is negligible at float32 resolution) + w64 = np.exp([20.0, 21.0] - np.float64(21.0)) + w64 = w64 / w64.sum() + ref = w64 * (np.array([1.0, 2.0]) - (w64 * [1.0, 2.0]).sum()) + np.testing.assert_allclose(data.grad.numpy(), ref, rtol=1e-4) + assert np.abs(ref).max() > 0.1 # nontrivial gradient, not all-zero + + # active-floor regime: signed denominator below the floor threshold + data = torch.tensor([-21.0, -21.0], dtype=torch.float32, requires_grad=True) + w = segment_softmax( + data, + ids, + 1, + phantom_count=torch.tensor([-1.0], dtype=torch.float32), + phantom_logit=-20.0, + slot_weight=torch.ones(2, dtype=torch.float32), + ) + (w * v).sum().backward() + assert torch.all(torch.isfinite(data.grad)) + + def test_float32_jax_backward_finite(self) -> None: + """Same float32 autodiff guarantee through JAX (the reviewer + reproduced the NaN gradient there as well). + """ + jax = pytest.importorskip("jax") + jnp = jax.numpy + + ids = np.array([0, 0], dtype=np.int64) + v = np.array([1.0, 2.0], dtype=np.float32) + + def loss(x): + w = segment_softmax( + x, + jnp.asarray(ids), + 1, + phantom_count=jnp.asarray([-1.0], dtype=jnp.float32), + phantom_logit=-20.0, + slot_weight=jnp.ones(2, dtype=jnp.float32), + ) + return (w * jnp.asarray(v)).sum() + + # inactive-floor regime (reviewer repro) + g = jax.grad(loss)(jnp.asarray([20.0, 21.0], dtype=jnp.float32)) + assert np.all(np.isfinite(np.asarray(g))) + w64 = np.exp([20.0, 21.0] - np.float64(21.0)) + w64 = w64 / w64.sum() + ref = w64 * (np.array([1.0, 2.0]) - (w64 * [1.0, 2.0]).sum()) + np.testing.assert_allclose(np.asarray(g), ref, rtol=1e-4) + # active-floor regime + g = jax.grad(loss)(jnp.asarray([-21.0, -21.0], dtype=jnp.float32)) + assert np.all(np.isfinite(np.asarray(g))) + + def test_floor_does_not_disturb_dense_parity_regime(self) -> None: + """In the in-design regime (all logits >= phantom_logit, count >= 0) + the floor never fires (its gate excludes non-negative counts), so + the fixed-width dense softmax is reproduced exactly. + """ + rng = np.random.default_rng(3) + data = rng.normal(size=7) # normal-scale logits, shift 20 below + seg = np.zeros(7, dtype=np.int64) + w_floor = segment_softmax( + data, + seg, + 1, + phantom_count=np.array([3.0]), + phantom_logit=-20.0, + slot_weight=np.ones(7), + ) + # reference: the exact fixed-width softmax with 3 phantom slots + full = np.concatenate([data, np.full(3, -20.0)]) + ref = np.exp(full - full.max()) + ref = ref / ref.sum() + np.testing.assert_allclose(w_floor, ref[:7], rtol=1e-12, atol=1e-15) + + +class TestSlotOccupancy: + """Capped water-filling occupancy (the phantom-denominator slot model).""" + + @staticmethod + def _h(t: np.ndarray) -> np.ndarray: + # C1 plateau saturator: t(2 - t) below the plateau, 1 above + return np.where(t >= 1.0, 1.0, t * (2.0 - t)) + + @classmethod + def _brute_force(cls, sw: np.ndarray, cap: float) -> np.ndarray: + # bisection on the inverse water level u: sum h(sw * u) == cap + if np.count_nonzero(sw) <= cap: + return (sw > 0).astype(np.float64) + lo, hi = 0.0, 2.0 * float(cap / max(sw.sum(), 1e-300)) + while cls._h(sw * hi).sum() < cap: + hi *= 2.0 + for _ in range(200): + mid = 0.5 * (lo + hi) + if cls._h(sw * mid).sum() < cap: + lo = mid + else: + hi = mid + return cls._h(sw * hi) + + def test_unbinding_capacity_is_bitwise_one(self) -> None: + from deepmd.dpmodel.utils.neighbor_graph.segment import ( + _slot_occupancy, + ) + + rng = np.random.default_rng(5) + sw = rng.uniform(0.01, 1.0, size=6) + seg = np.array([0, 0, 0, 1, 1, 1], dtype=np.int64) + theta = _slot_occupancy(sw, seg, 2, np.array([3.0, 5.0])) + np.testing.assert_array_equal(theta, np.ones(6)) # bitwise + + def test_binding_matches_bisection(self) -> None: + from deepmd.dpmodel.utils.neighbor_graph.segment import ( + _slot_occupancy, + ) + + rng = np.random.default_rng(7) + for cap0, cap1 in [(1.0, 2.0), (2.0, 3.0), (1.0, 1.0)]: + sw = rng.uniform(0.0, 1.0, size=9) + seg = np.array([0, 0, 0, 0, 1, 1, 1, 1, 1], dtype=np.int64) + theta = _slot_occupancy(sw, seg, 2, np.array([cap0, cap1])) + ref = np.concatenate( + [self._brute_force(sw[:4], cap0), self._brute_force(sw[4:], cap1)] + ) + np.testing.assert_allclose(theta, ref, atol=1e-10) + # capacity respected + assert theta[:4].sum() <= cap0 + 1e-10 + assert theta[4:].sum() <= cap1 + 1e-10 + + def test_ties_and_zeros(self) -> None: + from deepmd.dpmodel.utils.neighbor_graph.segment import ( + _slot_occupancy, + ) + + sw = np.array([1.0, 1.0, 1.0, 0.0]) + seg = np.zeros(4, dtype=np.int64) + theta = _slot_occupancy(sw, seg, 1, np.array([1.0])) + np.testing.assert_allclose(theta, [1 / 3, 1 / 3, 1 / 3, 0.0], rtol=1e-12) + + def test_empty_segment_no_nan(self) -> None: + from deepmd.dpmodel.utils.neighbor_graph.segment import ( + _slot_occupancy, + ) + + theta = _slot_occupancy( + np.array([0.5]), np.array([1], dtype=np.int64), 3, np.full(3, 2.0) + ) + assert np.all(np.isfinite(theta)) + np.testing.assert_array_equal(theta, [1.0]) + + def test_torch_matches_numpy(self) -> None: + import torch + + from deepmd.dpmodel.utils.neighbor_graph.segment import ( + _slot_occupancy, + ) + + rng = np.random.default_rng(11) + sw = rng.uniform(0.0, 1.0, size=8) + seg = np.array([0, 0, 0, 0, 0, 1, 1, 1], dtype=np.int64) + cap = np.array([2.0, 1.0]) + ref = _slot_occupancy(sw, seg, 2, cap) + out = _slot_occupancy( + torch.from_numpy(sw), torch.from_numpy(seg), 2, torch.from_numpy(cap) + ) + np.testing.assert_allclose(out.numpy(), ref, atol=1e-14) + + +class TestGradientContinuity: + """C1 regressions (OutisLi round 6): the denominator must carry no + first-derivative jump at the below-phantom surface ``e == P`` or at the + water-filling active-set transitions -- logits vary smoothly with + coordinates, so a weight-gradient kink is a force discontinuity. + """ + + @staticmethod + def _w(l0: float) -> np.ndarray: + return segment_softmax( + np.array([l0, -18.0]), + np.array([0, 0]), + 1, + phantom_count=np.array([-1.0]), # sel=1, n=2: binding, theta=1/2 + phantom_logit=-20.0, + slot_weight=np.ones(2), + ) + + def test_left_right_slope_match_at_phantom_logit(self) -> None: + """FD slopes of BOTH weights w.r.t. the crossing logit agree across + l = phantom_logit (the relu bracket jumped them by 1/theta = 2x). + """ + d = 1e-6 + left = (self._w(-20.0 - d) - self._w(-20.0 - 3.0 * d)) / (2.0 * d) + right = (self._w(-20.0 + 3.0 * d) - self._w(-20.0 + d)) / (2.0 * d) + assert np.all(np.abs(left) > 1e-4) # nontrivial slopes + np.testing.assert_allclose(left, right, rtol=5e-3) + + def test_theta_derivative_smooth_across_active_set_change(self) -> None: + """The occupancy's water-filling active-set transition (an entry + crossing the saturation plateau) is C1: the plateau's zero slope + makes the FD derivative of theta continuous. A pole or kink would + keep an O(1) derivative step under grid refinement. + """ + from deepmd.dpmodel.utils.neighbor_graph.segment import ( + _slot_occupancy, + ) + + seg = np.zeros(3, dtype=np.int64) + cap = np.array([2.0]) + + def dtheta(s: float, d: float = 1e-7) -> np.ndarray: + hi = _slot_occupancy(np.array([1.0, 0.2, s + d]), seg, 1, cap) + lo = _slot_occupancy(np.array([1.0, 0.2, s - d]), seg, 1, cap) + return (hi - lo) / (2.0 * d) + + # the k*=1 <-> k*=0 transition sits near s ~ 0.40 for sw=[1, 0.2, s] + def max_step(num: int) -> float: + grid = np.linspace(0.2, 0.6, num) + ds = np.stack([dtheta(float(s)) for s in grid]) + return float(np.abs(np.diff(ds, axis=0)).max()) + + coarse = max_step(41) + fine = max_step(401) + assert fine < 0.3 * coarse, ( + f"theta derivative steps do not shrink under refinement " + f"({coarse:.5f} -> {fine:.5f}): active-set transition is not C1" + ) + + +class TestFloat32HighLogitBackward: + """OutisLi round 7: in-design HIGH logits (count >= 0) make + ``ph_e = exp(phantom_logit - segment_max)`` subnormal in float32; a + ``log(ex_t / ph_e)`` tail -- even with the numerator where-substituted + -- overflows ``1 / ph_e`` (torch) or ``ph_e**2`` (jax) in the UNSELECTED + branch's backward, and 0 * inf poisons the selected gradient. The + log-space chi must keep float32 gradients finite and exactly dense. + """ + + def test_torch_high_logits_backward_finite(self) -> None: + import torch + + data = torch.tensor([68.0, 69.0], dtype=torch.float32, requires_grad=True) + ids = torch.tensor([0, 0], dtype=torch.int64) + w = segment_softmax( + data, + ids, + 1, + phantom_count=torch.tensor([0.0], dtype=torch.float32), + phantom_logit=-20.0, + slot_weight=torch.ones(2, dtype=torch.float32), + ) + v = torch.tensor([1.0, 2.0]) + (w * v).sum().backward() + assert torch.all(torch.isfinite(data.grad)) + w64 = np.exp([68.0, 69.0] - np.float64(69.0)) + w64 = w64 / w64.sum() + ref = w64 * (np.array([1.0, 2.0]) - (w64 * [1.0, 2.0]).sum()) + np.testing.assert_allclose(data.grad.numpy(), ref, rtol=1e-4) + + def test_jax_high_logits_backward_finite(self) -> None: + jax = pytest.importorskip("jax") + jnp = jax.numpy + + ids = np.array([0, 0], dtype=np.int64) + v = np.array([1.0, 2.0], dtype=np.float32) + + def loss(x): + w = segment_softmax( + x, + jnp.asarray(ids), + 1, + phantom_count=jnp.asarray([0.0], dtype=jnp.float32), + phantom_logit=-20.0, + slot_weight=jnp.ones(2, dtype=jnp.float32), + ) + return (w * jnp.asarray(v)).sum() + + # the reviewer's jax repro threshold: ph_e**2 underflows already here + g = jax.grad(loss)(jnp.asarray([24.0, 25.0], dtype=jnp.float32)) + assert np.all(np.isfinite(np.asarray(g))) + w64 = np.exp([24.0, 25.0] - np.float64(25.0)) + w64 = w64 / w64.sum() + ref = w64 * (np.array([1.0, 2.0]) - (w64 * [1.0, 2.0]).sum()) + np.testing.assert_allclose(np.asarray(g), ref, rtol=1e-4) + + +class TestFloat32WaterFillingStability: + """OutisLi round 8: the active-set selection must survive float32 weights + spanning ordinary decades. ``total - prefix`` suffix moments and the + ``S**2 - Q*room`` discriminant are ill-conditioned subtractions; in + float32 they rounded away a small positive discriminant, rejected the + valid saturated cut, and returned theta violating the defining capacity + equation (sum theta = 1.52 instead of 3.0 on the reviewer's repro). + The selection now runs in float64 internally (comparison-only, so no + gradient or export cost). + """ + + @staticmethod + def _bisect_theta(sw: np.ndarray, cap: float) -> np.ndarray: + def h(t): + return np.where(t >= 1.0, 1.0, t * (2.0 - t)) + + if np.count_nonzero(sw) <= cap: + return (sw > 0).astype(np.float64) + lo, hi = 0.0, 2.0 * cap / max(sw.sum(), 1e-300) + while h(sw * hi).sum() < cap: + hi *= 2.0 + for _ in range(200): + mid = 0.5 * (lo + hi) + if h(sw * mid).sum() < cap: + lo = mid + else: + hi = mid + return h(sw * hi) + + def test_reviewer_repro_float32(self) -> None: + from deepmd.dpmodel.utils.neighbor_graph.segment import ( + _slot_occupancy, + ) + + sw = np.array([1.0, 0.1, 0.01, 5e-7], dtype=np.float32) + seg = np.zeros(4, dtype=np.int64) + theta = _slot_occupancy(sw, seg, 1, np.array([3.0], dtype=np.float32)) + ref = self._bisect_theta(sw.astype(np.float64), 3.0) + # float64 bisection: [1, 1, 0.9999010, 0.0000990], sum 3.0; the old + # float32 selection returned sum 1.5208207 + np.testing.assert_allclose(theta.astype(np.float64), ref, atol=1e-5) + np.testing.assert_allclose(float(theta.sum()), 3.0, rtol=1e-5) + + def test_log_spaced_float32_weights_match_bisection(self) -> None: + from deepmd.dpmodel.utils.neighbor_graph.segment import ( + _slot_occupancy, + ) + + rng = np.random.default_rng(17) + for cap in (1.0, 2.0, 5.0): + # weights log-spaced over 8 decades, the reviewer's failure mode + sw = (10.0 ** rng.uniform(-8, 0, size=12)).astype(np.float32) + seg = np.zeros(12, dtype=np.int64) + theta = _slot_occupancy(sw, seg, 1, np.array([cap], dtype=np.float32)) + ref = self._bisect_theta(sw.astype(np.float64), cap) + np.testing.assert_allclose(theta.astype(np.float64), ref, atol=1e-4) + np.testing.assert_allclose(float(theta.sum()), cap, rtol=1e-4) + + def test_float32_forward_matches_float64_reference(self) -> None: + """The reviewer's forward repro: float32 vs float64 evaluation of the + same quantized inputs must agree (the broken selection gave 0.0612 + vs the 0.0792 reference through LocalAtten, ~23% error). + """ + data32 = np.full(4, -30.0, dtype=np.float32) + sw32 = np.array([1.0, 0.1, 0.01, 5e-7], dtype=np.float32) + seg = np.zeros(4, dtype=np.int64) + w32 = segment_softmax( + data32, + seg, + 1, + phantom_count=np.array([-1.0], dtype=np.float32), # sel=3, n=4 + phantom_logit=-20.0, + slot_weight=sw32, + ) + w64 = segment_softmax( + data32.astype(np.float64), + seg, + 1, + phantom_count=np.array([-1.0]), + phantom_logit=-20.0, + slot_weight=sw32.astype(np.float64), + ) + np.testing.assert_allclose(w32.astype(np.float64), w64, rtol=1e-3, atol=1e-6) diff --git a/source/tests/infer/deeppot_dpa3_spin.yaml b/source/tests/infer/deeppot_dpa3_spin.yaml new file mode 100644 index 0000000000..b192449624 --- /dev/null +++ b/source/tests/infer/deeppot_dpa3_spin.yaml @@ -0,0 +1,1884 @@ +backend: dpmodel +model: + backbone_model: + "@class": Model + "@variables": + out_bias: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - - 0.0 + - - 0.0 + - - 0.0 + - - 0.0 + out_std: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - - 1.0 + - - 1.0 + - - 1.0 + - - 1.0 + "@version": 2 + atom_exclude_types: &id002 + - 2 + - 3 + descriptor: + "@class": Descriptor + "@version": 2 + activation_function: silu + add_chg_spin_ebd: false + concat_output_tebd: false + default_chg_spin: null + env_protection: 1.0e-06 + exclude_types: &id003 + - - 3 + - 0 + - - 3 + - 1 + - - 3 + - 2 + - - 3 + - 3 + ntypes: 4 + precision: float64 + repflow_args: + a_compress_e_rate: 1 + a_compress_rate: 0 + a_compress_use_split: false + a_dim: 4 + a_rcut: 3.5 + a_rcut_smth: 0.5 + a_sel: 4 + axis_neuron: 4 + e_dim: 6 + e_rcut: 4.0 + e_rcut_smth: 0.5 + e_sel: 8 + edge_init_use_dist: false + fix_stat_std: 0.3 + n_dim: 8 + n_multi_edge_message: 1 + nlayers: 1 + optim_update: true + sel_reduce_factor: 10.0 + sequential_update: false + smooth_edge_update: false + update_angle: false + update_residual: 0.1 + update_residual_init: const + update_style: res_residual + use_dynamic_sel: false + use_exp_switch: false + repflow_variable: + "@variables": + davg: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + dstd: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + angle_embd: + "@class": Layer + "@variables": + b: null + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - 0.17649720801051316 + - 0.26111987625920485 + - -0.5130082451185703 + - -0.473906411761865 + "@version": 2 + activation_function: none + bias: false + precision: float64 + resnet: false + trainable: true + use_timestep: false + edge_embd: + "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - -0.9014522065120832 + - -0.2707576978619512 + - -0.11154097243018048 + - 0.5031862622181524 + - 1.443248998508462 + - -0.40919736174923005 + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - -0.41394371690540416 + - -0.020029235227998383 + - 0.7548439568852728 + - 0.18561320525199423 + - -0.1235585931191982 + - -0.6320668874586287 + "@version": 2 + activation_function: none + bias: true + precision: float64 + resnet: false + trainable: true + use_timestep: false + env_mat: + protection: 0.0 + rcut: 4.0 + rcut_smth: 0.5 + use_exp_switch: false + repflow_layers: + - "@class": RepFlowLayer + "@variables": + a_residual: [] + e_residual: + - "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + n_residual: + - "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + "@version": 2 + a_compress_e_rate: 1 + a_compress_rate: 0 + a_compress_use_split: false + a_dim: 4 + a_rcut: 3.5 + a_rcut_smth: 0.5 + a_sel: 4 + activation_function: silu + axis_neuron: 4 + e_dim: 6 + e_rcut: 4.0 + e_rcut_smth: 0.5 + e_sel: + - 8 + edge_self_linear: + "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.12186140180283762 + - -0.821800845268822 + - 1.5248690012827415 + - 0.2702504550116806 + - 1.383709381842487 + - -2.4456551147351098 + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - -0.05677878533841125 + - 0.08257140520321203 + - -0.17682541236689134 + - -0.06780335156677138 + - 0.2613225810769312 + - -0.2528499571680411 + - - -0.1060811611888437 + - 0.2834597688459181 + - 0.17021435817066793 + - -0.2119118384053945 + - 0.044257126661162875 + - -0.20348530980582044 + - - -0.17206345113221683 + - 0.19628652842871394 + - -0.03853877890775931 + - 0.0064469870425646406 + - -0.2035021228657865 + - 0.33893151231499635 + - - -0.1603541649096622 + - -0.07431170557593793 + - -0.1285051383598977 + - -0.09531516630926942 + - -0.10774430641710406 + - 0.10558546868439946 + - - -0.027408677366010784 + - 0.03171038951939523 + - -0.26649612755080526 + - 0.0749559135333121 + - 0.12753219377780048 + - -0.12375862279261161 + - - -0.3561807917324476 + - -0.028580689473013433 + - -0.2740045204725747 + - 0.12423725221263406 + - -0.0746927118825747 + - -0.16583458613892285 + - - -0.22497394079088218 + - -0.10329264538957945 + - -0.06015496745765555 + - -0.24047390264558702 + - -0.2470728805254821 + - -0.03091482236168157 + - - 0.313786674711663 + - -0.014345848137540798 + - -0.1446657411476756 + - -0.11134433415995286 + - -0.10957716367503506 + - 0.25318230359455945 + - - -0.11353216449019704 + - -0.24278855525542462 + - -0.0657328669264313 + - -0.08357873620530161 + - -0.19969579432068596 + - 0.0217399962565733 + - - -0.017346111123240478 + - -0.20460540022763518 + - 0.19580548714183002 + - 0.26081320850512657 + - -0.01937612111130992 + - 0.26782602217325135 + - - 0.169450738702664 + - 0.0007586409725729347 + - 0.2217946757639642 + - -0.08785618340632881 + - 0.08754553673729902 + - 0.0459550075486224 + - - -0.10347442434861592 + - -0.05665265742992996 + - -0.15657294594958857 + - 0.07518451488260593 + - -0.200469822163535 + - 0.008552407309104103 + - - -0.13418495292451688 + - -0.15007855071339413 + - -0.47561245640659827 + - 0.05519145026405931 + - 0.034426127944687676 + - -0.19833628864440492 + - - -0.061539077996517144 + - 0.140236735963936 + - 0.44907007382747016 + - -0.17502514002597466 + - -0.13141545313528988 + - -0.10225960767785013 + - - 0.15849623153238157 + - 0.14969793438965767 + - 0.05020396887825857 + - -0.42237393212574514 + - -0.43560848414739306 + - -0.34368434587411545 + - - -0.090558268807612 + - -0.10586479947976848 + - -0.17654116465986686 + - -0.17464251661717356 + - -0.17707748016637653 + - 0.4728011907076426 + - - 0.06839467741547146 + - 0.20172332403056187 + - -0.20761658659357723 + - -0.3179201458113386 + - 0.1570191398976539 + - 0.30829366728408747 + - - -0.07346831672915768 + - -0.01603422028135059 + - -0.2343121216044693 + - -0.09228986456390967 + - -0.12259985802096685 + - 0.13925704477109332 + - - 0.03112673892006412 + - -0.12259170091097643 + - -0.01873720650800844 + - -0.02825905483134531 + - -0.07410620360262994 + - -0.13890487670689447 + - - 0.2599426512954838 + - -0.030475413023044056 + - 0.04418102981236639 + - 0.14747916053674695 + - 0.11469436259489629 + - -0.12589465767715197 + - - 0.1534348683560527 + - -0.2598559665351654 + - 0.1691188844559884 + - -0.05815067519957393 + - -0.09922406205302857 + - -0.026111067214965193 + - - -0.09469687008709152 + - -0.25433614509748265 + - -0.3230603080176275 + - 0.10565697308598668 + - 0.11382397843310456 + - 0.12636033735242963 + "@version": 2 + activation_function: none + bias: true + precision: float64 + resnet: false + trainable: true + use_timestep: false + n_dim: 8 + n_multi_edge_message: 1 + node_edge_linear: + "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - -0.47607591249187536 + - -1.1996431006923678 + - -0.17281730905229956 + - -0.37252679074651174 + - -2.004295595595026 + - -1.3828864783672565 + - -1.6353313440832131 + - 0.22589145999360716 + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - -0.0696849625476975 + - 0.2297992430114777 + - 0.2016390404837622 + - -0.14268332516715398 + - 0.05104561028973491 + - -0.0047524771390660015 + - 0.2007647888532239 + - -0.13892443853282946 + - - -0.06269186432762 + - -0.3620323943560662 + - 0.1429598213093448 + - 0.32251103225335886 + - 0.03297508356302982 + - 0.09813020765990464 + - -0.2128967691985101 + - 0.1446653191521595 + - - 0.3408694676048958 + - 0.2108742996875453 + - 0.3708891734344055 + - 0.03603632207631642 + - -0.021861106604374982 + - 0.05885265981211556 + - -0.3850668694292795 + - -0.02565715855818745 + - - 0.2554067355579947 + - -0.12385934461529113 + - 0.20794804172087122 + - 0.34771760519493133 + - 0.0030775484903255083 + - 0.08033323613153229 + - 0.04547640227535313 + - 0.03256133523805088 + - - -0.20228475171413982 + - -0.40882303462099395 + - -0.13933573248982073 + - -0.09056898648309795 + - 0.06705102826758672 + - -0.10643998751725821 + - 0.3714789434592029 + - -0.15660714896565422 + - - 0.0620445405627027 + - -0.14981233984554174 + - 0.1377612457580642 + - -0.3264453797874674 + - 0.18886992363386892 + - -0.13120999191064697 + - 0.2639396300281778 + - 0.20744058178112204 + - - 0.0902283316504931 + - 0.2642720697422412 + - 0.11616051352480065 + - 0.33194344559115435 + - -0.07519119975054182 + - -0.05062288710700148 + - -0.0033899752949634763 + - 0.12074780296663348 + - - -0.14625494457636853 + - -0.39187048236106403 + - 0.005863181213654556 + - 0.08058988215606765 + - 0.229684677996952 + - 0.02491096095922189 + - -0.07923462148233366 + - 0.03463149425218323 + - - 0.1459761391053138 + - 0.1826916307305693 + - -0.24330282168960599 + - -0.15404338160080283 + - -0.15732528026070658 + - 0.10082194118502665 + - 0.10780094880007303 + - 0.10439459076027502 + - - 0.2967580322447816 + - 0.19310548831515634 + - 0.13271337197427788 + - -0.003964549207383962 + - -0.3053587625881187 + - 0.12883374510336365 + - 0.045960329737757634 + - 0.19761345822107057 + - - -0.13773367016862742 + - 0.09659775346412201 + - 0.13561552570758134 + - 0.07814681408507194 + - -0.28773064288403055 + - 0.07556744144481048 + - -0.09838713644355178 + - 0.009867107649393275 + - - -0.09655717020123253 + - -0.10871554819348611 + - -0.11670258568304015 + - 0.2177137774640066 + - -0.14817421356773255 + - -0.03606693672811542 + - -0.026029214690369364 + - -0.040666475049662726 + - - -0.0677385423671006 + - -0.12993597893178244 + - 0.1180039874263662 + - 0.1384604584579823 + - -0.024227421664540914 + - 0.1679245814762119 + - -0.19280274838451647 + - 0.0990223630355508 + - - 0.0758415027141385 + - 0.16215196433523008 + - 0.2767732385588474 + - 0.022163750613004355 + - -0.12254120989786124 + - -0.12391951174230557 + - -0.028791741351884195 + - -0.0595519969823867 + - - 0.22247449036902905 + - 0.07567917899966987 + - -0.18221068561029122 + - -0.1496346790319525 + - 0.01739141266484531 + - 0.03295277270138665 + - -0.27927822171693173 + - -0.13558030103477586 + - - 0.1712575942124072 + - 0.13705603104177683 + - 0.290608271870899 + - 0.25077636518593155 + - 0.06723740912894116 + - -0.29077479630216374 + - -0.25998108625190797 + - 0.15096707384533595 + - - -0.011258223444056591 + - 0.07940059884337107 + - 0.14539160696529732 + - -0.33401238196882443 + - 0.0359760729699335 + - -0.02226084988227022 + - 0.12276616178918343 + - 0.0439592954772777 + - - 0.07596667366672871 + - -0.11052600964268607 + - -0.13155071841622368 + - -0.07425437999013539 + - 0.00827734508288158 + - 0.07414300482320346 + - 0.052019022231599196 + - 0.16368644986528788 + - - 0.31022863799320216 + - 0.11380817934759249 + - 0.11671054675679823 + - 0.03833224311415518 + - 0.1545146635596559 + - 0.5283089690392868 + - -0.17235747525638992 + - -0.16802245441710034 + - - 0.19547575805994974 + - 0.03442738806627725 + - 0.035134165349037516 + - 0.1685202553837112 + - -0.13706885637245225 + - -0.09105484518308726 + - 0.24401116664356562 + - -0.042463896239058455 + - - 0.18293429344914702 + - -0.0797150153045118 + - 0.2837300628985514 + - -0.03290000697254011 + - 0.07484025269991934 + - 0.4486382833349405 + - 0.18215765586473062 + - 0.14222755521955213 + - - -0.054949228485595726 + - 0.2298266346316468 + - -0.13022437426681047 + - 0.31473958548227127 + - -0.16053599380138361 + - 0.12351036770696595 + - -0.2026640600757936 + - -0.3120452604960154 + "@version": 2 + activation_function: none + bias: true + precision: float64 + resnet: false + trainable: true + use_timestep: false + node_self_mlp: + "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 1.3358086643032931 + - 0.7847145686255231 + - 0.46740601094575585 + - -0.17204456063327273 + - -1.0586982191306735 + - -0.23498342309774128 + - 1.6521024027545508 + - -2.401400317086709 + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - -0.22061013087519665 + - 0.17161694901625085 + - 0.25079797681294247 + - -0.06984190636344022 + - 0.402412783105689 + - -0.13232509868240386 + - -0.12410592033624109 + - -0.5243896508356666 + - - -0.34531669337816745 + - 0.2590681097532894 + - -0.4170438578433154 + - 0.33209656716128205 + - 0.20907222698506978 + - 0.21026382825889875 + - -0.04125433055358784 + - -0.3362049950725693 + - - -0.02306669199993831 + - -0.27140136827851236 + - 0.08675906253383281 + - 0.20991982378397447 + - -0.20157467157772102 + - 0.10954533237221269 + - -0.30521247150866015 + - 0.1039196228402914 + - - 0.2927901959232568 + - -0.05686111266739088 + - -0.352867716741099 + - 0.06499009437306054 + - 0.2935084094905296 + - -0.5208455549268021 + - -0.06412894033597939 + - 0.2617524844957687 + - - -0.26859205166611555 + - -0.017740123512057532 + - -0.16973184286647353 + - -0.041497625408519805 + - -0.33848186563738925 + - -0.498133067071094 + - 0.06453515847241846 + - -0.28211046673410256 + - - -0.0031712540783364537 + - 0.14054927501098227 + - -0.16739625499774285 + - 0.02924799819668618 + - 0.19945724852581612 + - -0.07433092972702877 + - 0.33641837410477954 + - -0.1935354318143647 + - - -0.2896583115032089 + - -0.4291374752779325 + - 0.18521131755882006 + - 0.036186935403130116 + - 0.27669775576389155 + - -0.04763160274577408 + - 0.1400908330823242 + - 0.15697986928574623 + - - -0.45902865822845124 + - 0.33250108656046035 + - 0.0306169230429561 + - -0.035381192364331175 + - -0.0510947377580893 + - 0.03972955950151097 + - 0.6129808284962325 + - 0.027297205883797467 + "@version": 2 + activation_function: none + bias: true + precision: float64 + resnet: false + trainable: true + use_timestep: false + node_sym_linear: + "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.3521609009674381 + - 1.5631307030631036 + - 0.17947890305801448 + - -1.274607877122093 + - -1.1528521407168824 + - -1.6164563686220714 + - 0.056724431118804874 + - -1.6177337409601333 + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - 0.07307166266137728 + - 0.06127645860254937 + - -0.18492998679736772 + - 0.04613999102916452 + - 0.0071079000479441 + - -0.03731231022031221 + - -0.09483134725409749 + - -0.04779952388125717 + - - -0.06975495248291474 + - -0.19948091645922555 + - -0.19101500000694704 + - -0.0756612239190429 + - 0.18223713459959498 + - 0.004660326879702162 + - -0.07331926290215518 + - -0.11804049351864328 + - - -0.008643603296694117 + - -0.07891692454926386 + - -0.24683520896350286 + - -0.07498319216962253 + - 0.14604675984008694 + - 0.09601184516912262 + - -0.01561740011879576 + - 0.05490167651869453 + - - -0.019884970427089133 + - 0.0007666914047260165 + - -0.16505651916265357 + - -0.16723740821054547 + - 0.1234653183096876 + - -0.04403642108952563 + - -0.03727304005788303 + - -0.18190516409632088 + - - -0.09168767286099873 + - -0.1549419399425698 + - 0.08193144903871091 + - 0.15640675555750194 + - -0.06305034848986912 + - 0.16836512195133213 + - 0.009048302220263893 + - 0.05322075713280992 + - - -0.06468543813846328 + - 0.0948348631241292 + - 0.006867290444906741 + - -0.24931773871448817 + - 0.08788089155489308 + - -0.0739514302480491 + - 0.025288498321181765 + - -0.08305521153831118 + - - -0.07393598220040017 + - 0.07042478981157554 + - -0.07236047649200875 + - -0.04706083253081129 + - 0.011054293351306345 + - -0.08799610585856558 + - 0.1563680796185477 + - 0.04333789772104407 + - - -0.10653678039670528 + - 0.18112426500221723 + - 0.009186401470971654 + - 0.006153194931152504 + - -0.04989535662898608 + - -0.17876067282409114 + - -0.15602193322162777 + - 0.00781917954318876 + - - -0.06699569918753231 + - 0.30735630566871885 + - -0.016096279041795645 + - -0.22956358044083025 + - -0.0065529000816888765 + - -0.06902463180143781 + - 0.06768609922953323 + - 0.1187567871665586 + - - -0.03935798540152214 + - -0.10867329670955546 + - -0.04052094555163571 + - -0.04078630187590839 + - -0.0748763378601901 + - 0.01860182181594992 + - 0.20057959184872112 + - 0.16046549209905156 + - - 0.031478615338666395 + - 0.11567514563874234 + - 0.08294594125898624 + - -0.07089853590674343 + - 0.20101923186451937 + - -0.11766930025015115 + - 0.21570163379940235 + - 0.14563108587004206 + - - 0.07932986781606469 + - -0.19442968907969105 + - 0.05697454617840562 + - -0.19484656091831729 + - 0.04754566926156801 + - -0.12152155059832441 + - 0.08105546170302243 + - -0.09483406966077029 + - - -0.10943334690817784 + - 0.11702284889224986 + - 0.06551144399385757 + - -0.003108503735325857 + - -0.1466684268106551 + - 0.11582453333312602 + - -0.19609870968779317 + - -0.11809063481420465 + - - -0.11120967058944209 + - -0.07178289284260277 + - -0.07505138171189361 + - -0.17137771295621249 + - -0.012516091428859523 + - 0.056132912587423756 + - -0.011172736867909887 + - 0.0014926969164057145 + - - -0.1652803302650934 + - 0.08452449793427 + - -0.06260662069159101 + - -0.07909718643578055 + - 0.00574135469567161 + - -0.05691391300603163 + - 0.2457179942284785 + - -0.08037694862142311 + - - 0.1761032538671494 + - -0.15524353322856968 + - -0.20338260987738993 + - -0.09738847488694806 + - 0.05960295717975261 + - -0.0268406105291267 + - 0.19154482080963495 + - -0.05557739958347549 + - - -0.23162474155138468 + - 0.005428848189956548 + - 0.14498512403306713 + - 0.015859517797165032 + - 0.13342303538966063 + - 0.07757097608660568 + - 0.061885304048992174 + - -0.02774862502778554 + - - -0.099674682792698 + - 0.1743242267060875 + - 0.0565895993699819 + - -0.1431728246694354 + - -0.04572377374634247 + - 0.1932522842767088 + - -0.13605774184771868 + - -0.079596349847149 + - - 0.015159290423222593 + - 0.0741788473825365 + - -0.025111776236424455 + - 0.11728977172727281 + - -0.05246129405331076 + - -0.3560652693695576 + - 0.22489664505020285 + - -0.11322150427163667 + - - 0.1172685876179488 + - 0.015449206720498673 + - 0.11464505230123948 + - -0.13045379503420262 + - -0.18460226345307634 + - -0.0735660416536509 + - 0.02668836976483192 + - 0.009471901506209893 + - - -0.12415218588856815 + - -0.028427823628392242 + - -0.0726329032188482 + - 0.2205454016716484 + - -0.06981635935553832 + - -0.06914918285976224 + - -0.07547512647684368 + - 0.19585301943839276 + - - 0.02068794647278527 + - 0.11434955856950152 + - 0.04733548159377606 + - -0.0940771421180628 + - 0.106950218084799 + - 0.11995323224700441 + - 0.07016105028143815 + - -0.07349788842232614 + - - -0.028316732941958092 + - -0.006316920155388264 + - 0.014323448114816232 + - 0.07909510285143638 + - 0.08089223619428912 + - -0.1285448965066473 + - -0.02731037643994388 + - -0.048232324099890284 + - - -0.04229476466912251 + - -0.10545582133061814 + - -0.1399519987577358 + - -0.24859786794141928 + - 0.04555029533580089 + - -0.06637709714144181 + - -0.11891839416041088 + - -0.05608836594526548 + - - -0.1481671676394082 + - -0.11826472343612228 + - 0.18759449377634982 + - -0.0027813243183313764 + - -0.06187858233767373 + - -0.16870507895423517 + - 0.15432198341660605 + - 0.2442525725033602 + - - 0.11655618965628044 + - 0.16410614799338208 + - -0.15922334755571288 + - 0.05294100944284731 + - -0.042676438943807564 + - 0.05982722192738627 + - 0.08818007330306689 + - 0.08799006862019813 + - - -0.1816952674192488 + - 0.33018315199731113 + - 0.14825048237904745 + - -0.12977688627249692 + - -0.014039894202582361 + - 0.021698570605095405 + - -0.10536292700472008 + - -0.016298405526400214 + - - 0.18891280861168214 + - -0.037066320234429954 + - 0.051989201606798936 + - -0.33236261122879446 + - -0.2233240290736924 + - 0.17632501110044835 + - 0.02791043546786102 + - 0.08058616657592761 + - - -0.12416787825473675 + - -0.0018776550590277605 + - -0.1361594510955972 + - -0.031008628174283 + - -0.1510470016534144 + - -0.1968118582063139 + - 0.05923927005740039 + - 0.10906017525194028 + - - -0.01747528984400593 + - 0.043571037430286425 + - 0.09735765094593854 + - 0.038496104792229716 + - 0.021583898030338507 + - -0.11795161808253331 + - -0.11404406907043374 + - -0.06541831356900717 + - - 0.05781757062086345 + - 0.06545403068342133 + - 0.07182196888387801 + - 0.06571017380833269 + - 0.25549620850343796 + - -0.01712221435859817 + - 0.02746476505848508 + - -0.16813933068880024 + - - 0.15811742659496866 + - -0.10097487333290259 + - 0.0007478905750386516 + - 0.15986815657402492 + - 0.0879704571647486 + - 0.051839404360383305 + - 0.04773139180116972 + - -0.1562216704347126 + - - -0.00554177026701311 + - 0.026672084558123862 + - -0.026556168406945337 + - 0.017618135480540704 + - 0.04290442846891425 + - -0.16108845422437917 + - 0.03885069382837762 + - -0.08559226312134341 + - - -0.10984387513362157 + - 0.06020841962256015 + - 0.013439129456291792 + - -0.1211722988008539 + - 0.0321361577334442 + - 0.04742132269747014 + - -0.08371259477093888 + - -0.14250805695920574 + - - 0.04498243399350513 + - -0.03633434279549633 + - 0.17043129619564554 + - 0.13738977779076048 + - 0.03367329367643751 + - 0.13141345496526305 + - 0.14626062464255066 + - -0.087660426894852 + - - 0.13304548046946202 + - -0.02074921039690319 + - -0.19614199925540662 + - 0.09145888259449976 + - -0.16872056060024043 + - -0.057806035869808946 + - -0.012927002228426554 + - -0.18968555494779932 + - - 0.09056415309144267 + - 0.19579647713404205 + - 0.12419307551929215 + - 0.03068324855507999 + - 0.16324257199502792 + - 0.28864177653836015 + - -0.04884842530407823 + - 0.05243039778651716 + - - -0.1354513040660592 + - -0.0032083328727676315 + - 0.035763067000830435 + - -0.10752629467535854 + - -0.004527627068300205 + - -0.26678729966885645 + - 0.16095749546546945 + - -0.0768457166279081 + - - 0.24290534029168284 + - -0.19993818991295886 + - 0.05863500838014017 + - 0.1075745460176732 + - -0.2703641493668329 + - 0.022882752217475207 + - -0.18377439784177813 + - -0.02475991439750886 + - - -0.0970343883793403 + - 0.022190761521183 + - -0.31137609288015433 + - -0.12852583938411438 + - 0.06380585650762231 + - -0.05537350140183391 + - 0.009834307052428782 + - -0.18327381164681603 + - - -0.058720582106338445 + - 0.012207974777885133 + - 0.04906298973398652 + - -0.0252045636071624 + - 0.04064401311527239 + - -0.12030307623147056 + - -0.02607458251331658 + - -0.12104904963385374 + - - 0.14380149442345772 + - -0.08586187966457755 + - -0.0562380312253021 + - 0.1183995520092173 + - -0.008618891616010692 + - -0.30556252122213096 + - 0.157107693967395 + - -0.150824446001649 + - - -0.12986340463514554 + - -0.13953775800615473 + - 0.06688782609307184 + - 0.30709990962247197 + - -0.10057794483875744 + - -0.15572836837520085 + - 0.22240522808485344 + - 0.07486567450323982 + - - -0.00026497955681491453 + - -0.462148220797257 + - -0.04683339159019641 + - 0.10954858908660245 + - 0.048155719596331595 + - -0.08404934441388894 + - 0.15848474948089222 + - -0.029754000979091536 + - - -0.008795641657631076 + - -0.021341761230446545 + - -0.10489671109204046 + - 0.03213370243212562 + - -0.021792936100149974 + - -0.018371450392434912 + - 0.0007292277382723748 + - 0.07679112359755517 + - - -0.06130007400378907 + - -0.06581095863692285 + - 0.06501448048047738 + - -0.14197246804370967 + - 0.15983589537290877 + - -0.15693380789472725 + - -0.17963845906090375 + - 0.10204145028546817 + - - -0.07077050429398143 + - 0.1990098057969514 + - -0.2525111691805106 + - -0.22059894251537618 + - -0.27531410890875607 + - -0.0693243961021514 + - 0.03876302523241355 + - 0.12122101629786736 + - - -0.12820657692829063 + - -0.10772035941442479 + - 0.10829696580051636 + - 0.1493715060396245 + - 0.13488833866187872 + - 0.09022524867490032 + - 0.007332743974581279 + - 0.1529338321168549 + - - -0.22245363971842472 + - -0.08917661330105822 + - 0.10304564318043377 + - -0.07026805272160686 + - 0.016625750231852813 + - 0.23074109385732217 + - 0.053971407495566504 + - -0.15089059679319458 + - - 0.1294396068073317 + - -0.038487426453509534 + - 0.09393650831599386 + - 0.09638990927578407 + - 0.17905918157852316 + - 0.06760574587425355 + - 0.0639998107196389 + - -0.1587157815816586 + - - 0.06077231806824999 + - 0.006159909130812671 + - 0.15285274367932117 + - -0.026531120401424045 + - 0.06104797756042876 + - -0.174933801016035 + - 0.25284181425638513 + - -0.16931699181750984 + - - -0.09480440252644158 + - -0.11919995631753837 + - 0.1374865485894956 + - 0.03525829583245701 + - 0.055414318086174905 + - 0.039970825479268265 + - -0.028476173719310948 + - 0.007895110382084259 + - - -0.08849522170883828 + - 0.1556903658898126 + - -0.06942905817654972 + - 0.17917871676321492 + - -0.12839965901095401 + - -0.1457242708290995 + - 0.2073632537418445 + - -0.0033056633245595168 + - - -0.14321940581992326 + - 0.016216983383358995 + - -0.05603214608550905 + - 0.034067014410779244 + - -0.004165932252642813 + - 0.03579825379823718 + - 0.2274077472661256 + - 0.12282153328534674 + - - -0.17424677728325255 + - 0.03032450606197887 + - -0.3407467917235723 + - 0.08460871296272927 + - -0.21233509125037692 + - 0.038581785470083826 + - -0.1271081651221865 + - -0.05674635282930029 + - - -0.06365889303105148 + - -0.0346798442701684 + - 0.04178115473238202 + - -0.03570145798701077 + - 0.2255873927499116 + - -0.21936512330368732 + - -0.19469567244011848 + - -0.007461014512643234 + "@version": 2 + activation_function: none + bias: true + precision: float64 + resnet: false + trainable: true + use_timestep: false + ntypes: 4 + optim_update: true + precision: float64 + sel_reduce_factor: 10.0 + sequential_update: false + smooth_edge_update: false + update_angle: false + update_residual: 0.1 + update_residual_init: const + update_style: res_residual + use_dynamic_sel: false + trainable: true + type: dpa3 + type_embedding: + "@class": TypeEmbedNet + "@version": 2 + activation_function: Linear + embedding: + "@class": EmbeddingNetwork + "@version": 2 + activation_function: Linear + bias: false + in_dim: 4 + layers: + - "@class": Layer + "@variables": + b: null + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - 0.06913868355931278 + - -0.3276059448146492 + - -0.22478586008940918 + - -0.03129740042629991 + - -0.2511436154794455 + - -0.4760319710462916 + - 0.183856376649989 + - 0.220680920691283 + - - -0.1331166944050067 + - -0.2985446381663858 + - -0.1299144028716818 + - 0.12716526105014014 + - 0.24445281051361242 + - 0.052359417290304015 + - -0.06639194378815659 + - -0.0515428623822807 + - - -0.3302870133986425 + - 0.1177804767091647 + - 0.06915893387117533 + - -0.4204302050492702 + - -0.3161145657939801 + - 0.322920377419993 + - 0.19395457855721343 + - -0.11365337655752422 + - - -0.16993400446851198 + - -0.157416126804567 + - -0.08090448953478106 + - 0.20830555342316676 + - -0.11308079862243182 + - 0.044490575624147384 + - 0.28211395871639494 + - 0.07920112686609734 + "@version": 2 + activation_function: Linear + bias: false + precision: float64 + resnet: true + trainable: true + use_timestep: false + neuron: + - 8 + precision: float64 + resnet_dt: false + neuron: + - 8 + ntypes: 4 + padding: true + precision: float64 + resnet_dt: false + trainable: true + type_map: &id001 + - Ni + - O + - Ni_spin + - O_spin + use_econf_tebd: false + use_tebd_bias: false + type_map: *id001 + use_econf_tebd: false + use_loc_mapping: true + use_tebd_bias: false + fitting: + "@class": Fitting + "@variables": + aparam_avg: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.0 + aparam_inv_std: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 1.0 + bias_atom_e: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - 0.0 + - - 0.0 + - - 0.0 + - - 0.0 + case_embd: null + fparam_avg: null + fparam_inv_std: null + "@version": 4 + activation_function: tanh + atom_ener: null + default_fparam: null + dim_case_embd: 0 + dim_descrpt: 8 + dim_out: 1 + exclude_types: *id002 + layer_name: null + mixed_types: true + nets: + "@class": NetworkCollection + "@version": 1 + ndim: 0 + network_type: fitting_network + networks: + - "@class": FittingNetwork + "@version": 1 + activation_function: tanh + bias_out: true + in_dim: 9 + layers: + - "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - -1.95770695761676 + - -0.2476293590902135 + - 0.25443971144004524 + - 1.1986522321754531 + - 0.0030188217090154337 + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - 0.19224554841074107 + - 0.05429795979660471 + - -0.1389670848638964 + - -0.10462785043714116 + - 0.1089987221732007 + - - -0.2674528934854709 + - -0.34368813663218956 + - -0.32661387749828436 + - -0.054276053653440084 + - -0.4708180280377169 + - - -0.028453356966921094 + - -0.1818257915009699 + - 0.5172015902059852 + - 0.07298508659463093 + - -0.41150720205719726 + - - -0.009573595400007627 + - 0.36335859733193915 + - -0.42257433445595627 + - 0.04512922378046872 + - 0.01214874819312088 + - - 0.10523323982907387 + - -0.0813318495284912 + - -0.5972654039164439 + - 0.18665154460844918 + - -0.11162809779624402 + - - -0.6162981937526609 + - -0.3794599870226799 + - -0.13208161507478433 + - -0.07980535308944087 + - -0.13750862677493716 + - - 0.003549601048741757 + - -0.10261020888740208 + - 0.14583248443759755 + - -0.3660277389715169 + - -0.2670347033440862 + - - -0.2325831241229178 + - -0.1838381450084032 + - -0.11628330754042544 + - -0.055126134875232095 + - -0.04304460916326145 + - - -0.1699430653550485 + - -0.34364758746663254 + - -0.4487935046391692 + - 0.4739297137312101 + - 0.03146001518088704 + "@version": 2 + activation_function: tanh + bias: true + precision: float64 + resnet: true + trainable: true + use_timestep: false + - "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.3004481477996428 + - 0.5261420741708049 + - -1.6397400541209868 + - -2.447812539095814 + - -0.5930086265906365 + idt: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.0993080666106219 + - 0.09873045857916823 + - 0.10010274098022319 + - 0.10020061209981648 + - 0.09998711001511004 + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - -0.3141891418069332 + - 0.30132598326837057 + - -0.1868614701027005 + - -0.1853536726835805 + - -0.14904917553209618 + - - -0.4993776326714626 + - 0.2929711950476154 + - -0.3300253064210836 + - -0.4799775188835898 + - -0.12327559985245252 + - - 0.16627900477763782 + - 0.18281489789715116 + - -0.0796215789550366 + - 0.11637836794519682 + - 0.019126199990905587 + - - 0.47193798042526686 + - 0.3935489978037474 + - 0.1926588188573466 + - 0.11685532990383077 + - -0.3143759410105157 + - - 0.2619509948079511 + - 0.17134734041574828 + - 0.16467987243470003 + - -0.17768942725372738 + - 0.17196893072212313 + "@version": 2 + activation_function: tanh + bias: true + precision: float64 + resnet: true + trainable: true + use_timestep: true + - "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.4738132556710063 + - -0.32857031000381093 + - -2.2966637967768286 + - -0.47371878604557704 + - -1.1317456558916699 + idt: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.10104051010914285 + - 0.10073058893042686 + - 0.09780511012883421 + - 0.09938856388264454 + - 0.09707779684554663 + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - -0.3092290441156226 + - -0.496367611501348 + - -0.052492949379292775 + - 0.06663748312823926 + - 0.027714401468510886 + - - -0.10433141997317527 + - -0.323901631855259 + - -0.24739439873488192 + - 0.3076895568713741 + - 0.1593814472209255 + - - -0.07111829721069259 + - -0.27598680250101504 + - 0.16632764307325093 + - 0.1801382402999823 + - 0.3107523993064097 + - - -0.012140157566561928 + - 0.07469305237763302 + - 0.26428018852282276 + - -0.11500213881655802 + - -0.2731498304335624 + - - 0.29941998505510775 + - 0.39267279762211 + - 0.06586779164332648 + - 0.10010820203885952 + - -0.04143485413490972 + "@version": 2 + activation_function: tanh + bias: true + precision: float64 + resnet: true + trainable: true + use_timestep: true + - "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.10882142522066754 + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - 0.26432565710368733 + - - 0.17264367113482967 + - - -0.04729186377886323 + - - -0.08841444813809296 + - - 0.2969145415081517 + "@version": 2 + activation_function: none + bias: true + precision: float64 + resnet: false + trainable: true + use_timestep: false + neuron: + - 5 + - 5 + - 5 + out_dim: 1 + precision: float64 + resnet_dt: true + ntypes: 4 + neuron: + - 5 + - 5 + - 5 + ntypes: 4 + numb_aparam: 1 + numb_fparam: 0 + precision: float64 + rcond: null + resnet_dt: true + spin: null + tot_ener_zero: false + trainable: + - true + - true + - true + - true + type: ener + type_map: + - Ni + - O + - Ni_spin + - O_spin + use_aparam_as_mask: false + var_name: energy + pair_exclude_types: *id003 + preset_out_bias: null + rcond: null + type: standard + type_map: + - Ni + - O + - Ni_spin + - O_spin + spin: + use_spin: + - true + - false + virtual_scale: + - 0.314 + - 0.0 + type: spin_ener +model_def_script: + descriptor: + precision: float64 + repflow: + a_dim: 4 + a_rcut: 3.5 + a_rcut_smth: 0.5 + a_sel: 4 + axis_neuron: 4 + e_dim: 6 + e_rcut: 4.0 + e_rcut_smth: 0.5 + e_sel: 8 + n_dim: 8 + nlayers: 1 + update_angle: false + seed: 1 + type: dpa3 + use_loc_mapping: true + fitting_net: + neuron: + - 5 + - 5 + - 5 + numb_aparam: 1 + resnet_dt: true + seed: 1 + spin: + use_spin: + - true + - false + virtual_scale: + - 0.314 + - 0.0 + type_map: + - Ni + - O +software: deepmd-kit +time: "2026-06-04 01:43:22.415031+00:00" +version: 3.0.0 diff --git a/source/tests/infer/deeppot_dpa3_spin_mpi.yaml b/source/tests/infer/deeppot_dpa3_spin_mpi.yaml new file mode 100644 index 0000000000..d60084aac8 --- /dev/null +++ b/source/tests/infer/deeppot_dpa3_spin_mpi.yaml @@ -0,0 +1,1884 @@ +backend: dpmodel +model: + backbone_model: + "@class": Model + "@variables": + out_bias: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - - 0.0 + - - 0.0 + - - 0.0 + - - 0.0 + out_std: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - - 1.0 + - - 1.0 + - - 1.0 + - - 1.0 + "@version": 2 + atom_exclude_types: &id002 + - 2 + - 3 + descriptor: + "@class": Descriptor + "@version": 2 + activation_function: silu + add_chg_spin_ebd: false + concat_output_tebd: false + default_chg_spin: null + env_protection: 1.0e-06 + exclude_types: &id003 + - - 3 + - 0 + - - 3 + - 1 + - - 3 + - 2 + - - 3 + - 3 + ntypes: 4 + precision: float64 + repflow_args: + a_compress_e_rate: 1 + a_compress_rate: 0 + a_compress_use_split: false + a_dim: 4 + a_rcut: 3.5 + a_rcut_smth: 0.5 + a_sel: 4 + axis_neuron: 4 + e_dim: 6 + e_rcut: 4.0 + e_rcut_smth: 0.5 + e_sel: 8 + edge_init_use_dist: false + fix_stat_std: 0.3 + n_dim: 8 + n_multi_edge_message: 1 + nlayers: 1 + optim_update: true + sel_reduce_factor: 10.0 + sequential_update: false + smooth_edge_update: false + update_angle: false + update_residual: 0.1 + update_residual_init: const + update_style: res_residual + use_dynamic_sel: false + use_exp_switch: false + repflow_variable: + "@variables": + davg: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + - - 0.0 + - 0.0 + - 0.0 + - 0.0 + dstd: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + - - 0.3 + - 0.3 + - 0.3 + - 0.3 + angle_embd: + "@class": Layer + "@variables": + b: null + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - 0.17649720801051316 + - 0.26111987625920485 + - -0.5130082451185703 + - -0.473906411761865 + "@version": 2 + activation_function: none + bias: false + precision: float64 + resnet: false + trainable: true + use_timestep: false + edge_embd: + "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - -0.9014522065120832 + - -0.2707576978619512 + - -0.11154097243018048 + - 0.5031862622181524 + - 1.443248998508462 + - -0.40919736174923005 + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - -0.41394371690540416 + - -0.020029235227998383 + - 0.7548439568852728 + - 0.18561320525199423 + - -0.1235585931191982 + - -0.6320668874586287 + "@version": 2 + activation_function: none + bias: true + precision: float64 + resnet: false + trainable: true + use_timestep: false + env_mat: + protection: 0.0 + rcut: 4.0 + rcut_smth: 0.5 + use_exp_switch: false + repflow_layers: + - "@class": RepFlowLayer + "@variables": + a_residual: [] + e_residual: + - "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + n_residual: + - "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + - 0.1 + "@version": 2 + a_compress_e_rate: 1 + a_compress_rate: 0 + a_compress_use_split: false + a_dim: 4 + a_rcut: 3.5 + a_rcut_smth: 0.5 + a_sel: 4 + activation_function: silu + axis_neuron: 4 + e_dim: 6 + e_rcut: 4.0 + e_rcut_smth: 0.5 + e_sel: + - 8 + edge_self_linear: + "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.12186140180283762 + - -0.821800845268822 + - 1.5248690012827415 + - 0.2702504550116806 + - 1.383709381842487 + - -2.4456551147351098 + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - -0.05677878533841125 + - 0.08257140520321203 + - -0.17682541236689134 + - -0.06780335156677138 + - 0.2613225810769312 + - -0.2528499571680411 + - - -0.1060811611888437 + - 0.2834597688459181 + - 0.17021435817066793 + - -0.2119118384053945 + - 0.044257126661162875 + - -0.20348530980582044 + - - -0.17206345113221683 + - 0.19628652842871394 + - -0.03853877890775931 + - 0.0064469870425646406 + - -0.2035021228657865 + - 0.33893151231499635 + - - -0.1603541649096622 + - -0.07431170557593793 + - -0.1285051383598977 + - -0.09531516630926942 + - -0.10774430641710406 + - 0.10558546868439946 + - - -0.027408677366010784 + - 0.03171038951939523 + - -0.26649612755080526 + - 0.0749559135333121 + - 0.12753219377780048 + - -0.12375862279261161 + - - -0.3561807917324476 + - -0.028580689473013433 + - -0.2740045204725747 + - 0.12423725221263406 + - -0.0746927118825747 + - -0.16583458613892285 + - - -0.22497394079088218 + - -0.10329264538957945 + - -0.06015496745765555 + - -0.24047390264558702 + - -0.2470728805254821 + - -0.03091482236168157 + - - 0.313786674711663 + - -0.014345848137540798 + - -0.1446657411476756 + - -0.11134433415995286 + - -0.10957716367503506 + - 0.25318230359455945 + - - -0.11353216449019704 + - -0.24278855525542462 + - -0.0657328669264313 + - -0.08357873620530161 + - -0.19969579432068596 + - 0.0217399962565733 + - - -0.017346111123240478 + - -0.20460540022763518 + - 0.19580548714183002 + - 0.26081320850512657 + - -0.01937612111130992 + - 0.26782602217325135 + - - 0.169450738702664 + - 0.0007586409725729347 + - 0.2217946757639642 + - -0.08785618340632881 + - 0.08754553673729902 + - 0.0459550075486224 + - - -0.10347442434861592 + - -0.05665265742992996 + - -0.15657294594958857 + - 0.07518451488260593 + - -0.200469822163535 + - 0.008552407309104103 + - - -0.13418495292451688 + - -0.15007855071339413 + - -0.47561245640659827 + - 0.05519145026405931 + - 0.034426127944687676 + - -0.19833628864440492 + - - -0.061539077996517144 + - 0.140236735963936 + - 0.44907007382747016 + - -0.17502514002597466 + - -0.13141545313528988 + - -0.10225960767785013 + - - 0.15849623153238157 + - 0.14969793438965767 + - 0.05020396887825857 + - -0.42237393212574514 + - -0.43560848414739306 + - -0.34368434587411545 + - - -0.090558268807612 + - -0.10586479947976848 + - -0.17654116465986686 + - -0.17464251661717356 + - -0.17707748016637653 + - 0.4728011907076426 + - - 0.06839467741547146 + - 0.20172332403056187 + - -0.20761658659357723 + - -0.3179201458113386 + - 0.1570191398976539 + - 0.30829366728408747 + - - -0.07346831672915768 + - -0.01603422028135059 + - -0.2343121216044693 + - -0.09228986456390967 + - -0.12259985802096685 + - 0.13925704477109332 + - - 0.03112673892006412 + - -0.12259170091097643 + - -0.01873720650800844 + - -0.02825905483134531 + - -0.07410620360262994 + - -0.13890487670689447 + - - 0.2599426512954838 + - -0.030475413023044056 + - 0.04418102981236639 + - 0.14747916053674695 + - 0.11469436259489629 + - -0.12589465767715197 + - - 0.1534348683560527 + - -0.2598559665351654 + - 0.1691188844559884 + - -0.05815067519957393 + - -0.09922406205302857 + - -0.026111067214965193 + - - -0.09469687008709152 + - -0.25433614509748265 + - -0.3230603080176275 + - 0.10565697308598668 + - 0.11382397843310456 + - 0.12636033735242963 + "@version": 2 + activation_function: none + bias: true + precision: float64 + resnet: false + trainable: true + use_timestep: false + n_dim: 8 + n_multi_edge_message: 1 + node_edge_linear: + "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - -0.47607591249187536 + - -1.1996431006923678 + - -0.17281730905229956 + - -0.37252679074651174 + - -2.004295595595026 + - -1.3828864783672565 + - -1.6353313440832131 + - 0.22589145999360716 + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - -0.0696849625476975 + - 0.2297992430114777 + - 0.2016390404837622 + - -0.14268332516715398 + - 0.05104561028973491 + - -0.0047524771390660015 + - 0.2007647888532239 + - -0.13892443853282946 + - - -0.06269186432762 + - -0.3620323943560662 + - 0.1429598213093448 + - 0.32251103225335886 + - 0.03297508356302982 + - 0.09813020765990464 + - -0.2128967691985101 + - 0.1446653191521595 + - - 0.3408694676048958 + - 0.2108742996875453 + - 0.3708891734344055 + - 0.03603632207631642 + - -0.021861106604374982 + - 0.05885265981211556 + - -0.3850668694292795 + - -0.02565715855818745 + - - 0.2554067355579947 + - -0.12385934461529113 + - 0.20794804172087122 + - 0.34771760519493133 + - 0.0030775484903255083 + - 0.08033323613153229 + - 0.04547640227535313 + - 0.03256133523805088 + - - -0.20228475171413982 + - -0.40882303462099395 + - -0.13933573248982073 + - -0.09056898648309795 + - 0.06705102826758672 + - -0.10643998751725821 + - 0.3714789434592029 + - -0.15660714896565422 + - - 0.0620445405627027 + - -0.14981233984554174 + - 0.1377612457580642 + - -0.3264453797874674 + - 0.18886992363386892 + - -0.13120999191064697 + - 0.2639396300281778 + - 0.20744058178112204 + - - 0.0902283316504931 + - 0.2642720697422412 + - 0.11616051352480065 + - 0.33194344559115435 + - -0.07519119975054182 + - -0.05062288710700148 + - -0.0033899752949634763 + - 0.12074780296663348 + - - -0.14625494457636853 + - -0.39187048236106403 + - 0.005863181213654556 + - 0.08058988215606765 + - 0.229684677996952 + - 0.02491096095922189 + - -0.07923462148233366 + - 0.03463149425218323 + - - 0.1459761391053138 + - 0.1826916307305693 + - -0.24330282168960599 + - -0.15404338160080283 + - -0.15732528026070658 + - 0.10082194118502665 + - 0.10780094880007303 + - 0.10439459076027502 + - - 0.2967580322447816 + - 0.19310548831515634 + - 0.13271337197427788 + - -0.003964549207383962 + - -0.3053587625881187 + - 0.12883374510336365 + - 0.045960329737757634 + - 0.19761345822107057 + - - -0.13773367016862742 + - 0.09659775346412201 + - 0.13561552570758134 + - 0.07814681408507194 + - -0.28773064288403055 + - 0.07556744144481048 + - -0.09838713644355178 + - 0.009867107649393275 + - - -0.09655717020123253 + - -0.10871554819348611 + - -0.11670258568304015 + - 0.2177137774640066 + - -0.14817421356773255 + - -0.03606693672811542 + - -0.026029214690369364 + - -0.040666475049662726 + - - -0.0677385423671006 + - -0.12993597893178244 + - 0.1180039874263662 + - 0.1384604584579823 + - -0.024227421664540914 + - 0.1679245814762119 + - -0.19280274838451647 + - 0.0990223630355508 + - - 0.0758415027141385 + - 0.16215196433523008 + - 0.2767732385588474 + - 0.022163750613004355 + - -0.12254120989786124 + - -0.12391951174230557 + - -0.028791741351884195 + - -0.0595519969823867 + - - 0.22247449036902905 + - 0.07567917899966987 + - -0.18221068561029122 + - -0.1496346790319525 + - 0.01739141266484531 + - 0.03295277270138665 + - -0.27927822171693173 + - -0.13558030103477586 + - - 0.1712575942124072 + - 0.13705603104177683 + - 0.290608271870899 + - 0.25077636518593155 + - 0.06723740912894116 + - -0.29077479630216374 + - -0.25998108625190797 + - 0.15096707384533595 + - - -0.011258223444056591 + - 0.07940059884337107 + - 0.14539160696529732 + - -0.33401238196882443 + - 0.0359760729699335 + - -0.02226084988227022 + - 0.12276616178918343 + - 0.0439592954772777 + - - 0.07596667366672871 + - -0.11052600964268607 + - -0.13155071841622368 + - -0.07425437999013539 + - 0.00827734508288158 + - 0.07414300482320346 + - 0.052019022231599196 + - 0.16368644986528788 + - - 0.31022863799320216 + - 0.11380817934759249 + - 0.11671054675679823 + - 0.03833224311415518 + - 0.1545146635596559 + - 0.5283089690392868 + - -0.17235747525638992 + - -0.16802245441710034 + - - 0.19547575805994974 + - 0.03442738806627725 + - 0.035134165349037516 + - 0.1685202553837112 + - -0.13706885637245225 + - -0.09105484518308726 + - 0.24401116664356562 + - -0.042463896239058455 + - - 0.18293429344914702 + - -0.0797150153045118 + - 0.2837300628985514 + - -0.03290000697254011 + - 0.07484025269991934 + - 0.4486382833349405 + - 0.18215765586473062 + - 0.14222755521955213 + - - -0.054949228485595726 + - 0.2298266346316468 + - -0.13022437426681047 + - 0.31473958548227127 + - -0.16053599380138361 + - 0.12351036770696595 + - -0.2026640600757936 + - -0.3120452604960154 + "@version": 2 + activation_function: none + bias: true + precision: float64 + resnet: false + trainable: true + use_timestep: false + node_self_mlp: + "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 1.3358086643032931 + - 0.7847145686255231 + - 0.46740601094575585 + - -0.17204456063327273 + - -1.0586982191306735 + - -0.23498342309774128 + - 1.6521024027545508 + - -2.401400317086709 + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - -0.22061013087519665 + - 0.17161694901625085 + - 0.25079797681294247 + - -0.06984190636344022 + - 0.402412783105689 + - -0.13232509868240386 + - -0.12410592033624109 + - -0.5243896508356666 + - - -0.34531669337816745 + - 0.2590681097532894 + - -0.4170438578433154 + - 0.33209656716128205 + - 0.20907222698506978 + - 0.21026382825889875 + - -0.04125433055358784 + - -0.3362049950725693 + - - -0.02306669199993831 + - -0.27140136827851236 + - 0.08675906253383281 + - 0.20991982378397447 + - -0.20157467157772102 + - 0.10954533237221269 + - -0.30521247150866015 + - 0.1039196228402914 + - - 0.2927901959232568 + - -0.05686111266739088 + - -0.352867716741099 + - 0.06499009437306054 + - 0.2935084094905296 + - -0.5208455549268021 + - -0.06412894033597939 + - 0.2617524844957687 + - - -0.26859205166611555 + - -0.017740123512057532 + - -0.16973184286647353 + - -0.041497625408519805 + - -0.33848186563738925 + - -0.498133067071094 + - 0.06453515847241846 + - -0.28211046673410256 + - - -0.0031712540783364537 + - 0.14054927501098227 + - -0.16739625499774285 + - 0.02924799819668618 + - 0.19945724852581612 + - -0.07433092972702877 + - 0.33641837410477954 + - -0.1935354318143647 + - - -0.2896583115032089 + - -0.4291374752779325 + - 0.18521131755882006 + - 0.036186935403130116 + - 0.27669775576389155 + - -0.04763160274577408 + - 0.1400908330823242 + - 0.15697986928574623 + - - -0.45902865822845124 + - 0.33250108656046035 + - 0.0306169230429561 + - -0.035381192364331175 + - -0.0510947377580893 + - 0.03972955950151097 + - 0.6129808284962325 + - 0.027297205883797467 + "@version": 2 + activation_function: none + bias: true + precision: float64 + resnet: false + trainable: true + use_timestep: false + node_sym_linear: + "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.3521609009674381 + - 1.5631307030631036 + - 0.17947890305801448 + - -1.274607877122093 + - -1.1528521407168824 + - -1.6164563686220714 + - 0.056724431118804874 + - -1.6177337409601333 + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - 0.07307166266137728 + - 0.06127645860254937 + - -0.18492998679736772 + - 0.04613999102916452 + - 0.0071079000479441 + - -0.03731231022031221 + - -0.09483134725409749 + - -0.04779952388125717 + - - -0.06975495248291474 + - -0.19948091645922555 + - -0.19101500000694704 + - -0.0756612239190429 + - 0.18223713459959498 + - 0.004660326879702162 + - -0.07331926290215518 + - -0.11804049351864328 + - - -0.008643603296694117 + - -0.07891692454926386 + - -0.24683520896350286 + - -0.07498319216962253 + - 0.14604675984008694 + - 0.09601184516912262 + - -0.01561740011879576 + - 0.05490167651869453 + - - -0.019884970427089133 + - 0.0007666914047260165 + - -0.16505651916265357 + - -0.16723740821054547 + - 0.1234653183096876 + - -0.04403642108952563 + - -0.03727304005788303 + - -0.18190516409632088 + - - -0.09168767286099873 + - -0.1549419399425698 + - 0.08193144903871091 + - 0.15640675555750194 + - -0.06305034848986912 + - 0.16836512195133213 + - 0.009048302220263893 + - 0.05322075713280992 + - - -0.06468543813846328 + - 0.0948348631241292 + - 0.006867290444906741 + - -0.24931773871448817 + - 0.08788089155489308 + - -0.0739514302480491 + - 0.025288498321181765 + - -0.08305521153831118 + - - -0.07393598220040017 + - 0.07042478981157554 + - -0.07236047649200875 + - -0.04706083253081129 + - 0.011054293351306345 + - -0.08799610585856558 + - 0.1563680796185477 + - 0.04333789772104407 + - - -0.10653678039670528 + - 0.18112426500221723 + - 0.009186401470971654 + - 0.006153194931152504 + - -0.04989535662898608 + - -0.17876067282409114 + - -0.15602193322162777 + - 0.00781917954318876 + - - -0.06699569918753231 + - 0.30735630566871885 + - -0.016096279041795645 + - -0.22956358044083025 + - -0.0065529000816888765 + - -0.06902463180143781 + - 0.06768609922953323 + - 0.1187567871665586 + - - -0.03935798540152214 + - -0.10867329670955546 + - -0.04052094555163571 + - -0.04078630187590839 + - -0.0748763378601901 + - 0.01860182181594992 + - 0.20057959184872112 + - 0.16046549209905156 + - - 0.031478615338666395 + - 0.11567514563874234 + - 0.08294594125898624 + - -0.07089853590674343 + - 0.20101923186451937 + - -0.11766930025015115 + - 0.21570163379940235 + - 0.14563108587004206 + - - 0.07932986781606469 + - -0.19442968907969105 + - 0.05697454617840562 + - -0.19484656091831729 + - 0.04754566926156801 + - -0.12152155059832441 + - 0.08105546170302243 + - -0.09483406966077029 + - - -0.10943334690817784 + - 0.11702284889224986 + - 0.06551144399385757 + - -0.003108503735325857 + - -0.1466684268106551 + - 0.11582453333312602 + - -0.19609870968779317 + - -0.11809063481420465 + - - -0.11120967058944209 + - -0.07178289284260277 + - -0.07505138171189361 + - -0.17137771295621249 + - -0.012516091428859523 + - 0.056132912587423756 + - -0.011172736867909887 + - 0.0014926969164057145 + - - -0.1652803302650934 + - 0.08452449793427 + - -0.06260662069159101 + - -0.07909718643578055 + - 0.00574135469567161 + - -0.05691391300603163 + - 0.2457179942284785 + - -0.08037694862142311 + - - 0.1761032538671494 + - -0.15524353322856968 + - -0.20338260987738993 + - -0.09738847488694806 + - 0.05960295717975261 + - -0.0268406105291267 + - 0.19154482080963495 + - -0.05557739958347549 + - - -0.23162474155138468 + - 0.005428848189956548 + - 0.14498512403306713 + - 0.015859517797165032 + - 0.13342303538966063 + - 0.07757097608660568 + - 0.061885304048992174 + - -0.02774862502778554 + - - -0.099674682792698 + - 0.1743242267060875 + - 0.0565895993699819 + - -0.1431728246694354 + - -0.04572377374634247 + - 0.1932522842767088 + - -0.13605774184771868 + - -0.079596349847149 + - - 0.015159290423222593 + - 0.0741788473825365 + - -0.025111776236424455 + - 0.11728977172727281 + - -0.05246129405331076 + - -0.3560652693695576 + - 0.22489664505020285 + - -0.11322150427163667 + - - 0.1172685876179488 + - 0.015449206720498673 + - 0.11464505230123948 + - -0.13045379503420262 + - -0.18460226345307634 + - -0.0735660416536509 + - 0.02668836976483192 + - 0.009471901506209893 + - - -0.12415218588856815 + - -0.028427823628392242 + - -0.0726329032188482 + - 0.2205454016716484 + - -0.06981635935553832 + - -0.06914918285976224 + - -0.07547512647684368 + - 0.19585301943839276 + - - 0.02068794647278527 + - 0.11434955856950152 + - 0.04733548159377606 + - -0.0940771421180628 + - 0.106950218084799 + - 0.11995323224700441 + - 0.07016105028143815 + - -0.07349788842232614 + - - -0.028316732941958092 + - -0.006316920155388264 + - 0.014323448114816232 + - 0.07909510285143638 + - 0.08089223619428912 + - -0.1285448965066473 + - -0.02731037643994388 + - -0.048232324099890284 + - - -0.04229476466912251 + - -0.10545582133061814 + - -0.1399519987577358 + - -0.24859786794141928 + - 0.04555029533580089 + - -0.06637709714144181 + - -0.11891839416041088 + - -0.05608836594526548 + - - -0.1481671676394082 + - -0.11826472343612228 + - 0.18759449377634982 + - -0.0027813243183313764 + - -0.06187858233767373 + - -0.16870507895423517 + - 0.15432198341660605 + - 0.2442525725033602 + - - 0.11655618965628044 + - 0.16410614799338208 + - -0.15922334755571288 + - 0.05294100944284731 + - -0.042676438943807564 + - 0.05982722192738627 + - 0.08818007330306689 + - 0.08799006862019813 + - - -0.1816952674192488 + - 0.33018315199731113 + - 0.14825048237904745 + - -0.12977688627249692 + - -0.014039894202582361 + - 0.021698570605095405 + - -0.10536292700472008 + - -0.016298405526400214 + - - 0.18891280861168214 + - -0.037066320234429954 + - 0.051989201606798936 + - -0.33236261122879446 + - -0.2233240290736924 + - 0.17632501110044835 + - 0.02791043546786102 + - 0.08058616657592761 + - - -0.12416787825473675 + - -0.0018776550590277605 + - -0.1361594510955972 + - -0.031008628174283 + - -0.1510470016534144 + - -0.1968118582063139 + - 0.05923927005740039 + - 0.10906017525194028 + - - -0.01747528984400593 + - 0.043571037430286425 + - 0.09735765094593854 + - 0.038496104792229716 + - 0.021583898030338507 + - -0.11795161808253331 + - -0.11404406907043374 + - -0.06541831356900717 + - - 0.05781757062086345 + - 0.06545403068342133 + - 0.07182196888387801 + - 0.06571017380833269 + - 0.25549620850343796 + - -0.01712221435859817 + - 0.02746476505848508 + - -0.16813933068880024 + - - 0.15811742659496866 + - -0.10097487333290259 + - 0.0007478905750386516 + - 0.15986815657402492 + - 0.0879704571647486 + - 0.051839404360383305 + - 0.04773139180116972 + - -0.1562216704347126 + - - -0.00554177026701311 + - 0.026672084558123862 + - -0.026556168406945337 + - 0.017618135480540704 + - 0.04290442846891425 + - -0.16108845422437917 + - 0.03885069382837762 + - -0.08559226312134341 + - - -0.10984387513362157 + - 0.06020841962256015 + - 0.013439129456291792 + - -0.1211722988008539 + - 0.0321361577334442 + - 0.04742132269747014 + - -0.08371259477093888 + - -0.14250805695920574 + - - 0.04498243399350513 + - -0.03633434279549633 + - 0.17043129619564554 + - 0.13738977779076048 + - 0.03367329367643751 + - 0.13141345496526305 + - 0.14626062464255066 + - -0.087660426894852 + - - 0.13304548046946202 + - -0.02074921039690319 + - -0.19614199925540662 + - 0.09145888259449976 + - -0.16872056060024043 + - -0.057806035869808946 + - -0.012927002228426554 + - -0.18968555494779932 + - - 0.09056415309144267 + - 0.19579647713404205 + - 0.12419307551929215 + - 0.03068324855507999 + - 0.16324257199502792 + - 0.28864177653836015 + - -0.04884842530407823 + - 0.05243039778651716 + - - -0.1354513040660592 + - -0.0032083328727676315 + - 0.035763067000830435 + - -0.10752629467535854 + - -0.004527627068300205 + - -0.26678729966885645 + - 0.16095749546546945 + - -0.0768457166279081 + - - 0.24290534029168284 + - -0.19993818991295886 + - 0.05863500838014017 + - 0.1075745460176732 + - -0.2703641493668329 + - 0.022882752217475207 + - -0.18377439784177813 + - -0.02475991439750886 + - - -0.0970343883793403 + - 0.022190761521183 + - -0.31137609288015433 + - -0.12852583938411438 + - 0.06380585650762231 + - -0.05537350140183391 + - 0.009834307052428782 + - -0.18327381164681603 + - - -0.058720582106338445 + - 0.012207974777885133 + - 0.04906298973398652 + - -0.0252045636071624 + - 0.04064401311527239 + - -0.12030307623147056 + - -0.02607458251331658 + - -0.12104904963385374 + - - 0.14380149442345772 + - -0.08586187966457755 + - -0.0562380312253021 + - 0.1183995520092173 + - -0.008618891616010692 + - -0.30556252122213096 + - 0.157107693967395 + - -0.150824446001649 + - - -0.12986340463514554 + - -0.13953775800615473 + - 0.06688782609307184 + - 0.30709990962247197 + - -0.10057794483875744 + - -0.15572836837520085 + - 0.22240522808485344 + - 0.07486567450323982 + - - -0.00026497955681491453 + - -0.462148220797257 + - -0.04683339159019641 + - 0.10954858908660245 + - 0.048155719596331595 + - -0.08404934441388894 + - 0.15848474948089222 + - -0.029754000979091536 + - - -0.008795641657631076 + - -0.021341761230446545 + - -0.10489671109204046 + - 0.03213370243212562 + - -0.021792936100149974 + - -0.018371450392434912 + - 0.0007292277382723748 + - 0.07679112359755517 + - - -0.06130007400378907 + - -0.06581095863692285 + - 0.06501448048047738 + - -0.14197246804370967 + - 0.15983589537290877 + - -0.15693380789472725 + - -0.17963845906090375 + - 0.10204145028546817 + - - -0.07077050429398143 + - 0.1990098057969514 + - -0.2525111691805106 + - -0.22059894251537618 + - -0.27531410890875607 + - -0.0693243961021514 + - 0.03876302523241355 + - 0.12122101629786736 + - - -0.12820657692829063 + - -0.10772035941442479 + - 0.10829696580051636 + - 0.1493715060396245 + - 0.13488833866187872 + - 0.09022524867490032 + - 0.007332743974581279 + - 0.1529338321168549 + - - -0.22245363971842472 + - -0.08917661330105822 + - 0.10304564318043377 + - -0.07026805272160686 + - 0.016625750231852813 + - 0.23074109385732217 + - 0.053971407495566504 + - -0.15089059679319458 + - - 0.1294396068073317 + - -0.038487426453509534 + - 0.09393650831599386 + - 0.09638990927578407 + - 0.17905918157852316 + - 0.06760574587425355 + - 0.0639998107196389 + - -0.1587157815816586 + - - 0.06077231806824999 + - 0.006159909130812671 + - 0.15285274367932117 + - -0.026531120401424045 + - 0.06104797756042876 + - -0.174933801016035 + - 0.25284181425638513 + - -0.16931699181750984 + - - -0.09480440252644158 + - -0.11919995631753837 + - 0.1374865485894956 + - 0.03525829583245701 + - 0.055414318086174905 + - 0.039970825479268265 + - -0.028476173719310948 + - 0.007895110382084259 + - - -0.08849522170883828 + - 0.1556903658898126 + - -0.06942905817654972 + - 0.17917871676321492 + - -0.12839965901095401 + - -0.1457242708290995 + - 0.2073632537418445 + - -0.0033056633245595168 + - - -0.14321940581992326 + - 0.016216983383358995 + - -0.05603214608550905 + - 0.034067014410779244 + - -0.004165932252642813 + - 0.03579825379823718 + - 0.2274077472661256 + - 0.12282153328534674 + - - -0.17424677728325255 + - 0.03032450606197887 + - -0.3407467917235723 + - 0.08460871296272927 + - -0.21233509125037692 + - 0.038581785470083826 + - -0.1271081651221865 + - -0.05674635282930029 + - - -0.06365889303105148 + - -0.0346798442701684 + - 0.04178115473238202 + - -0.03570145798701077 + - 0.2255873927499116 + - -0.21936512330368732 + - -0.19469567244011848 + - -0.007461014512643234 + "@version": 2 + activation_function: none + bias: true + precision: float64 + resnet: false + trainable: true + use_timestep: false + ntypes: 4 + optim_update: true + precision: float64 + sel_reduce_factor: 10.0 + sequential_update: false + smooth_edge_update: false + update_angle: false + update_residual: 0.1 + update_residual_init: const + update_style: res_residual + use_dynamic_sel: false + trainable: true + type: dpa3 + type_embedding: + "@class": TypeEmbedNet + "@version": 2 + activation_function: Linear + embedding: + "@class": EmbeddingNetwork + "@version": 2 + activation_function: Linear + bias: false + in_dim: 4 + layers: + - "@class": Layer + "@variables": + b: null + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - 0.06913868355931278 + - -0.3276059448146492 + - -0.22478586008940918 + - -0.03129740042629991 + - -0.2511436154794455 + - -0.4760319710462916 + - 0.183856376649989 + - 0.220680920691283 + - - -0.1331166944050067 + - -0.2985446381663858 + - -0.1299144028716818 + - 0.12716526105014014 + - 0.24445281051361242 + - 0.052359417290304015 + - -0.06639194378815659 + - -0.0515428623822807 + - - -0.3302870133986425 + - 0.1177804767091647 + - 0.06915893387117533 + - -0.4204302050492702 + - -0.3161145657939801 + - 0.322920377419993 + - 0.19395457855721343 + - -0.11365337655752422 + - - -0.16993400446851198 + - -0.157416126804567 + - -0.08090448953478106 + - 0.20830555342316676 + - -0.11308079862243182 + - 0.044490575624147384 + - 0.28211395871639494 + - 0.07920112686609734 + "@version": 2 + activation_function: Linear + bias: false + precision: float64 + resnet: true + trainable: true + use_timestep: false + neuron: + - 8 + precision: float64 + resnet_dt: false + neuron: + - 8 + ntypes: 4 + padding: true + precision: float64 + resnet_dt: false + trainable: true + type_map: &id001 + - Ni + - O + - Ni_spin + - O_spin + use_econf_tebd: false + use_tebd_bias: false + type_map: *id001 + use_econf_tebd: false + use_loc_mapping: false + use_tebd_bias: false + fitting: + "@class": Fitting + "@variables": + aparam_avg: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.0 + aparam_inv_std: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 1.0 + bias_atom_e: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - 0.0 + - - 0.0 + - - 0.0 + - - 0.0 + case_embd: null + fparam_avg: null + fparam_inv_std: null + "@version": 4 + activation_function: tanh + atom_ener: null + default_fparam: null + dim_case_embd: 0 + dim_descrpt: 8 + dim_out: 1 + exclude_types: *id002 + layer_name: null + mixed_types: true + nets: + "@class": NetworkCollection + "@version": 1 + ndim: 0 + network_type: fitting_network + networks: + - "@class": FittingNetwork + "@version": 1 + activation_function: tanh + bias_out: true + in_dim: 9 + layers: + - "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - -1.95770695761676 + - -0.2476293590902135 + - 0.25443971144004524 + - 1.1986522321754531 + - 0.0030188217090154337 + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - 0.19224554841074107 + - 0.05429795979660471 + - -0.1389670848638964 + - -0.10462785043714116 + - 0.1089987221732007 + - - -0.2674528934854709 + - -0.34368813663218956 + - -0.32661387749828436 + - -0.054276053653440084 + - -0.4708180280377169 + - - -0.028453356966921094 + - -0.1818257915009699 + - 0.5172015902059852 + - 0.07298508659463093 + - -0.41150720205719726 + - - -0.009573595400007627 + - 0.36335859733193915 + - -0.42257433445595627 + - 0.04512922378046872 + - 0.01214874819312088 + - - 0.10523323982907387 + - -0.0813318495284912 + - -0.5972654039164439 + - 0.18665154460844918 + - -0.11162809779624402 + - - -0.6162981937526609 + - -0.3794599870226799 + - -0.13208161507478433 + - -0.07980535308944087 + - -0.13750862677493716 + - - 0.003549601048741757 + - -0.10261020888740208 + - 0.14583248443759755 + - -0.3660277389715169 + - -0.2670347033440862 + - - -0.2325831241229178 + - -0.1838381450084032 + - -0.11628330754042544 + - -0.055126134875232095 + - -0.04304460916326145 + - - -0.1699430653550485 + - -0.34364758746663254 + - -0.4487935046391692 + - 0.4739297137312101 + - 0.03146001518088704 + "@version": 2 + activation_function: tanh + bias: true + precision: float64 + resnet: true + trainable: true + use_timestep: false + - "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.3004481477996428 + - 0.5261420741708049 + - -1.6397400541209868 + - -2.447812539095814 + - -0.5930086265906365 + idt: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.0993080666106219 + - 0.09873045857916823 + - 0.10010274098022319 + - 0.10020061209981648 + - 0.09998711001511004 + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - -0.3141891418069332 + - 0.30132598326837057 + - -0.1868614701027005 + - -0.1853536726835805 + - -0.14904917553209618 + - - -0.4993776326714626 + - 0.2929711950476154 + - -0.3300253064210836 + - -0.4799775188835898 + - -0.12327559985245252 + - - 0.16627900477763782 + - 0.18281489789715116 + - -0.0796215789550366 + - 0.11637836794519682 + - 0.019126199990905587 + - - 0.47193798042526686 + - 0.3935489978037474 + - 0.1926588188573466 + - 0.11685532990383077 + - -0.3143759410105157 + - - 0.2619509948079511 + - 0.17134734041574828 + - 0.16467987243470003 + - -0.17768942725372738 + - 0.17196893072212313 + "@version": 2 + activation_function: tanh + bias: true + precision: float64 + resnet: true + trainable: true + use_timestep: true + - "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.4738132556710063 + - -0.32857031000381093 + - -2.2966637967768286 + - -0.47371878604557704 + - -1.1317456558916699 + idt: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.10104051010914285 + - 0.10073058893042686 + - 0.09780511012883421 + - 0.09938856388264454 + - 0.09707779684554663 + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - -0.3092290441156226 + - -0.496367611501348 + - -0.052492949379292775 + - 0.06663748312823926 + - 0.027714401468510886 + - - -0.10433141997317527 + - -0.323901631855259 + - -0.24739439873488192 + - 0.3076895568713741 + - 0.1593814472209255 + - - -0.07111829721069259 + - -0.27598680250101504 + - 0.16632764307325093 + - 0.1801382402999823 + - 0.3107523993064097 + - - -0.012140157566561928 + - 0.07469305237763302 + - 0.26428018852282276 + - -0.11500213881655802 + - -0.2731498304335624 + - - 0.29941998505510775 + - 0.39267279762211 + - 0.06586779164332648 + - 0.10010820203885952 + - -0.04143485413490972 + "@version": 2 + activation_function: tanh + bias: true + precision: float64 + resnet: true + trainable: true + use_timestep: true + - "@class": Layer + "@variables": + b: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - 0.10882142522066754 + idt: null + w: + "@class": np.ndarray + "@is_variable": true + "@version": 1 + dtype: float64 + value: + - - 0.26432565710368733 + - - 0.17264367113482967 + - - -0.04729186377886323 + - - -0.08841444813809296 + - - 0.2969145415081517 + "@version": 2 + activation_function: none + bias: true + precision: float64 + resnet: false + trainable: true + use_timestep: false + neuron: + - 5 + - 5 + - 5 + out_dim: 1 + precision: float64 + resnet_dt: true + ntypes: 4 + neuron: + - 5 + - 5 + - 5 + ntypes: 4 + numb_aparam: 1 + numb_fparam: 0 + precision: float64 + rcond: null + resnet_dt: true + spin: null + tot_ener_zero: false + trainable: + - true + - true + - true + - true + type: ener + type_map: + - Ni + - O + - Ni_spin + - O_spin + use_aparam_as_mask: false + var_name: energy + pair_exclude_types: *id003 + preset_out_bias: null + rcond: null + type: standard + type_map: + - Ni + - O + - Ni_spin + - O_spin + spin: + use_spin: + - true + - false + virtual_scale: + - 0.314 + - 0.0 + type: spin_ener +model_def_script: + descriptor: + precision: float64 + repflow: + a_dim: 4 + a_rcut: 3.5 + a_rcut_smth: 0.5 + a_sel: 4 + axis_neuron: 4 + e_dim: 6 + e_rcut: 4.0 + e_rcut_smth: 0.5 + e_sel: 8 + n_dim: 8 + nlayers: 1 + update_angle: false + seed: 1 + type: dpa3 + use_loc_mapping: false + fitting_net: + neuron: + - 5 + - 5 + - 5 + numb_aparam: 1 + resnet_dt: true + seed: 1 + spin: + use_spin: + - true + - false + virtual_scale: + - 0.314 + - 0.0 + type_map: + - Ni + - O +software: deepmd-kit +time: "2026-06-04 01:43:22.387456+00:00" +version: 3.0.0 diff --git a/source/tests/infer/gen_dpa1.py b/source/tests/infer/gen_dpa1.py index 9c743c9f88..55f68b44a0 100644 --- a/source/tests/infer/gen_dpa1.py +++ b/source/tests/infer/gen_dpa1.py @@ -333,6 +333,114 @@ def main(): ) print("\nAll graph sanity checks passed.") # noqa: T201 + + # ============================================================ + # Section C: graph-eligible DPA1 with numb_aparam=2 + # ============================================================ + # Distinct per-COMPONENT aparam values: the C++ graph route once + # under-sized the selected aparam buffer (select_real_atoms_coord + # missing the daparam factor) and the broadcasting ``copy_`` then + # smeared column 0 over all columns -- silent result corruption that a + # numb_aparam=1 fixture can never expose. The C++ gtest + # (test_deeppot_dpa1_graph_aparam_ptexpt.cc) feeds the same aparam + # values and compares against the independent nlist reference below. + aparam_config = copy.deepcopy(graph_config) + aparam_config["fitting_net"] = { + **graph_config["fitting_net"], + "numb_aparam": 2, + } + + print( # noqa: T201 + "\n---- Building graph-eligible DPA1 (attn_layer=0, numb_aparam=2) ----" + ) + model_a = get_model(copy.deepcopy(aparam_config)) + data_a = { + "model": model_a.serialize(), + "model_def_script": aparam_config, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + + # Distinct per-atom, per-component values -- MUST match the C++ gtest. + aparam_vals = (np.arange(1, 13, dtype=np.float64) * 0.1).reshape(1, 6, 2) + + # ---- C.1 Independent nlist reference (same weights, dense path) ---- + aparam_nlist_ref_pt2 = os.path.join( + base_dir, "deeppot_dpa1_graph_aparam_nlist_ref.pt2" + ) + print(f"Exporting reference nlist .pt2 to {aparam_nlist_ref_pt2} ...") # noqa: T201 + pt_expt_deserialize_to_file( + aparam_nlist_ref_pt2, + copy.deepcopy(data_a), + do_atomic_virial=True, + lower_kind="nlist", + ) + dp_ar = DeepPot(aparam_nlist_ref_pt2) + e_a1, f_a1, v_a1, ae_a1, av_a1 = dp_ar.eval( + coord, box, atype, atomic=True, aparam=aparam_vals + ) + e_anp, f_anp, v_anp, ae_anp, av_anp = dp_ar.eval( + coord, None, atype, atomic=True, aparam=aparam_vals + ) + print(f"Aparam nlist ref PBC energy: {e_a1[0, 0]:.18e}") # noqa: T201 + + # Anti-vacuity: swapping the two aparam COLUMNS must change the energy, + # otherwise this fixture cannot catch column-broadcast corruption. + aparam_swapped = aparam_vals[..., ::-1].copy() + e_sw = dp_ar.eval(coord, box, atype, atomic=False, aparam=aparam_swapped)[0] + col_sensitivity = abs(float(e_sw[0, 0]) - float(e_a1[0, 0])) + print(f"Aparam column-swap energy shift: {col_sensitivity:.6e}") # noqa: T201 + if col_sensitivity < 1e-10: + raise RuntimeError( + "aparam columns are degenerate (swap shift " + f"{col_sensitivity:.2e}); the fixture cannot expose " + "column-broadcast bugs." + ) + + # ---- C.2 Sidecar reference file ---- + aparam_ref_path = os.path.join(base_dir, "deeppot_dpa1_graph_aparam.expected") + write_expected_ref( + aparam_ref_path, + sections={ + "pbc": { + "expected_e": ae_a1[0, :, 0], + "expected_f": f_a1[0], + "expected_v": av_a1[0], + }, + "nopbc": { + "expected_e": ae_anp[0, :, 0], + "expected_f": f_anp[0], + "expected_v": av_anp[0], + }, + }, + source_script="source/tests/infer/gen_dpa1.py", + ) + print(f"Wrote {aparam_ref_path}") # noqa: T201 + + # ---- C.3 Export graph-form .pt2 and sanity-check vs nlist ref ---- + aparam_graph_pt2 = os.path.join(base_dir, "deeppot_dpa1_graph_aparam.pt2") + print(f"Exporting to {aparam_graph_pt2} (lower_kind='graph') ...") # noqa: T201 + pt_expt_deserialize_to_file( + aparam_graph_pt2, + copy.deepcopy(data_a), + do_atomic_virial=True, + lower_kind="graph", + ) + dp_ag = DeepPot(aparam_graph_pt2) + e_ag, f_ag, v_ag = dp_ag.eval(coord, box, atype, atomic=False, aparam=aparam_vals)[ + :3 + ] + aparam_force_diff = float(np.max(np.abs(f_ag[0] - f_a1[0]))) + print( # noqa: T201 + f"Aparam graph .pt2 vs nlist ref PBC force max diff: {aparam_force_diff:.2e}" + ) + if aparam_force_diff > 1e-5: + raise RuntimeError( + f"BLOCKED: aparam graph .pt2 PBC force differs from nlist " + f"reference by {aparam_force_diff:.2e} (threshold 1e-5)." + ) + print("\nDone!") # noqa: T201 diff --git a/source/tests/infer/gen_dpa2.py b/source/tests/infer/gen_dpa2.py index 093000bc59..be160f52fd 100644 --- a/source/tests/infer/gen_dpa2.py +++ b/source/tests/infer/gen_dpa2.py @@ -206,6 +206,257 @@ def main(): else: print("\n// Skipping .pth verification (file not generated).") # noqa: T201 + # ============================================================ + # Section B: graph-eligible DPA2 (use_three_body=False) model + # ============================================================ + # Three-body (repinit's optional se_t_tebd sub-block) is graph-ineligible, + # so the graph-form export needs a config with use_three_body=False. + # Everything else stays at the section-A defaults (attn toggles ON in + # repformer -- that's the point: repformer's own attention is graph- + # native, unlike DPA1's se_atten which requires attn_layer=0). + # + # Skip the whole graph section under LeakSanitizer. The C++ memleak matrix + # runs these gen scripts with ``LD_PRELOAD=liblsan`` (the sanitizer- + # instrumented deepmd op .so requires the LSAN runtime; see + # source/install/test_cc_local.sh). Evaluating the AOTInductor-compiled + # graph .pt2's BACKWARD (forces) under that runtime INTERMITTENTLY + # segfaults inside the repformer's fused backward kernel -- an + # AOTI-compiled-code vs LeakSanitizer allocator incompatibility, NOT a + # graph-code bug: the same backward is bit-identical and finite in eager, + # in the non-memleak C++ ctest, and in LAMMPS. dpa1's simpler graph .pt2 + # does not trip it. Leak-checking torch's own compiled kernels is + # meaningless anyway (the memleak build exists to leak-check deepmd's C++ + # ops). The dense DPA2 .pt2 (section A) is unaffected and still generated. + # The C++ dpa2_graph_ptexpt row GTEST_SKIPs when this artifact is absent + # (skip_if_artifact_missing). + # + # Detection is via the explicit DP_GEN_UNDER_SANITIZER flag set by + # test_cc_local.sh next to the preload: sniffing LD_PRELOAD here does NOT + # work -- the LSAN runtime removes its own entry from the process + # environment during startup, so os.environ never sees it (observed on + # CI). The LD_PRELOAD check is kept only as a belt-and-braces fallback for + # manual invocations where the runtime leaves the variable intact. + if ( + os.environ.get("DP_GEN_UNDER_SANITIZER", "") == "lsan" + or "lsan" in os.environ.get("LD_PRELOAD", "").lower() + ): + # remove any graph artifacts left by a previous non-LSAN run of a + # REUSED workspace: skipping regeneration alone leaves them present, + # and the C++ tests' skip_if_artifact_missing would then execute + # them under LSAN and hit the very crash this branch avoids + for name in ( + "deeppot_dpa2_graph_nlist_ref.pt2", + "deeppot_dpa2_graph.pt2", + "deeppot_dpa2_graph.expected", + "deeppot_dpa2_graph_aparam.pt2", + ): + stale = os.path.join(base_dir, name) + if os.path.exists(stale): + os.remove(stale) + print( # noqa: T201 + "\n// Skipping DPA2 graph section under LeakSanitizer " + "(AOTInductor .pt2 backward is incompatible with the LSAN runtime; " + "covered by the non-memleak C++/LAMMPS matrix)." + ) + return + + graph_config = copy.deepcopy(config) + graph_config["descriptor"]["repinit"]["use_three_body"] = False + + print("\n---- Building graph-eligible DPA2 (use_three_body=False) ----") # noqa: T201 + + # ---- B.1 Build dpmodel, serialize ---- + model_g = get_model(copy.deepcopy(graph_config)) + model_dict_g = model_g.serialize() + + data_g = { + "model": copy.deepcopy(model_dict_g), + "model_def_script": graph_config, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + + # ---- B.2 Independent cross-check via nlist .pt2 (dense-quartet) ---- + # Unlike the DPA1 graph fixture, deeppot_dpa2_graph.expected is NOT + # copied from the nlist artifact (see B.5): DPA2 is a message-passing + # model, so the graph path's per-edge full-to-source atomic-virial + # decomposition genuinely differs from the dense per-atom decomposition + # (only the SUM agrees), and cross-path energy/force agreement (~1e-10) + # sits at the universal gtest's double tolerance (1e-10). The nlist + # artifact is instead used here as the independent gen-time oracle: + # atomic energies, forces and the TOTAL virial of the graph .pt2 must + # match it (checked in B.4) or generation aborts. + # + # The nlist .pt2 is PERSISTED (deeppot_dpa2_graph_nlist_ref.pt2): a + # C++ gtest could load it alongside the graph .pt2 to cross-check + # graph≈dense on arbitrary system sizes without baking a second reference + # block into the .expected sidecar. Same weights as the graph model, so + # at non-binding sel the two paths must agree. + nlist_ref_pt2 = os.path.join(base_dir, "deeppot_dpa2_graph_nlist_ref.pt2") + print(f"Exporting reference nlist .pt2 to {nlist_ref_pt2} ...") # noqa: T201 + pt_expt_deserialize_to_file( + nlist_ref_pt2, + copy.deepcopy(data_g), + do_atomic_virial=True, + lower_kind="nlist", # independent: dense nlist, NOT graph + ) + dp_nlist_ref = DeepPot(nlist_ref_pt2) + + # PBC reference from nlist path + e_r1, f_r1, v_r1, ae_r1, av_r1 = dp_nlist_ref.eval(coord, box, atype, atomic=True) + # NoPBC reference from nlist path + e_rnp, f_rnp, v_rnp, ae_rnp, av_rnp = dp_nlist_ref.eval( + coord, None, atype, atomic=True + ) + + print(f"Nlist ref PBC energy: {e_r1[0, 0]:.18e}") # noqa: T201 + print(f"Nlist ref NoPBC energy: {e_rnp[0, 0]:.18e}") # noqa: T201 + max_ref_force_pbc = float(np.max(np.abs(f_r1))) + max_ref_force_nopbc = float(np.max(np.abs(f_rnp))) + print(f"Nlist ref PBC max |force|: {max_ref_force_pbc:.6e}") # noqa: T201 + print(f"Nlist ref NoPBC max |force|: {max_ref_force_nopbc:.6e}") # noqa: T201 + # ``not (x >= th)`` (rather than ``x < th``) so NaN forces -- e.g. from + # an inductor SIMD miscompile of the AOTI artifact -- fail the check + # instead of slipping through. Both the PBC and the NoPBC references are + # checked: each is baked into its own section of the .expected sidecar. + if ( + not (max_ref_force_pbc >= 1e-10) + or not (max_ref_force_nopbc >= 1e-10) + or not (np.all(np.isfinite(f_r1)) and np.all(np.isfinite(f_rnp))) + ): + raise RuntimeError( + f"Graph model nlist-ref forces are degenerate or non-finite " + f"(PBC max={max_ref_force_pbc:.2e}, " + f"NoPBC max={max_ref_force_nopbc:.2e}); weights may need " + f"perturbation, or the AOTI compile is broken (known inductor " + f"CPU-SIMD bug; workaround: torch._inductor.config.cpp.simdlen " + f"= 1)." + ) + + # ---- B.3 Export graph-form .pt2 ---- + # For DPA2, has_message_passing_across_ranks is True, so the + # lower_kind="graph" export automatically embeds the nested with-comm + # artifact (forward_lower_with_comm.pt2) alongside the graph forward -- + # no separate export call is needed. + graph_pt2_path = os.path.join(base_dir, "deeppot_dpa2_graph.pt2") + print(f"Exporting to {graph_pt2_path} (lower_kind='graph') ...") # noqa: T201 + pt_expt_deserialize_to_file( + graph_pt2_path, + copy.deepcopy(data_g), + do_atomic_virial=True, + lower_kind="graph", + ) + print("Graph .pt2 export done.") # noqa: T201 + + # ---- B.4 Cross-check: graph .pt2 vs independent nlist reference ---- + # Both use the SAME weights; at non-binding sel the math is equivalent. + # Atomic energies, forces and the TOTAL virial must agree. The per-atom + # virial is deliberately NOT compared: the graph path assigns each edge's + # virial contribution fully to the source atom, which for a + # message-passing model is a different (equally valid) decomposition + # than the dense path's -- only the sum is convention-independent. + dp_graph = DeepPot(graph_pt2_path) + + e_g1, f_g1, v_g1, ae_g1, av_g1 = dp_graph.eval(coord, box, atype, atomic=True) + e_gnp, f_gnp, v_gnp, ae_gnp, av_gnp = dp_graph.eval(coord, None, atype, atomic=True) + + cross_tol = 1e-8 + for label, (f_g, ae_g, v_g), (f_r, ae_r, v_r) in ( + ("PBC", (f_g1, ae_g1, v_g1), (f_r1, ae_r1, v_r1)), + ("NoPBC", (f_gnp, ae_gnp, v_gnp), (f_rnp, ae_rnp, v_rnp)), + ): + f_diff = float(np.max(np.abs(f_g[0] - f_r[0]))) + ae_diff = float(np.max(np.abs(ae_g[0] - ae_r[0]))) + v_diff = float(np.max(np.abs(v_g[0] - v_r[0]))) + print( # noqa: T201 + f"Graph .pt2 vs nlist ref {label}: ae {ae_diff:.2e}, " + f"f {f_diff:.2e}, total-virial {v_diff:.2e}" + ) + # NaN-safe: NaN fails ``<=`` + if not (f_diff <= cross_tol and ae_diff <= cross_tol and v_diff <= cross_tol): + raise RuntimeError( + f"BLOCKED: graph .pt2 {label} differs from nlist reference " + f"(ae {ae_diff:.2e}, f {f_diff:.2e}, v {v_diff:.2e}; " + f"threshold {cross_tol:.0e})." + ) + + # ---- B.5 Write sidecar reference file from the graph .pt2 eval ---- + # Self-referential like the dense fixtures' .expected (the C++ gtest is a + # regression test of the C++ inference path against the Python eval of + # the same artifact); independence from the graph path is enforced above + # in B.4. Sourcing e/f/v from the nlist artifact instead would break the + # per-atom virial comparison (convention, see B.4) and sit at the gtest's + # 1e-10 double tolerance for energies/forces (cross-path noise ~1e-10). + graph_ref_path = os.path.join(base_dir, "deeppot_dpa2_graph.expected") + write_expected_ref( + graph_ref_path, + sections={ + "pbc": { + "expected_e": ae_g1[0, :, 0], + "expected_f": f_g1[0], + "expected_v": av_g1[0], + }, + "nopbc": { + "expected_e": ae_gnp[0, :, 0], + "expected_f": f_gnp[0], + "expected_v": av_gnp[0], + }, + }, + source_script="source/tests/infer/gen_dpa2.py", + ) + print(f"Wrote {graph_ref_path}") # noqa: T201 + + print("\nAll graph sanity checks passed.") # noqa: T201 + + # ============================================================ + # Section C: graph-eligible DPA2 with numb_aparam=1 + # ============================================================ + # Consumed by test_lammps_dpa2_graph_pt2.py's empty-subdomain aparam + # regression: a ghost-only rank (nlocal == 0, nghost > 0) must + # SYNTHESIZE a zero aparam covering the graph's ghost nodes instead of + # feeding the empty owned tensor to the artifact -- the rank-local + # reshape failure would otherwise desynchronize the per-layer border_op + # collective sequence and hang the peers. No .expected sidecar: the + # LAMMPS test is an MP == SP parity check on this same artifact. + aparam_config = copy.deepcopy(graph_config) + aparam_config["fitting_net"] = { + **config["fitting_net"], + "numb_aparam": 1, + } + + print( # noqa: T201 + "\n---- Building graph-eligible DPA2 (use_three_body=False, numb_aparam=1) ----" + ) + model_a = get_model(copy.deepcopy(aparam_config)) + data_a = { + "model": model_a.serialize(), + "model_def_script": aparam_config, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + aparam_graph_pt2 = os.path.join(base_dir, "deeppot_dpa2_graph_aparam.pt2") + print(f"Exporting to {aparam_graph_pt2} (lower_kind='graph') ...") # noqa: T201 + pt_expt_deserialize_to_file( + aparam_graph_pt2, + copy.deepcopy(data_a), + do_atomic_virial=True, + lower_kind="graph", + ) + # Sanity: single-rank eval with a uniform aparam must be finite, and the + # aparam must genuinely reach the fitting. + dp_ag = DeepPot(aparam_graph_pt2) + natoms_c = len(atype) + aparam_c = np.full((1, natoms_c, 1), 0.25, dtype=np.float64) + e_c, f_c, v_c = dp_ag.eval(coord, box, atype, atomic=False, aparam=aparam_c)[:3] + e_c2 = dp_ag.eval(coord, box, atype, atomic=False, aparam=aparam_c + 1.0)[0] + if not (np.all(np.isfinite(f_c)) and abs(e_c2[0, 0] - e_c[0, 0]) > 1e-12): + raise RuntimeError( + "BLOCKED: dpa2 graph aparam artifact is degenerate (non-finite " + "forces or aparam-insensitive energy)." + ) + print("\nDone!") # noqa: T201 diff --git a/source/tests/pt/test_compile_compat.py b/source/tests/pt/test_compile_compat.py new file mode 100644 index 0000000000..f076a5cba5 --- /dev/null +++ b/source/tests/pt/test_compile_compat.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""forbidden_dims_from_model accessor handling (CodeRabbit #5779).""" + +import torch + +from deepmd.pt.utils.compile_compat import ( + forbidden_dims_from_model, +) + + +class _WithDims(torch.nn.Module): + def get_dim_fparam(self) -> int: + return 5 + + def get_dim_aparam(self) -> int: + return 3 + + +class TestForbiddenDimsFromModel: + def test_dims_collected_when_accessors_present(self) -> None: + forbidden = forbidden_dims_from_model(_WithDims(), []) + assert {3, 5} <= forbidden + + def test_missing_accessors_fall_through_best_effort(self) -> None: + # a bare Module lacks get_dim_fparam/get_dim_aparam: the lookup must + # happen inside the try (an eagerly-built accessor tuple raised + # AttributeError before the best-effort guard could catch it) + assert forbidden_dims_from_model(torch.nn.Module(), []) == set() diff --git a/source/tests/pt_expt/descriptor/test_dpa1_triton.py b/source/tests/pt_expt/descriptor/test_dpa1_triton.py index f3dbdea931..ecb08f2516 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1_triton.py +++ b/source/tests/pt_expt/descriptor/test_dpa1_triton.py @@ -53,6 +53,7 @@ def _build_dpa1_expt( precision="float64", tebd_input_mode="concat", smooth=True, + exclude_types=(), ): from deepmd.pt_expt.descriptor.dpa1 import ( DescrptDPA1, @@ -72,6 +73,7 @@ def _build_dpa1_expt( activation_function=act, precision=precision, smooth_type_embedding=smooth, + exclude_types=list(exclude_types), seed=1, ).to(device) des.eval() @@ -346,6 +348,21 @@ def test_parity_concat_identity(self) -> None: def test_parity_concat_nonpow2(self) -> None: self._assert_parity(_build_dpa1_expt(self.device, [12, 25, 50])) + def test_parity_concat_exclude_types(self) -> None: + # descriptor-level exclude_types is served by the dense fused path + # (the emask masks nlist / sw / rr in ``_env_mat``); the eligibility + # predicate must not force it back to the reference. + des_ex = _build_dpa1_expt(self.device, [8, 16, 32], exclude_types=[(0, 1)]) + self._assert_parity(des_ex) + # anti-vacuity: the exclusion visibly changes the descriptor (same + # seed as the unexcluded twin), so the parity above compared two + # implementations that BOTH applied it. + des_no = _build_dpa1_expt(self.device, [8, 16, 32]) + ec, ea, nl = self._dense_inputs(des_ex) + self.assertFalse( + torch.allclose(des_ex.call(ec, ea, nl)[0], des_no.call(ec, ea, nl)[0]) + ) + def test_make_fx_bakes_env_mat(self) -> None: des = _build_dpa1_expt(self.device, [8, 16, 32]) ec, ea, nl = self._dense_inputs(des) @@ -366,5 +383,123 @@ def fn(c): self.assertGreaterEqual(sum("env_mat.default" in t for t in targets), 1) +class TestDpa1TritonExcludeTypesRouting(unittest.TestCase): + """Descriptor-level ``exclude_types`` routing is per dispatch site. + + The eligibility predicate no longer carries an ``exclude_types`` veto: + the dense ``_call_triton`` applies the descriptor emask itself (in + ``_env_mat``) and stays fused, while the graph ``call_graph`` gates the + fused edge kernel out (``_call_graph_triton`` bypasses the reference + ``apply_pair_exclusion``) and serves exclusions from the dpmodel + reference. Routing is device-free, so this runs on CPU with + ``TRITON_AVAILABLE`` patched; kernel-level parity is covered by the + CUDA classes above. + """ + + def setUp(self) -> None: + # single import style for the module (CodeQL py/import-and-import-from) + from deepmd.pt_expt.descriptor import dpa1 as dpa1_mod + + self._dpa1_mod = dpa1_mod + self._saved_avail = dpa1_mod.TRITON_AVAILABLE + self._saved_level = os.environ.get("DP_TRITON_INFER") + dpa1_mod.TRITON_AVAILABLE = True + os.environ["DP_TRITON_INFER"] = "1" + self.device = torch.device("cpu") + gen = torch.Generator(device=self.device).manual_seed(3) + n = 24 + self.coord = (torch.rand(1, n, 3, generator=gen, device=self.device) * 8.0).to( + torch.float64 + ) + self.atype = torch.randint(0, 2, (1, n), generator=gen, device=self.device) + self.box = ( + (torch.eye(3, device=self.device) * 8.0).reshape(1, 9).to(torch.float64) + ) + + def tearDown(self) -> None: + self._dpa1_mod.TRITON_AVAILABLE = self._saved_avail + if self._saved_level is None: + os.environ.pop("DP_TRITON_INFER", None) + else: + os.environ["DP_TRITON_INFER"] = self._saved_level + + def _quartet(self, des): + return extend_input_and_build_neighbor_list( + self.coord, + self.atype, + des.get_rcut(), + des.get_sel(), + mixed_types=des.mixed_types(), + box=self.box, + ) + + def _graph(self, des): + from deepmd.dpmodel.utils.neighbor_graph import ( + from_dense_quartet, + ) + + ec, ea, mp, nl = self._quartet(des) + return from_dense_quartet(ec, nl, mp, compact=True), self.atype.reshape(-1) + + def test_predicate_admits_exclude_types(self) -> None: + des = _build_dpa1_expt(self.device, [8, 16, 32], exclude_types=[(0, 1)]) + self.assertTrue(des._fused_eligible("triton")) + + def test_dense_dispatch_serves_exclude_types(self) -> None: + from deepmd.dpmodel.descriptor.dpa1 import DescrptDPA1 as DescrptDPA1DP + + des = _build_dpa1_expt(self.device, [8, 16, 32], exclude_types=[(0, 1)]) + calls = [] + + def recording_triton(coord_ext, atype_ext, nlist): + calls.append(1) + # delegate to the reference body (no Triton on this box); the + # dispatch decision is what is under test + return DescrptDPA1DP.call.__wrapped__(des, coord_ext, atype_ext, nlist) + + des._call_triton = recording_triton + ec, ea, _mp, nl = self._quartet(des) + des.call(ec, ea, nl) + self.assertEqual(len(calls), 1) + + def test_graph_dispatch_exclude_types_stays_reference(self) -> None: + from deepmd.dpmodel.descriptor.dpa1 import DescrptDPA1 as DescrptDPA1DP + + des = _build_dpa1_expt(self.device, [8, 16, 32], exclude_types=[(0, 1)]) + + def forbidden_triton(graph, atype, type_embedding): + raise AssertionError( + "graph triton kernel must not serve descriptor exclude_types" + ) + + des._call_graph_triton = forbidden_triton + tebd = des.type_embedding.call() + graph, atype_local = self._graph(des) + grrg, _rot = des.call_graph(graph, atype_local, type_embedding=tebd) + # the reference path applied apply_pair_exclusion + ref, _rot_ref = DescrptDPA1DP.call_graph( + des, graph, atype_local, type_embedding=tebd + ) + torch.testing.assert_close(grrg, ref, atol=0.0, rtol=0.0) + + def test_graph_dispatch_no_exclusions_takes_triton(self) -> None: + from deepmd.dpmodel.descriptor.dpa1 import DescrptDPA1 as DescrptDPA1DP + + des = _build_dpa1_expt(self.device, [8, 16, 32]) + calls = [] + + def recording_triton(graph, atype, type_embedding): + calls.append(1) + return DescrptDPA1DP.call_graph( + des, graph, atype, type_embedding=type_embedding + ) + + des._call_graph_triton = recording_triton + tebd = des.type_embedding.call() + graph, atype_local = self._graph(des) + des.call_graph(graph, atype_local, type_embedding=tebd) + self.assertEqual(len(calls), 1) + + if __name__ == "__main__": unittest.main() diff --git a/source/tests/pt_expt/descriptor/test_dpa2.py b/source/tests/pt_expt/descriptor/test_dpa2.py index 217bcdb230..371d42655e 100644 --- a/source/tests/pt_expt/descriptor/test_dpa2.py +++ b/source/tests/pt_expt/descriptor/test_dpa2.py @@ -303,6 +303,15 @@ def test_compressed_forward(self, prec) -> None: dd0.repinit.stddev = torch.tensor(dstd, dtype=dtype, device=self.device) dd0.repformers.mean = torch.tensor(davg_2, dtype=dtype, device=self.device) dd0.repformers.stddev = torch.tensor(dstd_2, dtype=dtype, device=self.device) + # This test checks compression accuracy against the legacy DENSE + # body specifically (compression itself is graph-ineligible, see + # DescrptDPA2.uses_graph_lower). Force the dense route for the + # "uncompressed" reference forward too: this config is otherwise + # graph-eligible (strip tebd) and `mapping` is supplied, so it would + # silently take the graph-native adapter route instead, which is + # NOT what this test intends to exercise -- disable it explicitly + # rather than relying on incidental non-eligibility. + dd0.disable_graph_lower() coord_ext = torch.tensor(self.coord_ext, dtype=dtype, device=self.device) atype_ext = torch.tensor(self.atype_ext, dtype=int, device=self.device) diff --git a/source/tests/pt_expt/descriptor/test_repformer_parallel.py b/source/tests/pt_expt/descriptor/test_repformer_parallel.py index 1a6413d08f..4d4ab9d614 100644 --- a/source/tests/pt_expt/descriptor/test_repformer_parallel.py +++ b/source/tests/pt_expt/descriptor/test_repformer_parallel.py @@ -162,6 +162,14 @@ def test_parallel_matches_default( type_map=None, seed=GLOBAL_SEED, ).to(self.device) + # Pin the dense route: this test compares the dense comm_dict/border_op + # exchange against the dense mapping-gather default. Since the graph + # adapter became the default route for graph-eligible configs, the + # nonzero-davg fixture would hit the documented dense padding-residual + # divergence (see DescrptDPA2.call_graph Notes) instead of testing + # the parallel exchange. + dd.disable_graph_lower() + dd.repinit.mean = torch.tensor(davg, dtype=dtype, device=self.device) dd.repinit.stddev = torch.tensor(dstd, dtype=dtype, device=self.device) dd.repformers.mean = torch.tensor(davg_2, dtype=dtype, device=self.device) diff --git a/source/tests/pt_expt/infer/test_graph_deepeval.py b/source/tests/pt_expt/infer/test_graph_deepeval.py index 932fb25cfb..634f0dd61c 100644 --- a/source/tests/pt_expt/infer/test_graph_deepeval.py +++ b/source/tests/pt_expt/infer/test_graph_deepeval.py @@ -291,3 +291,66 @@ def test_graph_pt2_single_atom_no_edges(graph_pt2) -> None: e.reshape(-1), ref["energy"].reshape(-1), rtol=1e-10, atol=1e-10 ) np.testing.assert_allclose(f.reshape(-1), 0.0, atol=1e-12) + + +def test_graph_pt2_deepeval_aparam(tmp_path) -> None: + """Graph ``.pt2`` DeepEval with ``numb_aparam > 0``. + + DeepEval receives the user-facing rectangular ``(nf, natoms, nda)`` + aparam and must flatten it to the graph ABI's flat ``(N, nda)`` node + axis before feeding the artifact (regression for the aparam + graph-freeze review round). Checks parity vs the eager dense reference + with the same aparam, and that aparam genuinely reaches the fitting. + """ + from deepmd.pt_expt.model import ( + get_model, + ) + + config = copy.deepcopy(DPA1_CONFIG) + config["descriptor"]["smooth_type_embedding"] = False + config["fitting_net"] = {**config["fitting_net"], "numb_aparam": 1} + model = get_model(config).to(torch.float64) + model.eval() + pt2_path = str(tmp_path / "deeppot_dpa1_graph_aparam.pt2") + deserialize_to_file( + pt2_path, + {"model": model.serialize()}, + do_atomic_virial=True, + lower_kind="graph", + ) + + coords, cells, atype = _build_system(**_SYSTEMS["small_8"]) + natoms = atype.shape[0] + aparam = np.linspace(0.1, 0.9, natoms).reshape(1, natoms, 1) + + dp = DeepPot(pt2_path) + assert dp.deep_eval.metadata["lower_input_kind"] == "graph" + e, f, v = dp.eval(coords, cells, atype, atomic=False, aparam=aparam)[:3] + + # aparam must genuinely reach the fitting through the flat node axis + e_bump = dp.eval(coords, cells, atype, atomic=False, aparam=aparam + 1.0)[0] + assert not np.allclose(e_bump, e), "aparam bump must change the energy" + + # parity vs the eager dense (sel-capped, non-binding) reference + coord_t = torch.tensor( + coords.reshape(1, natoms, 3), dtype=torch.float64, device=DEVICE + ).requires_grad_(True) + atype_t = torch.tensor(atype.reshape(1, natoms), dtype=torch.int64, device=DEVICE) + box_t = torch.tensor(cells.reshape(1, 9), dtype=torch.float64, device=DEVICE) + ap_t = torch.tensor( + aparam.reshape(1, natoms, 1), dtype=torch.float64, device=DEVICE + ) + ret = model.call_common( + coord_t, + atype_t, + box_t, + aparam=ap_t, + neighbor_graph_method="legacy", + ) + np.testing.assert_allclose( + e.reshape(-1), + ret["energy_redu"].detach().cpu().numpy().reshape(-1), + rtol=1e-10, + atol=1e-10, + err_msg="energy (graph .pt2 + aparam vs eager dense)", + ) diff --git a/source/tests/pt_expt/model/test_dos_graph.py b/source/tests/pt_expt/model/test_dos_graph.py index dc8cfc0685..adfc71d465 100644 --- a/source/tests/pt_expt/model/test_dos_graph.py +++ b/source/tests/pt_expt/model/test_dos_graph.py @@ -5,7 +5,8 @@ routes into the graph path by default. Before the general output transform this KeyError'd on ``"energy"``; now every fitting (dos/dipole/polar/property/...) flows through :func:`fit_output_to_model_output_graph` with no change on the -fitting side. Each model's graph forward (default ``neighbor_graph_method``) +fitting side. Each model's graph forward (EXPLICIT ``neighbor_graph_method`` +opt-in; non-energy models default to dense since the energy gate) must match the dense path (``neighbor_graph_method="legacy"``) on every shared key (carry-all graph at non-binding ``sel`` reproduces the dense neighbor set). """ @@ -129,8 +130,14 @@ def test_dos_repro(self) -> None: [_make_dos, _make_dipole, _make_polar, _make_property], ) # one builder per fitting kind def test_graph_matches_dense(self, make_model) -> None: - """Graph (default) output matches the dense (``legacy``) path on every + """EXPLICIT graph opt-in matches the dense (``legacy``) path on every shared key, including derivatives for r/c-differentiable fittings. + + Non-energy models no longer DEFAULT-flip to the graph route (the + compiled-training trace is energy-specific; see the + ``_resolve_graph_method`` energy gate), but the eager graph lower + stays output-agnostic and available via an explicit + ``neighbor_graph_method`` -- which is what this parity test pins. """ tol = ( {"rtol": 1e-11, "atol": 1e-11} @@ -139,7 +146,13 @@ def test_graph_matches_dense(self, make_model) -> None: ) ds = _make_descriptor() m = make_model(ds) - graph = m.call_common(self.coord, self.atype, None, do_atomic_virial=True) + graph = m.call_common( + self.coord, + self.atype, + None, + do_atomic_virial=True, + neighbor_graph_method="dense", + ) # the dense path differentiates w.r.t. coord -> needs a coord leaf. dense = m.call_common( self.coord.detach().requires_grad_(True), diff --git a/source/tests/pt_expt/model/test_dpa2_graph_lower.py b/source/tests/pt_expt/model/test_dpa2_graph_lower.py new file mode 100644 index 0000000000..6e1a08481d --- /dev/null +++ b/source/tests/pt_expt/model/test_dpa2_graph_lower.py @@ -0,0 +1,1362 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""pt_expt DPA2 graph-lower Task 8 tests. + +The pt_expt model plumbing that routes ``forward_common`` through the +carry-all graph (default-flip ``_resolve_graph_method``, autograd force/ +virial via ``forward_common_lower_graph``, compiled training +``_trace_and_compile_graph``, the freeze gate) is GENERIC: it keys off +``descriptor.uses_graph_lower()`` and was implemented once (Task 7). DPA2 +inherits it with ZERO new plumbing code -- these tests PROVE that +inheritance (routing/parity), plus cover the one piece of real code this +task adds: ``DescrptBlockRepformers._exchange_ghosts_graph``, the pt_expt +override that performs the per-layer MPI ghost refresh via +``deepmd_export::border_op`` on the graph path (see +``deepmd/pt_expt/descriptor/repformers.py``). +""" + +import ctypes +import importlib + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.descriptor.dpa2 import ( + RepformerArgs, + RepinitArgs, +) +from deepmd.pt_expt.descriptor.dpa2 import ( + DescrptDPA2, +) +from deepmd.pt_expt.descriptor.repformers import ( + DescrptBlockRepformers, +) +from deepmd.pt_expt.fitting import ( + InvarFitting, +) +from deepmd.pt_expt.model import ( + EnergyModel, +) +from deepmd.pt_expt.utils import ( + env, +) + +from ...seed import ( + GLOBAL_SEED, +) + +# Trigger registration of the deepmd_export::border_op opaque wrapper +# (importlib form: the module is imported purely for its side effect). +importlib.import_module("deepmd.pt_expt.utils.comm") + +# --------------------------------------------------------------------------- +# Self-comm-dict helper (mirrors +# ``source/tests/pt_expt/descriptor/test_repflow_parallel.py``): builds a +# single-rank, self-only MPI ``comm_dict`` whose effect is a plain gather +# from the sendlist-indexed local rows into the ghost slots, so the +# ``border_op`` self-send branch (no MPI runtime needed) can be exercised +# eagerly. + + +def _addr_of(np_arr: np.ndarray) -> int: + """Return the raw int address of a numpy array's data buffer.""" + return np_arr.ctypes.data_as(ctypes.c_void_p).value + + +def _build_self_comm_dict( + *, + nloc: int, + nghost: int, + sendlist_indices: np.ndarray, + keepalive: list, +) -> dict: + """Build a comm_dict for a single-rank self-exchange. + + ``sendlist_indices`` (int32, length ``nghost``) gives the local row to + copy into each successive ghost slot ``[nloc, nloc + nghost)``. Control + tensors are forced to CPU: the C++ ``border_op`` host-side code + dereferences ``data_ptr()`` directly. + """ + sendlist_indices = np.ascontiguousarray(sendlist_indices, dtype=np.int32) + keepalive.append(sendlist_indices) + addr = _addr_of(sendlist_indices) + return { + "send_list": torch.tensor([addr], dtype=torch.int64, device="cpu"), + "send_proc": torch.zeros(1, dtype=torch.int32, device="cpu"), + "recv_proc": torch.zeros(1, dtype=torch.int32, device="cpu"), + "send_num": torch.tensor([nghost], dtype=torch.int32, device="cpu"), + "recv_num": torch.tensor([nghost], dtype=torch.int32, device="cpu"), + "communicator": torch.zeros(1, dtype=torch.int64, device="cpu"), + "nlocal": torch.tensor(nloc, dtype=torch.int32, device="cpu"), + "nghost": torch.tensor(nghost, dtype=torch.int32, device="cpu"), + } + + +def _make_dpa2_descriptor( + ntypes: int = 2, + repinit_rcut: float = 4.0, + repinit_nsel: int = 200, + repformer_rcut: float = 2.0, + repformer_nsel: int = 150, + repformer_attn: bool = False, + nlayers: int = 2, + g1_dim: int = 8, + g2_dim: int = 4, +) -> DescrptDPA2: + """Small graph-eligible DPA2 descriptor (use_three_body=False, concat + tebd) with configurable sel and attention toggles. + + ``repformer_attn=False`` (default) keeps ``update_g1_has_attn`` / + ``update_g2_has_attn`` off: with attention on, the carry-all graph's + attention is sel-independent BY DESIGN (NeighborGraph PR-D) and + diverges from the dense body even at non-binding sel, so exact + graph-vs-dense parity requires attention off (dpa1 + ``test_force_virial_parity_vs_legacy`` precedent, and the dpmodel + ``TestDPA2ModelEnergyCarryAll`` fixture). + """ + repinit = RepinitArgs( + rcut=repinit_rcut, + rcut_smth=0.5, + nsel=repinit_nsel, + neuron=[3, 6], + axis_neuron=2, + tebd_dim=4, + ) + repformer = RepformerArgs( + rcut=repformer_rcut, + rcut_smth=0.5, + nsel=repformer_nsel, + nlayers=nlayers, + g1_dim=g1_dim, + g2_dim=g2_dim, + axis_neuron=2, + update_g1_has_attn=repformer_attn, + update_g2_has_attn=repformer_attn, + attn1_hidden=4, + attn1_nhead=2, + attn2_hidden=4, + attn2_nhead=2, + ) + return DescrptDPA2( + ntypes=ntypes, + repinit=repinit, + repformer=repformer, + precision="float64", + seed=GLOBAL_SEED, + ) + + +class TestDpa2GraphLower: + def setup_method(self) -> None: + self.device = env.DEVICE + self.natoms = 6 + self.nt = 2 + self.type_map = ["foo", "bar"] + + generator = torch.Generator(device=self.device).manual_seed(GLOBAL_SEED) + cell = torch.rand( + [3, 3], dtype=torch.float64, device=self.device, generator=generator + ) + cell = (cell + cell.T) + 5.0 * torch.eye(3, device=self.device) + self.cell = cell.unsqueeze(0) # [1, 3, 3] + coord = torch.rand( + [self.natoms, 3], + dtype=torch.float64, + device=self.device, + generator=generator, + ) + coord = torch.matmul(coord, cell) + self.coord = coord.unsqueeze(0).to(self.device) # [1, natoms, 3] + self.atype = torch.tensor( + [[0, 0, 0, 1, 1, 1]], dtype=torch.int64, device=self.device + ) + + def _make_model( + self, + repformer_nsel: int = 150, + repinit_nsel: int = 200, + repformer_attn: bool = False, + numb_fparam: int = 0, + ) -> EnergyModel: + ds = _make_dpa2_descriptor( + ntypes=self.nt, + repinit_nsel=repinit_nsel, + repformer_nsel=repformer_nsel, + repformer_attn=repformer_attn, + ).to(self.device) + ft = InvarFitting( + "energy", + self.nt, + ds.get_dim_out(), + 1, + mixed_types=ds.mixed_types(), + numb_fparam=numb_fparam, + precision="float64", + seed=GLOBAL_SEED, + ).to(self.device) + return EnergyModel(ds, ft, type_map=self.type_map).to(self.device) + + # ------------------------------------------------------------------ + # 1. routing/parity: proves the Task-7 generic plumbing works for DPA2. + # ------------------------------------------------------------------ + def test_force_virial_parity_vs_legacy(self) -> None: + """Default-flip ``forward_common`` (graph) matches + ``neighbor_graph_method="legacy"`` (dense) bit-tight at non-binding + sel with repformer attention off. + """ + model = self._make_model() # non-binding sel, attention off + model.eval() + box = self.cell.reshape(1, 9) + + graph = model.forward_common( + self.coord.clone().requires_grad_(True), self.atype, box + ) + legacy = model.forward_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="legacy", + ) + tol = {"rtol": 1e-10, "atol": 1e-10} + torch.testing.assert_close(graph["energy_redu"], legacy["energy_redu"], **tol) + torch.testing.assert_close( + graph["energy_derv_r"], legacy["energy_derv_r"], **tol + ) + torch.testing.assert_close( + graph["energy_derv_c_redu"], legacy["energy_derv_c_redu"], **tol + ) + + def test_graph_native_hessian_matches_dense(self) -> None: + """The graph route computes the Hessian natively and it matches the + dense route bit-tight. + + A Hessian-enabled graph-eligible DPA2 model default-flips to the graph + route; ``_call_common_graph`` now runs its own ``_cal_hessian_ext_graph`` + loop (differentiating the reduced energy w.r.t. the LOCAL coords by + rebuilding the carry-all graph inside the autograd wrapper), so + ``energy_derv_r_derv_r`` is produced without falling back to dense. + The result must equal the dense route (forced via the + ``disable_graph_lower`` escape hatch on the same weights) at fp64 + parity, in the same ``(nf, 1, nloc*3, nloc*3)`` layout. + + A non-binding repformer sel is used so graph and dense agree on the + first derivatives too (attention off -- the fixture default); the + Hessian parity would otherwise inherit the binding-sel divergence. + """ + from deepmd.pt_expt.model.graph_lower import ( + model_uses_graph_lower as _model_uses_graph_lower, + ) + + model = self._make_model() # non-binding sel, attention off + model.eval() + assert model.atomic_model.descriptor.uses_graph_lower() is True + model.enable_hessian() + # a Hessian model stays graph-routed: the graph now produces the Hessian. + assert _model_uses_graph_lower(model) is True + + box = self.cell.reshape(1, 9) + # ``EnergyModel.forward`` exposes the Hessian under the ``"hessian"`` + # key (translated from ``energy_derv_r_derv_r``); it would KeyError if + # the graph route failed to produce it. + graph_out = model.forward( + self.coord.clone().requires_grad_(True), self.atype, box=box + ) + assert "hessian" in graph_out + assert graph_out["hessian"].shape == (1, self.natoms * 3, self.natoms * 3) + + # dense reference on the SAME weights via the escape hatch. + ref_model = self._make_model() + ref_model.eval() + ref_model.load_state_dict(model.state_dict(), strict=False) + ref_model.atomic_model.descriptor.disable_graph_lower() + ref_model.enable_hessian() + assert _model_uses_graph_lower(ref_model) is False + dense_out = ref_model.forward( + self.coord.clone().requires_grad_(True), self.atype, box=box + ) + torch.testing.assert_close( + graph_out["hessian"], dense_out["hessian"], rtol=1e-9, atol=1e-9 + ) + # Hessian symmetry (a genuine second derivative, not a shape artifact). + h = graph_out["hessian"][0] + torch.testing.assert_close(h, h.transpose(-1, -2), rtol=1e-9, atol=1e-9) + + def test_disable_graph_lower_escape_hatch(self) -> None: + """``descriptor.disable_graph_lower()`` is the documented legacy-dense + escape hatch: it flips ``uses_graph_lower()`` to ``False`` so the + default (``neighbor_graph_method=None``) eager forward takes the DENSE + route -- bit-identical to explicitly requesting ``"legacy"``. + + Uses a BINDING repformer sel so graph and dense genuinely differ: with + the hatch engaged, the default must match dense (not graph), which a + binding sel makes an unambiguous (non-bit-tight) distinction. + """ + generator = torch.Generator(device=self.device).manual_seed(GLOBAL_SEED) + nloc = 12 + box_size = 3.0 + coord = ( + torch.rand( + [nloc, 3], dtype=torch.float64, device=self.device, generator=generator + ) + * box_size + ).unsqueeze(0) + atype = torch.tensor( + [[ii % self.nt for ii in range(nloc)]], + dtype=torch.int64, + device=self.device, + ) + box = ( + torch.eye(3, dtype=torch.float64, device=self.device) * box_size + ).reshape(1, 9) + + model = self._make_model(repformer_nsel=3, repformer_attn=True) + model.eval() + assert model.atomic_model.descriptor.uses_graph_lower() is True + + # engage the escape hatch + model.atomic_model.descriptor.disable_graph_lower() + assert model.atomic_model.descriptor.uses_graph_lower() is False + + default_after = model.forward_common( + coord.clone().requires_grad_(True), atype, box + ) + legacy = model.forward_common( + coord.clone().requires_grad_(True), + atype, + box, + neighbor_graph_method="legacy", + ) + # With the hatch on, default (None) takes the same dense path as legacy, + # so the energy is bit-identical. The force goes through a scatter/index_add + # in the backward, which is non-deterministic on CUDA across two separate + # passes, so it matches only to reduction-order noise (~1e-18), far below + # any real route divergence (>=1e-8, cf. test_binding_sel_diverges). + torch.testing.assert_close( + default_after["energy_redu"], legacy["energy_redu"], rtol=0, atol=0 + ) + torch.testing.assert_close( + default_after["energy_derv_r"], + legacy["energy_derv_r"], + rtol=1e-10, + atol=1e-12, + ) + + def test_binding_sel_diverges(self) -> None: + """At binding repformer sel, the carry-all graph (sel-independent) + keeps neighbors the dense body truncates, so the two routes diverge. + + Uses a dense cluster (small box, more atoms) so each atom's real + neighbor count within ``repformer_rcut`` comfortably exceeds + ``repformer_nsel`` (dpmodel ``test_dpa2_call_graph.py`` binding-sel + precedent: box_size=3.0, nloc=12, repformer_nsel=3). + """ + generator = torch.Generator(device=self.device).manual_seed(GLOBAL_SEED) + nloc = 12 + box_size = 3.0 + coord = ( + torch.rand( + [nloc, 3], dtype=torch.float64, device=self.device, generator=generator + ) + * box_size + ) + coord = coord.unsqueeze(0) # [1, nloc, 3] + atype = torch.tensor( + [[ii % self.nt for ii in range(nloc)]], + dtype=torch.int64, + device=self.device, + ) + box = ( + torch.eye(3, dtype=torch.float64, device=self.device) * box_size + ).reshape(1, 9) + + model = self._make_model( + repformer_nsel=3 + ) # repinit_nsel stays 200 (non-binding) + model.eval() + + graph = model.forward_common(coord.clone().requires_grad_(True), atype, box) + legacy = model.forward_common( + coord.clone().requires_grad_(True), + atype, + box, + neighbor_graph_method="legacy", + ) + e_diff = (graph["energy_redu"] - legacy["energy_redu"]).abs().max().item() + e_scale = legacy["energy_redu"].abs().max().item() + # non-zero (proves the default routed to the GRAPH, not a silent dense + # fallback -- a fallback would make this exactly 0) AND bounded (the + # divergence is the carry-all's extra in-cutoff neighbors, not a blow-up) + assert e_diff > 1e-8, f"expected binding-sel divergence, got {e_diff:.3e}" + assert e_diff < e_scale, ( + f"binding-sel divergence {e_diff:.3e} must stay below the energy " + f"scale {e_scale:.3e} (bounded, not a blow-up)" + ) + + def test_binding_sel_diverges_with_attention(self) -> None: + """Same binding-sel graph-vs-dense divergence as + :meth:`test_binding_sel_diverges`, but with BOTH repformer attention + channels ON -- the config iProzd's review asks for. + + With attention enabled and the fixed-phantom-count compensation, the + two routes are bit-tight at NON-binding sel (see + ``test_neighbor_list.py::test_default_fallback[dpa2]``). At BINDING sel + they must still diverge, because the carry-all graph attends over + neighbors the dense body truncates -- a difference that is both + non-zero (the default IS the graph route) and bounded (below the energy + scale). If the default silently fell back to dense, the difference + would be exactly zero even with attention on. + """ + generator = torch.Generator(device=self.device).manual_seed(GLOBAL_SEED) + nloc = 12 + box_size = 3.0 + coord = ( + torch.rand( + [nloc, 3], dtype=torch.float64, device=self.device, generator=generator + ) + * box_size + ).unsqueeze(0) + atype = torch.tensor( + [[ii % self.nt for ii in range(nloc)]], + dtype=torch.int64, + device=self.device, + ) + box = ( + torch.eye(3, dtype=torch.float64, device=self.device) * box_size + ).reshape(1, 9) + + model = self._make_model(repformer_nsel=3, repformer_attn=True) + model.eval() + + graph = model.forward_common(coord.clone().requires_grad_(True), atype, box) + legacy = model.forward_common( + coord.clone().requires_grad_(True), + atype, + box, + neighbor_graph_method="legacy", + ) + for key in ("energy_redu", "energy_derv_r"): + diff = (graph[key] - legacy[key]).abs().max().item() + scale = legacy[key].abs().max().item() + assert diff > 1e-8, ( + f"expected binding-sel divergence in {key} with attention on, " + f"got {diff:.3e} (silent dense fallback?)" + ) + assert diff < max(scale, 1.0), ( + f"{key} binding-sel divergence {diff:.3e} must stay bounded " + f"below the scale {scale:.3e}" + ) + + def test_graph_lower_symbolic_trace(self) -> None: + """``make_fx`` symbolic trace of ``forward_common_lower_graph`` + reproduces the eager graph lower bit-tight. ``model.to("cpu")`` + before tracing mirrors the real ``.pt2`` export path and the dpa1 + CUDA lesson (traced inputs and params must share a device). + """ + from deepmd.pt_expt.utils.serialization import ( + build_synthetic_graph_inputs, + ) + + model = self._make_model().to("cpu") + model.eval() + sample = build_synthetic_graph_inputs( + model, + e_max=175, + nframes=2, + nloc=7, + dtype=torch.float64, + device=torch.device("cpu"), + ) + atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs = sample + traced = model.forward_lower_graph_exportable( + atype, + n_node, + nl, + ei, + ev, + em, + do, + drp, + so, + srp, + fparam=fp, + aparam=ap, + do_atomic_virial=True, + charge_spin=cs, + destination_sorted=True, + tracing_mode="symbolic", + _allow_non_fake_inputs=True, + ) + out = traced(atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs) + ref = model.forward_common_lower_graph( + atype, + n_node, + nl, + ei, + ev, + em, + do, + drp, + so, + srp, + destination_sorted=True, + fparam=fp, + aparam=ap, + do_atomic_virial=True, + ) + tol = {"rtol": 1e-12, "atol": 1e-12} + torch.testing.assert_close(out["energy"], ref["energy_redu"], **tol) + torch.testing.assert_close( + out["force"], ref["energy_derv_r"].reshape(out["force"].shape), **tol + ) + torch.testing.assert_close( + out["virial"], ref["energy_derv_c_redu"].reshape(out["virial"].shape), **tol + ) + + def test_graph_lower_symbolic_trace_with_attention(self) -> None: + """Same symbolic trace with the repformer attention ON (the CI + gen_dpa2 export config). The attention path runs the slot-occupancy + ``segment_softmax`` whose entry count is a DATA-DEPENDENT (unbacked) + symbol under export; a Python ``if n == 0`` early-out in + ``_slot_occupancy`` raised ``GuardOnDataDependentSymNode: Eq(u1, 0)`` + here while the attention-off trace above stayed green. + """ + from deepmd.pt_expt.utils.serialization import ( + build_synthetic_graph_inputs, + ) + + model = self._make_model(repformer_attn=True).to("cpu") + model.eval() + sample = build_synthetic_graph_inputs( + model, + e_max=175, + nframes=2, + nloc=7, + dtype=torch.float64, + device=torch.device("cpu"), + ) + atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs = sample + traced = model.forward_lower_graph_exportable( + atype, + n_node, + nl, + ei, + ev, + em, + do, + drp, + so, + srp, + fparam=fp, + aparam=ap, + do_atomic_virial=True, + charge_spin=cs, + destination_sorted=True, + tracing_mode="symbolic", + _allow_non_fake_inputs=True, + ) + out = traced(atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs) + ref = model.forward_common_lower_graph( + atype, + n_node, + nl, + ei, + ev, + em, + do, + drp, + so, + srp, + destination_sorted=True, + fparam=fp, + aparam=ap, + do_atomic_virial=True, + ) + tol = {"rtol": 1e-12, "atol": 1e-12} + torch.testing.assert_close(out["energy"], ref["energy_redu"], **tol) + torch.testing.assert_close( + out["force"], ref["energy_derv_r"].reshape(out["force"].shape), **tol + ) + + def test_compiled_training_graph_smoke(self) -> None: + """``_trace_and_compile_graph`` inductor-compiles the graph lower; + one synthetic batch matches the eager graph lower at fp64 1e-10. + + NOTE: torch 2.11.0+cpu on an AVX2-only box can make inductor's C++ + vectorizer assert on the per-frame virial scatter + (``AssertionError ... assert index.is_vec``, cpp.py:2980). + ``_trace_and_compile_graph`` already disables CPU SIMD for the + graph-lower compile internally (``extra_options={"cpp.simdlen": + 0}``) to route around this, so no test-side workaround is normally + needed; the accepted fallback (project memory: torch 2.11 AVX2 + inductor bug) is ``torch._inductor.config.cpp.simdlen = 1`` if it + ever resurfaces here. + """ + from deepmd.pt_expt.train.training import ( + _trace_and_compile_graph, + ) + from deepmd.pt_expt.utils.serialization import ( + build_synthetic_graph_inputs, + ) + + model = self._make_model().to("cpu") + model.eval() + + compiled_lower, buf_order = _trace_and_compile_graph(model, None, None, None) + assert isinstance(compiled_lower, torch.nn.Module) + assert buf_order == () + + sample = build_synthetic_graph_inputs( + model, + e_max=97, + nframes=3, + nloc=5, + dtype=torch.float64, + device=torch.device("cpu"), + want_fparam=False, + want_aparam=False, + want_charge_spin=False, + ) + atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs = sample + + compiled_out = compiled_lower( + atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs + ) + eager = model.forward_common_lower_graph( + atype, + n_node, + nl, + ei, + ev, + em, + do, + drp, + so, + srp, + destination_sorted=True, + do_atomic_virial=False, + fparam=fp, + aparam=ap, + charge_spin=cs, + ) + tol = {"rtol": 1e-10, "atol": 1e-10} + torch.testing.assert_close(compiled_out["energy"], eager["energy_redu"], **tol) + torch.testing.assert_close( + compiled_out["force"], + eager["energy_derv_r"].reshape(compiled_out["force"].shape), + **tol, + ) + torch.testing.assert_close( + compiled_out["virial"], + eager["energy_derv_c_redu"].reshape(compiled_out["virial"].shape), + **tol, + ) + + def test_graph_lower_hatch_state_dict_roundtrip_and_backcompat(self) -> None: + """``disable_graph_lower()`` is persisted state: it round-trips a + ``state_dict`` save/load (the ``graph_lower_disabled`` buffer), and a + PRE-knob checkpoint (no buffer key) still strict-loads, defaulting to + the graph route. + """ + model = self._make_model() + model.atomic_model.descriptor.disable_graph_lower() + sd = model.state_dict() + assert any("graph_lower_disabled" in k for k in sd), ( + "the hatch must ride the state_dict" + ) + + fresh = self._make_model() + assert fresh.atomic_model.descriptor.uses_graph_lower() is True + fresh.load_state_dict(sd) + assert fresh.atomic_model.descriptor.uses_graph_lower() is False, ( + "restored buffer must re-disable the graph route" + ) + + # back-compat: a checkpoint written before the knob was persisted + # lacks the buffer key -- the strict load must still succeed and + # keep the default (graph) route. + old_sd = { + k: v + for k, v in self._make_model().state_dict().items() + if "graph_lower_disabled" not in k + } + fresh2 = self._make_model() + fresh2.load_state_dict(old_sd) + assert fresh2.atomic_model.descriptor.uses_graph_lower() is True + + def test_nonzero_mean_stays_on_legacy_route(self) -> None: + """A ``set_davg_zero=False`` DPA2 with COMPUTED nonzero statistics is + gated off the default graph route. + + Regression (OutisLi review): the dense repinit accumulates the + ``(sel - n_real)`` padding slots' ``-davg/dstd`` environment rows, + which the sel-free carry-all lower has no phantom-row counterpart + for -- with normal ``compute_input_stats`` (nonzero davg), the graph + route differed from ``legacy`` by ~4e-2 in energy at NON-binding sel + with attention off, silently reinterpreting checkpoints trained + against the dense expression. ``uses_graph_lower()`` therefore + returns False for ``set_davg_zero=False``; the default forward is + bit-identical to ``legacy``. The anti-vacuity part transplants the + computed statistics into a ``set_davg_zero=True`` (graph-eligible) + twin and shows the graph-vs-dense divergence is real -- so this + gate cannot be removed without reproducing the dense padding math. + """ + from deepmd.dpmodel.descriptor.dpa2 import ( + RepformerArgs, + RepinitArgs, + ) + + def _build(set_davg_zero: bool) -> EnergyModel: + repinit = RepinitArgs( + rcut=4.0, + rcut_smth=0.5, + nsel=12, # NON-binding for the 6-atom fixture + neuron=[3, 6], + axis_neuron=2, + tebd_dim=4, + set_davg_zero=set_davg_zero, + ) + repformer = RepformerArgs( + rcut=2.0, + rcut_smth=0.5, + nsel=8, # NON-binding + nlayers=2, + g1_dim=8, + g2_dim=4, + axis_neuron=2, + update_g1_has_attn=False, + update_g2_has_attn=False, + ) + ds = DescrptDPA2( + ntypes=self.nt, + repinit=repinit, + repformer=repformer, + precision="float64", + seed=GLOBAL_SEED, + ).to(self.device) + ft = InvarFitting( + "energy", + self.nt, + ds.get_dim_out(), + 1, + mixed_types=ds.mixed_types(), + precision="float64", + seed=GLOBAL_SEED, + ).to(self.device) + return EnergyModel(ds, ft, type_map=self.type_map).to(self.device) + + model = _build(set_davg_zero=False) + # NORMAL computed statistics (not a hand-written mean). + sample = { + "coord": self.coord.clone(), + "atype": self.atype.clone(), + "box": self.cell.clone(), + } + model.atomic_model.descriptor.compute_input_stats([sample]) + davg = np.asarray(model.atomic_model.descriptor.repinit.mean.detach().cpu()) + assert np.abs(davg).max() > 0, "computed statistics must be nonzero" + + # 1. the gate: nonzero-mean configs stay on the legacy route. + assert model.atomic_model.descriptor.uses_graph_lower() is False + model.eval() + box = self.cell.reshape(1, 9) + out_default = model.forward_common( + self.coord.clone().requires_grad_(True), self.atype, box + ) + out_legacy = model.forward_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="legacy", + ) + torch.testing.assert_close( + out_default["energy_redu"], out_legacy["energy_redu"], rtol=0.0, atol=0.0 + ) + + # 2. anti-vacuity: transplant the computed stats into a graph-eligible + # twin -- the graph route genuinely diverges from dense at NON-binding + # sel once davg != 0, which is what the gate protects against. + twin = _build(set_davg_zero=True) + twin.load_state_dict(model.state_dict(), strict=False) + twin.atomic_model.descriptor.repinit.mean = ( + model.atomic_model.descriptor.repinit.mean + ) + twin.atomic_model.descriptor.repinit.stddev = ( + model.atomic_model.descriptor.repinit.stddev + ) + assert twin.atomic_model.descriptor.uses_graph_lower() is True + twin.eval() + out_graph = twin.forward_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="dense", + ) + out_dense = twin.forward_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="legacy", + ) + diff = float((out_graph["energy_redu"] - out_dense["energy_redu"]).abs().max()) + assert diff > 1e-8, ( + f"expected a genuine graph-vs-dense divergence for nonzero davg " + f"(the reason for the gate); got {diff:.3e}" + ) + + def test_non_energy_model_stays_off_graph_compiled_path(self) -> None: + """A graph-eligible DPA2 descriptor paired with a NON-energy fitting + (property) must stay OFF the graph compiled-training path AND off + the eager default-flip. + + ``_trace_and_compile_graph`` and the trainer's public-key + translation are energy-specific (``do_grad_r('energy')``, + ``_translate_energy_keys``), so a property model routed onto the + graph compiled path raised ``KeyError('energy')`` at its first + batch while its eager forward (the output-agnostic graph lower) + succeeded -- an eager != compiled split. The default-flip gate + keeps BOTH on dense for non-energy outputs: eager default must be + bit-identical to ``neighbor_graph_method="legacy"``, and the + compiled first batch must produce property outputs via the dense + compiler. + """ + from deepmd.pt_expt.fitting import ( + PropertyFittingNet, + ) + from deepmd.pt_expt.model import ( + PropertyModel, + ) + from deepmd.pt_expt.model.graph_lower import ( + model_uses_graph_lower as _model_uses_graph_lower, + ) + from deepmd.pt_expt.train.training import ( + _CompiledModel, + ) + + ds = _make_dpa2_descriptor(ntypes=self.nt).to(self.device) + ft = PropertyFittingNet( + self.nt, + ds.get_dim_out(), + task_dim=2, + mixed_types=ds.mixed_types(), + seed=GLOBAL_SEED, + ).to(self.device) + model = PropertyModel(ds, ft, type_map=self.type_map).to(self.device) + model.eval() + + # graph-eligible DESCRIPTOR, but non-energy MODEL -> gated off. + assert model.atomic_model.descriptor.uses_graph_lower() is True + assert _model_uses_graph_lower(model) is False + + # eager default-flip mirrors the gate: default == legacy (dense). + box = self.cell.reshape(1, 9) + out_default = model.forward_common(self.coord, self.atype, box) + out_legacy = model.forward_common( + self.coord, self.atype, box, neighbor_graph_method="legacy" + ) + torch.testing.assert_close( + out_default["property_redu"], + out_legacy["property_redu"], + rtol=0.0, + atol=0.0, + ) + + # compiled first batch: routed to the DENSE compiler; used to raise + # KeyError('energy') from the energy-specific graph trace. + cm = _CompiledModel(model, structure_key=("dpa2-property-regression",)) + out_c = cm(self.coord, self.atype, box=box) + assert any("property" in k for k in out_c), ( + f"compiled property forward must emit property keys; got {list(out_c)}" + ) + + def test_compiled_training_graph_small_sel(self) -> None: + """The compiled-training trace capacity derives from the synthetic + system's REAL edge count, not from ``sel``. + + The carry-all graph builder is sel-free (sel = normalization + constant only), so a sel-derived static trace capacity + (``ceil(1.25 * nloc * sum(sel))``) overflows whenever the synthetic + trace system's actual degree exceeds ``sel``: with repinit/repformer + ``nsel=10/6`` the trace used to raise ``edge overflow: 106 real + edges > edge_capacity 89``. The capacity is now probed from the + actual unpadded synthetic graph, so the trace must succeed and the + compiled lower must match the eager graph lower (compiled and eager + are the SAME route, so parity holds even at binding sel). + """ + from deepmd.pt_expt.train.training import ( + _trace_and_compile_graph, + ) + from deepmd.pt_expt.utils.serialization import ( + build_synthetic_graph_inputs, + ) + + model = self._make_model(repinit_nsel=10, repformer_nsel=6).to("cpu") + model.eval() + + compiled_lower, _ = _trace_and_compile_graph(model, None, None, None) + + sample = build_synthetic_graph_inputs( + model, + e_max=97, + nframes=3, + nloc=5, + dtype=torch.float64, + device=torch.device("cpu"), + want_fparam=False, + want_aparam=False, + want_charge_spin=False, + ) + atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs = sample + compiled_out = compiled_lower( + atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs + ) + eager = model.forward_common_lower_graph( + atype, + n_node, + nl, + ei, + ev, + em, + do, + drp, + so, + srp, + destination_sorted=True, + do_atomic_virial=False, + fparam=fp, + aparam=ap, + charge_spin=cs, + ) + tol = {"rtol": 1e-10, "atol": 1e-10} + torch.testing.assert_close(compiled_out["energy"], eager["energy_redu"], **tol) + torch.testing.assert_close( + compiled_out["force"], + eager["energy_derv_r"].reshape(compiled_out["force"].shape), + **tol, + ) + + def test_graph_lower_fparam_symbolic_trace_and_compile(self) -> None: + """A graph-eligible DPA2 model with ``numb_fparam > 0`` must export + (``make_fx`` symbolic) AND inductor-compile. + + Regression for the ``frame_id_from_n_node(graph.n_node)`` call in + ``dp_atomic_model.forward_atomic_graph``: without a STATIC ``n_total`` + it fell back to ``int(sum(n_node))``, which make_fx / torch.export + cannot evaluate on a traced tensor (``GuardOnDataDependentSymNode``), + so both the graph ``.pt2`` export and compiled training failed for any + fparam model. Both must now trace/compile and match eager bit-tight. + """ + from deepmd.pt_expt.train.training import ( + _trace_and_compile_graph, + ) + from deepmd.pt_expt.utils.serialization import ( + build_synthetic_graph_inputs, + ) + + model = self._make_model(numb_fparam=2).to("cpu") + model.eval() + assert model.atomic_model.descriptor.uses_graph_lower() is True + + sample = build_synthetic_graph_inputs( + model, + e_max=175, + nframes=2, + nloc=7, + dtype=torch.float64, + device=torch.device("cpu"), + ) + atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs = sample + assert fp is not None and fp.shape[-1] == 2, "fparam must be present" + + # (a) symbolic-trace export path + traced = model.forward_lower_graph_exportable( + atype, + n_node, + nl, + ei, + ev, + em, + do, + drp, + so, + srp, + fparam=fp, + aparam=ap, + do_atomic_virial=True, + charge_spin=cs, + destination_sorted=True, + tracing_mode="symbolic", + _allow_non_fake_inputs=True, + ) + out = traced(atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs) + ref = model.forward_common_lower_graph( + atype, + n_node, + nl, + ei, + ev, + em, + do, + drp, + so, + srp, + destination_sorted=True, + fparam=fp, + aparam=ap, + do_atomic_virial=True, + ) + tol = {"rtol": 1e-12, "atol": 1e-12} + torch.testing.assert_close(out["energy"], ref["energy_redu"], **tol) + torch.testing.assert_close( + out["force"], ref["energy_derv_r"].reshape(out["force"].shape), **tol + ) + + # (b) compiled-training path (fparam threaded through the compile) + compiled_lower, _ = _trace_and_compile_graph(model, fp, None, None) + compiled_out = compiled_lower( + atype, n_node, nl, ei, ev, em, do, drp, so, srp, fp, ap, cs + ) + ctol = {"rtol": 1e-10, "atol": 1e-10} + torch.testing.assert_close(compiled_out["energy"], ref["energy_redu"], **ctol) + + # ------------------------------------------------------------------ + # 2. the one piece of real code: the border_op graph exchange override. + # ------------------------------------------------------------------ + def test_graph_exchange_identity_when_no_comm(self) -> None: + """``comm_dict is None`` -> identity (ghost-free Python graphs / + extended single-process graphs need no exchange). + """ + block = _make_dpa2_descriptor().repformers.to(self.device) + assert isinstance(block, DescrptBlockRepformers) + g1 = torch.rand( + [self.natoms, block.get_dim_out()], + dtype=torch.float64, + device=self.device, + ) + out = block._exchange_ghosts_graph(g1, None, self.natoms) + assert out is g1 + + def test_graph_exchange_border_op_self_communication(self) -> None: + """A real (non-spin) ``comm_dict`` routes through ``border_op``: + local rows [0, nlocal) are untouched and ghost rows [nlocal, N) + are overwritten with the owner rows named by the sendlist -- the + actual new-code path this task adds (the identity/spin-reject + tests above already hold on the dpmodel base and do not exercise + ``border_op`` at all). + """ + block = _make_dpa2_descriptor().repformers.to(self.device) + nlocal, nghost = 4, 3 + n_total = nlocal + nghost + dim = block.get_dim_out() + g1_local = torch.rand([nlocal, dim], dtype=torch.float64, device=self.device) + g1_ghost_junk = torch.full( + [nghost, dim], -999.0, dtype=torch.float64, device=self.device + ) + g1 = torch.cat([g1_local, g1_ghost_junk], dim=0) + + # ghost slot ii (0-indexed within the ghost block) mirrors local + # owner (ii % nlocal). + owners = np.array([ii % nlocal for ii in range(nghost)], dtype=np.int32) + keepalive: list = [] + comm_dict = _build_self_comm_dict( + nloc=nlocal, + nghost=nghost, + sendlist_indices=owners, + keepalive=keepalive, + ) + + out = block._exchange_ghosts_graph(g1, comm_dict, n_total) + + torch.testing.assert_close(out[:nlocal], g1_local) + expected_ghosts = g1_local[torch.as_tensor(owners, dtype=torch.int64)] + torch.testing.assert_close(out[nlocal:], expected_ghosts) + + def test_graph_exchange_rejects_spin_comm(self) -> None: + """A ``comm_dict`` describing a spin system raises ``NotImplementedError`` + -- spin models never route the graph lower (``disable_graph_lower``), + so a ``has_spin`` comm_dict reaching the graph exchange seam is a + programming error, not a supported configuration. + """ + block = _make_dpa2_descriptor().repformers.to(self.device) + g1 = torch.rand( + [self.natoms, block.get_dim_out()], + dtype=torch.float64, + device=self.device, + ) + comm_dict = {"has_spin": torch.ones(1)} + with pytest.raises(NotImplementedError): + block._exchange_ghosts_graph(g1, comm_dict, self.natoms) + + +# --------------------------------------------------------------------------- +# Task 9: owned-node (``n_local``) energy mask in the graph output reduction. +# --------------------------------------------------------------------------- + + +class TestOwnedNodeMaskEnergyReduction: + """``n_local`` (owned-node mask) must exclude ghost rows from the + DIFFERENTIATED per-frame energy (each ghost atom is owned -- and counted + -- on another rank), while ``atom_energy`` itself stays FULL and the + mask is applied BEFORE ``edge_energy_deriv`` differentiates the summed + energy (so ``grad(energy, edge_vec)`` only carries owned-energy terms). + """ + + def setup_method(self) -> None: + self.device = env.DEVICE + self.natoms = 6 + self.nt = 2 + self.type_map = ["foo", "bar"] + + generator = torch.Generator(device=self.device).manual_seed(GLOBAL_SEED) + cell = torch.rand( + [3, 3], dtype=torch.float64, device=self.device, generator=generator + ) + cell = (cell + cell.T) + 5.0 * torch.eye(3, device=self.device) + self.cell = cell.unsqueeze(0) # [1, 3, 3] + coord = torch.rand( + [self.natoms, 3], + dtype=torch.float64, + device=self.device, + generator=generator, + ) + coord = torch.matmul(coord, cell) + self.coord = coord.unsqueeze(0).to(self.device) # [1, natoms, 3] + self.atype = torch.tensor( + [[0, 0, 0, 1, 1, 1]], dtype=torch.int64, device=self.device + ) + + def _make_model(self) -> EnergyModel: + ds = _make_dpa2_descriptor(ntypes=self.nt).to(self.device) + ft = InvarFitting( + "energy", + self.nt, + ds.get_dim_out(), + 1, + mixed_types=ds.mixed_types(), + precision="float64", + seed=GLOBAL_SEED, + ).to(self.device) + return EnergyModel(ds, ft, type_map=self.type_map).to(self.device) + + def _build_graph(self, model: EnergyModel): + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + ) + + box = self.cell.reshape(1, 9) + return build_neighbor_graph(self.coord, self.atype, box, model.get_rcut()) + + def test_owned_mask_energy_reduction(self) -> None: + """``energy_redu`` with ``n_local`` == sum(atom_energy[:n_local]); + ghost rows of the per-node output are masked to zero (#5758 + convention -- each ghost atom's fitting output is owned by another + rank); ghost rows [n_local:) of the assembled force are NOT all zero + (they still receive src-scatter contributions from OWNED centers' + edges). + """ + model = self._make_model() + model.eval() + ng = self._build_graph(model) + atype_flat = self.atype.reshape(-1) + n_local_val = 4 + + out_full = model.forward_common_lower_graph( + atype_flat, + ng.n_node, + ng.n_node.clone(), + ng.edge_index, + ng.edge_vec, + ng.edge_mask, + ) + n_local = torch.tensor([n_local_val], dtype=torch.int64, device=self.device) + out_masked = model.forward_common_lower_graph( + atype_flat, + ng.n_node, + n_local, + ng.edge_index, + ng.edge_vec, + ng.edge_mask, + ) + + # per-node output: owned rows unchanged, ghost rows masked to zero. + torch.testing.assert_close( + out_masked["energy"].reshape(-1)[:n_local_val], + out_full["energy"].reshape(-1)[:n_local_val], + ) + halo_energy = out_masked["energy"].reshape(-1)[n_local_val:] + torch.testing.assert_close(halo_energy, torch.zeros_like(halo_energy)) + + atom_energy = out_full["energy"].reshape(-1) + owned_sum = atom_energy[:n_local_val].sum() + torch.testing.assert_close( + out_masked["energy_redu"].reshape(-1), owned_sum.reshape(1) + ) + + # ghost-row partial forces survive (src-side scatter from owned edges). + force = out_masked["energy_derv_r"].reshape(self.natoms, 3) + halo_force = force[n_local_val:] + assert not torch.allclose(halo_force, torch.zeros_like(halo_force)), ( + "expected nonzero ghost-row partial force from owned-center edges" + ) + + # ordering check: the masked force must differ from the unmasked + # (full-graph) force -- if the mask were applied AFTER + # edge_energy_deriv (or not at all), both runs would produce the + # identical force since the autograd leaf (edge_vec) is unchanged. + force_full = out_full["energy_derv_r"].reshape(self.natoms, 3) + assert not torch.allclose(force, force_full), ( + "masked force must differ from the unmasked force -- the " + "owned-node mask must be applied BEFORE edge_energy_deriv" + ) + + def test_n_local_full_ownership_is_unmasked(self) -> None: + """``n_local == n_node`` (every node owned, the single-rank case) + reproduces the unmasked reduction exactly: ``energy_redu`` equals the + sum of the (unmasked) per-node energies. + """ + model = self._make_model() + model.eval() + ng = self._build_graph(model) + atype_flat = self.atype.reshape(-1) + + out = model.forward_common_lower_graph( + atype_flat, + ng.n_node, + ng.n_node.clone(), + ng.edge_index, + ng.edge_vec, + ng.edge_mask, + ) + torch.testing.assert_close( + out["energy_redu"].reshape(-1), + out["energy"].reshape(-1).sum().reshape(1), + ) + + +class TestGraphAparamExtendedAxis: + """aparam contract on the (extended-region) graph route. + + The graph ABI carries atomic parameters FLAT on the node axis -- + ``(N, nda)`` with ``N = sum(n_node)``, the same axis as ``atype``. On + extended-region graphs (the multi-rank C++ routes: ``N == nall_real``, + owned prefix first, then ghost) the caller must therefore supply an + extended-axis aparam with the ghost rows padded -- their fitting outputs + are discarded by the owned-node mask, so the padded values are inert. + This pins the contract the C++ runtime's ``extend_graph_aparam`` + (DeepPotPTExpt.cc) implements: an owned-only (nloc-row) aparam and a + rectangular ``(nf, nloc, nda)`` aparam must both fail loudly rather + than silently misalign parameter rows with nodes. + """ + + def setup_method(self) -> None: + self.device = env.DEVICE + self.natoms = 6 + self.n_local_val = 4 + self.nt = 2 + generator = torch.Generator(device=self.device).manual_seed(GLOBAL_SEED) + cell = torch.rand( + [3, 3], dtype=torch.float64, device=self.device, generator=generator + ) + cell = (cell + cell.T) + 5.0 * torch.eye(3, device=self.device) + self.cell = cell.unsqueeze(0) + coord = torch.rand( + [self.natoms, 3], + dtype=torch.float64, + device=self.device, + generator=generator, + ) + self.coord = torch.matmul(coord, cell).unsqueeze(0).to(self.device) + self.atype = torch.tensor( + [[0, 0, 0, 1, 1, 1]], dtype=torch.int64, device=self.device + ) + + def _make_model(self) -> EnergyModel: + ds = _make_dpa2_descriptor(ntypes=self.nt).to(self.device) + ft = InvarFitting( + "energy", + self.nt, + ds.get_dim_out(), + 1, + mixed_types=ds.mixed_types(), + numb_aparam=1, + precision="float64", + seed=GLOBAL_SEED, + ).to(self.device) + return EnergyModel(ds, ft, type_map=["foo", "bar"]).to(self.device) + + def _forward(self, model, aparam): + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + ) + + ng = build_neighbor_graph( + self.coord, self.atype, self.cell.reshape(1, 9), model.get_rcut() + ) + n_local = torch.tensor( + [self.n_local_val], dtype=torch.int64, device=self.device + ) + return model.forward_common_lower_graph( + self.atype.reshape(-1), + ng.n_node, + n_local, + ng.edge_index, + ng.edge_vec, + ng.edge_mask, + aparam=aparam, + ) + + def test_extended_axis_accepted_halo_rows_inert(self) -> None: + """Flat extended-axis aparam (``(N, nda)``) runs; ghost-row values are + inert for the owned (masked) energy, owned-row values are not. + """ + model = self._make_model() + model.eval() + base = torch.linspace( + 0.1, 0.6, self.natoms, dtype=torch.float64, device=self.device + ).reshape(self.natoms, 1) + out = self._forward(model, base) + + # ghost rows [n_local:) changed -> owned energy identical + halo_bump = base.clone() + halo_bump[self.n_local_val :, :] += 7.5 + out_ghost = self._forward(model, halo_bump) + torch.testing.assert_close(out_ghost["energy_redu"], out["energy_redu"]) + + # an OWNED row changed -> owned energy must change + owned_bump = base.clone() + owned_bump[0, :] += 7.5 + out_owned = self._forward(model, owned_bump) + assert not torch.allclose(out_owned["energy_redu"], out["energy_redu"]), ( + "owned-row aparam change must reach the owned energy" + ) + + def test_owned_only_axis_fails_loudly(self) -> None: + """A flat aparam with only the OWNED rows (``(nloc, nda)`` instead of + ``(N, nda)``) must raise -- silently misaligning parameter rows with + the N-node axis is the bug the extended-axis contract prevents. + """ + model = self._make_model() + model.eval() + owned_only = torch.linspace( + 0.1, 0.4, self.n_local_val, dtype=torch.float64, device=self.device + ).reshape(self.n_local_val, 1) + with pytest.raises((RuntimeError, ValueError)): + self._forward(model, owned_only) + + def test_rectangular_aparam_extends_consistently(self) -> None: + """A rectangular ``(nf, nloc, nda)`` aparam on the EXTENDED-region + route (``n_local`` present) is expanded by the ragged + ``_extend_graph_aparam`` gather (#5758) and must give the same owned + energy as the equivalent flat ``(N, nda)`` input; the flat layout + stays the canonical trace/export ABI (a rectangular sample would + hand ``torch.export`` an unprovable ``N == nf * nloc`` relation -- + the root cause of the aparam graph-freeze failures). + """ + model = self._make_model() + model.eval() + flat = torch.linspace( + 0.1, 0.6, self.natoms, dtype=torch.float64, device=self.device + ).reshape(self.natoms, 1) + rectangular = flat.reshape(1, self.natoms, 1) + out_rect = self._forward(model, rectangular) + out_flat = self._forward(model, flat) + torch.testing.assert_close(out_rect["energy_redu"], out_flat["energy_redu"]) diff --git a/source/tests/pt_expt/model/test_graph_export.py b/source/tests/pt_expt/model/test_graph_export.py index 1114b78848..fe35c66ad9 100644 --- a/source/tests/pt_expt/model/test_graph_export.py +++ b/source/tests/pt_expt/model/test_graph_export.py @@ -104,6 +104,127 @@ def test_graph_exportable_traces(): torch.testing.assert_close(te, eager["energy_redu"], rtol=1e-10, atol=1e-10) +def test_graph_export_aparam_flat_node_axis(): + """The regular graph export ABI carries ``aparam`` FLAT on the node axis, + sharing ``atype``'s dynamic ``N`` symbol. + + Regression (OutisLi review): the trace sample used to build ``aparam`` + rectangular ``(nf, nloc, nda)`` with independent ``nframes``/``nloc`` + dims while ``atype`` is flat ``(N,)``; the graph fitting views ``aparam`` + against the flat node count, so ``torch.export`` derived ``N == nf * + nloc`` and rejected every ``numb_aparam > 0`` regular graph freeze with + ``Constraints violated (n_node_total)`` (plus an ``N >= 4`` guard that + would reject valid 1-to-3-node systems). The flat ``(N, nda)`` ABI must + export, and the exported program must run BOTH a single-node ``nf=1, + N=1`` system and a multi-frame ``nf=3, N=15`` system, matching eager. + + ``numb_aparam=2`` also pins the trace-dim selection: traced at the old + fixed ``nframes=2``, make_fx duck-merged the STATIC ``nda == 2`` with + the dynamic ``nf`` symbol, baking outputs whose frame axis followed the + aparam width (silently wrong ``(2, 1)`` energy at ``nf=1``). The + production trace (``_trace_and_export``, exercised here) now picks + collision-free prime trace dims like compiled training does. + """ + from deepmd.pt_expt.model.get_model import ( + get_model, + ) + from deepmd.pt_expt.utils.serialization import ( + _trace_and_export, + build_synthetic_graph_inputs, + ) + + config = { + "type_map": ["A", "B"], + "descriptor": { + "type": "se_atten", + "sel": 20, + "rcut": _RCUT, + "rcut_smth": 0.5, + "neuron": [3, 6], + "axis_neuron": 2, + "attn_layer": 0, + "seed": GLOBAL_SEED, + }, + "fitting_net": { + "neuron": [16, 16], + "numb_aparam": 2, + "seed": GLOBAL_SEED, + }, + } + model = get_model(config).to("cpu") + model.eval() + + exported, _meta, _dj, _keys = _trace_and_export( + {"model": model.serialize()}, + model_json_override=None, + lower_kind="graph", + ) + loaded = exported.module() + # ``_trace_and_export`` moves the exported program to ``env.DEVICE``; run the + # eager reference there too so both sides use matching-device inputs. + model.to(env.DEVICE) + + # nf=1, N=1 (single node: zero real edges, guard rows only) and a + # multi-frame nf=3, N=15 system: both must pass the input guards and + # match the eager graph lower. + for nframes, nloc in ((1, 1), (3, 5)): + inputs = build_synthetic_graph_inputs( + model, + e_max=None, + nframes=nframes, + nloc=nloc, + dtype=torch.float64, + device=env.DEVICE, + ) + ( + a2, + nn2, + nl2, + ei2, + ev2, + em2, + do2, + drp2, + so2, + srp2, + fp2, + ap2, + cs2, + ) = inputs + # single-rank runtime: every node is owned + nl2 = nn2.clone() + # distinct per-row values so the sensitivity check below is real + ap2 = torch.linspace( + 0.1, 0.9, ap2.numel(), dtype=torch.float64, device=ap2.device + ).reshape(ap2.shape) + out = loaded(a2, nn2, nl2, ei2, ev2, em2, do2, drp2, so2, srp2, fp2, ap2, cs2) + ref = model.forward_common_lower_graph( + a2, + nn2, + nl2, + ei2, + ev2, + em2, + do2, + drp2, + so2, + srp2, + fparam=fp2, + aparam=ap2, + do_atomic_virial=False, + ) + torch.testing.assert_close( + out["energy"], ref["energy_redu"], rtol=1e-10, atol=1e-10 + ) + # aparam must actually reach the fitting: bump it -> energy changes + out_bump = loaded( + a2, nn2, nl2, ei2, ev2, em2, do2, drp2, so2, srp2, fp2, ap2 + 1.5, cs2 + ) + assert not torch.allclose(out_bump["energy"], out["energy"]), ( + f"aparam bump must change the energy (nf={nframes}, nloc={nloc})" + ) + + def test_graph_export_rejects_false_canonical_claim(): model = _model().eval() graph_inputs = list(_graph_inputs(model)) diff --git a/source/tests/pt_expt/model/test_graph_export_with_comm.py b/source/tests/pt_expt/model/test_graph_export_with_comm.py new file mode 100644 index 0000000000..d1c031074c --- /dev/null +++ b/source/tests/pt_expt/model/test_graph_export_with_comm.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Graph-form ``.pt2`` export: the with-comm nested artifact (task 10). + +Message-passing graph descriptors (dpa2's repformer block) need a SECOND +AOTI artifact -- ``model/extra/forward_lower_with_comm.pt2`` -- that accepts +8 extra positional comm tensors and drives per-layer cross-rank ghost-atom +exchange via ``deepmd_export::border_op``. This mirrors the dense +with-comm flow (``test_graph_pt2_metadata.py`` covers the plain-graph +``lower_input_kind`` branch; this file covers the *comm* branch on top of +it). +""" + +import copy +import json +import zipfile + +import pytest + +from deepmd.pt_expt.utils.env import ( + DEVICE, +) +from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, +) + +# Small graph-eligible dpa2 descriptor: tebd_input_mode defaults to +# "concat", use_three_body defaults to False -> uses_graph_lower() is True, +# and has_message_passing_across_ranks() is unconditionally True for any +# DPA2 (delegates to DescrptBlockRepformers, which always needs cross-rank +# g1 exchange) -- so this model always gets a with-comm artifact. +DPA2_CONFIG = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa2", + "repinit": { + "rcut": 4.0, + "rcut_smth": 0.5, + "nsel": 10, + "neuron": [4, 8], + "axis_neuron": 2, + }, + "repformer": { + "rcut": 3.0, + "rcut_smth": 0.5, + "nsel": 6, + "nlayers": 1, + "g1_dim": 8, + "g2_dim": 4, + }, + }, + "fitting_net": {"neuron": [8, 8], "seed": 1}, +} + +# dpa1 with attn_layer == 0: graph-eligible but NOT message-passing +# (has_message_passing_across_ranks() is False for a single se_atten +# descriptor) -- the regression pin that the with-comm artifact stays +# absent for non-GNN graph-eligible descriptors. +DPA1_CONFIG = { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_atten", + "sel": 10, + "rcut_smth": 2.0, + "rcut": 4.0, + "neuron": [2, 4], + "axis_neuron": 2, + "attn": 4, + "attn_layer": 0, + "attn_dotr": True, + "attn_mask": False, + "activation_function": "tanh", + "scaling_factor": 1.0, + "normalize": True, + "temperature": 1.0, + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "neuron": [4, 4], + "resnet_dt": True, + "seed": 1, + }, +} + + +def _build_data(config: dict) -> dict: + """Build a serialized dpmodel data dict (same shape as ``dp freeze`` input).""" + from deepmd.dpmodel.model.model import ( + get_model, + ) + + model = get_model(copy.deepcopy(config)) + return { + "model": model.serialize(), + "model_def_script": copy.deepcopy(config), + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + + +def _read_metadata(pt2_path: str) -> dict: + """Read ``model/extra/metadata.json`` from a ``.pt2`` ZIP archive.""" + with zipfile.ZipFile(pt2_path, "r") as zf: + raw = zf.read("model/extra/metadata.json").decode("utf-8") + return json.loads(raw) + + +@pytest.fixture(scope="module") +def dpa2_dpmodel_data() -> dict: + return _build_data(DPA2_CONFIG) + + +@pytest.fixture(scope="module") +def dpa1_dpmodel_data() -> dict: + return _build_data(DPA1_CONFIG) + + +def test_dpa2_graph_pt2_embeds_with_comm_artifact(dpa2_dpmodel_data, tmp_path) -> None: + """dpa2 (message-passing) graph ``.pt2`` embeds the nested with-comm artifact.""" + p = str(tmp_path / "m_dpa2_graph.pt2") + deserialize_to_file( + p, + copy.deepcopy(dpa2_dpmodel_data), + do_atomic_virial=True, + lower_kind="graph", + ) + with zipfile.ZipFile(p, "r") as zf: + names = zf.namelist() + assert "model/extra/forward_lower_with_comm.pt2" in names + + meta = _read_metadata(p) + assert meta["lower_input_kind"] == "graph" + assert meta["has_comm_artifact"] is True + assert meta["has_message_passing"] is True + + +def test_dpa2_graph_with_comm_aparam_freeze(tmp_path) -> None: + """A ``numb_aparam > 0`` message-passing model freezes to the graph + ``.pt2`` with the nested with-comm artifact. + + Regression (OutisLi review): the with-comm dynamic-shape spec gave + ``aparam``'s node axis an independent ``nloc`` Dim while the graph ABI + carries ``aparam`` on the same flat node axis as ``atype``; the graph + fitting views ``aparam`` against the flat node count, so + ``torch.export`` proved the two axes equal and rejected every + ``numb_aparam > 0`` with-comm graph freeze with ``Constraints violated + (nloc): aparam.size(1) and atype.size(0) must always be equal``. The + aparam node axis now IS ``atype``'s ``n_node_total`` axis (flat + ``(N, nda)`` ABI), so both the regular and the with-comm graph traces + of this freeze must succeed. + """ + cfg = copy.deepcopy(DPA2_CONFIG) + cfg["fitting_net"] = {**cfg["fitting_net"], "numb_aparam": 1} + data = _build_data(cfg) + p = str(tmp_path / "m_dpa2_graph_aparam.pt2") + deserialize_to_file( + p, + data, + do_atomic_virial=True, + lower_kind="graph", + ) + with zipfile.ZipFile(p, "r") as zf: + names = zf.namelist() + assert "model/extra/forward_lower_with_comm.pt2" in names + + meta = _read_metadata(p) + assert meta["lower_input_kind"] == "graph" + assert meta["has_comm_artifact"] is True + + +def test_dpa1_graph_pt2_has_no_comm_artifact(dpa1_dpmodel_data, tmp_path) -> None: + """dpa1 (non-message-passing) graph ``.pt2`` regression pin: no nested entry.""" + p = str(tmp_path / "m_dpa1_graph.pt2") + deserialize_to_file( + p, + copy.deepcopy(dpa1_dpmodel_data), + do_atomic_virial=True, + lower_kind="graph", + ) + with zipfile.ZipFile(p, "r") as zf: + names = zf.namelist() + assert "model/extra/forward_lower_with_comm.pt2" not in names + + meta = _read_metadata(p) + assert meta["lower_input_kind"] == "graph" + assert meta["has_comm_artifact"] is False + + +def test_graph_with_comm_n_local_is_separate_device_input( + dpa2_dpmodel_data, tmp_path +) -> None: + """The graph with-comm ABI splits the owned-count roles: the CPU + ``nlocal`` comm tensor is host control metadata for ``border_op``, and + the SEPARATE device ``n_local`` input (slot 2 of the graph base, #5758 + convention) feeds the in-graph owned-node mask. + + Regression for the device-scalar review finding: deriving the owned + mask from a device-placed ``nlocal`` comm tensor made every per-layer + ``border_op`` forward AND custom backward pull the scalar back with a + synchronizing D2H read (``4 * nlayers`` per MD step). With the split, + all 8 comm tensors stay on CPU (symmetric with the dense with-comm + artifact). This pins the ABI: the traced program exposes the extra + input (21 placeholders: 13 graph-base incl. slot-2 ``n_local`` and the + None-valued fparam/aparam/charge_spin slots + 8 comm), and the + owned-energy reduction follows ``n_local``, not the comm ``nlocal``. + """ + import numpy as np + import torch + + from deepmd.pt_expt.utils.serialization import ( + _trace_and_export, + ) + + exported, _meta, _dj, _keys = _trace_and_export( + copy.deepcopy(dpa2_dpmodel_data), + model_json_override=None, + with_comm_dict=True, + lower_kind="graph", + ) + loaded = exported.module() + placeholders = loaded.graph.find_nodes(op="placeholder") + assert len(placeholders) == 21, ( + f"graph with-comm program must accept 21 positional inputs " + f"(13 graph-base incl. n_local + 8 comm); got {len(placeholders)}" + ) + + # Build a single-rank self-comm system: 5 owned + 2 ghost nodes. + import ctypes + + from deepmd.dpmodel.model.model import get_model as get_dp_model + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + ) + + model = get_dp_model(copy.deepcopy(DPA2_CONFIG)) + rcut = model.get_rcut() + rng = np.random.default_rng(7) + n_total, nghost = 7, 2 + nlocal = n_total - nghost + coord = rng.random((1, n_total, 3)) * rcut * 0.6 + atype = np.array([[i % 2 for i in range(n_total)]]) + graph = build_neighbor_graph(coord, atype, None, rcut, canonicalize=True) + + # Graph inputs live on the device the exported program was moved to + # (env.DEVICE); the 8 comm tensors below stay on CPU (host control metadata). + atype_t = torch.tensor(atype.reshape(-1), dtype=torch.int64, device=DEVICE) + n_node_t = torch.as_tensor( + np.asarray(graph.n_node), dtype=torch.int64, device=DEVICE + ) + ei = torch.as_tensor(np.asarray(graph.edge_index), dtype=torch.int64, device=DEVICE) + ev = torch.as_tensor(np.asarray(graph.edge_vec), dtype=torch.float64, device=DEVICE) + em = torch.as_tensor(np.asarray(graph.edge_mask), dtype=torch.bool, device=DEVICE) + do_t = torch.as_tensor( + np.asarray(graph.destination_order), dtype=torch.int64, device=DEVICE + ) + drp_t = torch.as_tensor( + np.asarray(graph.destination_row_ptr), dtype=torch.int64, device=DEVICE + ) + so_t = torch.as_tensor( + np.asarray(graph.source_order), dtype=torch.int64, device=DEVICE + ) + srp_t = torch.as_tensor( + np.asarray(graph.source_row_ptr), dtype=torch.int64, device=DEVICE + ) + + sendlist_indices = np.ascontiguousarray( + np.arange(nghost, dtype=np.int32) + ) # keepalive below + addr = sendlist_indices.ctypes.data_as(ctypes.c_void_p).value + comm = ( + torch.tensor([addr], dtype=torch.int64), # send_list + torch.zeros(1, dtype=torch.int32), # send_proc + torch.zeros(1, dtype=torch.int32), # recv_proc + torch.tensor([nghost], dtype=torch.int32), # send_num + torch.tensor([nghost], dtype=torch.int32), # recv_num + torch.zeros(1, dtype=torch.int64), # communicator (null: self) + torch.tensor(nlocal, dtype=torch.int32), # nlocal (CPU, host role) + torch.tensor(nghost, dtype=torch.int32), # nghost (CPU, host role) + ) + + def run(n_local_val: int) -> torch.Tensor: + n_local_t = torch.tensor([n_local_val], dtype=torch.int64, device=DEVICE) + out = loaded( + atype_t, + n_node_t, + n_local_t, + ei, + ev, + em, + do_t, + drp_t, + so_t, + srp_t, + None, + None, + None, + *comm, + ) + return out["energy"] + + # Same comm nlocal, different n_local: the owned-energy reduction must + # follow the 17th input (fewer owned nodes -> different energy). + e_full = run(nlocal) + e_fewer = run(nlocal - 1) + assert torch.isfinite(e_full).all() and torch.isfinite(e_fewer).all() + assert not torch.allclose(e_full, e_fewer), ( + "owned-energy reduction must consume the dedicated n_local input" + ) + # keepalive: the raw pointer in ``comm`` must reference a live buffer + # through both ``run`` calls above (a real use, not a bare ``del``) + assert sendlist_indices.ctypes.data_as(ctypes.c_void_p).value == addr diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index c8a21a8a59..f945f14a75 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -792,6 +792,54 @@ def _train_and_get_ckpt(self, config: dict, tmpdir: str) -> str: self.assertTrue(os.path.exists(ckpt), "Checkpoint not created") return ckpt + def test_disable_graph_lower_persists_across_restart(self) -> None: + """The documented training opt-out survives a REAL save/restart. + + Regression (OutisLi review): ``disable_graph_lower()`` flipped only + a plain python bool, which a checkpoint restart silently reset -- + the fresh model is rebuilt from config before ``load_state_dict``, + and neither the state-dict keys nor ``_extra_state.model_params`` + carried the choice; on a binding-sel system the route-only switch + changed the training equation (energy by ~0.6) without warning. + The knob now lives in a persistent descriptor buffer + (``graph_lower_disabled``), so every state_dict round-trips it and + ``uses_graph_lower()`` re-syncs from the restored buffer. + """ + tmpdir = tempfile.mkdtemp(prefix="pt_expt_hatch_restart_") + try: + old_cwd = os.getcwd() + os.chdir(tmpdir) + try: + config = _make_config(self.data_dir, numb_steps=1) + config["model"]["descriptor"] = copy.deepcopy(_DESCRIPTOR_DPA1_NO_ATTN) + config = update_deepmd_input(config, warning=False) + config = normalize(config) + trainer = get_trainer(config) + desc = trainer.model.atomic_model.descriptor + self.assertTrue(desc.uses_graph_lower()) + desc.disable_graph_lower() + self.assertFalse(desc.uses_graph_lower()) + trainer.run() + + ckpt_path = os.path.join(tmpdir, "model.ckpt.pt") + self.assertTrue(os.path.exists(ckpt_path)) + + config2 = _make_config(self.data_dir, numb_steps=2) + config2["model"]["descriptor"] = copy.deepcopy(_DESCRIPTOR_DPA1_NO_ATTN) + config2 = update_deepmd_input(config2, warning=False) + config2 = normalize(config2) + trainer2 = get_trainer(config2, restart_model=ckpt_path) + desc2 = trainer2.model.atomic_model.descriptor + self.assertFalse( + desc2.uses_graph_lower(), + "disable_graph_lower() must survive a checkpoint restart " + "(route-only silent flip changes the training equation)", + ) + finally: + os.chdir(old_cwd) + finally: + shutil.rmtree(tmpdir, ignore_errors=True) + def test_restart(self) -> None: """Train 5 steps, restart from checkpoint, train 5 more.""" tmpdir = tempfile.mkdtemp(prefix="pt_expt_restart_") diff --git a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py index cadec99c89..8dd8af74c1 100644 --- a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py +++ b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py @@ -52,16 +52,24 @@ } -def _build_dpa1_data() -> dict: - """Build a serialized dpmodel data dict for a dpa1(attn_layer=0) energy model.""" +def _build_dpa1_data(config: dict | None = None) -> dict: + """Build a serialized dpmodel data dict for a dpa1(attn_layer=0) energy model. + + Parameters + ---------- + config : dict, optional + Model config to build from. Defaults to ``DPA1_CONFIG``. + """ from deepmd.dpmodel.model.model import ( get_model, ) - model = get_model(copy.deepcopy(DPA1_CONFIG)) + if config is None: + config = DPA1_CONFIG + model = get_model(copy.deepcopy(config)) return { "model": model.serialize(), - "model_def_script": copy.deepcopy(DPA1_CONFIG), + "model_def_script": copy.deepcopy(config), "backend": "dpmodel", "software": "deepmd-kit", "version": "3.0.0", @@ -97,6 +105,30 @@ def test_graph_pt2_has_lower_input_kind_graph(dpa1_dpmodel_data) -> None: assert "edge_capacity" not in meta +def test_graph_pt2_small_sel_exports() -> None: + """Graph-form ``.pt2`` export succeeds for a small-``sel`` model. + + The graph trace capacity derives from the synthetic trace system's + REAL edge count; the former sel-derived estimate + (``ceil(1.25 * nloc * sum(sel))``) overflowed the sel-free carry-all + builder whenever the actual degree exceeded ``sel`` (``edge overflow: + 36 real edges > edge_capacity 18`` at ``sel=2``). + """ + cfg = copy.deepcopy(DPA1_CONFIG) + cfg["descriptor"]["sel"] = 2 + data = _build_dpa1_data(cfg) + with tempfile.TemporaryDirectory() as d: + p = os.path.join(d, "m_graph_small_sel.pt2") + deserialize_to_file( + p, + data, + do_atomic_virial=True, + lower_kind="graph", + ) + meta = _read_metadata(p) + assert meta["lower_input_kind"] == "graph" + + @pytest.mark.parametrize( ("statistics_dtype", "expected"), [(torch.float32, "float32"), (torch.float64, "float64")], @@ -172,8 +204,9 @@ class _FakeDesc: def __init__(self, n_attn: int) -> None: self._n = n_attn - def get_numb_attn_layer(self) -> int: - return self._n + def uses_compact_edge_pairs(self) -> bool: + # mirrors dpa1's capability: attention rides center_edge_pairs + return self._n > 0 class _FakeAtomicModel: @@ -237,3 +270,85 @@ class _NoDesc: monkeypatch.setattr(torch, "__version__", "2.5.1") check_graph_trace_torch_version(_NoDesc()) + + +@pytest.mark.parametrize( + ("repformer_overrides", "should_raise"), + [ + # nlayers >= 2 so a non-last layer actually consumes compact pairs + # (the LAST layer is built with update_chnnl_2=False, which forces + # its g2/h2 updates off). + ({"nlayers": 2}, True), # default update_g2_has_attn=True + ({"nlayers": 2, "update_g2_has_attn": False, "update_h2": True}, True), + ( + {"nlayers": 2, "update_g2_has_attn": False, "update_h2": False}, + False, + ), # no pair consumers on any layer + # nlayers=1: the only layer is the last -> NO effective compact-pair + # consumer even with the arguments enabled; torch 2.5 stays usable. + ({"nlayers": 1}, False), + ], + ids=["g2_attn_2layers", "update_h2_2layers", "no_pair_consumers", "single_layer"], +) +def test_graph_trace_version_guard_dpa2_compact_pairs( + monkeypatch, repformer_overrides, should_raise +) -> None: + """A default graph-eligible DPA2 must trip the torch < 2.6 guard. + + Regression (OutisLi review): the guard keyed on dpa1's + ``get_numb_attn_layer``, which DPA2 does not implement, so every DPA2 + passed and compiled training / graph freeze failed deep inside + ``make_fx`` instead of the fast version error. The guard now keys on + the descriptor capability ``uses_compact_edge_pairs()``: DPA2's + ``update_g2_has_attn`` (default True) and ``update_h2`` both run the + compact ``center_edge_pairs`` realization; with both off the lower + traces backed symbols only and old torch stays usable. The capability + reads the EFFECTIVE per-layer flags: the last layer's g2/h2 updates + are structurally off (``update_chnnl_2=False``), so a single-layer + repformer never builds compact pairs and must NOT be rejected. + """ + import torch + + from deepmd.dpmodel.model.model import ( + get_model, + ) + from deepmd.pt_expt.utils.serialization import ( + check_graph_trace_torch_version, + ) + + cfg = copy.deepcopy(DPA2_GUARD_CONFIG) + cfg["descriptor"]["repformer"].update(repformer_overrides) + model = get_model(cfg) + assert model.atomic_model.descriptor.uses_graph_lower() is True + + monkeypatch.setattr(torch, "__version__", "2.5.1") + if should_raise: + with pytest.raises(RuntimeError, match=r"torch >= 2\.6"): + check_graph_trace_torch_version(model) + else: + check_graph_trace_torch_version(model) + + +# Small graph-eligible dpa2 for the version-guard regression above. +DPA2_GUARD_CONFIG = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa2", + "repinit": { + "rcut": 4.0, + "rcut_smth": 0.5, + "nsel": 10, + "neuron": [4, 8], + "axis_neuron": 2, + }, + "repformer": { + "rcut": 3.0, + "rcut_smth": 0.5, + "nsel": 6, + "nlayers": 1, + "g1_dim": 8, + "g2_dim": 4, + }, + }, + "fitting_net": {"neuron": [8, 8], "seed": 1}, +} diff --git a/source/tests/pt_expt/utils/test_neighbor_list.py b/source/tests/pt_expt/utils/test_neighbor_list.py index c52bdf03c2..e795f1eedc 100644 --- a/source/tests/pt_expt/utils/test_neighbor_list.py +++ b/source/tests/pt_expt/utils/test_neighbor_list.py @@ -289,6 +289,8 @@ # attention softmax mechanism (dpa1.py's `_graph_attention`, gated only on # `attn_layer > 0`, entered identically regardless of concat/strip), so one # tolerance covers both. +# dpa2: repformer smooth attention (sel-independent on carry-all graph, with +# sel-padding phantoms on dense) amplified through message passing. KNOWN_GRAPH_DENSE_DIVERGENT = {"dpa1_smooth", "se_atten_v2"} @@ -539,22 +541,34 @@ def test_default_fallback(name: str) -> None: the dense route (``call_common``'s ``neighbor_list is not None`` branch), while ``None`` lets pt_expt's default-flip (decision #17) route to the carry-all graph instead -- two genuinely different algorithms, not two - evaluations of one. Both hardcode/pin ``smooth_type_embedding=True`` - (unlike ``model_dpa1`` above, which pins it ``False`` for exactly this - reason), so graph and dense intentionally diverge (NeighborGraph PR-D: - dense keeps sel-padding phantom terms in the attention softmax - denominator, the graph route does not -- see + evaluations of one. ``dpa1_smooth`` and ``se_atten_v2`` hardcode/pin + ``smooth_type_embedding=True`` (unlike ``model_dpa1`` above, which pins it + ``False`` for exactly this reason), so graph and dense intentionally diverge + (NeighborGraph PR-D: dense keeps sel-padding phantom terms in the attention + softmax denominator, the DPA1 graph route does not -- see ``test_block_compact_graph_smooth_clean_divergence`` in - ``test_dpa1_graph_attention_parity.py`` for the same invariant at the - block level). At this test's non-binding ``sel=40`` (vs. <=5 real - neighbors), the gap is small but non-zero and deterministic (not CUDA - ULP-style non-determinism): energy differs by ~1e-7, force by ~1e-6, - virial (a derivative, so it amplifies the softmax-denominator - perturbation more than the value itself) by up to ~1.3e-5. We assert - BOUNDED closeness (atol=3e-5, rtol=1e-3 -- individual virial/force - components can be near-zero, so atol dominates) rather than bit-identity - for these two, keeping the check meaningful rather than silently dropping - coverage. + ``test_dpa1_graph_attention_parity.py`` for the same invariant at the block + level). At this test's non-binding ``sel=40`` (vs. <=5 real neighbors), the + gap is small but non-zero and deterministic (not CUDA ULP-style + non-determinism): energy differs by ~1e-7, force by ~1e-6, virial (a + derivative, so it amplifies the softmax-denominator perturbation more than + the value itself) by up to ~1.3e-5. + + ``dpa2`` is NOT in that set: its repformer attention uses the + fixed-phantom-count compensation (``segment_softmax(phantom_count=...)``, + ``Atten2Map``/``LocalAtten.call_graph``), so at non-binding ``sel`` the + carry-all graph softmax denominator equals the dense one term-for-term and + the two routes agree to the fp64 noise floor (measured 0 / 7e-18 / 3e-17 + on this fixture). It therefore takes the tight (``rtol=1e-10``) bound like + a non-divergent model. This test cannot by itself prove the graph route is + actually TAKEN for dpa2 (a silent dense fallback would also be bit-tight): + that positive check lives in + ``test_dpa2_graph_lower.py::TestDpa2GraphLower.test_binding_sel_diverges`` + and ``test_binding_sel_diverges_with_attention``, where a BINDING sel makes + the carry-all graph attend over more neighbors than the sel-truncated dense + body, so ``None`` (graph) differs from ``"legacy"`` (dense) by a bounded + non-zero amount -- a difference that would collapse to zero if the default + silently fell back to dense. """ coord_np, atype_np, box_np = _system() md = get_model(copy.deepcopy(ALL_MODELS[name])).to(env.DEVICE) @@ -570,11 +584,17 @@ def test_default_fallback(name: str) -> None: coord_np, dtype=torch.float64, device=env.DEVICE ).requires_grad_(True) outs[tag] = md.forward(coord_t, atype_t, box=box_t, do_atomic_virial=True, **kw) - tol = ( - {"rtol": 1e-3, "atol": 3e-5} - if name in KNOWN_GRAPH_DENSE_DIVERGENT - else {"rtol": 1e-10, "atol": 1e-12} - ) + if name in KNOWN_GRAPH_DENSE_DIVERGENT: + # per-model divergence envelopes: loosening dpa1's bound to dpa2's + # message-passing-amplified scale would let a real dpa1 graph bug + # three orders of magnitude above its documented divergence pass. + tol = ( + {"rtol": 1e-3, "atol": 2e-2} + if name == "dpa2" + else {"rtol": 1e-3, "atol": 3e-5} + ) + else: + tol = {"rtol": 1e-10, "atol": 1e-12} for k in ("energy", "force", "virial"): np.testing.assert_allclose( outs["none"][k].detach().cpu().numpy(),