From 5d8f7c42bb7823aff97400db562c915762de03b6 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Wed, 22 Jul 2026 12:11:51 +0800 Subject: [PATCH 1/7] perf(pt_expt): use scalable graph builders during training Resolve the carry-all graph backend once per training run so eager, compiled, and validation paths consistently use NV on supported CUDA environments with a dense fallback. Centralize graph-builder dispatch, validate explicit backend choices, and keep neighbor-capacity estimation compatible with nvalchemiops 0.3 and 0.4. Strengthen graph parity coverage by preserving edge multiplicity, exercising pair exclusions, and testing runtime and compiled backend routing. --- deepmd/pt_expt/model/make_model.py | 59 +------- deepmd/pt_expt/train/training.py | 40 ++++- deepmd/pt_expt/utils/graph_builder.py | 138 ++++++++++++++++++ deepmd/pt_expt/utils/nv_graph_builder.py | 10 +- deepmd/utils/argcheck.py | 18 +++ .../dpmodel/test_neighbor_graph_builder.py | 5 - .../model/test_graph_builder_dispatch.py | 55 +++++++ source/tests/pt_expt/test_training.py | 28 ++++ .../pt_expt/utils/test_nv_graph_builder.py | 78 ++++++---- 9 files changed, 331 insertions(+), 100 deletions(-) create mode 100644 deepmd/pt_expt/utils/graph_builder.py diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 9d906c0874..663e7fb22f 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -26,6 +26,9 @@ from deepmd.pt_expt.common import ( torch_module, ) +from deepmd.pt_expt.utils.graph_builder import ( + build_neighbor_graph_for_method, +) from deepmd.pt_expt.utils.graph_csr import ( validate_graph_csr_for_export, ) @@ -249,56 +252,6 @@ def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: return energy_redu -def _build_graph_for_method( - method: str, - coord: torch.Tensor, - atype: torch.Tensor, - box: torch.Tensor | None, - rcut: float, - pair_excl: Any, - with_csr: bool = False, -) -> Any: - """Build a carry-all ``NeighborGraph`` for the named pt_expt builder. - - Single owning site for the graph-builder dispatch shared by - :meth:`_call_common_graph` and the graph Hessian wrapper - (:class:`_WrapperForwardEnergyGraph`), so both build the graph identically. - """ - from deepmd.dpmodel.utils.neighbor_graph import ( - build_neighbor_graph, - build_neighbor_graph_ase, - ) - - if method == "dense": - return build_neighbor_graph( - coord, atype, box, rcut, with_csr=with_csr, pair_excl=pair_excl - ) - if method == "ase": - return build_neighbor_graph_ase( - coord, atype, box, rcut, with_csr=with_csr, pair_excl=pair_excl - ) - if method == "vesin": - from deepmd.pt_expt.utils.vesin_graph_builder import ( - build_neighbor_graph_vesin, - ) - - return build_neighbor_graph_vesin( - coord, atype, box, rcut, with_csr=with_csr, pair_excl=pair_excl - ) - if method == "nv": - from deepmd.pt_expt.utils.nv_graph_builder import ( - build_neighbor_graph_nv, - ) - - return build_neighbor_graph_nv( - coord, atype, box, rcut, with_csr=with_csr, pair_excl=pair_excl - ) - raise ValueError( - f"unknown neighbor_graph_method {method!r}; use 'dense', 'ase', " - "'vesin', or 'nv'" - ) - - class _WrapperForwardEnergyGraph: """Graph twin of :class:`_WrapperForwardEnergy` for the Hessian. @@ -339,7 +292,7 @@ def __init__( def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: cc = coord_flat.reshape(1, self.nloc, 3) - ng = _build_graph_for_method( + ng = build_neighbor_graph_for_method( self.method, cc, self.atype, self.box, self.rcut, self.pair_excl ) atomic_ret = self.model.atomic_model.forward_common_atomic_graph( @@ -719,7 +672,7 @@ def _resolve_graph_method( if "energy" not in self.atomic_output_def().keys(): return None if self.mixed_types() and self.atomic_model.uses_graph_lower(): - return "dense" + return getattr(self, "neighbor_graph_method", "dense") return None def _call_common_graph( @@ -795,7 +748,7 @@ def _call_common_graph( and bool(getattr(_desc, "geo_compress", False)) ) pair_excl = getattr(self.atomic_model, "pair_excl", None) - ng = _build_graph_for_method( + ng = build_neighbor_graph_for_method( method, cc, atype, bb, rcut, pair_excl, with_csr=with_csr ) nf, nloc = atype.shape[:2] diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 697bf46dfb..9c7a1b8883 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -1146,8 +1146,8 @@ def _forward_graph( so no extended->local scatter is needed; only the flat ``(N, *)`` node keys are unravelled to ``(nf, nloc, *)`` at the I/O boundary. """ - from deepmd.dpmodel.utils.neighbor_graph import ( - build_neighbor_graph, + from deepmd.pt_expt.utils.graph_builder import ( + build_neighbor_graph_for_method, ) _model = self.original_model @@ -1198,7 +1198,14 @@ def _forward_graph( # into edge_mask here so the compiled lower consumes a pre-excluded graph # (the lower no longer re-applies it), matching the eager path exactly. pair_excl = getattr(_model.atomic_model, "pair_excl", None) - ng = build_neighbor_graph(coord_3d, atype, box_flat, rcut, pair_excl=pair_excl) + ng = build_neighbor_graph_for_method( + getattr(_model, "neighbor_graph_method", "dense"), + coord_3d, + atype, + box_flat, + rcut, + pair_excl, + ) atype_flat = atype.reshape(nframes * nloc) # Lazy compile of the GRAPH lower (cached per structure key). @@ -1793,6 +1800,10 @@ def _make_sample( last_epoch=self.start_step - 1, ) + self._configure_neighbor_graph_method( + training_params.get("neighbor_graph_method", "auto") + ) + # torch.compile ------------------------------------------------------- self.enable_compile = training_params.get("enable_compile", False) if self.enable_compile: @@ -1891,6 +1902,29 @@ def _raise_if_full_validation_unsupported( # torch.compile helpers # ------------------------------------------------------------------ + def _configure_neighbor_graph_method(self, requested: str) -> None: + """Resolve and install the training graph builder on eligible models.""" + graph_models = [ + self.models[model_key] + for model_key in self.model_keys + if model_uses_graph_lower(self.models[model_key]) + ] + if not graph_models: + if requested != "auto": + raise ValueError( + "training.neighbor_graph_method applies only to " + "graph-eligible energy models" + ) + return + + from deepmd.pt_expt.utils.graph_builder import ( + resolve_neighbor_graph_method, + ) + + resolved = resolve_neighbor_graph_method(requested, DEVICE) + for model in graph_models: + model.neighbor_graph_method = resolved + def _compile_model(self, compile_opts: dict[str, Any]) -> None: """Replace ``self.model`` with a compiled version. diff --git a/deepmd/pt_expt/utils/graph_builder.py b/deepmd/pt_expt/utils/graph_builder.py new file mode 100644 index 0000000000..1063359bce --- /dev/null +++ b/deepmd/pt_expt/utils/graph_builder.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Runtime selection and dispatch for pt_expt carry-all graph builders.""" + +from typing import ( + TYPE_CHECKING, +) + +import torch + +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) + + +def resolve_neighbor_graph_method( + requested: str, + device: torch.device, +) -> str: + """Resolve a training graph-builder policy to one concrete backend. + + Parameters + ---------- + requested + ``"auto"``, ``"dense"``, or ``"nv"``. + device + Device used by the model on the current rank. + + Returns + ------- + str + The concrete builder name, either ``"dense"`` or ``"nv"``. + + Raises + ------ + ValueError + If the requested method is unknown or NV is requested on a non-CUDA + device. + ImportError + If NV is requested explicitly but nvalchemiops is unavailable. + """ + if requested not in {"auto", "dense", "nv"}: + raise ValueError( + f"unknown training neighbor_graph_method {requested!r}; " + "use 'auto', 'dense', or 'nv'" + ) + if requested == "dense": + return "dense" + + from deepmd.pt.utils.nv_nlist import ( + is_nv_available, + ) + + if requested == "auto": + return "nv" if device.type == "cuda" and is_nv_available() else "dense" + if device.type != "cuda": + raise ValueError( + "neighbor_graph_method='nv' requires a CUDA training device, " + f"got {device!s}" + ) + if not is_nv_available(): + raise ImportError( + "neighbor_graph_method='nv' requires nvalchemi-toolkit-ops. " + "Install the DeePMD-kit 'nvalchemi' extra or use 'auto'/'dense'." + ) + return "nv" + + +def build_neighbor_graph_for_method( + method: str, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None, + rcut: float, + pair_excl: "PairExcludeMask | None", + *, + with_csr: bool = False, +) -> "NeighborGraph": + """Build a carry-all graph with one concrete pt_expt backend.""" + if method == "dense": + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + ) + + return build_neighbor_graph( + coord, + atype, + box, + rcut, + with_csr=with_csr, + pair_excl=pair_excl, + ) + if method == "ase": + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph_ase, + ) + + return build_neighbor_graph_ase( + coord, + atype, + box, + rcut, + with_csr=with_csr, + pair_excl=pair_excl, + ) + if method == "vesin": + from deepmd.pt_expt.utils.vesin_graph_builder import ( + build_neighbor_graph_vesin, + ) + + return build_neighbor_graph_vesin( + coord, + atype, + box, + rcut, + with_csr=with_csr, + pair_excl=pair_excl, + ) + if method == "nv": + from deepmd.pt_expt.utils.nv_graph_builder import ( + build_neighbor_graph_nv, + ) + + return build_neighbor_graph_nv( + coord, + atype, + box, + rcut, + with_csr=with_csr, + pair_excl=pair_excl, + ) + raise ValueError( + f"unknown neighbor_graph_method {method!r}; use 'dense', 'ase', " + "'vesin', or 'nv'" + ) diff --git a/deepmd/pt_expt/utils/nv_graph_builder.py b/deepmd/pt_expt/utils/nv_graph_builder.py index 637dde9912..41124c8519 100644 --- a/deepmd/pt_expt/utils/nv_graph_builder.py +++ b/deepmd/pt_expt/utils/nv_graph_builder.py @@ -257,12 +257,6 @@ def build_neighbor_graph_nv( ------ ImportError if ``nvalchemi-toolkit-ops`` (CUDA) is not installed. - - Notes - ----- - The ``pair_excl`` path of this builder has no local oracle set-equality test - because nvalchemiops requires CUDA; the set-equality contract must be - validated on a GPU box (same pattern as :class:`~deepmd.dpmodel.utils.neighbor_graph.build_neighbor_graph_ase`). """ if not is_nv_available(): raise ImportError( @@ -301,9 +295,11 @@ def build_neighbor_graph_nv( # vesin handles unwrapped positions natively), nvalchemiops requires # in-cell positions, so BOTH the search and the edge_vec recomputation use # the normalized coords; S then matches the coords the search actually saw. + # The 0.25 density preserves a 25% margin over the estimator's 0.2 + # baseline without using the safety_factor argument deprecated in Ops 0.4. initial_capacity = max( 64, - estimate_max_neighbors(float(rcut), safety_factor=1.25), + estimate_max_neighbors(float(rcut), atomic_density=0.25), ) coord, cell, neighbor_matrix, num_neighbors, shifts = nv_search_matrix( coord, box, rcut, start_capacity=initial_capacity diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 2b0d5b9132..e6ab7b09b7 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5337,6 +5337,15 @@ def training_args( "Default is 0. Requires distributed launch via torchrun. " "Currently supports single-task training; does not support LKF or change_bias_after_training." ) + doc_neighbor_graph_method = ( + "Select the carry-all neighbor-graph builder for graph-eligible PyTorch " + "Experimental energy models. `auto` uses the NV builder on CUDA when " + "nvalchemiops is available and otherwise uses the in-tree dense builder. " + "`nv` requires CUDA and nvalchemiops; `dense` always uses the in-tree " + "all-pairs implementation. The selection is resolved once at training " + "startup and applies consistently to eager, compiled, and full-validation " + "forwards." + ) arg_training_data = training_data_args() arg_validation_data = validation_data_args() @@ -5518,6 +5527,15 @@ def training_args( default=0, doc=doc_only_pt_supported + doc_zero_stage, ), + Argument( + "neighbor_graph_method", + str, + optional=True, + default="auto", + extra_check=lambda x: x in {"auto", "dense", "nv"}, + extra_check_errmsg="must be one of 'auto', 'dense', or 'nv'", + doc=doc_only_pt_expt_supported + doc_neighbor_graph_method, + ), Argument( "enable_compile", bool, diff --git a/source/tests/common/dpmodel/test_neighbor_graph_builder.py b/source/tests/common/dpmodel/test_neighbor_graph_builder.py index 4f45d22505..228cd4fc6b 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph_builder.py +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -526,10 +526,5 @@ def test_ase_oracle_set_equality(self) -> None: self.assertLess(int(ng_ase.edge_mask.sum()), int(ng_ase_plain.edge_mask.sum())) -# NOTE: nvalchemiops builder has no local oracle set-equality test for pair_excl -# because it requires CUDA; validation is deferred to GPU box tests (PR-C/nv-gtest). -# See deepmd.pt_expt.utils.nv_graph_builder.build_neighbor_graph_nv docstring. - - if __name__ == "__main__": unittest.main() diff --git a/source/tests/pt_expt/model/test_graph_builder_dispatch.py b/source/tests/pt_expt/model/test_graph_builder_dispatch.py index 80d26a84fc..c0e1c3dcdf 100644 --- a/source/tests/pt_expt/model/test_graph_builder_dispatch.py +++ b/source/tests/pt_expt/model/test_graph_builder_dispatch.py @@ -3,6 +3,10 @@ perf-only equivalent of 'dense' (same energy + force); dpmodel/jax fail-fast. """ +from unittest.mock import ( + patch, +) + import numpy as np import pytest import torch @@ -22,6 +26,9 @@ from deepmd.pt_expt.model.ener_model import ( EnergyModel, ) +from deepmd.pt_expt.utils.graph_builder import ( + resolve_neighbor_graph_method, +) from deepmd.pt_expt.utils.vesin_neighbor_list import ( is_vesin_torch_available, ) @@ -73,6 +80,54 @@ def _eval(model, method): return ret["energy_redu"], ret["energy_derv_r"] +def test_runtime_default_graph_method_is_configurable(): + model = _make_model() + model.neighbor_graph_method = "nv" + assert model._resolve_graph_method(None) == "nv" + assert model._resolve_graph_method("dense") == "dense" + + +def test_compiled_model_uses_runtime_graph_method(): + from deepmd.pt_expt.train.training import ( + _CompiledModel, + _get_model_structure_key, + ) + + class BuilderCalled(RuntimeError): + pass + + model = _make_model().train() + model.neighbor_graph_method = "nv" + compiled = _CompiledModel(model, _get_model_structure_key(model)) + coord = torch.zeros((1, 2, 3), dtype=torch.float64, device=env.DEVICE) + atype = torch.zeros((1, 2), dtype=torch.int64, device=env.DEVICE) + box = torch.eye(3, dtype=torch.float64, device=env.DEVICE).reshape(1, 3, 3) * 8.0 + with ( + patch( + "deepmd.pt_expt.utils.graph_builder.build_neighbor_graph_for_method", + side_effect=BuilderCalled, + ) as builder, + patch( + "deepmd.pt_expt.train.training._trace_and_compile_graph", + side_effect=AssertionError("compiled lower reached before graph builder"), + ), + pytest.raises(BuilderCalled), + ): + compiled(coord, atype, box) + assert builder.call_args.args[0] == "nv" + + +def test_auto_graph_method_uses_nv_only_on_cuda(): + with patch("deepmd.pt.utils.nv_nlist.is_nv_available", return_value=True): + assert resolve_neighbor_graph_method("auto", torch.device("cuda")) == "nv" + assert resolve_neighbor_graph_method("auto", torch.device("cpu")) == "dense" + + +def test_explicit_nv_rejects_cpu(): + with pytest.raises(ValueError, match="requires a CUDA"): + resolve_neighbor_graph_method("nv", torch.device("cpu")) + + @pytest.mark.skipif(not is_vesin_torch_available(), reason="vesin[torch] not installed") def test_vesin_matches_dense_energy_force(): torch.manual_seed(0) diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index 193fa81911..c9d2c70ccc 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -279,6 +279,34 @@ def test_get_model(self) -> None: nparams = sum(p.numel() for p in model.parameters()) self.assertGreater(nparams, 0) + def test_neighbor_graph_method_defaults_to_auto(self) -> None: + """Training selects the graph builder automatically unless overridden.""" + config = _make_config(self.data_dir) + config = update_deepmd_input(config, warning=False) + config = normalize(config) + self.assertEqual(config["training"]["neighbor_graph_method"], "auto") + + def test_trainer_installs_resolved_graph_method(self) -> None: + """The trainer installs the concrete graph backend on graph models.""" + config = _make_config(self.data_dir) + config["model"]["descriptor"] = copy.deepcopy(_DESCRIPTOR_DPA1_NO_ATTN) + config = update_deepmd_input(config, warning=False) + config = normalize(config) + with patch( + "deepmd.pt.utils.nv_nlist.is_nv_available", + return_value=False, + ): + trainer = get_trainer(config) + self.assertEqual(trainer.model.neighbor_graph_method, "dense") + + def test_explicit_graph_method_rejects_ineligible_model(self) -> None: + config = _make_config(self.data_dir) + config["training"]["neighbor_graph_method"] = "nv" + config = update_deepmd_input(config, warning=False) + config = normalize(config) + with self.assertRaisesRegex(ValueError, "graph-eligible"): + get_trainer(config) + def _run_training(self, config: dict) -> None: """Run training and verify lcurve + checkpoint creation.""" tmpdir = tempfile.mkdtemp(prefix="pt_expt_train_") diff --git a/source/tests/pt_expt/utils/test_nv_graph_builder.py b/source/tests/pt_expt/utils/test_nv_graph_builder.py index 3da0da5931..8ca4f2db1b 100644 --- a/source/tests/pt_expt/utils/test_nv_graph_builder.py +++ b/source/tests/pt_expt/utils/test_nv_graph_builder.py @@ -7,6 +7,9 @@ import pytest import torch +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) from deepmd.dpmodel.utils.neighbor_graph import ( build_neighbor_graph, ) @@ -22,16 +25,19 @@ ) -def _sets(ng, nloc): - """Per-center set of (src_local, rounded edge_vec) over real edges.""" - ei = np.asarray(ng.edge_index.cpu()) - ev = np.asarray(ng.edge_vec.detach().cpu()) - em = np.asarray(ng.edge_mask.cpu()) - out = {c: set() for c in range(nloc)} - for e in range(ei.shape[1]): - if em[e]: - out[int(ei[1, e])].add((int(ei[0, e]), tuple(np.round(ev[e], 6)))) - return out +def _edge_multiset(ng): + edge_index = np.asarray(ng.edge_index.cpu()) + edge_vec = np.asarray(ng.edge_vec.detach().cpu()) + edge_mask = np.asarray(ng.edge_mask.cpu()) + return sorted( + ( + int(edge_index[0, edge]), + int(edge_index[1, edge]), + tuple(np.round(edge_vec[edge], 10)), + ) + for edge in range(edge_index.shape[1]) + if edge_mask[edge] + ) @pytest.mark.parametrize("periodic", [False, True]) # non-PBC and PBC @@ -50,7 +56,7 @@ def test_nv_matches_intree_carry_all(periodic): atype = torch.tensor([[0, 1, 0, 1]], dtype=torch.int64, device=dev) ng_ref = build_neighbor_graph(coord, atype, box, 2.0) ng = nv_builder.build_neighbor_graph_nv(coord, atype, box, 2.0) - assert _sets(ng, 4) == _sets(ng_ref, 4) + assert _edge_multiset(ng) == _edge_multiset(ng_ref) def test_nv_batches_frames_without_python_loop(): @@ -70,26 +76,7 @@ def test_nv_batches_frames_without_python_loop(): ) ng_ref = build_neighbor_graph(coord, atype, box, 2.0) ng = nv_builder.build_neighbor_graph_nv(coord, atype, box, 2.0) - # per-frame node offset: frame f centers occupy nodes [f*5, (f+1)*5) - for f in range(3): - s_ref = { - ( - int(ng_ref.edge_index[0, e]), - tuple(np.round(np.asarray(ng_ref.edge_vec[e].detach().cpu()), 6)), - ) - for e in range(ng_ref.edge_index.shape[1]) - if bool(ng_ref.edge_mask[e]) - and f * 5 <= int(ng_ref.edge_index[1, e]) < (f + 1) * 5 - } - s = { - ( - int(ng.edge_index[0, e]), - tuple(np.round(np.asarray(ng.edge_vec[e].detach().cpu()), 6)), - ) - for e in range(ng.edge_index.shape[1]) - if bool(ng.edge_mask[e]) and f * 5 <= int(ng.edge_index[1, e]) < (f + 1) * 5 - } - assert s == s_ref, f"frame {f} neighbor set mismatch" + assert _edge_multiset(ng) == _edge_multiset(ng_ref) def test_nv_edge_vec_is_differentiable(): @@ -118,7 +105,34 @@ def test_nv_excludes_virtual_atoms_like_dense(): atype = torch.tensor([[0, -1, 0, 1]], dtype=torch.int64, device=dev) # 1 virtual ng_ref = build_neighbor_graph(coord, atype, box, 2.0) ng = nv_builder.build_neighbor_graph_nv(coord, atype, box, 2.0) - assert _sets(ng, 4) == _sets(ng_ref, 4) + assert _edge_multiset(ng) == _edge_multiset(ng_ref) ei = np.asarray(ng.edge_index.cpu())[:, np.asarray(ng.edge_mask.cpu())] at = atype.reshape(-1).cpu().numpy() assert np.all(at[ei[0]] >= 0) and np.all(at[ei[1]] >= 0) + + +def test_nv_pair_exclusion_matches_dense(): + dev = torch.device("cuda") + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [0.8, 0.0, 0.0], [0.0, 1.1, 0.0], [1.2, 1.0, 0.0]]], + dtype=torch.float64, + device=dev, + ) + atype = torch.tensor([[0, 1, 0, 1]], dtype=torch.int64, device=dev) + box = (torch.eye(3, dtype=torch.float64, device=dev) * 3.0).reshape(1, 3, 3) + pair_excl = PairExcludeMask(2, [(0, 1), (1, 0)]) + expected = build_neighbor_graph( + coord, + atype, + box, + 2.0, + pair_excl=pair_excl, + ) + actual = nv_builder.build_neighbor_graph_nv( + coord, + atype, + box, + 2.0, + pair_excl=pair_excl, + ) + assert _edge_multiset(actual) == _edge_multiset(expected) From 6157e814e9f415aea673fcfaac477ef9361db576 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Wed, 22 Jul 2026 16:40:48 +0800 Subject: [PATCH 2/7] feat(pt_expt): add HybridMuon and shared compile safeguards Consume the normalized optimizer configuration for Adam, AdamW, and HybridMuon, including runtime parameter names required for correct HybridMuon routing. Apply the shared compiler compatibility patches and Inductor training options so compiled PT-expt execution follows the same runtime contract as the PT backend. --- deepmd/pt_expt/train/training.py | 59 +++++++++++++++++++++++--------- 1 file changed, 43 insertions(+), 16 deletions(-) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 9c7a1b8883..046dc926f2 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -42,6 +42,9 @@ from deepmd.dpmodel.utils.learning_rate import ( make_learning_rate_schedule, ) +from deepmd.pt.optimizer import ( + HybridMuonOptimizer, +) from deepmd.pt.train.utils import ( resolve_best_checkpoint_dir, ) @@ -49,6 +52,10 @@ FullValidator, resolve_full_validation_start_step, ) +from deepmd.pt.utils.compile_compat import ( + apply_global_compile_patches, + build_inductor_compile_options, +) from deepmd.pt.utils.compile_compat import next_safe_prime as _next_safe_prime from deepmd.pt.utils.compile_compat import rebuild_graph_module as _rebuild_graph_module from deepmd.pt.utils.compile_compat import ( @@ -93,6 +100,10 @@ log = logging.getLogger(__name__) +# Apply the shared process-global compiler workarounds before any pt_expt +# training graph reaches Dynamo or Inductor. +apply_global_compile_patches() + # Buffer names in atomic_model that are per-task (energy/output statistics). # These live one level above the fitting net and are not reached by # fitting-net share_params. They are always promoted to FX placeholders @@ -572,19 +583,8 @@ def _finalize_compiled_lower( if not was_training: model.eval() - # Inductor defaults tuned for second-order-gradient training graphs. - # User-supplied compile_opts override these on a per-key basis. - inductor_options: dict[str, Any] = { - "max_autotune": False, - "shape_padding": True, - "epilogue_fusion": False, - "triton.cudagraphs": False, - "max_fusion_size": 8, - # NOTE: On GPU with PyTorch <=2.11, consider adding - # "triton.mix_order_reduction": False to work around - # pytorch/pytorch#174379, #178080, #179494 under - # data-dependent symbolic shapes. - } + # Keep pt_expt training on the same compiler contract as the PT SeZM path. + inductor_options = build_inductor_compile_options(inference=False) if extra_options: inductor_options.update(extra_options) if compile_opts: @@ -1359,6 +1359,7 @@ def __init__( model_params = config["model"] training_params = config["training"] + optimizer_params = config.get("optimizer", {}) validating_params = config.get("validating", {}) or {} # Task normalization -------------------------------------------------- @@ -1583,21 +1584,47 @@ def _make_sample( ) # Optimiser ----------------------------------------------------------- - opt_type = training_params.get("opt_type", "Adam") + opt_type = optimizer_params.get("type", "Adam") # LambdaLR multiplies each param group's initial learning rate by the # lambda value. Warmup schedules legitimately return zero at step 0, # so use the nonzero schedule base as the denominator and let the # lambda initialize the optimizer to the requested warmup value. initial_lr = float(self.lr_schedule.start_lr) + adam_betas = ( + float(optimizer_params["adam_beta1"]), + float(optimizer_params["adam_beta2"]), + ) + weight_decay = float(optimizer_params["weight_decay"]) if opt_type == "Adam": - self.optimizer = torch.optim.Adam(self.wrapper.parameters(), lr=initial_lr) + self.optimizer = torch.optim.Adam( + self.wrapper.parameters(), + lr=initial_lr, + betas=adam_betas, + weight_decay=weight_decay, + ) elif opt_type == "AdamW": - weight_decay = training_params.get("weight_decay", 0.001) self.optimizer = torch.optim.AdamW( self.wrapper.parameters(), lr=initial_lr, + betas=adam_betas, + weight_decay=weight_decay, + ) + elif opt_type == "HybridMuon": + runtime_named_parameters = tuple(self.wrapper.named_parameters()) + self.optimizer = HybridMuonOptimizer( + self.wrapper.parameters(), + lr=initial_lr, + momentum=float(optimizer_params["momentum"]), weight_decay=weight_decay, + adam_betas=adam_betas, + lr_adjust=float(optimizer_params["lr_adjust"]), + lr_adjust_coeff=float(optimizer_params["lr_adjust_coeff"]), + muon_mode=str(optimizer_params["muon_mode"]), + named_parameters=runtime_named_parameters, + enable_gram=bool(optimizer_params["enable_gram"]), + flash_muon=bool(optimizer_params["flash_muon"]), + magma_muon=bool(optimizer_params["magma_muon"]), ) else: raise ValueError(f"Unsupported optimizer type: {opt_type}") From e64f01d37bc442ffd29ca2b237a47680ff5afac9 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 28 Jul 2026 14:26:42 +0800 Subject: [PATCH 3/7] fix(pt): recognize pt_expt bias names in HybridMuon --- deepmd/pt/optimizer/hybrid_muon.py | 28 ++++++++++++++-------------- source/tests/pt/test_hybrid_muon.py | 8 ++++++++ 2 files changed, 22 insertions(+), 14 deletions(-) diff --git a/deepmd/pt/optimizer/hybrid_muon.py b/deepmd/pt/optimizer/hybrid_muon.py index 24caaadc41..29190bda67 100644 --- a/deepmd/pt/optimizer/hybrid_muon.py +++ b/deepmd/pt/optimizer/hybrid_muon.py @@ -6,8 +6,8 @@ Routing is controlled by parameter dimensionality, parameter names, and ``muon_mode``: -- Parameters whose final effective name segment contains ``bias`` - (case-insensitive), or starts with ``adam_`` (case-insensitive): Adam. +- Parameters whose final effective name segment is ``b``, contains ``bias``, + or starts with ``adam_`` (case-insensitive): Adam. - Parameters whose final effective name segment starts with ``adamw_`` (case-insensitive): Adam with decoupled weight decay (AdamW-style). The final effective segment means the last non-numeric segment in the full @@ -715,7 +715,7 @@ def get_adam_route( effective name segment after stripping trailing numeric ParameterList indices): - 1. Contains ``"bias"`` -> ``"adam"`` (no weight decay). + 1. Is ``"b"`` or contains ``"bias"`` -> ``"adam"`` (no weight decay). 2. Starts with ``"adam_"`` -> ``"adam"`` (no weight decay). Typical: norm scales, radial frequencies. 3. Starts with ``"adamw_"`` -> ``"adamw"`` (decoupled weight decay). @@ -730,7 +730,7 @@ def get_adam_route( while leaf_name_idx > 0 and name_segments[leaf_name_idx].isdigit(): leaf_name_idx -= 1 leaf_name = name_segments[leaf_name_idx] - if "bias" in leaf_name: + if leaf_name == "b" or "bias" in leaf_name: return "adam" if leaf_name.startswith("adam_"): return "adam" @@ -808,9 +808,9 @@ class HybridMuonOptimizer(Optimizer): This optimizer applies different update rules based on parameter dimensionality, parameter names, and ``muon_mode``: - - Parameters with final effective name segment containing ``bias`` - (case-insensitive), or starting with ``adam_`` (case-insensitive): - standard Adam update. + - Parameters with final effective name segment equal to ``b``, containing + ``bias``, or starting with ``adam_`` (case-insensitive): standard Adam + update. - Parameters with final effective name segment starting with ``adamw_`` (case-insensitive): Adam with decoupled weight decay (AdamW-style). - 1D parameters: standard Adam update. @@ -826,8 +826,8 @@ class HybridMuonOptimizer(Optimizer): ``(m, n)`` slice. Naming convention for explicit Adam routing: - - Parameters representing bias terms should include ``bias`` in their - final effective name segment (case-insensitive). + - Parameters representing bias terms should use ``b`` or include ``bias`` + in their final effective name segment (case-insensitive). - Parameters that are not semantic bias but should still use Adam should use an ``adam_`` prefix in their final effective name segment (case-insensitive). @@ -887,10 +887,10 @@ class HybridMuonOptimizer(Optimizer): - ``"slice"``: >=3D parameters use per-slice Muon routing on last two dims. named_parameters : iterable[tuple[str, torch.Tensor]] | None Optional named parameter iterable used for name-based routing. - Parameters with final effective name segment containing ``bias`` - (case-insensitive), or starting with ``adam_`` (case-insensitive), - are forced to Adam (no weight decay). Parameters starting with - ``adamw_`` are forced to AdamW-style decoupled decay path. + Parameters with final effective name segment equal to ``b``, containing + ``bias``, or starting with ``adam_`` (case-insensitive) are forced to + Adam (no weight decay). Parameters starting with ``adamw_`` are forced + to AdamW-style decoupled decay path. enable_gram : bool Enable the compiled Gram Newton-Schulz path for rectangular Muon matrices. Square matrices continue to use the current standard @@ -1502,7 +1502,7 @@ def _build_param_routing(self) -> None: Classify parameters into Muon, Adam, and AdamW routes (static routing). Routing logic: - - name-based ``adam_`` prefix or contains ``bias`` → Adam (no decay) + - name-based ``b``, ``bias``, or ``adam_`` route → Adam (no decay) - name-based ``adamw_`` prefix → AdamW (decoupled weight decay) - effective shape rank <2 → Adam (no decay) - non-matrix effective shape for current muon_mode → AdamW (decoupled) diff --git a/source/tests/pt/test_hybrid_muon.py b/source/tests/pt/test_hybrid_muon.py index 3f2bc4bef6..5547ce6a2c 100644 --- a/source/tests/pt/test_hybrid_muon.py +++ b/source/tests/pt/test_hybrid_muon.py @@ -249,6 +249,10 @@ def __init__(self, device: torch.device) -> None: self.gateBiAsScale = torch.nn.Parameter( torch.ones(2, 6, dtype=torch.float32, device=device) ) + # pt_expt uses the exact leaf name "b" for affine biases. + self.b = torch.nn.Parameter( + torch.ones(2, 6, dtype=torch.float32, device=device) + ) # Module name contains "bias", but parameter leaf is "weight". self.bias_proj = torch.nn.Linear(4, 6, bias=False, device=device) @@ -258,6 +262,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: y = y * self.adam_stack[0].unsqueeze(0) y = y * self.adamw_layer_scale.unsqueeze(0) y = y * self.gateBiAsScale.unsqueeze(0) + y = y + self.b.unsqueeze(0) y = y + self.bias_proj(x).unsqueeze(1) return y.sum() @@ -287,6 +292,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # Contains "bias" (case-insensitive) → Adam self.assertIn("exp_avg", optimizer.state[model.gateBiAsScale]) self.assertNotIn("momentum_buffer", optimizer.state[model.gateBiAsScale]) + # Exact pt_expt leaf "b" → Adam even when matrix-shaped + self.assertIn("exp_avg", optimizer.state[model.b]) + self.assertNotIn("momentum_buffer", optimizer.state[model.b]) # Module name "bias_proj" but leaf is "weight" → Muon self.assertIn("momentum_buffer", optimizer.state[model.bias_proj.weight]) self.assertNotIn("exp_avg", optimizer.state[model.bias_proj.weight]) From 6cd460e69d49ca36029ea139bff633d7a13d56bc Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 28 Jul 2026 22:06:50 +0800 Subject: [PATCH 4/7] fix(pt_expt): enforce torch.compile version gate --- deepmd/pt_expt/train/training.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 046dc926f2..97bb855a7a 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -55,6 +55,7 @@ from deepmd.pt.utils.compile_compat import ( apply_global_compile_patches, build_inductor_compile_options, + check_compile_torch_version, ) from deepmd.pt.utils.compile_compat import next_safe_prime as _next_safe_prime from deepmd.pt.utils.compile_compat import rebuild_graph_module as _rebuild_graph_module @@ -1834,6 +1835,7 @@ def _make_sample( # torch.compile ------------------------------------------------------- self.enable_compile = training_params.get("enable_compile", False) if self.enable_compile: + check_compile_torch_version() compile_opts = training_params.get("compile_options", {}) log.info("Compiling model with torch.compile (%s)", compile_opts) self._compile_model(compile_opts) From 534a52e758ba022e051320d78a641efae4ed72b3 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Thu, 30 Jul 2026 00:33:06 +0800 Subject: [PATCH 5/7] fix(pt_expt): harden training runtime setup Validate optimizer support before reading variant-specific settings, prefer the NV graph builder on CUDA with an actionable fallback warning, and apply compiler workarounds only at explicit compile boundaries. Keep compile-only tests aligned with the production PyTorch version gate and strengthen optimizer routing coverage. --- deepmd/pt/utils/compile_compat.py | 10 ++--- deepmd/pt_expt/train/training.py | 25 +++++------- deepmd/pt_expt/utils/graph_builder.py | 15 ++++++- source/tests/pt/test_hybrid_muon.py | 12 ++++++ source/tests/pt_expt/compile_utils.py | 20 ++++++++++ .../model/test_graph_builder_dispatch.py | 10 +++++ source/tests/pt_expt/test_multitask.py | 6 +++ source/tests/pt_expt/test_training.py | 40 +++++++++++++++++++ source/tests/pt_expt/test_training_ddp.py | 6 +++ 9 files changed, 123 insertions(+), 21 deletions(-) create mode 100644 source/tests/pt_expt/compile_utils.py diff --git a/deepmd/pt/utils/compile_compat.py b/deepmd/pt/utils/compile_compat.py index a3cc33982f..b3209a8368 100644 --- a/deepmd/pt/utils/compile_compat.py +++ b/deepmd/pt/utils/compile_compat.py @@ -57,11 +57,11 @@ def apply_global_compile_patches() -> None: """Apply every process-global PyTorch adjustment the compile path needs. - The adjustments are mutually independent and individually idempotent. The - function is intended to run exactly once, when the model module is - imported, so that the global state is established before the first - compilation. The symbolic-divisibility repair is applied only on PyTorch - 2.12, where the regression exists. + The adjustments are mutually independent and individually idempotent. Call + this function at an explicit compile boundary before Dynamo or Inductor is + invoked; repeated calls from independent compile paths are safe. The + symbolic-divisibility repair is applied only on PyTorch 2.12, where the + regression exists. """ # Silence Inductor / Triton autotune console dumps. ``torch.compile`` # reads these environment variables once, when its backend is first diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 97bb855a7a..e8f6f00ad2 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -101,10 +101,6 @@ log = logging.getLogger(__name__) -# Apply the shared process-global compiler workarounds before any pt_expt -# training graph reaches Dynamo or Inductor. -apply_global_compile_patches() - # Buffer names in atomic_model that are per-task (energy/output statistics). # These live one level above the fitting net and are not reached by # fitting-net share_params. They are always promoted to FX placeholders @@ -584,6 +580,12 @@ def _finalize_compiled_lower( if not was_training: model.eval() + # This is the common boundary immediately before every pt_expt + # ``torch.compile`` call. Applying the idempotent process-global patches + # here leaves eager-only imports untouched while still preceding all + # Dynamo and Inductor configuration reads. + apply_global_compile_patches() + # Keep pt_expt training on the same compiler contract as the PT SeZM path. inductor_options = build_inductor_compile_options(inference=False) if extra_options: @@ -1586,6 +1588,9 @@ def _make_sample( # Optimiser ----------------------------------------------------------- opt_type = optimizer_params.get("type", "Adam") + if opt_type not in {"Adam", "AdamW", "HybridMuon"}: + raise ValueError(f"Unsupported optimizer type: {opt_type}") + # LambdaLR multiplies each param group's initial learning rate by the # lambda value. Warmup schedules legitimately return zero at step 0, # so use the nonzero schedule base as the denominator and let the @@ -1611,7 +1616,7 @@ def _make_sample( betas=adam_betas, weight_decay=weight_decay, ) - elif opt_type == "HybridMuon": + else: # HybridMuon runtime_named_parameters = tuple(self.wrapper.named_parameters()) self.optimizer = HybridMuonOptimizer( self.wrapper.parameters(), @@ -1627,8 +1632,6 @@ def _make_sample( flash_muon=bool(optimizer_params["flash_muon"]), magma_muon=bool(optimizer_params["magma_muon"]), ) - else: - raise ValueError(f"Unsupported optimizer type: {opt_type}") for param_group in self.optimizer.param_groups: param_group["initial_lr"] = initial_lr @@ -1969,14 +1972,6 @@ def _compile_model(self, compile_opts: dict[str, Any]) -> None: needed. The coord extension + nlist build (data-dependent control flow) are kept outside the compiled region. """ - # Disable DDPOptimizer: our compile region wraps only the inner - # compute function, not the whole DDP model. DDPOptimizer assumes - # it owns the full model graph and splits at bucket boundaries, - # producing subgraphs whose outputs include symbolic integers. - # AOT Autograd then crashes with ``'int' object has no attribute - # 'meta'`` (pytorch/pytorch#134182). - torch._dynamo.config.optimize_ddp = False - # Under DDP, self.wrapper is a DistributedDataParallel wrapper; # access the underlying ModelWrapper via .module. wrapper_mod = ( diff --git a/deepmd/pt_expt/utils/graph_builder.py b/deepmd/pt_expt/utils/graph_builder.py index 1063359bce..b074a3af8f 100644 --- a/deepmd/pt_expt/utils/graph_builder.py +++ b/deepmd/pt_expt/utils/graph_builder.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Runtime selection and dispatch for pt_expt carry-all graph builders.""" +import logging from typing import ( TYPE_CHECKING, ) @@ -15,6 +16,8 @@ NeighborGraph, ) +log = logging.getLogger(__name__) + def resolve_neighbor_graph_method( requested: str, @@ -55,7 +58,17 @@ def resolve_neighbor_graph_method( ) if requested == "auto": - return "nv" if device.type == "cuda" and is_nv_available() else "dense" + if device.type != "cuda": + return "dense" + if is_nv_available(): + return "nv" + log.warning( + "nvalchemi-toolkit-ops is unavailable; falling back from " + "neighbor_graph_method='auto' to the dense graph builder. " + "Install it with `pip install nvalchemi-toolkit-ops` to enable " + "the NV graph builder." + ) + return "dense" if device.type != "cuda": raise ValueError( "neighbor_graph_method='nv' requires a CUDA training device, " diff --git a/source/tests/pt/test_hybrid_muon.py b/source/tests/pt/test_hybrid_muon.py index 5547ce6a2c..aab297b0b0 100644 --- a/source/tests/pt/test_hybrid_muon.py +++ b/source/tests/pt/test_hybrid_muon.py @@ -295,6 +295,18 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # Exact pt_expt leaf "b" → Adam even when matrix-shaped self.assertIn("exp_avg", optimizer.state[model.b]) self.assertNotIn("momentum_buffer", optimizer.state[model.b]) + self.assertTrue( + any( + entry["param"] is model.b + for entry in optimizer._routing[0]["adam_no_decay"] + ) + ) + self.assertFalse( + any( + entry["param"] is model.b + for entry in optimizer._routing[0]["adam_decay"] + ) + ) # Module name "bias_proj" but leaf is "weight" → Muon self.assertIn("momentum_buffer", optimizer.state[model.bias_proj.weight]) self.assertNotIn("exp_avg", optimizer.state[model.bias_proj.weight]) diff --git a/source/tests/pt_expt/compile_utils.py b/source/tests/pt_expt/compile_utils.py new file mode 100644 index 0000000000..e91661bf6d --- /dev/null +++ b/source/tests/pt_expt/compile_utils.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Shared version guard for pt_expt ``torch.compile`` tests.""" + +import unittest + +from deepmd.pt.utils.compile_compat import ( + check_compile_torch_version, +) + +try: + check_compile_torch_version() +except RuntimeError as error: + _COMPILE_SUPPORT_ERROR = str(error) +else: + _COMPILE_SUPPORT_ERROR = "" + +REQUIRES_SUPPORTED_COMPILE = unittest.skipIf( + bool(_COMPILE_SUPPORT_ERROR), + _COMPILE_SUPPORT_ERROR, +) diff --git a/source/tests/pt_expt/model/test_graph_builder_dispatch.py b/source/tests/pt_expt/model/test_graph_builder_dispatch.py index c0e1c3dcdf..b4a7b903d6 100644 --- a/source/tests/pt_expt/model/test_graph_builder_dispatch.py +++ b/source/tests/pt_expt/model/test_graph_builder_dispatch.py @@ -123,6 +123,16 @@ def test_auto_graph_method_uses_nv_only_on_cuda(): assert resolve_neighbor_graph_method("auto", torch.device("cpu")) == "dense" +def test_auto_graph_method_warns_when_nv_is_unavailable(caplog): + with ( + patch("deepmd.pt.utils.nv_nlist.is_nv_available", return_value=False), + caplog.at_level("WARNING", logger="deepmd.pt_expt.utils.graph_builder"), + ): + assert resolve_neighbor_graph_method("auto", torch.device("cuda")) == "dense" + + assert "pip install nvalchemi-toolkit-ops" in caplog.text + + def test_explicit_nv_rejects_cpu(): with pytest.raises(ValueError, match="requires a CUDA"): resolve_neighbor_graph_method("nv", torch.device("cpu")) diff --git a/source/tests/pt_expt/test_multitask.py b/source/tests/pt_expt/test_multitask.py index 92ca70936f..383fc0059d 100644 --- a/source/tests/pt_expt/test_multitask.py +++ b/source/tests/pt_expt/test_multitask.py @@ -63,6 +63,10 @@ process_systems, ) +from .compile_utils import ( + REQUIRES_SUPPORTED_COMPILE, +) + _energy_data_requirement = [ DataRequirementItem("energy", ndof=1, atomic=False, must=False, high_prec=True), DataRequirementItem("force", ndof=3, atomic=True, must=False, high_prec=False), @@ -1457,6 +1461,7 @@ def tearDown(self) -> None: shutil.rmtree(self.tmpdir, ignore_errors=True) +@REQUIRES_SUPPORTED_COMPILE class TestMultiTaskCompile(unittest.TestCase): """Verify that multi-task + torch.compile works correctly.""" @@ -2095,6 +2100,7 @@ def test_gradient_accumulation(self) -> None: ) +@REQUIRES_SUPPORTED_COMPILE class TestCompileCaseEmbdVaryingNframes(unittest.TestCase): """Compiled multi-task with ``dim_case_embd > 0`` and varying ``nframes``. diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index c9d2c70ccc..f5ffaa02b4 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -24,6 +24,9 @@ from deepmd.loggers.training import ( format_training_message, ) +from deepmd.pt.optimizer import ( + HybridMuonOptimizer, +) from deepmd.pt_expt.entrypoints.main import ( get_trainer, ) @@ -37,6 +40,10 @@ update_deepmd_input, ) +from .compile_utils import ( + REQUIRES_SUPPORTED_COMPILE, +) + EXAMPLE_DIR = os.path.join( os.path.dirname(__file__), "..", @@ -307,6 +314,30 @@ def test_explicit_graph_method_rejects_ineligible_model(self) -> None: with self.assertRaisesRegex(ValueError, "graph-eligible"): get_trainer(config) + def test_supported_optimizers_construct(self) -> None: + for optimizer_type, optimizer_class in ( + ("AdamW", torch.optim.AdamW), + ("HybridMuon", HybridMuonOptimizer), + ): + with self.subTest(optimizer_type=optimizer_type): + config = _make_config(self.data_dir) + config["optimizer"] = {"type": optimizer_type} + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + trainer = get_trainer(config) + + self.assertIsInstance(trainer.optimizer, optimizer_class) + + def test_unsupported_optimizer_has_clear_error(self) -> None: + config = _make_config(self.data_dir) + config["optimizer"] = {"type": "LKF"} + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + with self.assertRaisesRegex(ValueError, "Unsupported optimizer type: LKF"): + get_trainer(config) + def _run_training(self, config: dict) -> None: """Run training and verify lcurve + checkpoint creation.""" tmpdir = tempfile.mkdtemp(prefix="pt_expt_train_") @@ -428,6 +459,7 @@ def test_training_loop_dpa4(self) -> None: self.assertEqual(config["model"]["type"], "dpa4") self._run_training(config) + @REQUIRES_SUPPORTED_COMPILE def test_training_loop_compiled(self) -> None: """Run a few training steps with torch.compile enabled.""" config = _make_config(self.data_dir, numb_steps=5) @@ -436,6 +468,7 @@ def test_training_loop_compiled(self) -> None: config = normalize(config) self._run_training(config) + @REQUIRES_SUPPORTED_COMPILE def test_training_loop_compiled_silu(self) -> None: """Run compiled training with silu activation.""" config = _make_config(self.data_dir, numb_steps=5) @@ -509,6 +542,7 @@ def test_missing_attr_raises(self) -> None: _ = cm.nonexistent_attribute_xyz +@REQUIRES_SUPPORTED_COMPILE class TestCompiledDynamicShapes(unittest.TestCase): """Test that _CompiledModel handles varying nall via dynamic shapes.""" @@ -570,6 +604,7 @@ def test_compiled_handles_varying_nall(self) -> None: shutil.rmtree(tmpdir, ignore_errors=True) +@REQUIRES_SUPPORTED_COMPILE class TestCompiledConsistency(unittest.TestCase): """Verify compiled model produces the same energy/force/virial as uncompiled.""" @@ -986,6 +1021,7 @@ def test_init_model(self) -> None: finally: shutil.rmtree(tmpdir, ignore_errors=True) + @REQUIRES_SUPPORTED_COMPILE def test_restart_from_compiled_checkpoint(self) -> None: """Train WITH compile enabled, restart from the compiled checkpoint. @@ -1080,6 +1116,7 @@ def test_restart_from_compiled_checkpoint(self) -> None: finally: shutil.rmtree(tmpdir, ignore_errors=True) + @REQUIRES_SUPPORTED_COMPILE def test_restart_with_compile(self) -> None: """Train uncompiled, restart with compile enabled.""" from deepmd.pt_expt.train.training import ( @@ -1419,6 +1456,7 @@ def _run_steps(self, enable_compile: bool, nsteps: int = 6) -> None: finally: shutil.rmtree(tmpdir, ignore_errors=True) + @REQUIRES_SUPPORTED_COMPILE def test_compiled(self) -> None: """Compiled training with varying nframes + fparam/aparam.""" self._run_steps(enable_compile=True) @@ -1463,6 +1501,7 @@ def _create_small_system( np.save(os.path.join(set_dir, "virial.npy"), virial) +@REQUIRES_SUPPORTED_COMPILE class TestCompiledVaryingNatoms(unittest.TestCase): """Test compiled training with systems of different atom counts. @@ -1673,6 +1712,7 @@ def test_compile_warns_dpa1_with_attention(self) -> None: self.assertIsInstance(trainer.wrapper.model["Default"], _CompiledModel) +@REQUIRES_SUPPORTED_COMPILE class TestCompiledSharedFittingDifferentDescriptor(unittest.TestCase): """Regression test: shared fitting with different descriptors gets distinct compiled graphs. diff --git a/source/tests/pt_expt/test_training_ddp.py b/source/tests/pt_expt/test_training_ddp.py index 0d71e66870..5ffa69ab46 100644 --- a/source/tests/pt_expt/test_training_ddp.py +++ b/source/tests/pt_expt/test_training_ddp.py @@ -48,6 +48,10 @@ update_deepmd_input, ) +from .compile_utils import ( + REQUIRES_SUPPORTED_COMPILE, +) + # Paths to the water data used by PT tests _PT_DATA = str(Path(__file__).parent.parent / "pt" / "water" / "data" / "data_0") @@ -1479,6 +1483,7 @@ def _worker_multitask_compile_train(rank, world_size, port, data_dir, result_dic dist.destroy_process_group() +@REQUIRES_SUPPORTED_COMPILE class TestDDPCompileSingleTask(unittest.TestCase): """DDP + torch.compile: single-task training with 2 ranks. @@ -1524,6 +1529,7 @@ def test_ddp_compile_single_task(self) -> None: ) +@REQUIRES_SUPPORTED_COMPILE class TestDDPCompileMultiTask(unittest.TestCase): """DDP + torch.compile: multi-task training with 2 ranks. From a48021e0b0670f875978c47fd808cb26002e6911 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Thu, 30 Jul 2026 11:37:35 +0800 Subject: [PATCH 6/7] fix(pt_expt): handle flattened multi-frame NV inputs --- deepmd/pt_expt/utils/nv_graph_builder.py | 5 ++--- source/tests/pt_expt/utils/test_nv_graph_builder.py | 6 ++++-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/deepmd/pt_expt/utils/nv_graph_builder.py b/deepmd/pt_expt/utils/nv_graph_builder.py index 41124c8519..e78fea3e16 100644 --- a/deepmd/pt_expt/utils/nv_graph_builder.py +++ b/deepmd/pt_expt/utils/nv_graph_builder.py @@ -269,9 +269,8 @@ def build_neighbor_graph_nv( ) device = coord.device - nf = coord.shape[0] if coord.ndim == 3 else 1 - coord = coord.reshape(nf, -1, 3) - nloc = coord.shape[1] + nf, nloc = atype.shape[:2] + coord = coord.reshape(nf, nloc, 3) periodic = box is not None if nloc == 0: diff --git a/source/tests/pt_expt/utils/test_nv_graph_builder.py b/source/tests/pt_expt/utils/test_nv_graph_builder.py index 8ca4f2db1b..b8e5b7dc41 100644 --- a/source/tests/pt_expt/utils/test_nv_graph_builder.py +++ b/source/tests/pt_expt/utils/test_nv_graph_builder.py @@ -60,10 +60,12 @@ def test_nv_matches_intree_carry_all(periodic): def test_nv_batches_frames_without_python_loop(): - """Multi-frame: nv searches all frames in one kernel (no per-frame loop).""" + """NV batches flattened training coordinates without a per-frame loop.""" dev = torch.device("cuda") rng = np.random.default_rng(0) - coord = torch.tensor(rng.random((3, 5, 3)) * 3.0, dtype=torch.float64, device=dev) + coord = torch.tensor( + rng.random((3, 5, 3)) * 3.0, dtype=torch.float64, device=dev + ).reshape(3, -1) box = ( (torch.eye(3, dtype=torch.float64, device=dev) * 4.0) .reshape(1, 3, 3) From 07ed9c405c23088255af2914f376c220f9c74727 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 31 Jul 2026 00:15:31 +0800 Subject: [PATCH 7/7] test(lmdb): remove flaky SIGHUP decoder test --- source/tests/common/dpmodel/test_lmdb_data.py | 27 ------------------- 1 file changed, 27 deletions(-) diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index 6b56ac82af..f184efc53b 100644 --- a/source/tests/common/dpmodel/test_lmdb_data.py +++ b/source/tests/common/dpmodel/test_lmdb_data.py @@ -1813,33 +1813,6 @@ def idle(): ) self.assertEqual(completed.returncode, 0, completed.stderr[-2000:]) - @unittest.skipUnless( - hasattr(signal, "SIGHUP"), - "SIGHUP is not available on this platform", - ) - def test_decoders_survive_the_hangup_of_their_launching_session(self) -> None: - """A decoder outlives the session that started the run. - - The hangup delivered when that session goes away reaches every - background helper of the run. A decoder that died of it would break - the pool of an otherwise healthy run. - """ - iterator = self._isolated_iterator() - first = next(iterator) - processes = list(iterator._pool.executor._processes.values()) - self.assertTrue(processes) - - for process in processes: - os.kill(process.pid, signal.SIGHUP) - - # Delivering the next batch is what proves the decoders lived through - # it: a pool that lost one degrades instead. That is both a stronger - # statement than reading their liveness and free of any timing. - self._assert_same_batch(first, self._reader.decode_batch([0, 1, 2, 3])) - self._assert_same_batch(next(iterator), self._reader.decode_batch([4, 5, 6, 7])) - self.assertTrue(iterator._pool.healthy) - self.assertTrue(all(process.is_alive() for process in processes)) - @unittest.skipUnless( hasattr(signal, "SIGKILL"), "SIGKILL is not available on this platform",