fix(pt_expt): keep use_amp out of serialization, fix the assembly boundary

use_amp is a runtime/training policy, not model state: revert the
dpa4/sezm serialize additions (a use_amp record leaks a torch runtime
option into cross-backend records -- the jax deserializer rejects
use_amp=true -- and deepmodeling#5963 established that checkpoints must not carry
the AMP switch).

The real pt_expt bug is in model assembly: make_model handed the raw
dpmodel atomic class to the dpmodel CM, so the constructed atomic model
was converted through the auto-wrap serialize()/deserialize() round-trip
and every runtime-only option on the live descriptor was reset to its
constructor default. Hand the CM the auto-wrapped atomic class instead:
the atomic model is constructed directly as a torch module and the live
(already wrapped) descriptor/fitting are kept as-is -- no round-trip.

Regression tests exercise the public construction path (get_model with
descriptor.use_amp=false) and pin that the portable record does not
carry use_amp.

Author: Han Wang

deepmd/dpmodel/descriptor/dpa4.py

@@ -2806,11 +2806,6 @@ def serialize(self) -> dict[str, Any]:
                 "mlp_bias": self.mlp_bias,
                 "exclude_types": self.exclude_types,
                 "eps": self.eps,
-                # Must round-trip: pt_expt rebuilds the descriptor from this
-                # dict, so omitting the key silently reset a configured
-                # ``use_amp: false`` to True and kept training in bfloat16.
-                # Older records without it still load (__init__ defaults it).
-                "use_amp": self.use_amp,
                 "trainable": self.trainable,
                 "seed": self.seed,
                 "inner_clamp_r_inner": self.inner_clamp_r_inner,

deepmd/pt/model/descriptor/sezm.py

@@ -2556,9 +2556,6 @@ def serialize(self) -> dict[str, Any]:
                 "mlp_bias": self.mlp_bias,
                 "exclude_types": self.exclude_types,
                 "eps": self.eps,
-                # Kept in step with the dpmodel serialize contract so both
-                # backends' records carry the same keys.
-                "use_amp": self.use_amp,
                 "trainable": self.trainable,
                 "seed": self.seed,
                 "inner_clamp_r_inner": self.inner_clamp_r_inner,

deepmd/pt_expt/common.py

@@ -136,32 +136,54 @@ def try_convert_module(value: Any) -> torch.nn.Module | None:
 _AUTO_WRAPPED_CLASSES: dict[type, type] = {}
 
 
-def _auto_wrap_native_op(value: NativeOP) -> torch.nn.Module:
-    """Auto-wrap any NativeOP as a torch.nn.Module via ``torch_module``.
+def auto_wrapped_class(cls: type) -> type:
+    """Return the cached ``torch_module`` auto-wrap of a dpmodel class.
 
     Creates a subclass with a generic ``forward`` that delegates to ``call``,
     then applies ``torch_module`` to get full ``__setattr__`` / post-init
     list conversion.  The wrapped class is cached per dpmodel type.
 
+    Constructing this class DIRECTLY (instead of building the raw dpmodel
+    class and converting the instance afterwards) is the lossless path:
+    instance conversion round-trips through ``serialize()``/``deserialize()``,
+    which by design drops runtime-only configuration (e.g. the DPA4
+    ``use_amp`` switch) from the portable record.
+
     Parameters
     ----------
-    value : NativeOP
-        The dpmodel object to wrap.
+    cls : type
+        The dpmodel NativeOP class to wrap.
 
     Returns
     -------
-    torch.nn.Module
-        The wrapped pt_expt module, deserialized from value's serialized state.
+    type
+        The ``torch_module``-wrapped subclass.
     """
-    cls = type(value)
     if cls not in _AUTO_WRAPPED_CLASSES:
         wrapped = type(
             cls.__name__,
             (cls,),
             {"forward": lambda self, *args, **kwargs: self.call(*args, **kwargs)},
         )
         _AUTO_WRAPPED_CLASSES[cls] = torch_module(wrapped)
-    wrapped_cls = _AUTO_WRAPPED_CLASSES[cls]
+    return _AUTO_WRAPPED_CLASSES[cls]
+
+
+def _auto_wrap_native_op(value: NativeOP) -> torch.nn.Module:
+    """Auto-wrap any NativeOP as a torch.nn.Module via ``torch_module``.
+
+    Parameters
+    ----------
+    value : NativeOP
+        The dpmodel object to wrap.
+
+    Returns
+    -------
+    torch.nn.Module
+        The wrapped pt_expt module, deserialized from value's serialized state.
+    """
+    cls = type(value)
+    wrapped_cls = auto_wrapped_class(cls)
     if not (hasattr(value, "serialize") and hasattr(wrapped_cls, "deserialize")):
         raise TypeError(
             f"Cannot auto-wrap {cls.__name__}: "

deepmd/pt_expt/model/make_model.py

@@ -24,6 +24,7 @@
     cuda_infer_level,
 )
 from deepmd.pt_expt.common import (
+    auto_wrapped_class,
     torch_module,
 )
 from deepmd.pt_expt.utils.graph_builder import (
@@ -422,7 +423,14 @@ def make_model(
         The model.
 
     """
-    DPModel = make_model_dp(T_AtomicModel)
+    # Hand the dpmodel CM the WRAPPED atomic class so `self.atomic_model =
+    # T_AtomicModel(...)` constructs the pt_expt module directly with the
+    # live (already wrapped) descriptor/fitting. Passing the raw dpmodel
+    # class instead would build a raw atomic model and convert the instance
+    # through a serialize()/deserialize() round-trip, which drops
+    # runtime-only configuration (e.g. the DPA4 `use_amp` switch) that the
+    # portable record deliberately does not carry.
+    DPModel = make_model_dp(auto_wrapped_class(T_AtomicModel))
 
     @torch_module
     class CM(DPModel, *T_Bases):

source/tests/common/dpmodel/test_descrpt_dpa4.py

@@ -265,26 +265,18 @@ def test_supported_feature_roundtrip(self, overrides) -> None:
         out2 = np.asarray(dd2.call(coord.reshape(nf, -1), atype, nlist)[0])
         np.testing.assert_array_equal(out1, out2)
 
-    @pytest.mark.parametrize(
-        "use_amp",
-        [
-            True,  # the constructor default; must not be clobbered either
-            False,  # the value that was silently lost, re-enabling autocast
-        ],
-    )
-    def test_use_amp_survives_roundtrip(self, use_amp) -> None:
-        """``use_amp`` must round-trip through serialize/deserialize.
-
-        The key was missing from the config, so a backend that rebuilds from
-        it (pt_expt does) reset ``use_amp: false`` to True and kept training in
-        bfloat16.  The forward-output round-trip test can't catch this --
-        dpmodel never autocasts, so outputs match either way.
+    def test_use_amp_stays_out_of_the_portable_record(self) -> None:
+        """``use_amp`` is a runtime/training policy, not model state.
+
+        The portable serialization must not carry it (a ``use_amp: true``
+        record would e.g. be rejected by the JAX deserializer); a fresh
+        deserialize falls back to the constructor default. Construction-time
+        survival is pinned at the pt_expt assembly boundary instead
+        (``test_get_model_dpa4.py``).
         """
-        dd = make_descriptor(use_amp=use_amp)
-        assert dd.use_amp is use_amp
-        assert dd.serialize()["config"]["use_amp"] is use_amp
-        dd2 = DescrptDPA4.deserialize(dd.serialize())
-        assert dd2.use_amp is use_amp
+        dd = make_descriptor(use_amp=False)
+        assert dd.use_amp is False
+        assert "use_amp" not in dd.serialize()["config"]
 
     def test_legacy_spin_gate_is_squared_on_deserialize(self) -> None:
         """Version 1.2 stores the env-seed spin gate after the quadratic form.

source/tests/pt_expt/model/test_get_model_dpa4.py

@@ -285,3 +285,61 @@ def test_unrelated_construction_error_propagates(self) -> None:
 
 if __name__ == "__main__":
     unittest.main()
+
+
+class TestUseAmpSurvivesAssembly(unittest.TestCase):
+    """``use_amp`` is runtime policy: it must survive model ASSEMBLY without
+    entering the portable serialization record (which the JAX deserializer
+    rejects for ``use_amp: true``). The wrapping of the atomic model must
+    therefore not round-trip the constructed descriptor through
+    ``serialize()``/``deserialize()``.
+    """
+
+    def test_get_model_keeps_use_amp_false(self) -> None:
+        model = get_model(
+            _make_raw_model_config(
+                descriptor={
+                    "sel": 20,
+                    "rcut": 4.0,
+                    "channels": 8,
+                    "n_radial": 4,
+                    "lmax": 1,
+                    "mmax": 1,
+                    "n_blocks": 1,
+                    "precision": "float64",
+                    "seed": 1,
+                    "use_amp": False,
+                }
+            )
+        )
+        assert model.atomic_model.descriptor.use_amp is False
+
+    def test_get_model_keeps_the_use_amp_default(self) -> None:
+        model = get_model(_make_raw_model_config())
+        assert model.atomic_model.descriptor.use_amp is True
+
+    def test_standard_type_keeps_use_amp_false(self) -> None:
+        """The plain `standard` route wraps through the same boundary."""
+        cfg = _make_raw_model_config(
+            descriptor={
+                "type": "dpa4",
+                "sel": 20,
+                "rcut": 4.0,
+                "channels": 8,
+                "n_radial": 4,
+                "lmax": 1,
+                "mmax": 1,
+                "n_blocks": 1,
+                "precision": "float64",
+                "seed": 1,
+                "use_amp": False,
+            },
+            fitting_net={
+                "type": "dpa4_ener",
+                "precision": "float64",
+                "seed": 1,
+            },
+        )
+        del cfg["type"]
+        model = get_model(cfg)
+        assert model.atomic_model.descriptor.use_amp is False