perf(pt_expt): use scalable graph builders during training

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]

deepmd/pt_expt/train/training.py

@@ -42,13 +42,20 @@
 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,
 )
 from deepmd.pt.train.validation import (
     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:
@@ -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).
@@ -1352,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 --------------------------------------------------
@@ -1576,22 +1584,48 @@ 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}")
 
@@ -1793,6 +1827,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 +1929,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.