feat(pt_expt): honor enable_tf32 / DP_TF32_INFER like the pt backend

pt_expt accepted `model.enable_tf32` and threw it away with a warning, so
DPA4/SeZM training always ran at "highest" matmul precision while the pt
backend -- reading the same input.json -- ran its training forwards under
`set_float32_matmul_precision("high")`.  On Ampere and later that is the
difference between TF32 tensor cores and fp32 CUDA cores for every matmul,
and GEMM is ~60% of compiled device time on this workload, so the two
backends were not comparable on that hardware at all.

Mirror pt's policy exactly: TRAINING forwards follow `enable_tf32`
(argcheck default True), EVAL forwards follow `DP_TF32_INFER` (0/1/2 ->
highest/high/medium, invalid values rejected).  Scope matches pt, where
argcheck declares the knob inside the dpa4 model arg block and only the
sezm builders wire it: pt_expt attaches it in `get_sezm_model` and
`get_native_spin_model`, and every other model keeps class defaults that
select full fp32 in both modes.

Ownership: `call_common` is the single owner for eager forwards -- every
pt_expt model's `forward` reaches the backbone through it, and the export
trace roots at `call_common_lower`, so the precision switch never enters an
exported graph.  The compiled path needs its own application because
`_CompiledModel.forward` bypasses `call_common` entirely; placing the
context only on the model would have left it dead on exactly the path this
is meant to speed up.  The context spans the lazy compile there, since
Inductor picks its GEMM backend while lowering.

Gating on `self.training` is what keeps the existing 1e-12 parity tests
valid: eval and export stay at "highest" unless DP_TF32_INFER asks
otherwise.

Author: Han Wang

deepmd/pt_expt/model/get_model.py

@@ -8,6 +8,7 @@
 
 import copy
 import logging
+import os
 
 from deepmd.dpmodel.atomic_model.dp_atomic_model import (
     DPAtomicModel,
@@ -52,8 +53,49 @@
 
 log = logging.getLogger(__name__)
 
-# Warn at most once per process for backend-ignored switches (keyed by name).
-_WARNED_ONCE: set[str] = set()
+#: ``DP_TF32_INFER`` -> eval-time matmul precision, copied from the pt backend's
+#: ``deepmd.pt.model.model.sezm_model._TF32_INFER_PRECISION_CHOICES`` so the two
+#: backends read the same environment variable the same way.
+_TF32_INFER_PRECISION_CHOICES = {
+    "0": "highest",
+    "1": "high",
+    "2": "medium",
+}
+
+
+def _apply_tf32_policy(model: BaseModel, data: dict) -> BaseModel:
+    """Attach the DPA4/SeZM TF32 matmul-precision policy to a built model.
+
+    Mirrors the pt backend: TRAINING forwards follow ``model.enable_tf32``
+    (default ``True``) and EVAL forwards follow ``DP_TF32_INFER``.  The model
+    layer applies the policy -- see ``call_common`` in
+    :func:`deepmd.pt_expt.model.make_model.make_model` and
+    ``_CompiledModel.forward`` for the compiled path.
+
+    Parameters
+    ----------
+    model : BaseModel
+        The freshly built model to configure.
+    data : dict
+        The model config section, read for ``enable_tf32``.
+
+    Returns
+    -------
+    BaseModel
+        The same model, with the precision policy attached.
+
+    Raises
+    ------
+    ValueError
+        If ``DP_TF32_INFER`` is set to anything other than ``0``, ``1``, or
+        ``2``.
+    """
+    model.enable_tf32 = bool(data.get("enable_tf32", True))
+    tf32_infer_env = os.environ.get("DP_TF32_INFER", "0").strip().lower()
+    if tf32_infer_env not in _TF32_INFER_PRECISION_CHOICES:
+        raise ValueError(f"DP_TF32_INFER must be one of 0/1/2, got {tf32_infer_env!r}")
+    model.tf32_infer_precision = _TF32_INFER_PRECISION_CHOICES[tf32_infer_env]
+    return model
 
 
 _model_factory = BackendModelFactory(
@@ -82,17 +124,11 @@ def get_sezm_model(data: dict) -> EnergyModel:
 
     Notes
     -----
-    ``enable_tf32`` is accepted but ignored: the pt backend uses it to toggle
-    TF32 matmul precision, while the pt_expt backend always runs at full
-    ("highest") matmul precision, which is numerically conservative.
+    ``enable_tf32`` follows the pt backend: TRAINING forwards run at TF32
+    ("high") matmul precision when it is true (the default), while EVAL
+    forwards follow ``DP_TF32_INFER``.  See :func:`_apply_tf32_policy`.
     """
     data = copy.deepcopy(data)
-    if bool(data.get("enable_tf32", True)) and "enable_tf32" not in _WARNED_ONCE:
-        log.warning(
-            "`enable_tf32` has no effect on the pt_expt backend, which "
-            "always runs at full ('highest') matmul precision; ignoring it."
-        )
-        _WARNED_ONCE.add("enable_tf32")
     if "spin" in data:
         if str(data["spin"].get("scheme", "deepspin")) != "native":
             raise NotImplementedError(
@@ -192,8 +228,8 @@ def get_sezm_model(data: dict) -> EnergyModel:
             atom_exclude_types=data.get("atom_exclude_types", []),
             pair_exclude_types=pair_exclude_types,
         )
-        return LinearEnergyModel(atomic_model_=composed)
-    return model
+        return _apply_tf32_policy(LinearEnergyModel(atomic_model_=composed), data)
+    return _apply_tf32_policy(model, data)
 
 
 def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
@@ -257,7 +293,10 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
             "spin scheme 'native' requires an atomic model declaring "
             "supports_native_spin()"
         )
-    return NativeSpinEnergyModel(atomic_model_=backbone_model.atomic_model, spin=spin)
+    return _apply_tf32_policy(
+        NativeSpinEnergyModel(atomic_model_=backbone_model.atomic_model, spin=spin),
+        data,
+    )
 
 
 def get_linear_model(model_params: dict) -> BaseModel:

deepmd/pt_expt/model/make_model.py

@@ -1,6 +1,10 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import contextlib
 import math
 import types
+from collections.abc import (
+    Generator,
+)
 from typing import (
     Any,
 )
@@ -435,6 +439,59 @@ def get_min_nbor_dist(self) -> float | None:
             """Get the minimum distance between two atoms."""
             return self.min_nbor_dist
 
+        # === TF32 matmul precision ===================================
+        # Mirrors the pt backend's ``SeZMModel`` policy (see
+        # ``deepmd.pt.model.model.sezm_model.SeZMModel.tf32_precision_ctx``):
+        # TRAINING forwards follow ``model.enable_tf32`` and EVAL forwards
+        # follow ``DP_TF32_INFER``.  Both attributes are set by the DPA4/SeZM
+        # builders in ``deepmd.pt_expt.model.get_model``; every other pt_expt
+        # model keeps these defaults, which select full fp32 in both modes and
+        # therefore leave its numerics untouched.
+        enable_tf32: bool = False
+        tf32_infer_precision: str = "highest"
+
+        @contextlib.contextmanager
+        def tf32_precision_ctx(self) -> Generator[None, None, None]:
+            """Select the matmul precision for one forward, then restore it.
+
+            Yields
+            ------
+            None
+                With ``torch.set_float32_matmul_precision`` set for the
+                duration of the block.
+            """
+            if not torch.cuda.is_available():
+                yield
+                return
+            prev_precision = torch.get_float32_matmul_precision()
+            try:
+                if self.training:
+                    precision = "high" if self.enable_tf32 else "highest"
+                else:
+                    precision = self.tf32_infer_precision
+                torch.set_float32_matmul_precision(precision)
+                yield
+            finally:
+                torch.set_float32_matmul_precision(prev_precision)
+
+        def call_common(self, *args: Any, **kwargs: Any) -> dict[str, torch.Tensor]:
+            """Run the shared dense/graph forward under the TF32 policy.
+
+            This is the ONE owner of matmul precision for eager forwards: every
+            pt_expt model's ``forward`` reaches the backbone through here, and
+            the export trace roots at ``call_common_lower`` instead, so the
+            precision switch never enters an exported graph.  The compiled
+            training path bypasses this method entirely and applies the same
+            policy at its own entry point (``_CompiledModel.forward``).
+
+            Returns
+            -------
+            dict[str, torch.Tensor]
+                The backbone's output dict, unchanged.
+            """
+            with self.tf32_precision_ctx():
+                return super().call_common(*args, **kwargs)
+
         def forward(self, *args: Any, **kwargs: Any) -> dict[str, torch.Tensor]:
             """Default forward delegates to call().
 

deepmd/pt_expt/train/training.py

@@ -919,7 +919,24 @@ def __getattr__(self, name: str) -> Any:
         except AttributeError:
             return getattr(self.original_model, name)
 
-    def forward(
+    def forward(self, *args: Any, **kwargs: Any) -> dict[str, torch.Tensor]:
+        """Run the compiled forward under the wrapped model's TF32 policy.
+
+        The compiled path never reaches ``call_common`` -- which owns matmul
+        precision for eager forwards -- so it applies the same policy here, at
+        its own entry point.  The context also spans the LAZY compile below:
+        Inductor selects its GEMM backend while lowering, so a precision set
+        only around the call would never reach the generated kernels.
+
+        Returns
+        -------
+        dict[str, torch.Tensor]
+            The model prediction dict.
+        """
+        with self.original_model.tf32_precision_ctx():
+            return self._forward_dispatch(*args, **kwargs)
+
+    def _forward_dispatch(
         self,
         coord: torch.Tensor,
         atype: torch.Tensor,

source/tests/pt_expt/model/test_get_model_dpa4.py

@@ -271,51 +271,113 @@ def test_default_unsupported_values_pass(self) -> None:
         self.assertIsInstance(model, EnergyModel)
 
 
-# `enable_tf32` toggles TF32 matmul precision in pt but is ignored by pt_expt
-# (always "highest" precision); a truthy value must emit a warn-once message.
-@pytest.mark.parametrize("enable_tf32", [True, False])  # truthy warns, falsy silent
-def test_enable_tf32_warns_once(enable_tf32, monkeypatch) -> None:
-    import importlib
-
-    # the package __init__ rebinds the name ``get_model`` to the function, so
-    # ``import ...get_model as`` would shadow the submodule; load it explicitly
-    gm_mod = importlib.import_module("deepmd.pt_expt.model.get_model")
-
-    # reset the warn-once set so the assertion is deterministic regardless of
-    # test ordering (other get_sezm_model calls may have already warned)
-    monkeypatch.setattr(gm_mod, "_WARNED_ONCE", set())
-
-    # Count emissions on the EMITTING logger with our own handler rather than
-    # through caplog: caplog reads a root handler, so whatever global logging
-    # state earlier tests left behind (set_log_handles flips the ``deepmd``
-    # logger's propagate off and installs its own handlers) changes how many
-    # records reach it -- zero when propagation is off, more than one when the
-    # record is seen through several attached handlers.  A handler on the
-    # emitting logger sees exactly one record per ``log.warning`` call.
-    records: list[logging.LogRecord] = []
-
-    class _Collect(logging.Handler):
-        def emit(self, record: logging.LogRecord) -> None:
-            records.append(record)
-
-    handler = _Collect(level=logging.WARNING)
-    old_level = gm_mod.log.level
-    gm_mod.log.setLevel(logging.WARNING)
-    gm_mod.log.addHandler(handler)
+# === TF32 matmul precision ==================================================
+# pt_expt mirrors the pt backend (``SeZMModel.tf32_precision_ctx``): TRAINING
+# forwards follow ``model.enable_tf32`` (default True) and EVAL forwards follow
+# ``DP_TF32_INFER``.  The knob is DPA4/SeZM-scoped, matching pt, where argcheck
+# declares it inside the dpa4 model arg block.
+
+
+@pytest.mark.parametrize(
+    "enable_tf32",
+    [
+        True,  # the argcheck default; training must select TF32 ("high")
+        False,  # opt-out; training must stay at full fp32
+    ],
+)
+def test_enable_tf32_is_stored(enable_tf32) -> None:
+    """The config knob reaches the model instead of being warned away."""
+    model = get_model(_make_raw_model_config(enable_tf32=enable_tf32))
+    assert model.enable_tf32 is enable_tf32
+
+
+def test_enable_tf32_defaults_true() -> None:
+    """An absent key follows pt's ``default=True`` (argcheck.py `enable_tf32`)."""
+    raw = _make_raw_model_config()
+    assert "enable_tf32" not in raw
+    assert get_model(raw).enable_tf32 is True
+
+
+@pytest.mark.parametrize(
+    ("env_value", "expected"),
+    [
+        (None, "highest"),  # unset -> pt's "0" default, full fp32
+        ("0", "highest"),
+        ("1", "high"),
+        ("2", "medium"),
+    ],
+)
+def test_tf32_infer_precision_from_env(env_value, expected, monkeypatch) -> None:
+    """Eval precision follows ``DP_TF32_INFER``, as in the pt backend."""
+    if env_value is None:
+        monkeypatch.delenv("DP_TF32_INFER", raising=False)
+    else:
+        monkeypatch.setenv("DP_TF32_INFER", env_value)
+    assert get_model(_make_raw_model_config()).tf32_infer_precision == expected
+
+
+def test_tf32_infer_precision_rejects_garbage(monkeypatch) -> None:
+    """An unusable ``DP_TF32_INFER`` fails fast rather than silently defaulting."""
+    monkeypatch.setenv("DP_TF32_INFER", "yes")
+    with pytest.raises(ValueError, match="DP_TF32_INFER"):
+        get_model(_make_raw_model_config())
+
+
+@pytest.mark.parametrize(
+    ("enable_tf32", "training", "expected"),
+    [
+        (True, True, "high"),  # the only combination that selects TF32
+        (False, True, "highest"),  # opt-out keeps training at full fp32
+        (True, False, "highest"),  # eval ignores enable_tf32 (uses DP_TF32_INFER)
+        (False, False, "highest"),
+    ],
+)
+def test_tf32_precision_ctx_selects_and_restores(
+    enable_tf32, training, expected, monkeypatch
+) -> None:
+    """The context selects pt's precision for the mode and restores the old one.
+
+    ``torch.set_float32_matmul_precision`` is a process global, so a forward
+    that leaked its setting would silently change every later matmul in the
+    process; the restore is as much of the contract as the selection.
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("tf32_precision_ctx is a no-op without CUDA")
+    monkeypatch.delenv("DP_TF32_INFER", raising=False)
+    model = get_model(_make_raw_model_config(enable_tf32=enable_tf32))
+    model.train(training)
+
+    torch.set_float32_matmul_precision("highest")
     try:
-        gm_mod.get_sezm_model(_make_raw_model_config(enable_tf32=enable_tf32))
-        matches = [r for r in records if "enable_tf32" in r.getMessage()]
-        if enable_tf32:
-            assert len(matches) == 1, [r.getMessage() for r in records]
-            # a second call must NOT warn again (warn-once per process)
-            records.clear()
-            gm_mod.get_sezm_model(_make_raw_model_config(enable_tf32=enable_tf32))
-            assert not [r for r in records if "enable_tf32" in r.getMessage()]
-        else:
-            assert not matches, [r.getMessage() for r in records]
+        with model.tf32_precision_ctx():
+            assert torch.get_float32_matmul_precision() == expected
+        assert torch.get_float32_matmul_precision() == "highest"
     finally:
-        gm_mod.log.removeHandler(handler)
-        gm_mod.log.setLevel(old_level)
+        torch.set_float32_matmul_precision("highest")
+
+
+def test_non_sezm_model_keeps_full_precision() -> None:
+    """The knob is DPA4/SeZM-scoped: other pt_expt models are untouched.
+
+    pt declares ``enable_tf32`` inside the dpa4 model arg block and wires it
+    only in its sezm builders, so a plain se_e2_a model must keep the class
+    defaults -- full fp32 in both train and eval.
+    """
+    model = get_model(
+        {
+            "type_map": ["O", "H"],
+            "descriptor": {
+                "type": "se_e2_a",
+                "sel": [4, 4],
+                "rcut": 4.0,
+                "rcut_smth": 3.5,
+                "seed": 1,
+            },
+            "fitting_net": {"seed": 1},
+        }
+    )
+    assert model.enable_tf32 is False
+    assert model.tf32_infer_precision == "highest"
 
 
 class TestNativeSpinErrorTranslation(unittest.TestCase):