fix(pt): use strict=False when loading .pt checkpoints for inference

[wanghan-iapcm]

Summary

• Use strict=False in load_state_dict when loading .pt checkpoints in the no_jit=True inference path
• Log warnings for missing/unexpected keys instead of crashing

Root cause

PR #5323 and #5057 added new buffers/parameters to DescrptBlockSeAtten.__init__:

• type_embd_data (buffer, for type embedding compression)
• compress_info (ParameterList, for geometric compression)
• compress_data (ParameterList, for geometric compression)

These are initialized to zero-size tensors (uncompressed state). Old checkpoints saved before these PRs don't contain these keys, so load_state_dict(strict=True) raises RuntimeError: Missing key(s).

Fix

Use strict=False with warning logs, matching the pattern already used in training.py:740. The zero-initialized defaults correctly represent uncompressed state.

Test plan

• Pre-existing tests pass (no behavioral change for new checkpoints)
• Fixes loading of old DPA-2 checkpoints (dpa-2.4-7M.pt) with no_jit=True

Reported by @njzjz .

Summary by CodeRabbit

• Bug Fixes
  • Improved checkpoint loading compatibility to support models with missing or unexpected parameters, with detailed warning messages logged for troubleshooting.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: d5ec8938b7

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[coderabbitai]

📝 Walkthrough

Walkthrough

When loading PyTorch checkpoints in the inference module, load_state_dict now accepts mismatched keys by using strict=False, capturing and logging warnings for any missing or unexpected parameters instead of throwing an error.

Changes

Cohort / File(s)	Summary
Checkpoint Loading deepmd/pt/infer/deep_eval.py	Modified load_state_dict call to use strict=False, capturing missing and unexpected keys; added warning logs to report missing keys (attributed to older checkpoint versions) and unexpected keys encountered during model weight loading.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• feat(pt): allow --init-frz-model for pt model converted from tf model #5091: Implements the same checkpoint loading strategy with strict=False and warning logs for missing/unexpected keys during model initialization.

Suggested labels

Python

Suggested reviewers

• iProzd

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

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

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately describes the main change: using strict=False when loading .pt checkpoints, which is the core fix addressed in this PR.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

_{✏️ 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 60.00000% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.26%. Comparing base (2a82988) to head (d5ec893).
⚠️ Report is 3 commits behind head on master.

Files with missing lines	Patch %	Lines
deepmd/pt/infer/deep_eval.py	60.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5353      +/-   ##
==========================================
- Coverage   82.26%   82.26%   -0.01%     
==========================================
  Files         799      799              
  Lines       82563    82567       +4     
  Branches     4066     4066              
==========================================
+ Hits        67924    67926       +2     
- Misses      13424    13426       +2     
  Partials     1215     1215              

☔ View full report in Codecov by Sentry.
📢 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.