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.47% 79.23% -0.24%
==========================================
Files 1072 1073 +1
Lines 125041 125299 +258
Branches 4536 4564 +28
==========================================
- Hits 99373 99286 -87
- Misses 24044 24371 +327
- Partials 1624 1642 +18 ☔ 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.
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 (2)
doc/model/train-energy-hessian.md:98
- The warning overstates PyTorch-Exportable limitations: pt_expt implements Hessian computation paths internally (e.g. eager-only), so saying it "cannot construct a Hessian model" is misleading. It would be more accurate to state that the PyTorch-Exportable freeze/runtime path does not currently preserve/expose Hessian outputs in exported models, while JAX can via
--hessian.
> 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
- This warning uses the same over-strong wording as
train-energy-hessian.md("cannot construct a Hessian model"). To avoid implying a hard capability gap, consider phrasing this as a limitation of the PyTorch-Exportable freeze/runtime path (Hessian outputs not preserved/exposed in exported models), while JAX can retain Hessians viadp --jax freeze --hessian.
> 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).
Resolve the Hessian documentation overlap by preserving the new JAX workflow together with the PyTorch-TorchScript naming and PyTorch-Exportable limitation. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
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 (4)
doc/model/change-bias.md:35
- The PyTorch section documents only the TorchScript route (
dp --pt ...), butdp --pt-expt change-biasalso supports.pte/.pt2frozen models (and multi-task is explicitly unsupported there). Adding an explicit pt-expt example prevents users from assuming change-bias is unavailable for PyTorch-Exportable.
:::{tab-item} PyTorch-TorchScript Backend {{ pytorch_icon }}
**Changing bias using provided systems for trained `.pt`/`.pth` models:**
```sh
doc/model/train-se-e2-r.md:79
- This backend-specific limitation also applies to PyTorch-Exportable: the pt_expt
se_e2_rdescriptor inherits the dpmodelse_rimplementation, which raisesNotImplementedErrorwhentype_one_sideisfalse. The docs should include PyTorch-Exportable in this list to avoid implying thattype_one_side=falseworks there.
In the PyTorch-TorchScript, JAX, and DP backends, {ref}`type_one_side <model[standard]/descriptor[se_e2_r]/type_one_side>` cannot be set to `false`.
doc/model/change-bias.md:4
dp change-biasis also implemented for the PyTorch-Exportable backend (dp --pt-expt change-bias ...indeepmd/pt_expt/entrypoints/main.py). The supported-backends note should include PyTorch-Exportable to stay consistent with the actual CLI capability and with the backend overview docs.
This issue also appears on line 31 of the same file.
> [!NOTE]
> **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch-TorchScript {{ pytorch_icon }}
doc/freeze/compress.md:113
- The PR description/linked #5923 calls out the StableHLO
.hlocompression route (dp --jax compress -i frozen_model.hlo -o compressed_model.hlo), but the JAX tab example only shows the.jaxroute. Including the.hloexample (and the.jaxalternative) better matches the “compression for deployment” workflow and the issue being closed.
```bash
dp --jax compress -i frozen_model.jax -o compressed_model.jax
JAX compression accepts .jax and .hlo inputs. Use .jax for the general,
</details>
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Both points are addressed, and I checked the files at 3bc4ff06a rather than taking the replies at their word.
On the pairtab page: Paddle is gone from the title and from both notes, and DP is back on the linear-combination note and on linear.md. Since the fix introduced a claim I had not checked in the first round, I verified it too — the new "PyTorch-Exportable" on both lines holds, because deepmd/pt_expt/model/ carries both dp_zbl_model.py and dp_linear_model.py. The asymmetry between the two notes is also right rather than an oversight: JAX and TF2 appear on the short-range interpolation line, where each has pairtab_atomic_model.py and dp_zbl_model.py, but not on the linear-combination line, where neither has a dp_linear_model.py — only linear_atomic_model.py. And Paddle genuinely has none of them.
On the compression page: switching the headline example to -i frozen_model.jax -o compressed_model.jax removes the case that would have raised. The prose now points at .jax as the general lossless path and defers the StableHLO question to the descriptor pages; I confirmed that deferral is not a dead pointer, since train-se-e2-a.md and train-se-atten.md both state the type_one_side requirement for .hlo export. One small thing for the record only: your reply said the type_one_side requirement was documented for compressed .hlo export, and in compress.md itself it is referred out rather than stated. That is the better placement given the requirement is descriptor-specific, so I am not asking for a change — just noting the wording difference in case you expected it inline.
The Hessian point from my earlier review body is resolved as well. train-energy-hessian.md now says "Supported backends: PyTorch-TorchScript, JAX" and labels the tabs the same way, instead of the blanket PyTorch icon. That matches where hessian_mode is actually honored — deepmd/pt/, deepmd/jax/, and deepmd/dpmodel/model/make_hessian_model.py, with nothing on the pt_expt side — so the page no longer implies coverage that does not exist.
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Both points are addressed, and I checked the files at 3bc4ff06a rather than taking the replies at their word.
On the pairtab page: Paddle is gone from the title and from both notes, and DP is back on the linear-combination note and on linear.md. Since the fix introduced a claim I had not checked in the first round, I verified it too — the new "PyTorch-Exportable" on both lines holds, because deepmd/pt_expt/model/ carries both dp_zbl_model.py and dp_linear_model.py. The asymmetry between the two notes is also right rather than an oversight: JAX and TF2 appear on the short-range interpolation line, where each has pairtab_atomic_model.py and dp_zbl_model.py, but not on the linear-combination line, where neither has a dp_linear_model.py — only linear_atomic_model.py. And Paddle genuinely has none of them.
On the compression page: switching the headline example to -i frozen_model.jax -o compressed_model.jax removes the case that would have raised. The prose now points at .jax as the general lossless path and defers the StableHLO question to the descriptor pages; I confirmed that deferral is not a dead pointer, since train-se-e2-a.md and train-se-atten.md both state the type_one_side requirement for .hlo export. One small thing for the record only: your reply said the type_one_side requirement was documented for compressed .hlo export, and in compress.md itself it is referred out rather than stated. That is the better placement given the requirement is descriptor-specific, so I am not asking for a change — just noting the wording difference in case you expected it inline.
The Hessian point from my earlier review body is resolved as well. train-energy-hessian.md now says "Supported backends: PyTorch-TorchScript, JAX" and labels the tabs the same way, instead of the blanket PyTorch icon. That matches where hessian_mode is actually honored — deepmd/pt/, deepmd/jax/, and deepmd/dpmodel/model/make_hessian_model.py, with nothing on the pt_expt side — so the page no longer implies coverage that does not exist.
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Both points are addressed, and I checked the files at 3bc4ff06a rather than taking the replies at their word.
On the pairtab page: Paddle is gone from the title and from both notes, and DP is back on the linear-combination note and on linear.md. Since the fix introduced a claim I had not checked in the first round, I verified it too — the new "PyTorch-Exportable" on both lines holds, because deepmd/pt_expt/model/ carries both dp_zbl_model.py and dp_linear_model.py. The asymmetry between the two notes is also right rather than an oversight: JAX and TF2 appear on the short-range interpolation line, where each has pairtab_atomic_model.py and dp_zbl_model.py, but not on the linear-combination line, where neither has a dp_linear_model.py — only linear_atomic_model.py. And Paddle genuinely has none of them.
On the compression page: switching the headline example to -i frozen_model.jax -o compressed_model.jax removes the case that would have raised. The prose now points at .jax as the general lossless path and defers the StableHLO question to the descriptor pages; I confirmed that deferral is not a dead pointer, since train-se-e2-a.md and train-se-atten.md both state the type_one_side requirement for .hlo export. One small thing for the record only: your reply said the type_one_side requirement was documented for compressed .hlo export, and in compress.md itself it is referred out rather than stated. That is the better placement given the requirement is descriptor-specific, so I am not asking for a change — just noting the wording difference in case you expected it inline.
The Hessian point from my earlier review body is resolved as well. train-energy-hessian.md now says "Supported backends: PyTorch-TorchScript, JAX" and labels the tabs the same way, instead of the blanket PyTorch icon. That matches where hessian_mode is actually honored — deepmd/pt/, deepmd/jax/, and deepmd/dpmodel/model/make_hessian_model.py, with nothing on the pt_expt side — so the page no longer implies coverage that does not exist.
Duplicate submission on my side; superseded by the identical approval below.
|
|
||
| DPA4 is a PyTorch-only model. Train it with the standard `dp --pt` workflow: | ||
| DPA4 is available only through the PyTorch framework backends. Train a | ||
| PyTorch-TorchScript model with the standard `dp --pt` workflow: |
There was a problem hiding this comment.
DPA4 is not a pytorch-torchscript model
There was a problem hiding this comment.
Fixed in 30858ac. The quick-start section no longer calls DPA4 a "PyTorch-TorchScript model": it now says DPA4/SeZM is trained with dp --pt but does not use TorchScript, exporting through the AOTInductor .pt2 path instead.
Coding agent: opencode
opencode version: 1.18.9
Model: ustc/deepseek-v4-flash
Reasoning effort: max
| ### Freeze to `.pt2` | ||
|
|
||
| DPA4/SeZM checkpoints use the PyTorch `.pt2` (AOTInductor) export path; the | ||
| DPA4/SeZM checkpoints use the PyTorch-TorchScript `.pt2` (AOTInductor) export path; the |
There was a problem hiding this comment.
Same here. SeZM/DPA4 is realized in pt backend, but the deployment path is built on the pt export backend
There was a problem hiding this comment.
Fixed in 30858ac. The "Freeze to .pt2" section now says DPA4/SeZM checkpoints are exported through the AOTInductor .pt2 path of the PyTorch backend and that the ordinary TorchScript freeze path is not used; the "PyTorch-TorchScript backend detects DPA4/SeZM" wording was dropped. This was collateral from the earlier mechanical PyTorch→PyTorch-TorchScript rename, which should not have applied to DPA4/SeZM given its non-TorchScript export.
Coding agent: opencode
opencode version: 1.18.9
Model: ustc/deepseek-v4-flash
Reasoning effort: max
DPA4/SeZM is realized in the PyTorch backend but exports through the AOTInductor .pt2 path, not TorchScript. The backend rename to PyTorch-TorchScript was mechanically applied to this page earlier; restore accurate wording in the quick-start and freeze-to-.pt2 sections. Coding-Agent: opencode opencode-Version: 1.18.9 Model: ustc/deepseek-v4-flash Reasoning-Effort: max
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/troubleshooting/howtoset_num_nodes.md:20
- This section attributes the torchrun-based distributed training workflow only to the PyTorch-TorchScript backend, but the PyTorch-Exportable backend also initializes distributed training when launched via torchrun/srun (pt_expt entrypoint). Consider wording that covers both backends to avoid implying
dp --pt-expt traincannot be launched with torchrun.
doc/train/parallel-training.md:98 - This section title and opening sentence scope DDP/torchrun parallel training to the PyTorch-TorchScript backend, but the PyTorch-Exportable backend also supports torchrun-launched distributed training. Updating the heading/intro to include both helps keep the backend documentation consistent.
## PyTorch-TorchScript Implementation {{ pytorch_icon }}
Currently, parallel training in pytorch version is implemented in the form of PyTorch Distributed Data Parallelism [DDP](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html).
DeePMD-kit will decide whether to launch the training in parallel (distributed) mode or in serial mode depending on your execution command.
doc/model/train-se-e2-r.md:79
type_one_side=falseis also unsupported in the PyTorch-Exportable backend forse_e2_r(pt_expt uses the dpmodelse_rimplementation, which raises NotImplementedError whentype_one_sideis false). The backend list here should include PyTorch-Exportable as well.
In the TensorFlow backend, {ref}`env_protection <model[standard]/descriptor[se_e2_r]/env_protection>` cannot be set to a non-zero value.
In the PyTorch-TorchScript, JAX, and DP backends, {ref}`type_one_side <model[standard]/descriptor[se_e2_r]/type_one_side>` cannot be set to `false`.
iProzd
left a comment
iProzd
left a comment
There was a problem hiding this comment.
Re-reviewed the current head after the DPA4/SeZM follow-up. The documentation now correctly separates the dp --pt training route from the AOTInductor .pt2 export/runtime path and explicitly states that DPA4/SeZM does not use ordinary TorchScript. The prior backend-contract concerns are resolved, and I found no new blocking issue.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
One remaining backend-classification error is noted inline.
| PyTorch-Exportable runtime, which reads this metadata to select the correct | ||
| input path. The `--lower-kind` option controls only the PyTorch-Exportable | ||
| freeze route; see the [DPA4 export documentation](model/dpa4.md#freeze-to-pt2) | ||
| for the separate PyTorch-TorchScript route. |
There was a problem hiding this comment.
[P2] Do not call the DPA4 export route TorchScript
The latest commit correctly states in dpa4.md that DPA4/SeZM bypasses TorchScript and exports through AOTInductor, but this overview still labels the same dp --pt freeze path as the “PyTorch-TorchScript route.” That leaves the documentation internally contradictory and misclassifies the artifact. Please describe it as the separate PyTorch AOTInductor export route (or equivalent).
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
One additional backend-workflow omission is noted inline.
|
|
||
| > [!NOTE] | ||
| > **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }} | ||
| > **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch-TorchScript {{ pytorch_icon }} |
There was a problem hiding this comment.
[P2] Document the supported PyTorch-Exportable change-bias route
This support list and the examples below omit PyTorch-Exportable, even though doc/backend.md explicitly advertises change-bias for that backend and deepmd/pt_expt/entrypoints/main.py implements dp --pt-expt change-bias for .pt, .pte, and .pt2 inputs, with dedicated tests. Please document that route here so the backend overview and workflow page agree, and state its current limitation that multi-task checkpoints are not supported.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
One additional PyTorch-Exportable training-workflow omission is noted inline.
| ``` | ||
|
|
||
| ## PyTorch Implementation {{ pytorch_icon }} | ||
| ## PyTorch-TorchScript Implementation {{ pytorch_icon }} |
There was a problem hiding this comment.
[P2] Document the existing PyTorch-Exportable DDP route separately
This rename leaves the parallel-training guide with only the PyTorch-TorchScript route, although deepmd/pt_expt/entrypoints/main.py::setup_run() initializes torch.distributed under torchrun/srun and the pt_expt trainer wraps both single-task and multi-task models in DistributedDataParallel. Please add a separate PyTorch-Exportable DDP subsection and a torchrun ... dp --pt-expt train input.json example. It should remain separate from the ZeRO/FSDP2 subsection below, because those options are specific to the PyTorch-TorchScript trainer rather than pt_expt.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
One additional graph-freeze behavior mismatch is noted inline.
|
|
||
| > [!NOTE] | ||
| > **Default route change in pt_expt (eager & training).** For a graph-eligible DPA-2 descriptor, the pt_expt backend now defaults to the carry-all graph route not only for `--lower-kind graph` freezing but also in **eager inference/evaluation and in (compiled) training** (`neighbor_graph_method=None` resolves to the graph). This changes the numerical behavior of existing pt_expt configurations relative to the dense neighbor-list route (by the amounts described above — negligible at non-binding `sel`). The other backends (dpmodel/PyTorch/Paddle/TensorFlow/JAX) are unaffected: they keep the dense route as their only path. | ||
| > **Default route change in pt_expt (eager & training).** For a graph-eligible DPA-2 descriptor, the pt_expt backend now defaults to the carry-all graph route not only for `--lower-kind graph` freezing but also in **eager inference/evaluation and in (compiled) training** (`neighbor_graph_method=None` resolves to the graph). This changes the numerical behavior of existing pt_expt configurations relative to the dense neighbor-list route (by the amounts described above — negligible at non-binding `sel`). The other backends (dpmodel/PyTorch-TorchScript/Paddle/TensorFlow/JAX) are unaffected: they keep the dense route as their only path. |
There was a problem hiding this comment.
[P1] Remove the unavailable dense fallback later in this section
This new default-route description conflicts with the final paragraph, which tells users to avoid the empty-rank restriction by using “the dense (nlist, the default) artifact.” For the graph-eligible DPA2 discussed here, deepmd/pt_expt/entrypoints/main.py::freeze() overrides every non-graph lower request—including an explicit --lower-kind nlist—to graph whenever model_uses_graph_lower(m) is true. Thus nlist is neither the default nor an alternative artifact selectable for the same model; a dense artifact requires changing the model into a non-graph-eligible configuration. Please remove that fallback advice and describe the forced graph behavior consistently.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
One additional backend-workflow documentation gap is noted inline.
| Select this backend with `dp --pt-expt` (alias | ||
| `dp --pytorch-exportable`). It uses PyTorch with the backend-independent model | ||
| implementation and supports training, including multi-task training and | ||
| fine-tuning, freezing, change-bias, and testing. Compression support and export |
There was a problem hiding this comment.
[P2] Add the advertised TF2 and pt_expt fine-tuning workflows
This overview now advertises fine-tuning for TensorFlow 2 and PyTorch-Exportable, but the dedicated doc/train/finetuning.md page still contains only the legacy TensorFlow, dp --pt, and Paddle workflows. Both new backends have implemented fine-tuning entrypoints, and pt_expt has coverage for single-task, multi-task, DDP, and .pt/.pte/.pt2 sources. Please update the dedicated guide with executable dp --tf2 train --finetune ... and dp --pt-expt train --finetune ... routes and their backend-specific accepted formats/limitations, so this capability claim leads to a complete public contract.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
One additional spin-scheme semantic mismatch is noted inline.
| > - In the TensorFlow backend, the `se_e2_a` descriptor will treat those atom types with spin as new (virtual) types, | ||
| > and duplicate their corresponding selected numbers of neighbors ({ref}`sel <model[standard]/descriptor[se_e2_a]/sel>`) from their real atom types. | ||
| > - In the PyTorch backend, if spin settings are added, all the types (with or without spin) will have their virtual types. | ||
| > - In the PyTorch-TorchScript and PyTorch-Exportable backends, if spin settings are added, all the types (with or without spin) will have their virtual types. |
There was a problem hiding this comment.
[P1] Limit the virtual-type statement to the deepspin scheme
This is true for the default deepspin virtual-atom scheme, but not for DPA4/SeZM with model.spin.scheme: native. The native path passes the per-atom spin vector directly into the descriptor and introduces neither virtual atoms nor a doubled type map. Moreover, native DPA4 is graph-only, and graph neighbor construction is completely sel-free: every edge within rcut is retained, so sel neither truncates nor sizes its neighbor set. The advice to extend sel therefore applies only to dense virtual-atom routes, not to native DPA4. The DPA4 page documents this distinction and links back to this page as the general spin-training workflow. Please qualify this paragraph as deepspin-specific, mention the DPA4 native exception, and include DPA4 in the scheme-aware supported-descriptor discussion below so the two pages expose one consistent spin contract.
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