docs(backend): document backend support and workflows - #5929
Conversation
Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughDocumentation now covers seven backends, including TensorFlow 2 and PyTorch-Exportable, across backend selection, training, freezing, compression, model formats, inference detection, and Hessian handling. ChangesBackend documentation
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@doc/train/training.md`:
- Line 23: Remove the shell prompt marker from each command-only example, or add
corresponding output, to eliminate MD014 warnings: update the TensorFlow 2,
PyTorch, and PyTorch-Exportable examples at doc/train/training.md lines 23, 30,
and 37, and the TensorFlow 2 and PyTorch-Exportable examples at
doc/freeze/freeze.md lines 21 and 52.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 16cb4b9c-60c4-416c-bbc6-ec382d110acd
📒 Files selected for processing (4)
doc/backend.mddoc/freeze/compress.mddoc/freeze/freeze.mddoc/train/training.md
There was a problem hiding this comment.
Pull request overview
This PR adds user-facing documentation for selecting and using the TensorFlow 2 (eager) and PyTorch-Exportable (pt-expt) backends, covering backend selection plus train, freeze, and compress CLI usage across the relevant docs pages.
Changes:
- Add TF2 and PyTorch-Exportable tabs to the training instructions.
- Add TF2 and PyTorch-Exportable tabs to the freeze and compress docs (including
--lower-kindmention for pt-expt). - Extend
doc/backend.mdto include the two backends, their model/checkpoint formats, and inference suffix detection notes.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| doc/backend.md | Adds TF2 and PyTorch-Exportable backend documentation (selection, file formats, detection). |
| doc/train/training.md | Adds TF2 and PyTorch-Exportable training command tabs. |
| doc/freeze/freeze.md | Adds TF2 and PyTorch-Exportable freeze command tabs and notes. |
| doc/freeze/compress.md | Updates supported-backends note and adds TF2 / PyTorch-Exportable compress tabs. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Document TF2, PyTorch-Exportable, and JAX availability across model and descriptor guides, and add DP/JAX compression routes. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
njzjz
changed the title
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
doc/backend.md (1)
114-117: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftDocument the DPA4
.pt2ABI, not just the filename classification.
doc/model/dpa4.mddocumentsdp --pt freezefor DPA4 anddoc/backend.mdsays.pt2belongs to PyTorch-Exportable. The code docs show thesept.pt2artifacts are AOTInductor, but both ABIs can be classified as.pt2, while the PT and pt_expt code paths/ABI are distinct. Make the inference guidance explicit about “which.pt2” or rename/correct one contract to avoid loading the wrong graph type.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@doc/backend.md` around lines 114 - 117, Update the `.pt2` inference guidance in the backend documentation to distinguish DPA4 `dp --pt freeze` artifacts from PyTorch-Exportable/AOTInductor artifacts, explicitly identifying which ABI each `.pt2` file uses. Align the wording with the separate PT and pt_expt code paths so users do not load a `.pt2` graph through the wrong backend contract.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@doc/model/train-fitting-dos.md`:
- Around line 4-11: Update the DOS training documentation to cover TensorFlow 2
and PyTorch-Exportable routes alongside the existing dp --tf train and dp --pt
train examples, adding dp --tf2 train and dp --pt-expt train tabs. If either
backend cannot train DOS models, instead revise the supported-backends
description to explicitly mark it as conversion/inference-only, consistent with
JAX.
---
Outside diff comments:
In `@doc/backend.md`:
- Around line 114-117: Update the `.pt2` inference guidance in the backend
documentation to distinguish DPA4 `dp --pt freeze` artifacts from
PyTorch-Exportable/AOTInductor artifacts, explicitly identifying which ABI each
`.pt2` file uses. Align the wording with the separate PT and pt_expt code paths
so users do not load a `.pt2` graph through the wrong backend contract.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 2b1880b6-d0b4-460d-b54d-d31775212ee3
📒 Files selected for processing (24)
doc/backend.mddoc/freeze/compress.mddoc/freeze/freeze.mddoc/model/dpa2.mddoc/model/dpa3.mddoc/model/dpa4.mddoc/model/dprc.mddoc/model/linear.mddoc/model/overall.mddoc/model/pairtab.mddoc/model/train-energy-hessian.mddoc/model/train-energy-spin.mddoc/model/train-energy.mddoc/model/train-fitting-dos.mddoc/model/train-fitting-property.mddoc/model/train-fitting-tensor.mddoc/model/train-hybrid.mddoc/model/train-se-atten.mddoc/model/train-se-e2-a-tebd.mddoc/model/train-se-e2-a.mddoc/model/train-se-e2-r.mddoc/model/train-se-e3-tebd.mddoc/model/train-se-e3.mddoc/train/training.md
🚧 Files skipped from review as they are similar to previous changes (1)
- doc/freeze/compress.md
Document the distinct .pt2 lower-input ABIs and add the missing DOS training commands for TensorFlow 2 and PyTorch-Exportable. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Document the dense nlist exception used by PT-frozen deepspin DPA4 models. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5929 +/- ##
==========================================
- Coverage 79.64% 79.39% -0.25%
==========================================
Files 1085 1085
Lines 126583 126583
Branches 4593 4599 +6
==========================================
- Hits 100813 100501 -312
- Misses 24117 24430 +313
+ Partials 1653 1652 -1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
This is a large and mostly accurate sweep -- I checked the great majority of the new support claims against backend registration at head and they hold up, including the ones that are easy to get wrong. Specifically confirmed correct: every dp --tf2 / dp --pt-expt / dp --jax flag against the aliases in deepmd/backend/*.py; the dp --pt_expt to dp --pt-expt and .xlo to .hlo corrections (both were real doc bugs); all the frozen-model suffixes against each backend's suffixes; dp --jax freeze --hessian; the JAX energy-loss-only and no---init-frz-model caveats; the TF2 and Paddle additions on dpa2/dpa3/se_e3_tebd; and the DPA4 note, where JAX really does have deepmd/jax/descriptor/dpa4.py but no jax/model/dpa4_model.py, so "descriptor-only" is exactly right.
Two claims do not hold. Both are inline.
One smaller thing I did not raise inline. The Hessian rewording in overall.md and train-energy-hessian.md is not wrong -- it names dp --pt freeze and the TorchScript route specifically -- but hessian_mode is honored only at deepmd/pt/model/model/__init__.py#L371-L374, and neither pt_expt/model/get_model.py nor dpmodel/model/model.py mention it, so PyTorch-Exportable cannot construct a Hessian model at all. Since this PR establishes that {{ pytorch_icon }} covers both, the page is now silent on pt_expt where it used to be blanket. Worth a clause if you touch these lines again; not worth a round trip on its own.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The current DPA4 native-spin multi-rank guidance contradicts both the implementation and itself. The graph-route/native-spin section of doc/model/dpa4.md says that native-spin .pt2 archives carry no with-comm artifact, have has_comm_artifact=false, and are single-rank only. At this head, _needs_with_comm_artifact explicitly admits a graph-lowered NativeSpinModelKind whose descriptor communicates across ranks; test_native_spin_needs_with_comm_on_the_graph_lower_only asserts that behavior; and DeepSpinPTExpt loads and dispatches forward_lower_with_comm.pt2 for multi-rank graph native spin. The same page also correctly states near the earlier multi-rank overview and the Limitations section that pt_expt native spin supports multi-rank. Please remove or update the obsolete single-rank paragraphs so this page documents one contract that matches the current export and C++ runtime.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The documented PyTorch-Exportable compression workflow does not actually put a compressed model into the executable artifact for the ordinary .pte example. Except for the special fused DPA1 graph .pt2 branch, enable_compression passes the original uncompressed model dictionary to deserialize_to_file and stores the compressed dictionary only as model_json_override. DeepEval reconstructs that model.json object for metadata and introspection, but energy inference runs exported_module or _pt2_runner, so the tabulated model is not executed. I reproduced the documented se_e2_a-style .pte workflow locally: the input and output exported graphs were identical (711 nodes each), the output executable contained no tabulate_fusion operator, while only model.json contained the compression state. Therefore model-compress.pte provides no compression speed or memory benefit even though this page and backend.md present pt_expt compression as generally supported. Please either export the genuinely compressed executable or document the narrow supported case instead of advertising the ordinary .pte workflow as compressed.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The newly documented JAX .jax frozen-model route is not usable by the normal DeePMD inference/test interface. Please correct the format contract or implement the missing loader support.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The newly documented TensorFlow 2 compression workflow fails for the normalized default se_e2_a configuration. The implementation or the documented support contract needs to be corrected.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The earlier correction that distinguishes JAX .jax serialization from DeepEval-compatible frozen formats is still incomplete in the Hessian documentation. train-energy-hessian.md says backend checkpoints can be tested directly and gives dp --jax test -m model.ckpt.jax, but test is a common entrypoint and the JAX DeepEval constructor accepts only .hlo and .savedmodel; a .jax path raises ValueError(Unsupported file extension). Please remove the direct .jax test workflow and use the already documented freeze --hessian to .hlo/.savedmodel route instead, unless .jax DeepEval loading is implemented.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The .pt2 lower-input ABI overview is incomplete: it enumerates edge_vec, nlist, and graph and then describes these as all variants, but the documented pt_expt DPA1 compression workflow can produce a fourth production ABI, lower_input_kind=dpa1_canonical, whenever canonical_model_eligible(model) is true. DeepPotPTExpt also has a distinct lower_input_is_canonical_ dispatch branch for it. Please document dpa1_canonical as the compact compressed-DPA1 ABI in both the ABI overview and the later extension-based inference-routing summary instead of subsuming it under NeighborGraph.
Address the remaining CHANGES_REQUESTED round on deepmodeling#5929: - dpa4.md: remove the nonexistent --lower-kind auto option; state that the CLI accepts only nlist (default) and graph, that a graph-capable DPA4/SeZM model is always frozen to the graph lower (nlist is overridden with a warning), and that a native-spin descriptor has no dense lower at all. - train-energy-hessian.md: drop the direct dp --jax test of .jax checkpoints, which the JAX DeepEval route does not accept; JAX models must be frozen to .hlo/.savedmodel with dp --jax freeze --hessian before testing, as already documented above. - backend.md: document the compact dpa1_canonical ABI produced by dp --pt-expt compress for eligible compressed DPA-1 in the .pt2 lower-input ABI overview and the inference-routing summary. Coding-Agent: opencode opencode-Version: 1.18.13 Model: ustc/deepseek-v4-flash Reasoning-Effort: max
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
All three findings from the latest CHANGES_REQUESTED round are addressed in commit e095aa7e6:
-
dpa4.md (
--lower-kind auto) — Removed the nonexistentautooption. The text now states--lower-kindaccepts onlynlist(the default) andgraph, that a graph-capable DPA4/SeZM descriptor (plain-energy or native-spin) is always frozen to the graph lower because the freeze entrypoint overrides annlistrequest with a warning, and that a native-spin descriptor has no dense (nlist) lower at all, so--lower-kindcannot produce a dense artifact for it. -
train-energy-hessian.md (
dp --jax test -m model.ckpt.jax) — Removed the direct.jaxtest workflow (JAXdp testaccepts only frozen.hlo/.savedmodel). The page now keeps the PyTorch-TorchScript direct-checkpoint test and states that a JAX Hessian model must be frozen withdp --jax freeze --hessianfirst, pointing to the freeze route documented above. -
backend.md (
.pt2lower-input ABI overview) — Documented the fourth production ABI:dp --pt-expt compresscan export an eligible compressed DPA-1 model through the compact canonicaldpa1_canonicalABI (lower_input_kind: dpa1_canonical). It is now listed both in the.pt2ABI overview and in the extension-based inference-routing summary alongsideedge_vec, densenlist, and NeighborGraph.
Coding agent: opencode
opencode version: 1.18.13
Model: ustc/deepseek-v4-flash
Reasoning effort: max
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 34 out of 34 changed files in this pull request and generated no new comments.
Suppressed comments (3)
doc/train/finetuning.md:5
- The supported-backends note omits PyTorch-Exportable even though this page now documents a PyTorch-Exportable fine-tuning workflow below. This is misleading because it suggests pt-expt fine-tuning is out of scope for this page.
> [!NOTE]
> **Supported backends**: TensorFlow {{ tensorflow_icon }}, TensorFlow 2
> {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, Paddle {{ paddle_icon }}
doc/model/train-energy-hessian.md:165
- This claims
.jaxcheckpoints cannot be tested directly, but the JAX backend is registered for.jaxmodel paths and the DeepEval implementation can load.jaxvia Orbax (serialize_from_fileindeepmd/jax/utils/serialization.py). As written, it incorrectly forbids a workflow that should work (dp --jax test -m model.ckpt.jax ...).
The JAX training checkpoint (`.jax`) cannot be tested directly: the JAX
`dp test` route accepts only frozen `.hlo` or `.savedmodel` artifacts, so a
JAX Hessian model must be frozen with `dp --jax freeze --hessian` (as shown
above) before testing.
doc/backend.md:107
- This says
dp test/DeepPotdoes not load.jaxserializations, but.jaxis a registered JAX model suffix (deepmd/backend/jax.py) andDeepEvalBackenddispatches by suffix, sodp test -m *.jax/DeepPot('*.jax')should work. If the intent is to discourage.jaxfor deployment, it’s more accurate to scope the limitation to the C++ interface.
JAX supports training with `dp --jax train`; training checkpoints use the
`.jax` extension. Freezing can write a DeepEval-compatible `.hlo` or
`.savedmodel` model, or a lossless `.jax` serialization for checkpoint
round-tripping and JAX-MD. The normal `dp test`/`DeepPot` route does not load
`.jax` serializations.
# Conflicts: # doc/model/dpa4.md
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 33 out of 33 changed files in this pull request and generated no new comments.
Suppressed comments (3)
doc/backend.md:52
- This sentence asserts a future deprecation plan ("DeePMD-kit will deprecate this backend") without any pointer to a deprecation notice, release note, or tracking issue. In user docs this reads like a guarantee/timeline and can become stale quickly. Consider rephrasing to describe the current technical difference and, if desired, recommend pt-expt without promising deprecation.
[PyTorch](https://pytorch.org/) 2.1 or above is required. Select this backend
with `dp --pt`. It uses TorchScript for most frozen models; DPA4/SeZM uses a
separate AOTInductor export path. Because PyTorch has deprecated TorchScript,
DeePMD-kit will deprecate this backend and replace it with PyTorch-Exportable.
doc/model/train-energy-hessian.md:118
- The warning says “PyTorch-Exportable cannot construct a Hessian model for freezing or inference”, but the pt-expt implementation does support Hessian computation in eager mode (it is just not exportable / not preserved in frozen artifacts). To avoid contradicting the implementation, scope this statement to exported
.pte/.pt2models (torch.export/AOTI) and/or todp teston frozen artifacts.
> The PyTorch-TorchScript freeze route does not preserve Hessian output. A
> PyTorch-TorchScript model frozen with `dp --pt freeze` is treated as a
> standard energy model. PyTorch-Exportable cannot construct a Hessian model
> for freezing or inference. The JAX backend can preserve Hessian output in a
> frozen model with `dp --jax freeze --hessian`.
doc/model/overall.md:64
- Same issue as in
train-energy-hessian.md: the current wording is too broad (“PyTorch-Exportable cannot construct a Hessian model for freezing or inference”). The pt-expt backend can compute Hessians in eager mode, but Hessians are not preserved in exported.pte/.pt2artifacts. Consider scoping the warning accordingly so it matches the implementation.
> The PyTorch-TorchScript freeze route cannot output Hessians, and
> PyTorch-Exportable cannot construct a Hessian model for freezing or
> inference. The JAX backend can retain Hessian output in a frozen model with
> `dp --jax freeze --hessian`; see [Freeze a model](../freeze/freeze.md).
|
Rechecked the three requested backend-documentation corrections against the current implementation. The findings were valid when raised and are already addressed by e095aa7, retained in the current head 7760c8d:
Source inspection confirms the forced-graph branch, the JAX DeepEval extension rejection, and the dedicated Coding agent: Codex |
OutisLi-Bot
left a comment
OutisLi-Bot
left a comment
There was a problem hiding this comment.
Review
Docs-only backend-support sweep. Rechecked this head against the CLI, backend registrations, and the previously requested contract fixes. Those fixes are in the files, not only in the reply threads.
Approve.
Previously requested corrections that hold at 7760c8d6:
- DPA4/SeZM
--lower-kindisnlist/graphonly; graph-capable models overridenlist. There is noautoCLI value. - JAX Hessian testing goes through
freeze --hessianto.hlo/.savedmodel..jaxis not a DeepEval format. .pt2ABI docs include compressed-DPA1dpa1_canonical.- DPA4 export is the AOTInductor route, not TorchScript.
- PyTorch-Exportable executable compression is limited to the fused graph DPA-1
.pt2path; other outputs only refreshmodel.json. - TF2
se_e2_aSavedModel export requirestype_one_side: true. - TF2 DOS frozen inference is caveated; PyTorch-Exportable DOS was removed and pointed at #5949.
No blocking false-command issues found at this head.
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
## 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 a
strictly 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:
- a backend-neutral DPA4C descriptor and a native `pt_expt`
implementation;
- graph-native training, serialization, export, compression, and
calibration;
- fused CUDA descriptor, fitting, force, virial, and magnetic-force
paths;
- native-spin conditioning from the descriptor through Python, C, C++,
and
LAMMPS/Kokkos interfaces;
- frame-level charge and spin-multiplicity conditioning, including
runtime
re-specialization of compressed artifacts;
- function-preserving fine-tuning from a spin-free checkpoint;
- ragged mixed-size training batches without exposing phantom atoms to
the
network; and
- user documentation plus non-spin and native-spin examples.
## 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 `sel` parameter, no capacity derived
from
the densest training frame, and no neighbor truncation. Its persistent
per-atom state is determined by `channels` and `lmax`, not by the number
of
neighbors.
## Descriptor architecture
### Edge representation
For every directed edge `j -> i`, DPA4C combines:
- the DPA4 Bessel or Gaussian radial basis;
- a bias-free one-hidden-layer SwiGLU radial network;
- ordered PairFiLM scale and shift terms for `(type_i, type_j)`;
- optional pair-conditioned shared radial modes; and
- a C3 cutoff envelope whose value and first three radial derivatives
join
continuously to zero at `rcut`.
`radial_modes` increases chemical/radial resolution without widening the
per-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:
- exact aligned Gram matrices within each degree;
- normalized low-rank bispectrum contractions across allowed degree
triples;
- the projected `Qv` quartic; and
- the two neighborhood-mass coordinates.
Only 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:
- `channels` in `{8, 16, 32, 64, 128}`;
- `lmax` in `{2, 3, 4}`;
- `basis_type` in `{bessel, gaussian}`;
- `n_radial`;
- `radial_modes`; and
- `use_amp`, which applies bf16 autocast only to the edge-dominated
stage and
restores descriptor precision before reduction and invariant
contraction.
## Frame charge-state conditioning
When `add_chg_spin_ebd` is enabled, DPA4C accepts one frame-level
`[charge, multiplicity]` condition. This condition is independent of the
per-atom native-spin vector. It enters at two finite locations:
1. a shift of the center type embedding; and
2. a bias of the ordered-pair encoder hidden state.
The portable graph path keeps the condition per frame, so one batch may
contain
different charge states. `default_chg_spin` supplies the fallback state
when an
input 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. The
compiled descriptor supports:
```text
channels in {8, 16, 32, 64, 128}
lmax in {2, 3, 4}
radial_modes in {0, 2, 4, 8}
precision = float32
```
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=1` enables the fused descriptor/fitting path with
autograd force
assembly. `DP_CUDA_INFER=2` additionally uses the compact canonical
fused
energy/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
- Registers `descriptor.type: dpa4c` for the PyTorch Exportable backend
and
documents its arguments in `argcheck`.
- Adds model serialization, graph export, compression routing, inference
metadata, and evaluation inputs for both charge state and native spin.
- Extends C and C++ energy/spin interfaces with charge-state dimensions,
setters, and per-call inputs.
- Adds non-spin water and native-spin NiO examples and a full user
guide.
- Adds backend-neutral, PyTorch, CUDA, graph-lower, export, fine-tuning,
symmetry, derivative, serialization, compression, and deployment tests.
- Adapts the DPA1 shared graph-kernel helpers without changing DPA1's
public
descriptor contract.
The final integration commit also replaces the removed
`doc_only_pt_expt_supported` symbol with the current
`supported_backends("pt_expt")` registry introduced on `master` by
deepmodeling#5929.
This is the only modification made after cherry-picking the four DPA4C
commits.
## Current scope and limitations
- DPA4C is implemented for `pt_expt`; other backends are not added here.
- Compressed inference is float32-only and restricted to the structural
profiles listed above. Unsupported profiles continue to use the portable
path or are rejected by explicit compression validation.
- Descriptor-level excluded type pairs are not supported by the fused
compact
kernel.
- Native spin requires `scheme: native`; the virtual-atom `deepspin`
scheme is
not used by DPA4C.
- The symmetric spin invariant basis does not represent the
antisymmetric
Dzyaloshinskii-Moriya interaction.
- The provided LAMMPS example covers evaluation and spin minimization.
Spin
dynamics through stock `fix nve/spin` additionally depends on that fix
recognizing the new pair style.
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Added the DPA4C descriptor with native-spin, charge-state
conditioning, compressed CUDA inference, and canonical graph support.
* Added native-spin LAMMPS pair styles and expanded C/C++ APIs for spin,
charge-state configuration, and GPU graph inference.
* Added compression capability detection and support for analytically
bounded compression domains.
* **Bug Fixes**
* Improved force, virial, magnetic-force, charge-state, and loss
handling consistency.
* **Documentation**
* Added DPA4C guides, training configurations, and spin-enabled LAMMPS
examples.
* **Tests**
* Expanded coverage for DPA4C, CUDA compression, export, validation,
spin, and charge-state behavior.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Summary
Closes #5919.
Closes #5920.
Closes #5923.
Supersedes #5930.
Validation
ruff format .ruff check .uvx pre-commit run mdformat --files ...for all changed Markdown filesdp --dp compress --helpdp --jax compress --helpgit diff --checkCoding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit