fix(pd): preserve fparam/aparam inputs when freezing models

[wanghan-iapcm]

Problem

The Paddle freeze entrypoint (deepmd/pd/utils/serialization.py::deserialize_to_file) converts model.forward and model.forward_lower to static graphs with fparam and aparam hardcoded to None in the input_spec. For models trained with required frame parameters (get_dim_fparam() > 0) or atomic parameters (get_dim_aparam() > 0), the exported static signature therefore bakes both values as None, so inference through the frozen Paddle model cannot supply the required fparam/aparam.

Fix

Build the fparam/aparam InputSpec from the model's get_dim_fparam()/get_dim_aparam() — a spec with shape [-1, dim_fparam] / [-1, -1, dim_aparam] when the dim is nonzero, and None otherwise — and use it in both the forward and forward_lower static signatures. Models that do not use these inputs keep the None placeholders as before.

Test

source/tests/pd/model/test_serialization_fparam.py covers the spec builder for the unused case (both None), the used case (correct shapes and names), and the fparam-only case.

Note on verification

Verified locally with paddlepaddle==3.3.1. The spec-builder test is version-independent; the end-to-end paddle.jit.save export (heavier and sensitive to the exact Paddle build) remains covered by CI's pinned nightly (paddlepaddle==3.4.0.dev20260310).

Fix #5687

Summary by CodeRabbit

• Bug Fixes

  • Improved static export so optional input parameters are included only when a model actually uses them.
  • Reduced export issues by avoiding unnecessary placeholder inputs when these parameters are not needed.

• Tests

  • Added coverage for export input signatures to confirm optional parameters are omitted, included, or partially included as expected based on model settings.

[coderabbitai]

📝 Walkthrough

Walkthrough

Modifies Paddle model serialization to conditionally include fparam/aparam InputSpec entries in JIT static export signatures based on whether the model uses nonzero fparam/aparam dimensions, replacing prior placeholder None, None entries, and adds corresponding unit tests.

Changes

Optional fparam/aparam JIT export specs

Layer / File(s)	Summary
Input spec helper and export wiring deepmd/pd/utils/serialization.py	Adds _fparam_aparam_input_specs(model) to compute InputSpec objects for fparam/aparam only when their dimensions are nonzero, and threads fparam_spec/aparam_spec into the input_spec lists for model.forward and model.forward_lower to_static calls, replacing prior None, None placeholders.
Unit tests for input spec helper source/tests/pd/model/test_serialization_fparam.py	Adds a stub model and TestFparamAparamInputSpecs test cases validating that specs are None when unused, both present when used, and only fparam present when aparam is unused.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related issues

• [Code scan] Preserve Paddle fparam/aparam inputs when freezing models #5687: Directly addresses the Paddle freeze issue of preserving fparam/aparam InputSpecs in serialization.py described in this issue.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main fix: preserving fparam/aparam inputs during model freezing.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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]

🧹 Nitpick comments (1)

source/tests/pd/model/test_serialization_fparam.py (1)

23-41: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add a symmetric aparam-only test case.

Coverage includes both-unused, both-used, and fparam-only, but not the mirror case (aparam used, fparam unused). Adding it would fully cover the branch logic in _fparam_aparam_input_specs.

✅ Suggested additional test

     def test_only_fparam(self) -> None:
         fparam_spec, aparam_spec = _fparam_aparam_input_specs(_StubModel(2, 0))
         self.assertIsNotNone(fparam_spec)
         self.assertIsNone(aparam_spec)
+
+    def test_only_aparam(self) -> None:
+        fparam_spec, aparam_spec = _fparam_aparam_input_specs(_StubModel(0, 3))
+        self.assertIsNone(fparam_spec)
+        self.assertIsNotNone(aparam_spec)

🤖 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 `@source/tests/pd/model/test_serialization_fparam.py` around lines 23 - 41, Add
the missing mirror coverage for _fparam_aparam_input_specs by introducing an
aparam-only test in TestFparamAparamInputSpecs, analogous to test_only_fparam.
Verify that the returned fparam_spec is None, aparam_spec is present, and the
aparam_spec shape matches the stub model case where fparam is unused and aparam
is used.

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

Nitpick comments:
In `@source/tests/pd/model/test_serialization_fparam.py`:
- Around line 23-41: Add the missing mirror coverage for
_fparam_aparam_input_specs by introducing an aparam-only test in
TestFparamAparamInputSpecs, analogous to test_only_fparam. Verify that the
returned fparam_spec is None, aparam_spec is present, and the aparam_spec shape
matches the stub model case where fparam is unused and aparam is used.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 8f597054-126d-4164-9258-aef26f69a559

📥 Commits

Reviewing files that changed from the base of the PR and between 9bd1f16 and 239bb4e.

📒 Files selected for processing (2)

• deepmd/pd/utils/serialization.py
• source/tests/pd/model/test_serialization_fparam.py

[codecov]

Codecov Report

❌ Patch coverage is 87.50000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 81.11%. Comparing base (0e5c170) to head (239bb4e).
⚠️ Report is 15 commits behind head on master.

Files with missing lines	Patch %	Lines
deepmd/pd/utils/serialization.py	87.50%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5713      +/-   ##
==========================================
- Coverage   81.97%   81.11%   -0.86%     
==========================================
  Files         959      981      +22     
  Lines      105748   109867    +4119     
  Branches     4102     4234     +132     
==========================================
+ Hits        86684    89117    +2433     
- Misses      17573    19222    +1649     
- Partials     1491     1528      +37     

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