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/deepmd/pt/utils/compile_compat.py b/deepmd/pt/utils/compile_compat.py index 2cf3a64f35..9ff51db003 100644 --- a/deepmd/pt/utils/compile_compat.py +++ b/deepmd/pt/utils/compile_compat.py @@ -84,11 +84,11 @@ def _torch_release() -> tuple[int, int]: 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 the - releases where the regression exists. + The adjustments are mutually independent and individually idempotent. + Invoke this function before the first Dynamo or Inductor compilation in + each entry path; repeated calls from independent compile paths are safe. + The symbolic-divisibility repair is applied only on releases 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/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 195f1f3e7e..1d443e97e6 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -50,6 +50,9 @@ from deepmd.dpmodel.utils.training_utils import ( compute_total_numb_batch, ) +from deepmd.pt.optimizer import ( + HybridMuonOptimizer, +) from deepmd.pt.train.utils import ( resolve_best_checkpoint_dir, ) @@ -57,6 +60,11 @@ FullValidator, resolve_full_validation_start_step, ) +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 from deepmd.pt.utils.compile_compat import ( @@ -583,19 +591,14 @@ 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. - } + # 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: inductor_options.update(extra_options) if compile_opts: @@ -1157,8 +1160,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 @@ -1209,7 +1212,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). @@ -1363,6 +1373,7 @@ def __init__( model_params = config["model"] training_params = config["training"] + optimizer_params = config.get("optimizer", {}) validating_params = config.get("validating", {}) or {} # Task normalization -------------------------------------------------- @@ -1614,24 +1625,51 @@ def initialize_statistics( ) # Optimiser ----------------------------------------------------------- - opt_type = training_params.get("opt_type", "Adam") + 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 # 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, ) - else: - raise ValueError(f"Unsupported optimizer type: {opt_type}") + else: # 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"]), + ) for param_group in self.optimizer.param_groups: param_group["initial_lr"] = initial_lr @@ -1840,9 +1878,14 @@ def update_finetune_bias( 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: + 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) @@ -1938,6 +1981,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. @@ -1953,14 +2019,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 new file mode 100644 index 0000000000..b074a3af8f --- /dev/null +++ b/deepmd/pt_expt/utils/graph_builder.py @@ -0,0 +1,151 @@ +# 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, +) + +import torch + +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) + +log = logging.getLogger(__name__) + + +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": + 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, " + 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..e78fea3e16 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( @@ -275,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: @@ -301,9 +294,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 2987f46bbf..f8f3eaf146 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5501,6 +5501,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() @@ -5696,6 +5705,15 @@ def training_args( default=0, doc=supported_backends("pt") + 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=supported_backends("pt_expt") + doc_neighbor_graph_method, + ), Argument( "enable_compile", bool, 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", 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/test_hybrid_muon.py b/source/tests/pt/test_hybrid_muon.py index 3f2bc4bef6..aab297b0b0 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,21 @@ 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]) + 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 80d26a84fc..b4a7b903d6 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,64 @@ 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_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")) + + @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_multitask.py b/source/tests/pt_expt/test_multitask.py index 2abda5bb02..478f0cf6aa 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), @@ -1493,6 +1497,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.""" @@ -2131,6 +2136,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 0db0df4977..f32f860825 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -30,6 +30,9 @@ from deepmd.loggers.training import ( format_training_message, ) +from deepmd.pt.optimizer import ( + HybridMuonOptimizer, +) from deepmd.pt_expt.entrypoints.main import ( get_trainer, ) @@ -47,6 +50,9 @@ assert_energy_stat_cache_round_trip, energy_model_params, ) +from .compile_utils import ( + REQUIRES_SUPPORTED_COMPILE, +) EXAMPLE_DIR = os.path.join( os.path.dirname(__file__), @@ -324,6 +330,58 @@ 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 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_") @@ -445,6 +503,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) @@ -453,6 +512,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) @@ -526,6 +586,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.""" @@ -587,6 +648,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.""" @@ -1039,6 +1101,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. @@ -1133,6 +1196,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 ( @@ -1472,6 +1536,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) @@ -1516,6 +1581,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. @@ -1726,6 +1792,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 b866db4f38..3f723dcfa8 100644 --- a/source/tests/pt_expt/test_training_ddp.py +++ b/source/tests/pt_expt/test_training_ddp.py @@ -54,6 +54,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") @@ -1568,6 +1572,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. @@ -1613,6 +1618,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. 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..b8e5b7dc41 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,14 +56,16 @@ 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(): - """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) @@ -70,26 +78,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 +107,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)