feat(pt-expt): add compact descriptor DPA4C 🎉🎉🎉 - #5972
Conversation
Introduce DPA4C as a graph-native descriptor built from degree-scaled Cartesian moments, exact invariant readouts, and pair-conditioned radial modes. - support backend-neutral training, serialization, graph export, calibration, and mixed-precision execution - add compressed CUDA and canonical inference for the supported channel and angular profiles - expose neighborhood masses, remove the fixed-capacity path, and tile compressed evaluation to bound memory - cover parity, gradients, compression, export, and end-to-end energy/force/virial behavior
Add per-atom native-spin conditioning to DPA4C from descriptor training through compressed deployment. - implement spin-aware invariant channels, statistics, serialization, evaluation, and validation - expose magnetic outputs through Python, C/C++, and LAMMPS/Kokkos interfaces - extend the compressed CUDA path and fused reductions for magnetic forces - document the model contract and provide a non-spin-dynamics LAMMPS example
…ding Condition DPA4C on frame charge and multiplicity while keeping it independent of per-atom native spin. - inject charge-state features into the type and ordered-pair routes - rebuild compressed constants once per runtime state and expose the setting across evaluation and LAMMPS interfaces - preserve unconditioned behavior and validate portable/compressed parity - reject graph folding without valid ghost-owner mappings instead of silently dropping halo edges
Naming a magnetic type on a pretraining that declared none must leave the predicted energy untouched, because the spin routes the activation releases never received a gradient. On FeC it otherwise moves the energy by several eV per atom with a configuration-dependent sign, which is what forces the output-bias regression to solve for a per-type constant of several keV. A single scalar gate on the whole spin branch makes the activation exact. It multiplies the block after the calibration, so a closed gate feeds the fitting network exactly zero whatever preconditioner was measured, and the invariants are linear in it, so zero is a starting point whose gradient is the branch itself rather than a stationary point. No weight inside the branch can play that role: the families reach the fitting network by several routes and at two spin orders, and a factor on the conditioned moment would enter the degree-one Grams squared and the quadrupole Grams to the fourth power. Constructing the gate closed is the whole mechanism. A transfer either copies a closed gate or keeps the freshly constructed one, so no reset hook takes part, and a checkpoint predating the gate carries no value for it that the runtime would invent. Fine-tuning such a corpus needs batches of unequal atom count. The graph lower already reads a flat node axis, so a ragged batch feeds it directly while a rectangular one has its phantom padding compacted away before the network sees it, and the loss reads each frame's own atom count from the graph. Two corrections ride along. Calibration accepts every scale the storage precision can represent instead of rejecting representable extremes, and the node-backward group width follows the occupancy the running device reports rather than a compiled-in constant, which the launch bounds of newer architectures ignore.
for more information, see https://pre-commit.ci
OutisLi
changed the title
OutisLi
changed the title
iProzd
left a comment
iProzd
left a comment
There was a problem hiding this comment.
Thanks—the two previously reported blockers are resolved at cc81d7b: batch normalization now validates every present training/statistics state, including LMDB data, before the embedding gathers, and C++ now checks the range count. One new correctness blocker remains: the eager range-count check uses dim_chg_spin == 0 before loading the charge-state fold, so valid compressed DPA4C archives carrying two ranges cannot be loaded by DeepPotPTExpt or NativeSpinPTExpt. Please validate against the effective settable width after fold detection and add compressed C++ loading coverage. CI is still in progress.
…nown A compressed archive reports dim_chg_spin == 0 because its compiled lower carries no conditioning input, yet it still records the ranges of the tables its charge-state fold indexes. The loader reads the ranges eagerly, before fold detection, so the count check added with them compared two ranges against a width of zero and refused every compressed charge-conditioned archive. The reader now returns empty for an unknown width, mirroring read_default_chg_spin, which the two call sites already invoke in pair with it. The fold block then reads both again with the width the fold names, so the ranges are still checked wherever a charge state can be set.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Three independent reviews were consolidated against the current head. The inline findings below are limited to issues not already covered by existing review threads.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
iProzd
left a comment
iProzd
left a comment
There was a problem hiding this comment.
The loader-order blocker from my previous review is fixed correctly at 2fc9121: zero-width eager reads now defer range parsing, while the post-fold read still validates the effective width. I still see three merge blockers already captured in the current unresolved threads: (1) source/api_cc/include/DeepPot.h:223 lets input-driven PT/JAX set_charge_spin calls succeed without updating their stored default, so later evaluations without a per-call condition silently use the old state; (2) deepmd/pt_expt/infer/deep_eval.py:1043 applies a folded state only to the AOTI runner, while eval_descriptor/eval_fitting_last_layer evaluate the separately deserialized DPModel at the archived default; and (3) source/api_cc/src/NativeSpinPTExpt.cc:822 hard-codes nframes=1 although the public C/C++ DeepSpin API supplies multi-frame buffers. Please resolve those threads and add regressions. The loader fix also needs a helper-level range-reader regression; that can cover the zero-width/after-fold sequence without a CUDA artifact. CI is still in progress.
…model A charge/spin condition reached its consumers through the width of the conditioning input a compiled forward reads, which a compressed model sets to zero because the condition is folded into frozen tables instead. Three places took that zero to mean the model carries no condition at all. The atomic model derived `add_chg_spin_ebd` from it, so a compressed model lost the flag on a serialization round trip and was frozen without the rebuild that lets a deployment serve any other state. Descriptors now declare the condition through `has_chg_spin_ebd`, which is about the model rather than about one forward. The introspection methods took the compiled forward's condition, which is none, while evaluating the model deserialized beside it, which reads one as an argument; `eval_descriptor` therefore answered for the archive default while `eval` answered for the requested state. They now build the condition the model they evaluate reads. `DeepPotPT` and `DeepPotJAX` inherited a `set_charge_spin` that reports success without storing anything, while both fall back to their own default whenever a call omits the condition, so a new state was accepted and then ignored. Both now validate and persist it.
…n inference `DeepSpinBackend::computew` takes coordinates and spins sized nframes x natoms x 3 and a cell per frame, and `DP_DeepSpinCompute2` builds exactly those. The standalone native-spin path read them as one frame, so a multi-frame call was answered from its first frame alone. The evaluation of one frame is unchanged and becomes `compute_frame`. The entry point derives the frame count from the coordinates, validates every input against it, and evaluates the frames in turn: each carries its own cell and therefore its own ghost set, which no batched forward would share.
A compressed archive reports a charge-state width of zero at load and names its width only once the fold is read, so the loader reads the ranges twice. The first read must yield nothing rather than judge a count it cannot yet know; before that ordering was in place it rejected every compressed charge-conditioned archive. The cases cover the widths either side of that boundary and the counts either side of the width, and need no model.
iProzd
left a comment
iProzd
left a comment
There was a problem hiding this comment.
Thanks for the follow-up. I re-audited the full current head rather than only the latest delta. The previously raised blockers are materially addressed, but four merge gates remain: the standalone native-spin charge-aware wrapper still validates multi-frame input as one frame; the new frame slicing rejects the documented single-block fparam/aparam broadcast forms; and the new PT/JAX persistent setters do not validate categorical charge-state domains. In addition, the new dpa4spin host integration and device-resident dpa4spin/kk path have no runtime test (source/lmp/pair_dpa4spin.cpp:451); all existing native-spin LAMMPS tests exercise pair_style deepspin instead. Compilation does not verify the new unit conversion, magnetic-force scaling, virial mapping, or reverse-communication contracts. Please add at least a single-rank runtime comparison for energy, force, magnetic force, and virial, plus a CUDA/Kokkos smoke test for /kk if that path remains part of this release. Please address these three contract breaks and add the focused API and LAMMPS regressions before merge.
…s to Reading the frames out of the coordinates left three inputs behind. `computew` accepts the parameter inputs in two layouts: one block per frame, or the single block every frame is to be evaluated with. Only the first was accepted, so a caller using the documented broadcast form was refused. The two layouts are now resolved once, at validation, and the frame loop reads any input the same way. The charge/spin condition was still checked against one frame, so a multi-frame call naming the state it serves was refused before reaching that loop. The check now sees the frames the call carries. `set_charge_spin` persisted any pair of the right width, but charge and multiplicity index one row each of an embedding table, and neither the gather nor the kernel bounds-checks the row: a fractional value lands on a neighbouring row and one out of range reads past the table. Both are now refused, against the same bounds `deepmd/utils/charge_state.py` holds for the Python boundaries. The native-spin archives the C++ suite carries all ship a with-comm artifact and are served by `DeepSpinPTExpt`, so nothing exercised `NativeSpinPTExpt`. A DPA4C fixture, whose compact descriptor keeps its messaging local, reaches it, and carries parameters and a charge state so that one archive covers every input the frame loop divides.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Flooring leaves an integer where it is and moves everything else down, so a strict drop is exactly a fractional part, and asking the value to be inside its bounds says the same as asking it not to be outside them for every value a table can address. Both forms are exact, and neither is an equality between floating-point values, which CodeQL flags because it is usually an epsilon mistake. The range test also gains a NaN, which compares false against every bound and so now falls out of it rather than through it.
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Approving. I re-verified every resolved thread against e3480cf rather than taking the resolution as evidence, and the substantive items are genuinely fixed.
On the @version bump -- you are right and I withdraw it. I checked your trace independently rather than just reading it: dpa4.py:2572 keeps only serialize()["@variables"] and discards the versioned wrapper; pt/model/descriptor/sezm.py never calls radial_basis.serialize() at all; pt_expt/descriptor/dpa4.py does not override serialize; and dpa4c has zero occurrences under deepmd/pt/. I also confirmed the converse, which is the part that settles it: across the whole tree, dpa4c.py:1763 is the only site that stores the full nested serialize() dict, so the versioned RadialBasis block genuinely comes into existence with DPA4C. A pre-PR reader therefore fails at descriptor dispatch on type: "dpa4c" long before RadialBasis.deserialize is reached, and the silent envelope-doubling I described cannot happen. Bumping would record a v1 that no archive ever contained, which is a worse artifact than the ambiguity it would remove. Leave it at 1.
The residual you identified -- hand-lifting the nested block into a v3.2.0 RadialBasis.deserialize -- is real but requires reading an archive that the same build cannot load as a model, and config.get("apply_envelope", True) covers the direction that matters. Not worth a version ladder. Thanks for pushing back with the call-site trace instead of just applying the change; that was the right response, and the conclusion is better for it.
The rest, all verified at HEAD:
disable_graph_lowernow refuses withNotImplementedErrorrather than inheriting the no-op, andtest_descriptor_dpa4c.pyassertsuses_graph_lower()is stillTrueafter the refusal. Choosing to raise rather than to add a dead_graph_lower_disabledflag is the more honest reading of the base contract for a descriptor with no dense form._DPA4_SEZM_DESCRIPTOR_TYPESnow listsdpa4c/DPA4C, covered bytest_virtual_atom_spin_scheme_is_refused.- The energy profile carries both the stress and virial families, so
v:mae/v:rmsekeep working alongside the news:keys. apply_envelopeis parametrised over both branches, withtest_radial_basis_envelope_modespinningbasis_env(r) == basis_raw(r) * E(r). I confirmed the round-trip assertion is non-vacuous by reading the code path rather than by executing it: with the key stripped,deserializereturnsTrueand theis apply_envelopeassertion fails on theFalsecase.- The example README freezes before compressing, and
build_compression_artifactsnow raises on a non-emptyexclude_types, so the documented "explicit error" is real.
One note rather than a request: Test Python on CUDA and Test C++ on CUDA are still in progress on this head. My approval is on the code; please let those land green before merging, since the fused kernels and the DPA1 canonical uint32 switch are only exercised there.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
The previous charge-state blockers are addressed, and the real CUDA C++ job is now green. I found one remaining API-contract bug in the new native-spin backend: uncompressed graph artifacts have a real per-call charge-state input, but the charge-aware overload rejects any state other than the installed default and then discards the argument. This should be fixed before merge.
— Agent: ChatGPT; Model: GPT-5.6 Sol
…rved The two artifacts this backend serves carry the condition differently, and the charge-aware entry point treated both as the folded one: it refused any state but the installed default and then dropped the argument, leaving run_graph_payload to build the input from that default. For a compressed artifact this is the contract. Its condition lives in frozen tables that set_charge_spin rebuilds and that stand for the whole run, and its compact lower takes no condition at all. Nothing there changes. An uncompressed one keeps the condition in the argument list of its compiled forward, which was already being filled on every evaluation, just from the default rather than from the caller. The condition now travels with the frames it belongs to, divided by the layout the other per-frame inputs use, so successive calls may each name their own state. The width of the forward's argument is what tells the two apart, and the charge-aware overload becomes the implementation the charge-unaware one delegates to, so the condition is an ordinary input rather than a second entry point.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Re-reviewed the latest head after the per-call charge-state fix. The previous P1 is addressed: uncompressed native-spin graph artifacts now validate and forward explicit per-call/per-frame states through the compiled input, while folded artifacts retain the installed-state contract. The new regression tests cover successive calls with different explicit states, multi-frame state propagation, and installed-state fallback.
I did not find another code-level correctness blocker in this update. The current PyPI workflow failure is an infrastructure timeout while restoring the macOS x86_64 setup-uv cache (read ETIMEDOUT), before wheel building begins; the other wheel jobs succeed. Python and C++ test workflows for this head are still running at the time of this review.
— Agent: ChatGPT; Model: GPT-5.6 Sol
Summary
This PR introduces DPA4C, the compact and compressible degree-wise member of
the DPA4 family, as a PyTorch Exportable (
pt_expt) descriptor. DPA4C is astrictly local, one-hop model intended for high-throughput molecular dynamics:
it reads each directed neighbor edge once, performs one destination reduction,
and converts the resulting degree-wise moments into a fixed invariant vector
without cross-atom message passing.
The PR includes the complete path from training to deployment:
pt_exptimplementation;LAMMPS/Kokkos interfaces;
re-specialization of compressed artifacts;
network; and
Why DPA4C
DPA4/SeZM uses equivariant message passing to target the accuracy frontier.
DPA4C targets a different operating point: a compact local student whose
radial dependence can be tabulated and whose angular computation can be fused
into bounded per-edge and per-node CUDA kernels.
The descriptor consumes a carry-all cutoff graph rather than a fixed-capacity
neighbor list. It therefore has no
selparameter, no capacity derived fromthe densest training frame, and no neighbor truncation. Its persistent
per-atom state is determined by
channelsandlmax, not by the number ofneighbors.
Descriptor architecture
Edge representation
For every directed edge
j -> i, DPA4C combines:(type_i, type_j);continuously to zero at
rcut.radial_modesincreases chemical/radial resolution without widening theper-atom moment state. The portable implementation accepts any non-negative
mode count; the compressed CUDA path specializes the production profiles
listed below.
One-reduction degree-wise moments
The edge direction is expanded in real Cartesian harmonics through
lmax.All scalar masses and all angular moments are packed into one edge payload and
accumulated with one destination segment reduction. Two smooth neighborhood
masses normalize the scalar and non-scalar blocks and are also emitted as
descriptor coordinates so the fitting network retains effective coordination
information.
The channel schedule keeps degree 0 wide, retains several channels for degrees
1 and 2, and uses one channel for degrees 3 and 4. This bounds the node state
while preserving the low-degree angular information that dominates the model.
Fixed invariant readout
The node-local readout combines:
Qvquartic; andOnly O(3)-even invariant scalars reach the standard energy fitting network.
Energy is therefore invariant under rotations, reflections, and neighbor
permutations, while force and virial remain conservative derivatives of the
same total energy.
The public structural controls are:
channelsin{8, 16, 32, 64, 128};lmaxin{2, 3, 4};basis_typein{bessel, gaussian};n_radial;radial_modes; anduse_amp, which applies bf16 autocast only to the edge-dominated stage andrestores descriptor precision before reduction and invariant contraction.
Frame charge-state conditioning
When
add_chg_spin_ebdis enabled, DPA4C accepts one frame-level[charge, multiplicity]condition. This condition is independent of theper-atom native-spin vector. It enters at two finite locations:
The portable graph path keeps the condition per frame, so one batch may contain
different charge states.
default_chg_spinsupplies the fallback state when aninput does not provide one.
Compression folds a single state into the finite type table and ordered-pair
caches, leaving the radial table, angular equations, and CUDA kernel layout
unchanged. The exported artifact carries a charge-state fold that rebuilds only
the affected constants when the evaluator, C/C++ API, or LAMMPS pair style
selects another state. This keeps the compact canonical inference ABI free of a
per-edge runtime condition while avoiding a permanently baked-in charge state.
Compression and deployment
Compression tabulates the distance-only radial network with quintic Hermite
splines on
[0, rcut]and snapshots the finite ordered-type-pair tables. Thecompiled descriptor supports:
The fused implementation includes forward and backward descriptor operators,
compact canonical graph operators, fitting-network kernels, and force/virial
assembly. The backward saves the minimum node moment state and recomputes the
edge-local radial and angular terms, avoiding a persistent per-edge moment
tensor. Evaluation is tiled so temporary memory stays bounded for large edge
sets.
DP_CUDA_INFER=1enables the fused descriptor/fitting path with autograd forceassembly.
DP_CUDA_INFER=2additionally uses the compact canonical fusedenergy/force/virial composition. The export metadata records the graph ABI and
dtype contract used by the C++ and LAMMPS loaders.
Graph folding now fails explicitly when a topology requests local-owner folding
but does not provide a valid owner for every ghost. This prevents a malformed
standalone C++ call from silently dropping halo-edge contributions. Extended
multi-rank paths keep ghosts as distinct nodes and use reverse communication as
their force-folding contract.
Integration surface
descriptor.type: dpa4cfor the PyTorch Exportable backend anddocuments its arguments in
argcheck.metadata, and evaluation inputs for both charge state and native spin.
setters, and per-call inputs.
symmetry, derivative, serialization, compression, and deployment tests.
descriptor contract.
The final integration commit also replaces the removed
doc_only_pt_expt_supportedsymbol with the currentsupported_backends("pt_expt")registry introduced onmasterby #5929.This is the only modification made after cherry-picking the four DPA4C commits.
Current scope and limitations
pt_expt; other backends are not added here.profiles listed above. Unsupported profiles continue to use the portable
path or are rejected by explicit compression validation.
kernel.
scheme: native; the virtual-atomdeepspinscheme isnot used by DPA4C.
Dzyaloshinskii-Moriya interaction.
dynamics through stock
fix nve/spinadditionally depends on that fixrecognizing the new pair style.
Summary by CodeRabbit