fix(dpa4): stabilize envelope-gated attention normalization

[OutisLi]

Summary by CodeRabbit

• Bug Fixes

  • Improved gated attention numerical stability for tiny edge/source weights, including smoother behavior at cutoff boundaries.
  • Refined attention normalization using explicit “physical null mass” and improved mixed-precision handling.
  • Updated cutoff envelope computation to a cancellation-free, factorized form for more reliable non-negative envelopes.

• Tests

  • Added regression coverage for cutoff-boundary behavior, vanishing-weight continuity, and second-derivative (Hessian) correctness.

• Documentation

  • Updated DPA4 example READMEs and input configs to use DPA4-Mini-style parameter sets.
  • Added an informational message when freezing SeZM models about setting the Triton inference level.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[coderabbitai]

📝 Walkthrough

Walkthrough

The change refactors C3 cutoff envelopes into factorized series evaluations, updates gated attention normalization with log-domain logits and physical null mass handling, adds numerical regression tests, and refreshes DPA4 example configurations and freeze guidance.

Changes

Numerical stability and DPA4 updates

Layer / File(s)	Summary
Factorized cutoff-envelope evaluation deepmd/dpmodel/descriptor/dpa4_nn/radial.py, deepmd/pt/model/descriptor/sezm_nn/radial.py	C3 cutoff envelopes now use precomputed binomial series coefficients, Horner evaluation, and a factorized u^4 form instead of explicit polynomial coefficients.
Log-domain gated attention normalization deepmd/dpmodel/descriptor/dpa4_nn/attention.py, deepmd/pt/model/descriptor/sezm_nn/attention.py	Attention uses active-edge effective logits, destination maxima containing physical null mass, shared shifted exponentials, null-mass normalization, and promoted computation for low-precision inputs.
Numerical regression coverage source/tests/pt/model/test_descriptor_sezm.py, source/tests/pt/model/test_dpa4_dpmodel_parity.py	Tests cover vanishing-edge continuity, cutoff-adjacent distances, and Hessian behavior for extremely small source weights across implementations.
DPA4 examples and freezing guidance examples/water/dpa4/*, deepmd/pt/entrypoints/freeze_pt2.py	Example configurations and documentation now describe DPA4-Mini settings and ratio-based warmup; freezing logs Triton inference configuration guidance.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related PRs

• deepmodeling/deepmd-kit#5599: Both changes modify destination-indexed segment_envelope_gated_softmax behavior in deepmd/dpmodel/descriptor/dpa4_nn/attention.py.

Suggested reviewers: wanghan-iapcm, iprozd, njzjz

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 57.14% 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: stabilizing DPA4 envelope-gated attention normalization.
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.

✨ 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.}

[coderabbitai]

Actionable comments posted: 2

🤖 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 `@source/tests/pt/model/test_descriptor_sezm.py`:
- Around line 1937-1943: Strengthen the Hessian validation in the test around
attention_sum and physical_sum by adding a separate assertion for the
reference[0, 0] component, using a meaningful tolerance appropriate for its
approximately -0.25 value. Keep the existing full-matrix assert_close for the
remaining Hessian comparison.
- Around line 1857-1868: Extend the assertions in the test around
evaluate(1.0e-12) to verify that the affected tiny-positive edge remains
nonzero, not just the unaffected zero edge. Use the existing near tensor and
device/dtype conventions, while preserving the current expected_stable and
zero-edge assertions.

🪄 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

Run ID: dd0e7c59-320e-43df-ad6e-378a7a98568b

📥 Commits

Reviewing files that changed from the base of the PR and between 1c5ff68 and 426bb6b.

📒 Files selected for processing (5)

• deepmd/dpmodel/descriptor/dpa4_nn/attention.py
• deepmd/dpmodel/descriptor/dpa4_nn/radial.py
• deepmd/pt/model/descriptor/sezm_nn/attention.py
• deepmd/pt/model/descriptor/sezm_nn/radial.py
• source/tests/pt/model/test_descriptor_sezm.py

[codecov]

Codecov Report

❌ Patch coverage is 95.18072% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.31%. Comparing base (1c5ff68) to head (d1c877d).
⚠️ Report is 5 commits behind head on master.

Files with missing lines	Patch %	Lines
deepmd/dpmodel/descriptor/dpa4_nn/attention.py	92.10%	3 Missing ⚠️
deepmd/pt/model/descriptor/sezm_nn/attention.py	96.55%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5808      +/-   ##
==========================================
- Coverage   78.57%   78.31%   -0.26%     
==========================================
  Files        1049     1050       +1     
  Lines      120650   120659       +9     
  Branches     4348     4358      +10     
==========================================
- Hits        94806    94500     -306     
- Misses      24285    24588     +303     
- Partials     1559     1571      +12     

☔ 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.

[wanghan-iapcm]

The numerics check out: the log-space softmax is algebraically equivalent with a correct log-sum-exp shift (the max-shift cancels through the null term, so autograd stays correct), the all-vanish case gives a well-defined alpha=0, the removed isfinite/eps guards were the original stabilization rather than a targeted patch and are provably unnecessary under the new null-mass formulation, and the C3 envelope factorization is exact (verified for p=1..7) and reaches exactly 0 at the cutoff. dpmodel and pt are faithful mirrors. LGTM.

One non-blocking follow-up on test coverage: the new extreme-regime tests (TestCutoffNumerics — nextafter-cutoff, src_weight=1e-30, high-logit vanishing, tiny-source-weight Hessian) are written only against the pt implementation. The dpmodel implementation is a separate rewrite and is the source of truth for the pt_expt and jax backends (pt_expt's descriptor subclasses this C3CutoffEnvelope), yet the existing cross-backend parity test (test_dpa4_dpmodel_parity.py) only feeds moderate rng.normal() values and never reaches these underflow/near-cutoff extremes. So nothing currently proves the dpmodel branch is equally stable in the exact regime this PR targets. Consider mirroring these edge cases into the parity file or adding a dpmodel-side equivalent, so the stability guarantee holds for all backends that consume dpmodel.

Minor: the Step 1 comment in deepmd/dpmodel/descriptor/dpa4_nn/attention.py lists edge_mask in the "edge_env**2 * src_weight * edge_mask underflows" product, but edge_mask is a binary 0/1 gate handled purely via the active predicate (never folded into log_weight), as the adjacent comment and the pt version correctly state. Worth aligning the wording.

[coderabbitai]

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 `@source/tests/pt/model/test_dpa4_dpmodel_parity.py`:
- Around line 2339-2361: Update the cutoff-attention test around dp_softmax and
pt_softmax to use a larger logit that keeps the expected attention weight in the
normal float32 range. Replace the current range-only assertion with an explicit
comparison against the expected physical value, while preserving the existing
DP/PyTorch parity check.

🪄 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

Run ID: 5aa0366e-ab43-479a-8bcd-bcf800caaa01

📥 Commits

Reviewing files that changed from the base of the PR and between d3f2fa3 and d1c877d.

📒 Files selected for processing (2)

• deepmd/dpmodel/descriptor/dpa4_nn/attention.py
• source/tests/pt/model/test_dpa4_dpmodel_parity.py

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

• deepmd/dpmodel/descriptor/dpa4_nn/attention.py