opt dpa1

Author: OutisLi

deepmd/kernels/cuda/dpa1/graph_compress.py

@@ -107,6 +107,7 @@ def _forward_fake(
     gate_table: torch.Tensor,
     type_one_side: int,
     concat_tebd: int,
+    write_rotation: int,
     smooth: int,
     axis: int,
     canonical: bool,
@@ -127,7 +128,13 @@ def _forward_fake(
     dev = edge_vec.device
     return (
         torch.empty(n_node, out_dim, dtype=torch.float32, device=dev),
-        torch.empty(n_node, ng, 3, dtype=torch.float32, device=dev),
+        torch.empty(
+            n_node if write_rotation else 0,
+            ng,
+            3,
+            dtype=torch.float32,
+            device=dev,
+        ),
         torch.empty(n_node, 4, ng, dtype=torch.float32, device=dev),
     )
 
@@ -160,6 +167,7 @@ def _setup_context(ctx: Any, inputs: tuple, output: tuple) -> None:
         gate_table,
         type_one_side,
         concat_tebd,
+        _write_rotation,
         smooth,
         axis,
         canonical,
@@ -269,7 +277,7 @@ def _backward(
         protection,
         nnei,
     )
-    return (d_edge_vec,) + (None,) * 24
+    return (d_edge_vec,) + (None,) * 25
 
 
 # ======================================================================
@@ -493,6 +501,7 @@ def _cpu_forward(
     gate_table: torch.Tensor,
     type_one_side: int,
     concat_tebd: int,
+    write_rotation: int,
     smooth: int,
     axis: int,
     canonical: bool,
@@ -540,6 +549,8 @@ def _cpu_forward(
         nnei,
         n_node,
     )
+    if not write_rotation:
+        rot_mat = rot_mat.new_empty(0, ng, 3)
     return grrg, rot_mat, gr
 
 
@@ -717,6 +728,7 @@ def dpa1_graph_compress(
         gate_table,
         int(se.type_one_side),
         int(desc.concat_output_tebd),
+        1,
         int(se.smooth),
         int(se.axis_neuron),
         bool(graph.destination_sorted),
@@ -761,7 +773,9 @@ def dpa1_graph_compress_energy_force(
     retained. Destination/source CSR tensors are part of the graph contract and
     are reused by both descriptor and force kernels. A graph marked
     ``destination_sorted`` selects direct destination-major addressing; edge
-    masks remain authoritative in both addressing modes.
+    masks remain authoritative in both addressing modes. This inference-only
+    path suppresses the unused rotation output, and fitting backward does not
+    retain the descriptor after its last value use.
 
     Parameters
     ----------
@@ -850,6 +864,7 @@ def dpa1_graph_compress_energy_force(
         gate_table,
         int(se.type_one_side),
         int(desc.concat_output_tebd),
+        0,
         int(se.smooth),
         int(se.axis_neuron),
         bool(graph.destination_sorted),
@@ -911,12 +926,12 @@ def dpa1_graph_compress_energy_force(
     ).index_add_(0, frame_index, atom_energy)
     descriptor_gradient = torch.ops.deepmd.graph_fitting_backward(
         torch.ones_like(atom_energy),
-        descriptor,
         fitting_saved,
         weights,
         residuals,
         head_weight,
     )
+    del descriptor, fitting_saved
     edge_gradient = torch.ops.deepmd.dpa1_graph_compress_backward(
         descriptor_gradient,
         None,

deepmd/kernels/cuda/dpa1/graph_energy_force.py

@@ -216,8 +216,9 @@ def _cpu(
     energy = torch.zeros(nf, 1, dtype=atom_e.dtype, device=atom_e.device)
     energy = energy.index_add(0, frame_id, atom_e)
     d_grrg = torch.ops.deepmd.graph_fitting_backward(
-        torch.ones_like(atom_e), grrg, fit_saved, fit_ws, fit_resnets, w_head
+        torch.ones_like(atom_e), fit_saved, fit_ws, fit_resnets, w_head
     )
+    del grrg, fit_saved
     g_e = torch.ops.deepmd.dpa1_graph_descriptor_backward(
         d_grrg,
         None,

deepmd/kernels/cuda/graph_fitting.py

@@ -20,6 +20,10 @@
   flat buffer of ``adot chunks`` (the activations themselves are a forward-only
   transient); it is an autograd save, never a user-facing value, and receives
   no gradient (``set_materialize_grads(False)``).
+* The backward infers the node count and descriptor width from the saved
+  derivative buffer and first weight. It deliberately does not retain the
+  descriptor tensor, allowing inference memory planners to reuse descriptor
+  storage for its gradient after the fitting forward.
 * The head bias is passed as a device tensor, not a Python float: reading a
   value host-side (``.item()``) inside the dispatch path would fail under
   symbolic tracing (``GuardOnDataDependentSymNode``) and force a GPU sync per
@@ -121,28 +125,29 @@ def _forward_fake(
 
 def _backward_fake(
     d_e: torch.Tensor,
-    x: torch.Tensor,
     saved: torch.Tensor,
     ws: list[torch.Tensor],
     resnets: list[int],
     w_head: torch.Tensor,
 ) -> torch.Tensor:
-    return torch.empty_like(x)
+    total_width = sum(int(w.shape[1]) for w in ws)
+    n_node = saved.shape[0] // total_width
+    return saved.new_empty(n_node, ws[0].shape[0])
 
 
 def _setup_context(ctx: Any, inputs: tuple, output: tuple) -> None:
     x, atype, ws, bs, idts, resnets, w_head, b_head, bias_atom_e, act = inputs
     _e, saved = output
-    ctx.save_for_backward(x, saved, w_head, *ws)
+    ctx.save_for_backward(saved, w_head, *ws)
     ctx.n_layers = len(ws)
     ctx.resnets = resnets
     ctx.set_materialize_grads(False)
 
 
 def _backward(ctx: Any, d_e: torch.Tensor, d_saved: Any) -> tuple:
-    x, saved, w_head, *ws = ctx.saved_tensors
+    saved, w_head, *ws = ctx.saved_tensors
     d_x = torch.ops.deepmd.graph_fitting_backward(
-        d_e, x, saved, list(ws), ctx.resnets, w_head
+        d_e, saved, list(ws), ctx.resnets, w_head
     )
     none_list = [None] * ctx.n_layers
     return (d_x, None, none_list, none_list, none_list, None, None, None, None, None)
@@ -192,13 +197,13 @@ def _cpu_forward(
 
 def _cpu_backward(
     d_e: torch.Tensor,
-    x: torch.Tensor,
     saved: torch.Tensor,
     ws: list[torch.Tensor],
     resnets: list[int],
     w_head: torch.Tensor,
 ) -> torch.Tensor:
-    n_node = x.shape[0]
+    total_width = sum(int(w.shape[1]) for w in ws)
+    n_node = saved.shape[0] // total_width
     offset = [0]
     for w in ws:
         offset.append(offset[-1] + int(w.shape[1]))

deepmd/pt_expt/infer/deep_eval.py

@@ -1752,8 +1752,10 @@ def _eval_model_graph(
         destination_order, destination_row_ptr, source_row_ptr, source_order,
         fparam, aparam, charge_spin)`` to the exported forward. The AOTI artifact's edge axis
         is DYNAMIC (B2.0), so no ``edge_capacity`` padding is needed.  The
-        forward returns the LOCAL public keys directly, so results are reshaped
-        without ``communicate_extended_output``.
+        ``graph_edge_dtype`` metadata selects float32 geometry for compressed
+        DPA1 and float64 for generic graph descriptors. The forward returns the
+        LOCAL public keys directly, so results are reshaped without
+        ``communicate_extended_output``.
         """
         from deepmd.pt_expt.utils.env import (
             DEVICE,
@@ -1781,7 +1783,16 @@ def _eval_model_graph(
         edge_index_t = torch.as_tensor(
             graph.edge_index, dtype=torch.int64, device=DEVICE
         )
-        edge_vec_t = torch.as_tensor(graph.edge_vec, dtype=torch.float64, device=DEVICE)
+        edge_dtype = (
+            torch.float32
+            if self.metadata.get("graph_edge_dtype") == "float32"
+            else torch.float64
+        )
+        edge_vec_t = torch.as_tensor(
+            graph.edge_vec,
+            dtype=edge_dtype,
+            device=DEVICE,
+        )
         edge_mask_t = torch.as_tensor(graph.edge_mask, dtype=torch.bool, device=DEVICE)
         destination_order_t = torch.as_tensor(
             graph.destination_order,

deepmd/pt_expt/utils/serialization.py

@@ -375,6 +375,7 @@ def build_synthetic_graph_inputs(
     nloc: int = 7,
     *,
     dtype: torch.dtype,
+    edge_dtype: torch.dtype | None = None,
     device: torch.device | None = None,
     want_fparam: bool = True,
     want_aparam: bool = True,
@@ -412,9 +413,12 @@ def build_synthetic_graph_inputs(
     nloc : int
         Number of local atoms per frame (``N == nframes * nloc``).
     dtype : torch.dtype
-        Float precision of ``coord``/``edge_vec``/``fparam``/... .  The exported
-        ``.pt2`` is float64-only (C++ ABI); training passes
-        ``GLOBAL_PT_FLOAT_PRECISION``.
+        Float precision of ``coord``/``fparam`` and other conditioning inputs.
+    edge_dtype : torch.dtype, optional
+        Precision of ``edge_vec``. Defaults to ``dtype``. A compressed DPA1
+        graph artifact uses float32 because its descriptor and analytical
+        backward both compute in float32; generic graph artifacts preserve the
+        model input precision.
     device : torch.device, optional
         Target device.  Defaults to ``deepmd.pt_expt.utils.env.DEVICE``; the
         export path passes ``cpu`` explicitly (make_fx traces on CPU).
@@ -431,6 +435,8 @@ def build_synthetic_graph_inputs(
 
     if device is None:
         device = _env.DEVICE
+    if edge_dtype is None:
+        edge_dtype = dtype
 
     rcut = model.get_rcut()
     ntypes = len(model.get_type_map())
@@ -479,7 +485,7 @@ def build_synthetic_graph_inputs(
         atype_t.reshape(-1),
         graph.n_node,
         graph.edge_index,
-        graph.edge_vec,
+        graph.edge_vec.to(edge_dtype),
         graph.edge_mask,
         graph.destination_order,
         graph.destination_row_ptr,
@@ -644,6 +650,20 @@ def _build_dynamic_shapes(
     return (*base, None, None, None, None, None, None, None, None)
 
 
+def _graph_edge_dtype(model: torch.nn.Module, lower_kind: str) -> str:
+    """Return the graph edge-vector dtype encoded by the deployment artifact.
+
+    Geometrically compressed DPA1 evaluates both descriptor directions in
+    float32 and therefore accepts float32 geometry directly. Other graph
+    descriptors retain the model-agnostic float64 geometry ABI.
+    """
+    atomic_model = getattr(model, "atomic_model", None)
+    descriptor = getattr(atomic_model, "descriptor", None)
+    if lower_kind == "graph" and bool(getattr(descriptor, "geo_compress", False)):
+        return "float32"
+    return "float64"
+
+
 def _collect_metadata(
     model: torch.nn.Module, is_spin: bool = False, lower_kind: str = "nlist"
 ) -> dict:
@@ -764,6 +784,7 @@ def _probe_has_message_passing(obj: object) -> bool | None:
     #   "graph" → NeighborGraph (atype, n_node, edge_index, edge_vec, edge_mask)
     # The C++ loader branches on this to build the matching inputs.
     meta["lower_input_kind"] = "graph" if lower_kind == "graph" else "nlist"
+    meta["graph_edge_dtype"] = _graph_edge_dtype(model, lower_kind)
     return meta
 
 
@@ -1066,14 +1087,21 @@ def _trace_and_export(
         nnei = sum(model.get_sel())
         e_sample = math.ceil(1.25 * nloc_sample * nnei)
 
-        # make_fx traces on CPU; the .pt2 C++ ABI is float64-only.  Pass device
-        # and dtype explicitly instead of mutating the module-level env.DEVICE.
+        # make_fx traces on CPU. Conditioning inputs retain the model-agnostic
+        # float64 ABI; compressed DPA1 graph geometry enters directly in its
+        # float32 compute precision, as recorded in metadata.
+        edge_dtype = (
+            torch.float32
+            if metadata["graph_edge_dtype"] == "float32"
+            else torch.float64
+        )
         sample_inputs = build_synthetic_graph_inputs(
             model,
             e_max=e_sample,
             nframes=2,
             nloc=nloc_sample,
             dtype=torch.float64,
+            edge_dtype=edge_dtype,
             device=torch.device("cpu"),
         )
 

source/api_cc/include/DeepPot.h

@@ -330,6 +330,20 @@ class DeepPotBackend : public DeepBaseModelBackend {
                                  const std::vector<double>& aparam,
                                  const int nall_nodes,
                                  const InputNlist* comm_nlist);
+  virtual void compute_edges_gpu(double* d_atom_energy,
+                                 double* d_force,
+                                 double* d_atom_virial,
+                                 const double* d_coord,
+                                 const int* d_atype,
+                                 const int* d_edge_index,
+                                 const float* d_edge_vec,
+                                 const int nloc,
+                                 const int nedge,
+                                 const std::vector<double>& fparam,
+                                 const std::vector<double>& aparam,
+                                 const int nall_nodes,
+                                 const InputNlist* comm_nlist);
+  virtual bool uses_fp32_edge_vectors() const;
 };
 
 /**
@@ -751,6 +765,30 @@ class DeepPot : public DeepBaseModel {
                          const int nall_nodes = 0,
                          const InputNlist* comm_nlist = nullptr);
 
+  /**
+   * @brief Device-edge inference with FP32 edge vectors.
+   *
+   * Call this overload only when ``uses_fp32_edge_vectors()`` is true.
+   */
+  void compute_edges_gpu(double* d_atom_energy,
+                         double* d_force,
+                         double* d_atom_virial,
+                         const double* d_coord,
+                         const int* d_atype,
+                         const int* d_edge_index,
+                         const float* d_edge_vec,
+                         const int nloc,
+                         const int nedge,
+                         const std::vector<double>& fparam,
+                         const std::vector<double>& aparam,
+                         const int nall_nodes = 0,
+                         const InputNlist* comm_nlist = nullptr);
+
+  /**
+   * @brief Whether the loaded artifact expects FP32 device edge vectors.
+   */
+  bool uses_fp32_edge_vectors() const;
+
   int dim_chg_spin() const;
 
  protected: