Merge upstream/master into feat-sfpg-multirank

Conflict in deepmd/pt_expt/utils/serialization.py: upstream deepmodeling#5913
relocated the graph edge-dtype helper to deepmd/pt_expt/model/graph_lower.py
as a public graph_edge_dtype(), while this branch had rewritten the same
helper in place to answer from the atomic-model capability instead of
reaching through a single .descriptor.

Resolution keeps BOTH: upstream's location (serialization.py imports it
from graph_lower) with this branch's implementation. The reach-in version
is wrong for compositions -- a LinearEnergyAtomicModel has no .descriptor,
so a bridged/ZBL model would silently report float64 regardless of its
children -- which is exactly what the capability aggregation fixes.
Dropped the now-unused torch import from graph_lower.py.

Author: Han Wang

deepmd/backend/pt_expt.py

@@ -13,6 +13,9 @@
 from deepmd.backend.backend import (
     Backend,
 )
+from deepmd.utils.pt_checkpoint import (
+    detect_pt_checkpoint_backend,
+)
 
 if TYPE_CHECKING:
     from argparse import (
@@ -51,11 +54,8 @@ def match_filename(cls, filename: str) -> int:
         Returns
         -------
         - 1 for the regular `.pte` / `.pt2` suffixes (default behaviour).
-        - 2 for `.pt` files whose state-dict uses pt_expt's dpmodel
-            parameter naming (`.w`/`.b`); this outranks the legacy pt
-            backend's default suffix score (1) so pt_expt-trained `.pt`
-            checkpoints route here, while genuine pt-trained `.pt` files
-            (which use `.matrix`/`.bias`) keep going to the pt backend.
+        - 2 for `.pt` files whose state dictionary uses the pt_expt parameter
+            dialect. This outranks the pt backend's default suffix score (1).
         - 0 otherwise.
         """
         score = super().match_filename(filename)
@@ -69,21 +69,14 @@ def match_filename(cls, filename: str) -> int:
 
             # weights_only=True avoids unpickling arbitrary code from an
             # untrusted .pt — sniffing only needs the dict keys.
-            sd = torch.load(filename, map_location="cpu", weights_only=True)
+            checkpoint = torch.load(filename, map_location="cpu", weights_only=True)
         except Exception:
             # Not a valid torch archive (corrupt file, wrong format, or a
             # weights_only=True restriction trip).  Surrender the claim so
             # the dispatcher falls back to the default suffix match — pt's
             # default score (1) will pick up the file under `dp --pt`.
             return 0
-        if isinstance(sd, dict) and "model" in sd:
-            sd = sd["model"]
-        keys = list(sd.keys()) if hasattr(sd, "keys") else []
-        has_pt_expt = any(k.endswith(".w") or k.endswith(".b") for k in keys)
-        has_pt = any(k.endswith(".matrix") or k.endswith(".bias") for k in keys)
-        if has_pt_expt and not has_pt:
-            return 2
-        return 0
+        return 2 if detect_pt_checkpoint_backend(checkpoint) == "pt-expt" else 0
 
     def is_available(self) -> bool:
         """Check if the backend is available.

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)

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