feat(pt_expt): DPA4 inference (.pt2 freeze, DeepEval, C++, LAMMPS)

[wanghan-iapcm]

PR-3 (final) of the DPA4/SeZM porting series — pt_expt inference: freeze to .pt2, Python DeepEval, pt→pt_expt checkpoint interop, C++ single-rank, and LAMMPS single-rank. Follows PR-1 (#5515, dpmodel core) and PR-2 (#5522, pt_expt training/export).

What's included

• Model freeze to .pt2 (deepmd/pt_expt/model/ener_model.py, deepmd/dpmodel/.../ener_model.py, deepmd/dpmodel/descriptor/dpa4.py): register EnergyModel under sezm_ener/dpa4_ener model-type aliases so BaseModel.deserialize resolves a standard DPA4 energy model (whose fitting type is sezm_ener). Fixed a torch.export specialization where int() on symbolic shapes baked nf*nloc (embedding/so2/attention).
• NoPBC export fix (deepmd/dpmodel/descriptor/dpa4.py): the atype_ext[:, :nloc] slice emitted a spurious Ne(nall, nloc) shape guard that crashed the compiled artifact when nall==nloc (no ghosts); replaced with xp_take_first_n (index_select). NoPBC now matches PBC.
• pt→pt_expt checkpoint interop (deepmd/pt_expt/model/model.py): BaseModel.deserialize unwraps pt's bespoke SeZMModel serialization (type:"SeZM", nested sezm_atomic atomic model with the pt-only dens head), validates versions, rejects unsupported features (bridging/lora/dens/active_mode) with NotImplementedError, and delegates to the standard path.
• Warn on silently-ignored flags (use_amp descriptor, enable_tf32 model): warn-once instead of silent drop.

Tests

• Model freeze source/tests/pt_expt/model/test_dpa4_export.py: dual-artifact .pt2, metadata, AOTI load, artifact-vs-eager parity (1e-10). (CI-skipped — AOTI is slow; run locally.)
• DeepEval parity vs pt source/tests/pt_expt/infer/test_dpa4_deep_eval.py: pt .pt vs pt_expt .pt2 energy/force/global-virial/atom-energy at fp64 1e-10, PBC and NoPBC; doubles as the checkpoint-interop proof. Per-atom virial compared by sum (pt's edge-scatter from feat(dpa4): use edge force and atomic virial #5518 redistributes it; global virial matches). (CI-skipped — AOTI.)
• Interop unit tests source/tests/pt_expt/model/test_dpa4_interop.py (CI-runnable, no AOTI): happy-path pt-checkpoint→pt_expt round-trip + every guard branch + version validation + @variables filtering.
• Alias deserialize guard + use_amp/enable_tf32 warn-once tests (CI-runnable).
• Fixture generator source/tests/infer/gen_dpa4.py (+ wired into source/install/test_cc_local.sh).
• C++ single-rank source/api_cc/tests/test_deeppot_dpa4_ptexpt.cc: 20 tests (double+float), dpa3-matched tolerances. Validated locally.
• LAMMPS single-rank source/lmp/tests/test_lammps_dpa4_pt2.py: parity + atom_modify map yes + the fix(pt_expt): fail-fast on .pt2 GNN inference without LAMMPS atom-map #5450 no-atom-map fail-fast. Validated on a GPU box (7 passed).

PR-1 parity suites stay green; the small dpmodel edits are parity-revalidated.

Known limitations

• Single-rank only. Multi-rank/MPI LAMMPS for DPA4 is deferred (no live multi-rank cell; the with-comm artifact compiles but its runtime is not exercised). DPA4 is a message-passing descriptor, so multi-rank follows the existing fix(pt_expt): fail-fast on .pt2 GNN inference without LAMMPS atom-map #5450/feat(pt_expt): multi-rank LAMMPS support for GNN models (DPA3 / DPA2 / spin) #5430 machinery in a later PR.
• No .pth (torch.jit) DPA4 — the pt backend has no sezm_ener model registration, so .pth freeze of a standard DPA4 energy model isn't available; not needed for the pt_expt inference path.
• Per-atom virial is not compared element-wise pt-vs-pt_expt (only its global sum) — feat(dpa4): use edge force and atomic virial #5518 changed pt's edge-scatter distribution; both are correct, the distribution differs.
• AOTI tests are CI-skipped (multi-minute compile) — the freeze/DeepEval paths are validated locally, not in CI; the interop/alias/warn tests give CI coverage of the non-AOTI logic.
• fp64 only; fp32 freeze untested. CUDA validated at LAMMPS level on a GPU box; the AOTI parity numbers are from CPU fp64.
• use_amp/enable_tf32 remain functionally ignored (now warned) — by design for this series.
• pt SeZM features out of scope (guarded NotImplementedError): spin, ZBL bridging, LoRA, dens/direct-force/denoising heads, SO3 grid projection, GridMLP, SO(2) attention extensions.

Summary by CodeRabbit

Release Notes

• New Features

  • Enabled DPA4 model inference via the pt_expt backend using dual-artifact compilation.
  • Registered the EnergyModel under additional aliases: sezm_ener and dpa4_ener.

• Improvements

  • Improved dynamic/symbolic shape handling across DPA4 components for export/tracing stability.
  • Enhanced pt SeZM/DPA4 checkpoint deserialization and normalization for interoperability.
  • Added one-time warnings when use_amp or enable_tf32 settings are ineffective.

• Tests

  • Added C++ and Python coverage for pt2 inference, LAMMPS integration, model export/freeze, parity, interop, and warning behavior.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 4dd7cc76-cb22-4612-8794-d38319bc26cf

📥 Commits

Reviewing files that changed from the base of the PR and between 5ed93a1 and dfddcd5.

📒 Files selected for processing (5)

• deepmd/dpmodel/descriptor/dpa4.py
• deepmd/pt_expt/model/get_model.py
• source/tests/pt_expt/descriptor/test_dpa4.py
• source/tests/pt_expt/model/test_dpa4_export.py
• source/tests/pt_expt/model/test_get_model_dpa4.py

🚧 Files skipped from review as they are similar to previous changes (5)

• source/tests/pt_expt/descriptor/test_dpa4.py
• source/tests/pt_expt/model/test_get_model_dpa4.py
• deepmd/dpmodel/descriptor/dpa4.py
• deepmd/pt_expt/model/get_model.py
• source/tests/pt_expt/model/test_dpa4_export.py

---

📝 Walkthrough

Walkthrough

Adds DPA4/SeZM model support to the pt_expt backend: fixes symbolic shape handling in descriptor submodules to enable torch.export, registers sezm_ener/dpa4_ener model type aliases, implements a pt→pt_expt checkpoint deserialization shim, generates reference test fixtures, and adds C++, LAMMPS, and Python test coverage.

Changes

DPA4/SeZM pt_expt backend support

Layer / File(s)	Summary
Symbolic shape fixes for DPA4 descriptor export deepmd/dpmodel/descriptor/dpa4.py, deepmd/dpmodel/descriptor/dpa4_nn/attention.py, deepmd/dpmodel/descriptor/dpa4_nn/embedding.py, deepmd/dpmodel/descriptor/dpa4_nn/so2.py	Removes int() casts on n_nodes/n_edge across SO2Convolution, segment_envelope_gated_softmax, GeometricInitialEmbedding, and EnvironmentInitialEmbedding to keep shapes symbolic at trace time; reshape node axes switch to -1; call() uses xp_take_first_n instead of a direct slice to avoid export contiguity guards; adds one-time use_amp warning infrastructure via a module-level logger and process-global _WARNED_ONCE tracking set.
Model registry aliases and pt deserialization compatibility deepmd/dpmodel/model/ener_model.py, deepmd/pt_expt/model/ener_model.py, deepmd/pt_expt/model/get_model.py, deepmd/pt_expt/model/model.py	Registers EnergyModel under sezm_ener and dpa4_ener in both dpmodel and pt_expt registries; adds one-time enable_tf32-ignored warning in get_sezm_model via module logger and process-global tracking; overrides BaseModel.deserialize in pt_expt to detect pt SeZM/DPA4 wrapper types, validate @version/bridging_method/lora constraints, extract the inner atomic_model, and normalize sezm_atomic dicts by stripping pt-only fields and rewriting type/@version to the dpmodel standard schema.
Reference fixture generator and CI wiring source/tests/infer/gen_dpa4.py, source/install/test_cc_local.sh	Adds gen_dpa4.py which builds a pt_expt DPA4 model, seeds residual parameters to non-zero values via deterministic pseudo-random weights, exports .pt2 (required) and .pth (best-effort), evaluates PBC and NoPbc inference via DeepPot, writes deeppot_dpa4.expected reference file with per-atom energy/force/virial for both boundary modes, and cross-validates .pth/.pt2 parity; hooks gen_dpa4.py into the parallel model-generation section of the C++ test setup script.
C++ GTest inference coverage source/api_cc/tests/test_deeppot_dpa4_ptexpt.cc	Adds typed test fixtures TestInferDeepPotDpa4PtExpt and TestInferDeepPotDpa4PtExptNoPbc, covering direct compute with reference tolerance matching, EnergyModelTest numfv sensitivity sweeps, atomic output paths, LAMMPS-style InputNlist neighbor lists with ghost atom mapping and fold-back, 2× cutoff variants, extended vir-atom type-selection scenarios, and summary printing; guarded under BUILD_PYTORCH && BUILD_PT_EXPT.
Python pt_expt model unit and interop tests source/tests/pt_expt/descriptor/test_dpa4.py, source/tests/pt_expt/model/test_dpa4_interop.py, source/tests/pt_expt/model/test_dpa4_export.py, source/tests/pt_expt/model/test_get_model_dpa4.py	Adds test_use_amp_warns_once to descriptor tests; introduces test_dpa4_interop.py with pt→pt_expt deserialization happy path, atomic variable filtering, parametrized guard-branch exception testing, and version range acceptance; adds test_dpa4_export.py to freeze DPA4 to .pt2, validate dual-artifact ZIP structure/metadata, load AOTI artifact, and assert AOTI-vs-eager parity; extends test_get_model_dpa4.py with serialize/deserialize alias round-trip and enable_tf32 warn-once tests.
Deep eval parity and LAMMPS integration tests source/tests/pt_expt/infer/test_dpa4_deep_eval.py, source/lmp/tests/test_lammps_dpa4_pt2.py	Adds test_dpa4_deep_eval_parity/test_dpa4_deep_eval_metadata for fp64-close pt vs pt_expt energy/force/virial/per-atom-energy comparison and metadata consistency; adds full LAMMPS test suite for the DPA4 .pt2 artifact across metal/real/SI units, type-map variants, per-atom virial output, and atom-map-disabled fail-fast validation against deeppot_dpa4.expected reference.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Suggested reviewers

• njzjz
• iProzd

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 35.82% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately reflects the main change: adding DPA4 inference capabilities to pt_expt with .pt2 freeze, DeepEval support, C++, and LAMMPS testing.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[codecov]

Codecov Report

❌ Patch coverage is 86.44501% with 53 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.23%. Comparing base (16df629) to head (dfddcd5).
⚠️ Report is 2 commits behind head on master.

Files with missing lines	Patch %	Lines
source/api_cc/tests/test_deeppot_dpa4_ptexpt.cc	83.79%	53 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5540      +/-   ##
==========================================
+ Coverage   82.21%   82.23%   +0.01%     
==========================================
  Files         892      893       +1     
  Lines      101506   101912     +406     
  Branches     4242     4270      +28     
==========================================
+ Hits        83456    83807     +351     
- Misses      16749    16801      +52     
- Partials     1301     1304       +3     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.