feat(pt_expt): enhance per-module export tests + fix nframes symbol collision (#5367)

## Summary
- Enhance existing `test_make_fx` (9 descriptors, 6 fittings) and
`test_forward_lower_exportable` (8 models) to cover the full `dp freeze`
pipeline: symbolic tracing, `torch.export` with dynamic shapes, `.pte`
save/load round-trip, and inference with different `nframes`
- Add `**make_fx_kwargs` forwarding to 6 model
`forward_lower_exportable` methods (dipole, dos, polar, property,
dp_zbl, dp_linear)
- Create `source/tests/pt_expt/export_helpers.py` with shared test
helpers
- Add new `test_make_fx` to `test_ener_fitting.py` and
`test_invar_fitting.py` with fparam/aparam parametrization
- Fix nframes symbol collision: change tracing `nframes` from 2 to 5 in
`_trace_and_export` to avoid PyTorch's symbolic tracer merging dim
symbols when `nframes == numb_fparam`

## Test plan
- [x] 16 descriptor `test_make_fx` pass
- [x] 6 fitting `test_make_fx` pass
- [x] 8 model `test_forward_lower_exportable` pass
- [x] `test_export_pipeline` with `fparam=True` passes for se_e2_a and
dpa1
- [ ] CI

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* FX tracing for export is now caller-configurable for exportable model
methods.
* Increased sample frames used for export traces to avoid unwanted
specialization/collisions.

* **Tests**
* Added comprehensive export → save → load round-trip validations across
descriptor, fitting, and model tests.
* Introduced shared test helpers/utilities for dynamic-shape export
verification.
* Added export-pipeline and linear-model test coverage; many tests
refactored to use the new helpers and one obsolete test module removed.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: wanghan-iapcm

deepmd/pt_expt/model/dipole_model.py

@@ -124,6 +124,7 @@ def forward_lower_exportable(
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
         do_atomic_virial: bool = False,
+        **make_fx_kwargs: Any,
     ) -> torch.nn.Module:
         model = self
 
@@ -146,6 +147,6 @@ def fn(
                 do_atomic_virial=do_atomic_virial,
             )
 
-        return make_fx(fn)(
+        return make_fx(fn, **make_fx_kwargs)(
             extended_coord, extended_atype, nlist, mapping, fparam, aparam
         )

deepmd/pt_expt/model/dos_model.py

@@ -104,6 +104,7 @@ def forward_lower_exportable(
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
         do_atomic_virial: bool = False,
+        **make_fx_kwargs: Any,
     ) -> torch.nn.Module:
         model = self
 
@@ -126,6 +127,6 @@ def fn(
                 do_atomic_virial=do_atomic_virial,
             )
 
-        return make_fx(fn)(
+        return make_fx(fn, **make_fx_kwargs)(
             extended_coord, extended_atype, nlist, mapping, fparam, aparam
         )

deepmd/pt_expt/model/dp_linear_model.py

@@ -129,6 +129,7 @@ def forward_lower_exportable(
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
         do_atomic_virial: bool = False,
+        **make_fx_kwargs: Any,
     ) -> torch.nn.Module:
         model = self
 
@@ -151,7 +152,7 @@ def fn(
                 do_atomic_virial=do_atomic_virial,
             )
 
-        return make_fx(fn)(
+        return make_fx(fn, **make_fx_kwargs)(
             extended_coord, extended_atype, nlist, mapping, fparam, aparam
         )
 

deepmd/pt_expt/model/dp_zbl_model.py

@@ -126,6 +126,7 @@ def forward_lower_exportable(
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
         do_atomic_virial: bool = False,
+        **make_fx_kwargs: Any,
     ) -> torch.nn.Module:
         model = self
 
@@ -148,6 +149,6 @@ def fn(
                 do_atomic_virial=do_atomic_virial,
             )
 
-        return make_fx(fn)(
+        return make_fx(fn, **make_fx_kwargs)(
             extended_coord, extended_atype, nlist, mapping, fparam, aparam
         )

deepmd/pt_expt/model/polar_model.py

@@ -104,6 +104,7 @@ def forward_lower_exportable(
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
         do_atomic_virial: bool = False,
+        **make_fx_kwargs: Any,
     ) -> torch.nn.Module:
         model = self
 
@@ -126,6 +127,6 @@ def fn(
                 do_atomic_virial=do_atomic_virial,
             )
 
-        return make_fx(fn)(
+        return make_fx(fn, **make_fx_kwargs)(
             extended_coord, extended_atype, nlist, mapping, fparam, aparam
         )

deepmd/pt_expt/model/property_model.py

@@ -111,6 +111,7 @@ def forward_lower_exportable(
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
         do_atomic_virial: bool = False,
+        **make_fx_kwargs: Any,
     ) -> torch.nn.Module:
         model = self
 
@@ -133,6 +134,6 @@ def fn(
                 do_atomic_virial=do_atomic_virial,
             )
 
-        return make_fx(fn)(
+        return make_fx(fn, **make_fx_kwargs)(
             extended_coord, extended_atype, nlist, mapping, fparam, aparam
         )

deepmd/pt_expt/utils/serialization.py

@@ -330,12 +330,17 @@ def _trace_and_export(
     metadata = _collect_metadata(model)
 
     # 3. Create sample inputs on CPU for tracing
-    # Use nframes=2 so make_fx doesn't specialize on nframes=1
+    # Use nframes=5 to avoid two specialization traps:
+    #   - nframes=1 causes make_fx to specialize on the scalar case
+    #   - nframes=N where N == numb_fparam or numb_aparam causes PyTorch's
+    #     symbolic tracer to merge symbols (e.g. fparam.shape=(2,2) when
+    #     nframes=2 and numb_fparam=2), so a guard on one dim constrains
+    #     the other.  5 is unlikely to collide with typical param counts.
     _orig_device = _env.DEVICE
     _env.DEVICE = torch.device("cpu")
     try:
         ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam = _make_sample_inputs(
-            model, nframes=2
+            model, nframes=5
         )
     finally:
         _env.DEVICE = _orig_device

source/tests/pt_expt/descriptor/test_dpa1.py

@@ -25,6 +25,10 @@
 from ...seed import (
     GLOBAL_SEED,
 )
+from ..export_helpers import (
+    export_save_load_and_compare,
+    make_descriptor_dynamic_shapes,
+)
 
 
 class TestDescrptDPA1(TestCaseSingleFrameWithNlist):
@@ -223,3 +227,15 @@ def fn(coord_ext, atype_ext, nlist):
             rtol=rtol,
             atol=atol,
         )
+
+        # --- symbolic trace + export + .pte round-trip ---
+        dynamic_shapes = make_descriptor_dynamic_shapes(has_mapping=False)
+        inputs = (coord_ext, atype_ext, nlist)
+        export_save_load_and_compare(
+            fn,
+            inputs,
+            (rd_eager, grad_eager),
+            dynamic_shapes,
+            rtol=rtol,
+            atol=atol,
+        )

source/tests/pt_expt/descriptor/test_dpa2.py

@@ -29,6 +29,10 @@
 from ...seed import (
     GLOBAL_SEED,
 )
+from ..export_helpers import (
+    export_save_load_and_compare,
+    make_descriptor_dynamic_shapes,
+)
 
 
 class TestDescrptDPA2(TestCaseSingleFrameWithNlist):
@@ -410,3 +414,15 @@ def fn(coord_ext, atype_ext, nlist, mapping):
             rtol=rtol,
             atol=atol,
         )
+
+        # --- symbolic trace + export + .pte round-trip ---
+        dynamic_shapes = make_descriptor_dynamic_shapes(has_mapping=True)
+        inputs = (coord_ext, atype_ext, nlist, mapping)
+        export_save_load_and_compare(
+            fn,
+            inputs,
+            (rd_eager, grad_eager),
+            dynamic_shapes,
+            rtol=rtol,
+            atol=atol,
+        )

source/tests/pt_expt/descriptor/test_dpa3.py

@@ -28,6 +28,10 @@
 from ...seed import (
     GLOBAL_SEED,
 )
+from ..export_helpers import (
+    export_save_load_and_compare,
+    make_descriptor_dynamic_shapes,
+)
 
 
 class TestDescrptDPA3(TestCaseSingleFrameWithNlist):
@@ -244,3 +248,15 @@ def fn(coord_ext, atype_ext, nlist, mapping):
             rtol=rtol,
             atol=atol,
         )
+
+        # --- symbolic trace + export + .pte round-trip ---
+        dynamic_shapes = make_descriptor_dynamic_shapes(has_mapping=True)
+        inputs = (coord_ext, atype_ext, nlist, mapping)
+        export_save_load_and_compare(
+            fn,
+            inputs,
+            (rd_eager, grad_eager),
+            dynamic_shapes,
+            rtol=rtol,
+            atol=atol,
+        )