feat(dpa4): DPA4/SeZM on the NeighborGraph lower — graph-native port, multi-rank, native spin, charge-spin/bridging#5884
Conversation
…erLayer.call_graph parity Add parametrized test case 'no_chnnl_2_no_gg1' to test_repformer_layer_call_graph_parity that exercises two previously untested branches: 1. update_chnnl_2=False: Tests the short-circuit path where g2/h2 updates are skipped and returned unchanged (lines 2397-2401 in repformers.py call_graph). This is the critical production path for the LAST layer of every DPA2 model. 2. cal_gg1=False (gg1=None): Tests the graph path when no gg1-dependent updates are needed (all of update_g1_has_drrd, update_g1_has_conv, update_g1_has_attn, update_g2_has_g1g1 are False). With g1_out_mlp=False, g1_mlp is seeded with the identity [g1] for xp.concat. Test passes all 30 cases (29 baseline + 1 new).
gen_dpa2.py section B builds a graph-eligible dpa2 (use_three_body=False; three-body is graph-ineligible), exports the graph-form .pt2 (which embeds the nested with-comm artifact automatically for message-passing models) plus an independent dense-nlist .pt2 of the same weights, cross-checks atomic energies/forces/total virial between the two at gen time (NaN-safe), and writes deeppot_dpa2_graph.expected from the graph artifact eval. Unlike the dpa1 precedent the .expected is not copied from the nlist oracle: for a message-passing model the graph path's full-to-source per-edge atomic-virial decomposition legitimately differs from the dense per-atom decomposition (only the sum agrees), and cross-path e/f noise (~1e-10) sits at the universal gtest's 1e-10 double tolerance. The dpa2_graph_ptexpt VariantDeepPotCase row mirrors dpa1_graph_ptexpt (same tolerances and capability flags).
- Add nall_real == 0 guard at the top of graph with-comm arm (lines 922-929) to throw a clear error instead of silently crashing with Dim violation - Fix stale comment (lines 1050-1051) to reflect that ghost forces fold back via both plain and with-comm graph routes
…h-comm route The with-comm GRAPH artifact derives n_local from the nlocal comm tensor IN-GRAPH (owned-node energy mask); after move_to_device_pass that derivation is a device kernel, so feeding it a CPU tensor makes the kernel read a host pointer as device memory -> CUDA illegal memory access on every multi-rank run (first live exercise of the route, T4). The other six comm tensors are consumed only by the opaque border_op whose host code dereferences their data_ptr, so they stay on CPU -- same placement the dense with-comm route uses for all eight. Verified on Tesla T4: minimal AOTI repro crashes with CPU nlocal and matches the plain graph artifact bit-for-bit with device nlocal.
… DEVICE
_trace_and_compile_graph built its synthetic NeighborGraph trace inputs on
the module-level DEVICE constant instead of the device the model's own
parameters/buffers actually live on. In real training these always match
(the trainer moves the model to DEVICE before compiling), but any caller
that intentionally holds the model on a different device (e.g. a test
pinning the model to CPU while running on a CUDA host, mirroring the .pt2
export's model.to("cpu") pattern) hits a device split between the traced
graph tensors and the model params.
DPA2.call_graph's tebd gather (dpmodel/descriptor/dpa2.py:1159) surfaced
this first: xp.take(tebd_table, atype_local, axis=0) requires tebd_table
(a model param, left un-asarray'd to avoid detaching its gradient -- the
documented dpa1 lesson) and atype_local (moved to device(graph.edge_vec))
to share a device. Under test_compiled_training_graph_smoke the model is
pinned to CPU but the trace sample was still built on the global DEVICE
(cuda:0 on a GPU box), so tebd_table stayed CPU while atype_local was CUDA.
DPA1's call_graph has the identical gather pattern and would fail the same
way, but no existing dpa1 test calls _trace_and_compile_graph directly with
a device-pinned model, so the bug was latent there since PR-B (deepmodeling#5604).
Fix: add _model_trace_device(), reading the device off the model's own
parameters/buffers (falling back to DEVICE only if the model exposes
neither), and use it to build the trace sample instead of the global
DEVICE. No change to the gather itself -- gradient flow through
type_embedding is untouched.
…back test dpa2 became graph-eligible (uses_graph_lower=True), so neighbor_list=None now routes to the carry-all graph while explicit DefaultNeighborList() forces the dense route. The smooth repformer attention diverges intentionally between these paths (sel-independent on graph, sel-padded on dense), amplified through message passing. Observed divergence at this fixture: - energy rel ~1e-4, abs ~9.5e-4 - force abs ~5e-3 - virial abs ~1.3e-2 Set atol=2e-2, rtol=1e-3 to accommodate dpa2's larger divergence compared to dpa1_smooth/se_atten_v2.
for more information, see https://pre-commit.ci
…adapter DPA2's repinit is hardwired to attn_layer=0, where the dense se_atten body leaks a phantom padding-neighbor -davg/dstd residual that the graph path deliberately omits (the physically-correct behavior, accepted as KNOWN_GRAPH_DENSE_DIVERGENT). So _call_graph_adapter is bit-exact vs _call_dense only in the trivial-statistics regime (davg=0, dstd=1); for any model with non-trivial stats it diverges. The T7 call gate routed the dense .call() through _call_graph_adapter for graph-eligible configs, which made the cross-backend consistency reference diverge from the leaky tf/pt/pd/jax dense bodies (source/tests/pt|pd/model/ test_dpa2.py::test_consistency, 71% mismatch) and crashed the universal zero_forward on empty systems. The graph-native route is reached exclusively through call_graph (pt_expt forward_atomic_graph + C++ graph .pt2), never through call, so .call() now always runs _call_dense. _call_graph_adapter is retained as the bit-exact-regime reference exercised by TestDPA2AdapterBitExact. Updates TestDPA2CallRouting.test_call_routing_graph_eligible to assert the corrected routing (call == _call_dense) and corrects the _call_graph_adapter bit-exactness docstring.
The C++ memleak matrix runs gen scripts with LD_PRELOAD=liblsan (the sanitizer-instrumented deepmd op .so needs the LSAN runtime). Evaluating the AOTInductor-compiled deeppot_dpa2_graph.pt2 BACKWARD (forces) under that runtime segfaults inside the repformer's fused backward kernel (cpp_fused_..._slice_backward_...): an AOTI-compiled-code vs LeakSanitizer allocator incompatibility, NOT a graph-code bug. The identical backward is finite and correct 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 (the memleak build exists to leak-check deepmd's C++ ops), so excluding the dpa2 graph artifact from the memleak build loses no real coverage. - gen_dpa2.py: skip the graph section (Section B) when LD_PRELOAD contains liblsan; the dense DPA2 .pt2 (section A) is still generated. - test_deeppot_universal.cc: add skip_if_artifact_missing to VariantDeepPotCase and set it on the dpa2_graph_ptexpt row, so SetUp GTEST_SKIPs (instead of failing) when the artifact is absent under memleak. Every other case still hard-fails on a missing artifact. LAMMPS/ipi tests already run only when !check_memleak, so no other memleak consumer of the artifact remains.
for more information, see https://pre-commit.ci
…red trunk Review: PR deepmodeling#5884 (OutisLi, P1). The shared dpmodel _run_graph hard-coded random_gamma=False, so the default random_gamma=True was silently ignored for all pt_expt DPA4 training (graph route AND dense adapter), diverging from pt's 'self.random_gamma and self.training'. - DescrptDPA4 (dpmodel) gains an _in_training_mode() runtime hook (False: dpmodel is an inference/reference backend), and the edge-cache call site becomes 'random_gamma=self.random_gamma and self._in_training_mode()'. - The pt_expt wrapper overrides the hook with the torch module's training flag: train-mode forwards draw a fresh gamma per call, eval/export forwards stay fixed (the export path already calls model.eval() before tracing, serialization.py:1440). - Tests: pt_expt spy test pins train->True / eval->False / dpmodel->False at the edge-cache seam plus eval-mode bit-determinism (spy-based because the model is roll-equivariant -- output differences have no deterministic teeth); dpmodel test pins the hook False + bit-identical repeated calls with random_gamma=True.
…ing_label in the native-spin builders Review: PR deepmodeling#5884 (OutisLi, P1 x2). The DPA4 native-spin builders did a direct boolean conversion of spin.use_spin, breaking the public schema's index and symbol forms (['Ni'] became [True] and failed the length check; [0, 1] became [False, True]), and dropped spin.allow_missing_label (so datasets without spin.npy were rejected despite the documented zero default). - New pure helper deepmd.utils.spin.normalize_spin_use_spin(use_spin, type_map) -> list[bool]: single owner of the three-form contract pt implements in its frozen tree (bool passthrough / index scatter / symbol lookup with ValueError on unknown symbols); direct unit tests cover all forms incl. the empty list and input purity. - Both builder twins (dpmodel model.py get_dpa4_native_spin_model and pt_expt get_model._get_dpa4_native_spin_model) normalize via the helper and forward allow_missing_label into Spin(...). - get_additional_data_requirement drops its getattr has_spin/spin dances for direct calls (has_spin is base-declared; spin models carry .spin). - Tests: config-form parametrization (symbols/indices/booleans/unknown symbol) on both builders; the pt_expt data-requirement regression pins must=False + default=0.0 with allow_missing_label=True and must=True without.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| the model is roll-equivariant, so an output-difference check would | ||
| have no deterministic teeth. | ||
| """ | ||
| import deepmd.dpmodel.descriptor.dpa4 as dpa4_mod |
…calls Twin of the pt_expt cleanup: has_spin() is declared on the shared make_base_model base with a concrete False default, so the getattr(model, 'has_spin', False) + callable() probes in pd's train wrapper, data-requirement helper, and DeepEval are dead-defensive. DeepEval's variant was additionally self-defeating: it forced _has_spin=False whenever the attribute WAS callable, which only worked because pd has no spin model classes -- with the base-declared method that branch became always-taken. Direct calls preserve behavior (False for every current pd model) and raise on a typo instead of silently degrading. Validated locally by py_compile+ruff only: paddle is not installed in this venv (verified), so pd unit tests must be exercised by CI.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The native-spin change-bias path still diverges from the PT reference and can silently leave the output bias unchanged; see the inline comment.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The pt_expt native-spin migration still rejects a public conditioning combination supported by the PT reference; see the inline comment.
| raise NotImplementedError( | ||
| "spin scheme 'native' requires the DPA4/SeZM descriptor" | ||
| ) | ||
| if data["descriptor"].get("add_chg_spin_ebd", False): |
There was a problem hiding this comment.
[P1] Preserve PT native spin combined with charge-spin FiLM
The PT reference does not reject this configuration: SeZMNativeSpinModel.forward and its exportable lower both accept charge_spin alongside the native spin input. I built the same add_chg_spin_ebd=True native-spin model at this HEAD; PT returns energy, force, and force_mag, and changing only charge_spin changes the energy (5.97e-05 in the probe), whereas pt_expt stops here with NotImplementedError. This is not a limitation of the migrated graph math or eager wrapper: bypassing this builder guard and constructing the current DPA4NativeSpinModel over the same descriptor produces both a nonzero charge-spin response and a nonzero magnetic force. Please port the combined public configuration rather than dropping PT behavior, including a charge_spin slot in the native-spin graph export/metadata and the corresponding serialization, DeepEval/C++ inference, training, and freeze regressions.
There was a problem hiding this comment.
Ported in c4ddea343, 57a008e30, 9cf9f36fc, d1dac261e — the combined public configuration is now supported end to end rather than rejected:
- Builders: both guards lifted;
add_chg_spin_ebd+spin.scheme='native'constructs the combined model in dpmodel and pt_expt. - Eager: weight-copied fp64 parity vs pt's
SeZMNativeSpinModelwith the samecharge_spinat rtol/atol 1e-12 (energy/force/force_mag, two conditioning values), plus nonzero charge-spin AND spin responses in the same model (your probe, integer-valued: theChargeSpinEmbeddinglookup is categorical). - Graph export ABI:
charge_spinoccupies a conditional slot-13 tail (spin stays pinned at index 10); symbolic-trace parity at 1e-12 with a LIVE slot (an integer-valued change moves the traced energy — pins that it is not baked), dynamic-shape spec + synthetic sample + freeze metadata (has_chg_spin_ebd/dim_chg_spin) included. - Serialization/freeze/DeepEval: a combined model freezes to a graph
.pt2and DeepEval evaluates it end to end through the compiled artifact — slot live, parity vs eager at 1e-10;DeepPot.eval(..., spin=..., charge_spin=...). - Training: 2-step trainer smoke on the NiO spin data through the
default_chg_spinmetadata path (this also exposed and fixed a latent dpmodel bug: the default-branch converted the numpy default with the numpy namespace against a torch device). - C++:
DeepSpinPTExpt::run_model_graphappends the conditional charge_spin tail fromdefault_chg_spinmetadata, mirroring its nlist/edge siblings (metadata-default is the C++ production contract, as elsewhere; compiles clean).
Known residual: no dedicated C++ gtest row for a combined fixture yet (the Python DeepEval e2e covers the identical artifact); multi-rank stays out of scope (native spin has no with-comm ABI by contract).
There was a problem hiding this comment.
The implementation itself now passes my targeted PT-parity, graph-export, training, and compiled-artifact checks, but the current HEAD still leaves source/tests/pt_expt/model/test_get_model_dpa4.py::TestGetModelDPA4::test_native_spin_add_chg_spin_ebd_raises asserting that this configuration must raise NotImplementedError. I reproduced that failure at f51ca5e3, and the current Python CI shard fails on the same stale negative regression. Please remove or replace it with a positive construction assertion (the new end-to-end coverage can remain the main behavioral regression). I am leaving this thread unresolved until the branch is internally consistent and CI is green.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The new graph tests cover only the SFPG descriptor gate, while the PT model-level analytical ZBL half of bridging is still absent and rejected in pt_expt; see the inline comment.
| return np.asarray(out).reshape(n_points, 3, -1) | ||
|
|
||
|
|
||
| def test_call_graph_bridging_frozen_sphere_invariance() -> None: |
There was a problem hiding this comment.
[P1] Implement and test the complete PT bridging model
This regression proves only the descriptor-side InnerClamp/SFPG invariance. In the PT reference, bridging_method="ZBL" also constructs a model-level InterPotential and adds its analytical ZBL atomic energy before the shared edge autograd, which is what supplies the short-range repulsive energy, force, and virial. That half is absent here: pt_expt still rejects every non-none bridging_method, its PT-checkpoint deserializer rejects bridging models, and there is no InterPotential equivalent in dpmodel/pt_expt. The current PT regressions confirm that ZBL creates the potential and adds positive energy, while the pt_expt regression confirms only NotImplementedError. Consequently users cannot construct the claimed single-rank bridging model, and manually enabling these descriptor switches would omit the physical ZBL term. The separate DP-ZBL model is not equivalent to SeZM's native SFPG plus additive ZBL path. Please port the model-level potential and add PT-parity training, freeze, DeepEval/C++ inference tests rather than treating the descriptor-only gate as complete bridging support.
There was a problem hiding this comment.
Ported in 0c388bec8, 32514115e, 412cc2129, 69a2697f0 — the complete model-level ZBL term is now implemented, not just the descriptor-side SFPG/InnerClamp gate:
InterPotential(backend-agnostic): exact array-API port of pt's universal-ZBL math (constants/screening verbatim; per-edge half energies scattered per atom; virtual-type wrap+mask; fp64). Unit tests port pt's known-value OO/OH cases, masking, and error paths, plus a torch-namespace autograd smoke.- Injection: owned by
DPAtomicModel— added to the per-atom energy after fitting+masking, before the model layer's edge autograd (pt's sezm Step 5).graph.edge_vecthere is pt_expt's autograd leaf, so the ZBL force/virial flow through the sameedge_energy_derivbackward as the learned energy. The dense (nlist) route raises for bridging models (pt has no dense conservative path; fail-closed beats silently dropping the term). - Builders/serialization/interop:
bridging_method/bridging_r_inner/bridging_r_outerbuild working models in dpmodel and pt_expt (radii feed the descriptor's InnerClamp/BridgingSwitch, as in pt);bridging_methodrides the atomic dict (emitted only when active); the pt-checkpoint deserializer now maps pt's wrapper key instead of rejecting — a pt-serialized ZBL checkpoint loads into pt_expt with 1e-12 energy parity. - Regressions: weight-copied fp64 parity vs pt's
SeZMModelwith ZBL at 1e-12 (energy/force/virial); ZBL-vs-plain positive repulsion on identical weights; central-FD force; graph.pt2freeze + DeepEval through the compiled artifact at 1e-10; 2-step training smoke; native-spin × bridging composition (pt supports it — ours composes via the atomic layer);bridging_r_*defaults; with-comm gate off (pt'ssupports_edge_parallel()==Falsecontract →has_comm_artifact=False, single-rank only). argcheck'sdoc_only_pt_supportedflags dropped from the three keys.
Known residuals: no dedicated C++ gtest/LAMMPS fixture row (the term is baked into the AOTI artifact — C++ needs no code change; the Python DeepEval e2e exercises the identical compiled artifact), and multi-rank stays off by the pt contract.
There was a problem hiding this comment.
The ZBL implementation now passes my formula-level, PT-parity, finite-difference, multi-frame, pair-exclusion, training, and compiled-artifact checks. However, the current HEAD still contains two stale negative regressions that make CI fail: TestGetModelDPA4.test_unsupported_keys_raise still lists bridging_method="ZBL" as unsupported, and TestDPA4Interop.test_guard_branches_raise still requires a PT ZBL wrapper to raise NotImplementedError. I reproduced both failures at f51ca5e3, matching the current Python CI shards. Please remove the bridging entry from the unsupported-key table and replace/remove the obsolete interop guard (the new positive ZBL checkpoint-interoperability test is the appropriate behavior). I am leaving this thread unresolved until those expectations and CI are updated.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Requesting changes for the backend-breaking native-spin BaseModel deserialization dispatch described inline.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Requesting changes for the broken default native-spin freeze route described inline.
| # implementation. ``is_native_spin`` distinguishes them at every seam | ||
| # below (model rebuild, graph rejection, graph sample-input/dynamic-shape | ||
| # ABI, trace call site). | ||
| is_native_spin = data["model"].get("type") in ( |
There was a problem hiding this comment.
[P1] Route native-spin freeze to graph before choosing the export ABI
The public freeze path still defaults --lower-kind to nlist and chooses a .pte suffix from that value, while this native-spin wrapper has no dense/nlist lower. I reproduced the ordinary default path at this HEAD with a freshly built native-spin model: deserialize_to_file(..., lower_kind="nlist") reaches the dense-spin branch and fails with TypeError: make_model..CM.forward_common_lower_exportable() got multiple values for argument 'fparam'. Existing native-spin export tests all pass lower_kind="graph" explicitly, so they do not cover the user-facing default. This also contradicts the new documentation saying native-spin models are always graph-frozen regardless of --lower-kind.
Please resolve native-spin to graph at the public freeze layer before selecting both the export ABI and default output suffix (or reject nlist explicitly and make the default auto/graph), and add a regression that freezes a native-spin checkpoint through the default public CLI/function path without supplying --lower-kind, asserting a graph-kind .pt2 artifact.
There was a problem hiding this comment.
Rechecked at f51ca5e: this is still unresolved. The public freeze() and deserialize_to_file() defaults remain lower_kind="nlist", and _resolve_lower_kind() only reroutes when the caller explicitly passes "auto". The native-spin export tests still pass lower_kind="graph" explicitly, so the ordinary default path is not covered and still selects the dense/nlist ABI (and .pte default suffix) for a graph-only native-spin model. Please resolve native-spin to graph before choosing both the export ABI and default suffix, and add a regression through the default public freeze path with no lower-kind argument.
…NativeSpinEnergyModel Native spin is a model input/output/autograd contract, not a DPA4-specific model type (PR deepmodeling#5884 review 3638137290). Replace the has-a DPA4NativeSpinModel wrapper with the make_model-aligned is-a pattern: - make_native_spin_model(T_Model) class factory (mirrors make_model and pt's SeZMNativeSpinModel(SeZMModel)): spin conditioning, spin-aware output defs, public-key call, and serialize = parent flat dict + 'spin' field. The deserialize closure over T_Model rebuilds through the CALLING backend's model class. - NativeSpinEnergyModel = make_native_spin_model(EnergyModel), registered in the BaseModel plugin registry under the single descriptor-agnostic wire type 'native_spin'; the hard-coded branch in BaseBaseModel.deserialize is deleted (registry dispatch is backend-aware by construction). The descriptor-specific legacy strings 'dpa4_native_spin'/'sezm_native_spin' are NOT aliased to the generic class -- they fail fast at the registry (pinned by test; nothing is released, and pt-checkpoint conversion was never claimed). - get_native_spin_model builder gates on the built descriptor's supports_native_spin() capability instead of a descriptor-type list: a future native-spin descriptor needs zero model/dispatch/builder changes. - The wrapper's 15 delegation methods and __getattr__ fallback are gone -- inheritance provides them.
…registry dispatch replaces native-spin seams The pt_expt twin instantiates the shared factory on THIS backend's EnergyModel and registers under 'native_spin' in the PT_EXPT BaseModel registry, so the public BaseModel.deserialize returns the torch class (backend-aware, no hard-coded branches). Consequences: - The bespoke deserialize override and the __getattr__ delegation shim are deleted: the factory's deserialize closure rebuilds through the pt_expt EnergyModel, and inheritance replaces delegation. forward and forward_lower_graph_exportable (spin@index-10 ABI unchanged) now call the inherited call_common/forward_common_lower_graph_exportable on self instead of a backbone attribute. - get_native_spin_model (public, matching its sibling builders) routes the backbone through get_sezm_model for the DPA4/SeZM family (keeping its guards) or get_standard_model otherwise, gates on descriptor.supports_native_spin(), and re-classes via atomic_model_. The standard-typed get_model path now dispatches scheme='native' too, so a future native-spin descriptor works from a plain standard config. - deep_eval and deserialize_to_file drop their native-spin special cases (registry + has_spin() cover them); the eval_typeebd unwrap is isinstance-precise on the virtual-atom SpinModel only. - NEW NativeSpinModelKind marker base in the factory module: each backend's concrete class is a parallel factory product with NO subclass relation, so the old isinstance-against-dpmodel-class with-comm gate went silently dead -- the freeze then compiled a with-comm artifact for the single-rank-only spin ABI (caught by test_native_spin_graph_freeze cold-cache A/B). Seams test the shared marker instead. Validated locally: native-spin+spin batteries 25, export file 4/5 (test_dpa4_freeze_to_pt2[auto-graph] fails identically at the pre-refactor commit with a cold inductor cache -- pre-existing local inductor codegen bug, not a regression), deep_eval 93, dpmodel 19.
Asserts the reviewer's exact contract (PR deepmodeling#5884, 3638137290): the public pt_expt BaseModel.deserialize(model.serialize()) returns the pt_expt NativeSpinEnergyModel, a torch.nn.Module accepting .to(DEVICE) (the concrete finetune failure), carrying the graph-export machinery, with forward parity at 1e-12 on energy/force/force_mag.
for more information, see https://pre-commit.ci
| def call( | ||
| self, | ||
| coord: np.ndarray, | ||
| atype: np.ndarray, | ||
| spin: np.ndarray, | ||
| box: np.ndarray | None = None, | ||
| fparam: np.ndarray | None = None, | ||
| aparam: np.ndarray | None = None, | ||
| do_atomic_virial: bool = False, | ||
| ) -> dict[str, np.ndarray]: |
| def forward( | ||
| self, | ||
| coord: torch.Tensor, | ||
| atype: torch.Tensor, | ||
| spin: torch.Tensor, | ||
| box: torch.Tensor | None = None, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, | ||
| do_atomic_virial: bool = False, | ||
| charge_spin: torch.Tensor | None = None, | ||
| ) -> dict[str, torch.Tensor]: |
…8047227, task 1) pt's SeZMNativeSpinModel accepts charge_spin alongside the native spin input; the dpmodel builder rejected the combination. Lift the guard and thread charge_spin through the native-spin factory call() into call_common (the graph lower already carries both: forward_atomic_graph forwards each kwarg per the descriptor's declared capabilities). Test pins the combined model: builds with has_chg_spin_ebd() AND has_spin(), charge_spin conditions the energy, and spin still conditions it in the same model. Note the categorical contract discovered while writing it: ChargeSpinEmbedding casts the (charge, spin) frame pair to int64 lookup indices, so only integer-valued changes condition the model.
…view 3638047227, task 2) - Lift the pt_expt builder rejection; the eager forward already threads charge_spin through call_common (kept for signature compatibility, now live). - Weight-copied fp64 parity vs pt's SeZMNativeSpinModel with the SAME charge_spin input at rtol/atol 1e-12 (both cs=0 and an integer-valued probe), plus nonzero charge_spin AND spin responses in the same combined model. - Trainer smoke on the NiO spin data (no charge_spin file) exercises the default_chg_spin metadata path: get_additional_data_requirement marks the label optional with the descriptor default; 2 steps, finite loss. - Fix a latent dpmodel bug the smoke exposed (first torch-path exercise of the default): _canonicalize_charge_spin derived its namespace from the numpy default_chg_spin attribute, so numpy's asarray met a torch device object and raised. The caller's xp is now threaded in (array- API pitfall: convert numpy attributes into the caller's namespace).
…38047227, task 3) The graph-spin positional ABI gains a conditional charge_spin tail at slot 13 (spin stays pinned at slot 10; fparam/aparam at 11/12), mirroring the energy model's convention: - NativeSpinEnergyModel.forward_lower_graph_exportable threads charge_spin through both trace layers instead of pinning None. - build_synthetic_graph_inputs emits the native-spin sample with the slot-13 charge_spin (None when the model has no add_chg_spin_ebd); _build_graph_dynamic_shapes marks it nframes-dynamic; the freeze trace call passes it through. Metadata already carries dim_chg_spin/has_default_chg_spin generically. - Tests: combined symbolic-trace parity vs eager forward_common_lower_graph at 1e-12 with a LIVE slot-13 (an integer-valued charge_spin change moves the traced energy -- pins that the slot is not baked); existing exportable/torch-export tests extended to the 14-slot ABI; combined graph freeze produces a .pt2 whose metadata pins has_chg_spin_ebd/dim_chg_spin. Non-combined native freeze re-validated (slot 13 = None).
…C++ (review 3638047227, task 4) - DeepEval: _eval_model_spin threads charge_spin into the graph-native fast path; _eval_model_graph_spin builds the slot-13 input via _make_charge_spin_input (explicit value, metadata default, or None when the model has no add_chg_spin_ebd -- the exported signature drops the slot in that case). - C++ DeepSpinPTExpt::run_model_graph appends the conditional charge_spin tail from the default_chg_spin metadata when dim_chg_spin > 0, mirroring its nlist/edge/with-comm siblings (the C++ interface's production contract is the metadata default; explicit per-call charge_spin remains a Python-API feature). Compiles clean (deepmd_cc target). - The combined freeze test now runs DeepEval end to end through the compiled artifact: slot-13 is LIVE (integer-valued charge_spin change moves the energy) and matches the eager forward at 1e-10.
for more information, see https://pre-commit.ci
…ng term) Array-API port of pt's InterPotential (review 3638077323, task 1 of the bridging plan): universal ZBL screened nuclear repulsion on the edge form -- per-edge half energies scattered into per-atom energies via segment_sum, virtual-type wrap+mask, r clamped at 1e-10, fp64 math cast back to edge_vec dtype. Constants and ELEMENT_TO_Z copied verbatim from the frozen pt reference. Tests port pt's known-value OO/OH cases (formula reproduced in-test; half-split sums to the full analytic pair energy at 1e-5), virtual-type masking, edge-mask zeroing, unknown element/mode ValueErrors, and a torch-namespace smoke with autograd gradient through edge_vec (1e-12 vs numpy).
…ew 3638077323, task 2)
- InterPotential moves to deepmd/dpmodel/atomic_model/ (the atomic layer
owns per-atom energy assembly; also breaks an import cycle).
- BaseAtomicModel gains bridging_method ('none' default): builds the
InterPotential from (type_map, mode); forward_common_atomic_graph adds
the term to the per-atom energy AFTER fitting+masking, BEFORE the
model layer's edge autograd (pt's sezm Step 5) -- graph.edge_vec is
pt_expt's autograd leaf, so force/virial ride the shared edge backward
with no extra autograd code. The dense (nlist) route RAISES for
bridging models (pt has no dense conservative path; silently dropping
the term is the dangerous direction).
- Serialization: 'bridging_method' emitted in the atomic dict only when
active (absent == disabled keeps existing dicts byte-stable); the
InterPotential is stateless and rebuilt on deserialize, like pt.
- get_standard_model wires the config: bridging_r_inner/r_outer feed the
descriptor's InnerClamp/BridgingSwitch (mirrors pt's builder);
bridging_method reaches the atomic model through the CM kwargs.
- Tests: ZBL model vs byte-identical plain twin shows positive repulsion
on a close pair; round-trip energy identical at 1e-12; dense-route
raise pinned; plain dicts carry no new key. Broad atomic-model battery
62 passed.
…review 3638077323, task 3) Ownership correction from review feedback: the generic BaseAtomicModel must not carry the bridging feature. All bridging state and behavior move to DPAtomicModel (the descriptor+fitting atomic model that assembles the per-atom energy): ctor kwarg, InterPotential build, conditional 'bridging_method' serialize key, the graph-route injection override (after fitting+masking, before the model layer's edge autograd -- graph.edge_vec is pt_expt's leaf, so ZBL force/virial ride the shared edge backward), and the dense-route rejection. The base gains ONLY the concrete-default capability method has_analytical_bridging() (the supports_native_spin pattern; no state, no imports, implementer- agnostic), which the with-comm export gate calls -- bridging models are single-rank only (pt's supports_edge_parallel contract). pt_expt wiring: - get_sezm_model builds bridging models (radii -> descriptor InnerClamp keys, method -> atomic-model kwarg) instead of rejecting. - pt-checkpoint interop: _unwrap_pt_sezm_model maps pt's model-level wrapper key onto our atomic-dict key instead of raising (the wrapper radii are redundant -- pt's descriptor dict carries the InnerClamp). - InnerClamp/BridgingSwitch/InterPotential get register_dpmodel_mapping converters (parameter-free rebuilds) -- first pt_expt instantiation. Tests (5 pt_expt + 4 dpmodel): weight-copied fp64 parity vs pt's SeZMModel with ZBL at 1e-12 (energy/force/virial), positive-repulsion twin, central-FD force at 1e-6, with-comm gate off, and a pt-serialized ZBL checkpoint deserializing into pt_expt with 1e-12 energy parity.
…38077323, task 4) - Graph .pt2 freeze of a ZBL model: the analytical term survives make_fx/AOTI tracing; DeepPot through the compiled artifact matches eager forward at 1e-10 (energy/force); metadata pins has_comm_artifact=False (single-rank contract). - 2-step trainer smoke on the NiO data with finite loss. - Native-spin x bridging composition pinned (pt's SeZMNativeSpinModel inherits the term; our atomic-layer injection composes with the native-spin factory for free): both the close-pair ZBL repulsion and force_mag live in one model. - bridging_r_inner/r_outer defaults (0.5/0.8) pinned. - argcheck: drop doc_only_pt_supported from the three bridging keys (supported by dpmodel/pt_expt now).
for more information, see https://pre-commit.ci
| def call( | ||
| self, | ||
| edge_vec: Array, | ||
| edge_index: Array, | ||
| atype_flat: Array, | ||
| edge_mask: Array, | ||
| n_node: int, | ||
| real_type_count: int | None = None, | ||
| ) -> Array: |
| "aparam", _model.get_dim_aparam(), atomic=True, must=True | ||
| ) | ||
| ) | ||
| if _model.has_spin(): |
There was a problem hiding this comment.
[P1] Preserve the capability-safe non-spin fallback
This unconditional call regresses the existing contract of get_additional_data_requirement and diverges from the PT reference, which uses getattr(_model, "has_spin", False) and only reads spin.allow_missing_label after the capability is present. At this HEAD, both existing TestAdditionalDataRequirement::test_has_default_fparam and test_no_default_fparam fail here with AttributeError: _M has no attribute has_spin; these are also current CI failures. Forwarding allow_missing_label does not require narrowing the helper to models that implement the new spin API. Please retain the previous capability-safe has_spin check and the nested spin.allow_missing_label fallback, while still producing must=False, default=0.0 for the native-spin model.
| back to a NaN-filled placeholder for those, same as any other | ||
| unavailable output. | ||
| """ | ||
| if odef.name in ("mask", "mask_mag"): |
There was a problem hiding this comment.
[P1] Return the real magnetic mask for graph-native-spin artifacts
This branch guarantees an invalid public result: DeepPot.eval() always returns results["mask_mag"] for a spin model, but mapping mask_mag to None makes _eval_model_graph_spin fill the entire requested output with NaNs. For use_spin=[True, False] and atom types [0,0,0,1,1,1], the documented result is [True,True,True,False,False,False], not an unavailable-output placeholder. The new end-to-end test currently receives _mm and deliberately does not assert it, so it codifies the gap instead of catching it. The adapter already has atom_types and the artifact metadata contains use_spin, so it can construct the Boolean mask without adding neural-network work (alternatively, export mask_mag from the graph forward). Please restore the PT/public DeepPot semantics and assert the exact mask through the graph .pt2 path.
5 participants
Summary
Brings the DPA4/SeZM descriptor fully onto the NeighborGraph ("graph") lower — the sel-free, edge-native route that dpa1 (#5583) and dpa2 (#5779) already use — and adds the conditioning inputs DPA4 needs on that route. Stacked on the merged dpa2 graph work (#5779); this PR is DPA4's turn on the same infrastructure.
Four layers, built and reviewed incrementally:
callbecomes a thin adapter over one graph-native math owner (_call_graph_impl); the model-level graph seam drives DPA4 throughcall_graph. Parity vs the dense route at fp64 1e-12; parity vs the pt reference at ~6e-15..pt2— per-blockborder_opghost-feature exchange on the graph lower; graph-kind.pt2embeds the nested with-comm artifact. Dense (nlist) lower stays comm-less and fails fast on multi-rank; bridging models also fail fast (a rank cannot observe a ghost owner's full outgoing-edge set).(nf, nloc, 3)conditions the descriptor and is a second autograd leaf givingforce_mag = -dE/dspin. Rides the graph route only; frozen to a graph-kind.pt2(spin at positional ABI index 10), evaluated in Python (DeepEval) and C++ (DeepSpinPTExptgraph route), run in LAMMPS (pair_style deepspin). Single-rank.DPA4NativeSpinModelin dpmodel + pt_expt.charge_spinand flips the capability gates).Validation
force_magvs finite difference: 5.5e-10 (atol 1e-6); pt weight-copied parity 1e-12.deeppot_universalDPA4 dense + graph rows 38 passed / 8 skipped (NoPbc profile); native-spinDeepSpinPTExptgraph gtest 16 passed (double/float × PBC/NoPbc).deepspinenergy/force/force_mag; 2-rank multi-rank fail-fast aborts cleanly.index_addatomics introduce 1-2 fp64 ULP run-to-run nondeterminism).Known limitations
.pt2carrieshas_comm_artifact=falseand C++ fails fast on multi-rank. Bridging is intrinsically single-rank (per-node freeze fold over the full outgoing-edge set).add_chg_spin_ebd+ native spin) — a small follow-up.sezm_native_spincheckpoints is not claimed — the deserialize alias covers the type string with dpmodel structure only.edge_vec.pt2lower schema (DPA4's original inference rail) is quarantined withDeprecationWarning, not removed — its removal is a follow-up gated on repointing the pt-SeZM freeze (and, for its spin part, on graph-spin support, now landed here).Summary by CodeRabbit
force_mag) and spin masking..pt2export/inference support on the graph route (with appropriate single-rank/MPI behavior).edge_vecdeprecation.