fix(jax): add hessian energy loss - #5607
Conversation
b56a253 to
ae453b0
Compare
|
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:
📝 WalkthroughWalkthrough
ChangesHessian Loss Feature
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Trainer
participant EnergyModel
participant DataSystem
participant EnergyLoss
Trainer->>EnergyLoss: read Hessian data requirement
Trainer->>EnergyModel: enable Hessian outputs
DataSystem->>EnergyLoss: provide padded Hessian labels
EnergyModel->>EnergyLoss: provide Hessian predictions
EnergyLoss->>EnergyLoss: compute weighted masked loss and metrics
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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: 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 `@deepmd/dpmodel/loss/ener.py`:
- Around line 565-574: The Hessian entry in `DPTrainer.data_requirements` is
advertising the wrong tensor shape for `label_requirement`. Update the
`DataRequirementItem("hessian", ...)` definition in `ener.py` so it matches the
real on-disk Hessian layout used by the new loss path and tests, rather than the
current atomic `ndof=1` schema. If the dataset loader cannot yet consume the
full Hessian tensor, add the loader support first and keep `has_h` gated until
the contract is consistent.
In `@deepmd/jax/train/trainer.py`:
- Around line 120-121: The Hessian enablement in Trainer setup is unguarded, so
`self.model.enable_hessian()` can break for model/loss combinations like the JAX
`zbl` path that do not implement it. Update the `Trainer` logic to check that
the model actually exposes `enable_hessian` before calling it, using the
existing `self.loss.has_h` condition as the trigger and keeping the current
hessian tensor handling unchanged.
🪄 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: 23778160-e75a-4d2b-be1f-2babb0b6710f
📒 Files selected for processing (3)
deepmd/dpmodel/loss/ener.pydeepmd/jax/train/trainer.pysource/tests/common/dpmodel/test_loss_ener.py
|
Possible reviewers based on changed lines, exact file history, and exact-file review history:
No review request was made automatically. Coding agent: Codex |
Adapt Hessian loss support to the multi-task JAX trainer by enabling Hessian outputs per task and exposing the transformed Hessian through the shared model-evaluation path. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Merge the latest master changes and retain the existing serialization and data-requirement schemas when Hessian loss is disabled. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5607 +/- ##
==========================================
- Coverage 79.31% 79.21% -0.11%
==========================================
Files 1070 1072 +2
Lines 124601 125020 +419
Branches 4532 4536 +4
==========================================
+ Hits 98831 99038 +207
- Misses 24150 24359 +209
- Partials 1620 1623 +3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Exclude Hessian rows and columns associated with placeholder atoms from the training MSE and the MAE/RMSE denominators while preserving the unmasked reduction. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Keep the energy loss as the canonical schema for Hessian prefactors, normalize the legacy ener_hess type to ener, and update documentation links to the canonical loss fields. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Move Hessian supervision into EnergyStdLoss, keep EnergyHessianStdLoss and ener_hess as compatibility aliases, enable Hessian when either prefactor is active, and apply the shared placeholder-pair mask to PT metrics and training loss. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
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.
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 `@deepmd/pt/loss/ener.py`:
- Around line 789-819: Update the Hessian block in the loss computation to
branch on self.loss_func like the energy, force, virial, and atom-energy terms:
use the squared residual for "mse", absolute residual for "mae", and raise
NotImplementedError for unsupported values. Apply the selected objective to loss
+= pref_h * ... while preserving the existing display metrics and masking
behavior.
🪄 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: 5bc53c5d-25cd-410a-af3b-84cdcb879872
📒 Files selected for processing (7)
deepmd/dpmodel/loss/ener.pydeepmd/dpmodel/loss/reduction.pydeepmd/dpmodel/model/ener_model.pydeepmd/jax/train/trainer.pydeepmd/pt/loss/ener.pydeepmd/pt/model/model/ener_model.pydeepmd/pt/train/training.py
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/dpmodel/loss/ener.py
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Folding Hessian supervision into the shared EnergyLoss is the right call, and inverting the control flow so the loss's label_requirement drives enable_hessian() is a much better contract than the trainer sniffing start_pref_h. I checked masked_pair_mean numerically rather than by eye: the pair mask really is the outer product of the per-atom mask, the denominator is (3*n_real)**2, an all-padding frame returns a neutral 0.0, and it reduces exactly to mean under an all-ones mask. The enable_hessian() idempotency guards are correct on both backends, and with both prefs left at 0 the whole thing is inert.
Two things I'd like resolved before this goes in, both inline.
Non-blocking notes:
-
@versionstays at 4 in bothserialize()implementations while two new fields are added, andcheck_version_compatibility(version, 4, 1)is unchanged. Because both__init__s end in**kwargs, an older deepmd-kit reading a Hessian-enabled dict passes the version gate and then silently discardsstart_pref_h/limit_pref_h- the model loads with no Hessian term and no error. Worth a bump plus an explicit 4-to-5 branch. -
has_hpicks upor inference, which is consistent with howhas_e/has_f/has_v/has_ae/has_pfalready behave - but every one of those labels is linear in N, whereas the Hessian is9N**2.dp change-biasbuilds the loss withinference=Truepurely to harvestlabel_requirement, so it now requests a Hessian it never uses, and_load_single_datazero-fills it at roughly 72 MB per frame for a 1000-atom system (vs 24 KB for force). The convention is fine; it just doesn't extend safely to a quadratic label. Excluding the Hessian item whenhas_his true only because ofinferencewould fix it. -
deepmd/pdis untouched, so after this lands the same input.json trains a Hessian term on pt/jax/dpmodel but not on pd, which keeps the old AND semantics. Related:alias=["ener_hess"]makestype: ener_hessa valid input for the first time, and it now normalizes toenerand reaches pd's un-audited path with no pd test. -
doc/model/train-energy-hessian.mdstill advertises PyTorch only, though this PR wires and tests JAX end-to-end. -
The
and->orchange means an existingstart_pref_h: 100, limit_pref_h: 0config now trains with an annealed Hessian term where it previously trained with none. I think the new behaviour is right - argcheck's sharedstart_pref()text has always said the term is ignored only when both are zero, so the old AND was the outlier - but it changes results for existing configs silently and deserves a release note.
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Both blocking items are fixed, and I verified them at a1ed51d2a rather than from the replies.
The Hessian term now dispatches on loss_func in both backends (dpmodel, pt), with has_h added to the Huber guard in both constructors. The _merge_batch_data branch is in the right place -- ahead of the atomic/non-atomic split, embedding each frame's 3N_i x 3N_i block in the top-left of the padded square (data_system.py).
The regression tests genuinely fail pre-fix, which is what I wanted to see. test_mae_loss_func_uses_l1_for_the_hessian zeroes every prefactor except pref_h and uses a constant residual of 10, so the old code contributes 100 where it asserts 10, and it cross-checks that the mse config still gives 100. test_data_system_hessian.py goes through the real loader with 2- and 3-atom systems, so the old np.concatenate hits (1, 36) against (1, 81). The unknown-loss_func and use_huber cases are covered in both backends. Good coverage of the exact intersections that were missing.
Two of my earlier non-blocking notes are still open. Neither blocks this, but flagging that they were not answered either way:
-
@versionis still 4 in bothserialize()implementations (dpmodel, pt) whilestart_pref_h/limit_pref_hare now emitted conditionally, andcheck_version_compatibility(version, 4, 1)is unchanged. Since both__init__s end in**kwargs, an older deepmd-kit reading a Hessian-enabled dict passes the version gate and then silently drops both fields -- the model loads with no Hessian term and no error. This is the case the repo's serialization-version discipline is meant to catch, so a bump plus an explicit 4-to-5 branch would be worth doing before or shortly after this lands. -
pt/loss/ener.py#L166still readsself.has_h = (start_pref_h != 0.0 or limit_pref_h != 0.0) or inference.dp change-biasconstructs the loss withinference=Trueonly to harvestlabel_requirement, so it now requests a Hessian it never uses and zero-fills roughly 9N^2 per frame -- about 72 MB per frame at 1000 atoms, against 24 KB for force. The convention matches the otherhas_*flags; it just does not extend safely to a quadratic label.
Thanks for the doc cleanup to the loss[ener] references as well.
Address Hessian serialization compatibility, avoid quadratic change-bias labels, align Paddle activation, and document JAX support. 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.
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 `@deepmd/pd/loss/ener.py`:
- Around line 652-654: In the initializer after assigning `self.has_h`, reject
the unsupported combination by checking `self.use_huber` together with
`self.has_h` and raising the established `RuntimeError` used by the dpmodel and
PyTorch implementations. Keep the existing endpoint activation and inference
logic unchanged.
- Around line 706-710: Update EnergyHessianStdLoss.forward() to mask padded
Hessian atom pairs along both atom axes before flattening and reducing. Use the
Paddle equivalent of masked_pair_mean, matching the dpmodel and PyTorch
implementations, so only valid Hessian entries contribute to the loss and
rmse_h.
🪄 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: 6fdcf074-92d1-4643-95f6-54a66925dc31
📒 Files selected for processing (8)
deepmd/dpmodel/loss/ener.pydeepmd/pd/loss/ener.pydeepmd/pd/train/training.pydeepmd/pt/loss/ener.pydoc/model/train-energy-hessian.mdsource/tests/common/dpmodel/test_loss_ener.pysource/tests/pd/test_loss.pysource/tests/pt/test_loss.py
🚧 Files skipped from review as they are similar to previous changes (1)
- doc/model/train-energy-hessian.md
Resolve the energy-loss argument conflict by retaining the ener_hess alias and master’s backend support documentation. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
@wanghan-iapcm Addressed the remaining non-blocking notes in 4d667b2 and resolved the latest master conflict in eb42bec:
Validation after merging current master: 20 focused loss, serialization, argument-schema, and JAX tests passed; the complete PyTorch change-bias scenario passed; ruff format . and ruff check . passed. A Paddle regression test is included, but Paddle is not installed in the local environment. Coding agent: Codex |
Reject unsupported Huber Hessian supervision and normalize Hessian MSE/MAE over real-real Cartesian pairs in mixed-system batches. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Construct PyTorch Hessian test inputs on the configured backend device so CUDA assertions compare colocated tensors. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
for more information, see https://pre-commit.ci
Summary
EnergyLossinstead of introducing a separate loss classTests
source venv/bin/activate && pytest source/tests/common/dpmodel/test_loss_ener.py -qsource venv/bin/activate && ruff check .source venv/bin/activate && ruff format .Summary by CodeRabbit
ener/legacyener_hessalias behavior.