feat(pt_expt): use inductor+dynamic for training compile

Replace aot_eager+padding+manual recompile with symbolic make_fx +
torch.compile(backend="inductor", dynamic=True). The compiled graph
natively handles varying nframes/nloc/nall so the per-batch padding
and runtime _recompile pass can be removed.

Use a trace-time nframes of 7 (prime) and reshape with -1 in dpmodel
(general_fitting, env_mat) to prevent PyTorch's symbolic tracer from
unifying the batch dim with numb_fparam / numb_aparam / ntypes /
dim_case_embd. Add TestCompiledVaryingNframesWithParams covering
collisions with fparam/aparam, and TestCompileCaseEmbdVaryingNframes
covering dim_case_embd > 0 with runtime nframes matching the embed dim.

Author: Han Wang

deepmd/dpmodel/fitting/general_fitting.py

@@ -674,15 +674,20 @@ def _call_common(
         if self.numb_fparam > 0:
             assert fparam is not None, "fparam should not be None"
             try:
-                fparam = xp.reshape(fparam, (nf, self.numb_fparam))
+                # Use -1 for nframes so the shape is inferred from the total
+                # size.  Passing the concrete symbol `nf` here would let
+                # torch.fx's symbolic tracer specialise when `nf` happens to
+                # equal another tensor dim (e.g. numb_fparam), baking the
+                # batch size into the compiled graph.
+                fparam = xp.reshape(fparam, (-1, self.numb_fparam))
             except (ValueError, RuntimeError) as e:
                 raise ValueError(
                     f"input fparam: cannot reshape {fparam.shape} "
                     f"into ({nf}, {self.numb_fparam})."
                 ) from e
             fparam = (fparam - self.fparam_avg[...]) * self.fparam_inv_std[...]
             fparam = xp.tile(
-                xp.reshape(fparam, (nf, 1, self.numb_fparam)), (1, nloc, 1)
+                xp.reshape(fparam, (-1, 1, self.numb_fparam)), (1, nloc, 1)
             )
             xx = xp.concat(
                 [xx, fparam],
@@ -697,7 +702,9 @@ def _call_common(
         if self.numb_aparam > 0 and not self.use_aparam_as_mask:
             assert aparam is not None, "aparam should not be None"
             try:
-                aparam = xp.reshape(aparam, (nf, nloc, self.numb_aparam))
+                # Use -1 for nframes so the shape is inferred from the total
+                # size; see the fparam branch above for rationale.
+                aparam = xp.reshape(aparam, (-1, nloc, self.numb_aparam))
             except (ValueError, RuntimeError) as e:
                 raise ValueError(
                     f"input aparam: cannot reshape {aparam.shape} "
@@ -744,8 +751,12 @@ def _call_common(
                     device=array_api_compat.device(descriptor),
                 )
             for type_i in range(self.ntypes):
+                # Use -1 for nframes so the shape is inferred; see the fparam
+                # branch above for rationale (avoid symbolic-dim collision
+                # with numb_fparam / other dims during symbolic tracing).
                 mask = xp.tile(
-                    xp.reshape((atype == type_i), (nf, nloc, 1)), (1, 1, net_dim_out)
+                    xp.reshape((atype == type_i), (-1, nloc, 1)),
+                    (1, 1, net_dim_out),
                 )
                 atom_property = self.nets[(type_i,)](xx)
                 if self.remove_vaccum_contribution is not None and not (
@@ -761,7 +772,7 @@ def _call_common(
                 if self.eval_return_middle_output and len(self.neuron) > 0:
                     mid = self.nets[(type_i,)].call_until_last(xx)
                     mid_mask = xp.tile(
-                        xp.reshape((atype == type_i), (nf, nloc, 1)),
+                        xp.reshape((atype == type_i), (-1, nloc, 1)),
                         (1, 1, self.neuron[-1]),
                     )
                     mid = xp.where(mid_mask, mid, xp.zeros_like(mid))
@@ -778,7 +789,7 @@ def _call_common(
                 xp.reshape(atype, (-1,)),
                 axis=0,
             ),
-            (nf, nloc, net_dim_out),
+            (-1, nloc, net_dim_out),
         )
         # nf x nloc
         exclude_mask = self.emask.build_type_exclude_mask(atype)

deepmd/dpmodel/utils/env_mat.py

@@ -68,16 +68,22 @@ def _make_env_mat(
     xp = array_api_compat.array_namespace(nlist)
     nf, nloc, nnei = nlist.shape
     # nf x nall x 3
-    coord = xp.reshape(coord, (nf, -1, 3))
+    # Callers may pass either (nf, nall*3) or (nf, nall, 3); normalise
+    # both to (nf, nall, 3) using -1 for nframes so the shape is inferred
+    # from the total size.  Passing the symbolic nf here can trigger
+    # torch.fx symbolic-tracer specialisation when nf happens to collide
+    # with another dim (e.g. numb_fparam) during training compile.
+    if coord.ndim == 2:
+        coord = xp.reshape(coord, (-1, coord.shape[1] // 3, 3))
     mask = nlist >= 0
     nlist = nlist * xp.astype(mask, nlist.dtype)
     # nf x (nloc x nnei) x 3
-    index = xp.tile(xp.reshape(nlist, (nf, -1, 1)), (1, 1, 3))
+    index = xp.tile(xp.reshape(nlist, (-1, nloc * nnei, 1)), (1, 1, 3))
     coord_r = xp_take_along_axis(coord, index, 1)
     # nf x nloc x nnei x 3
-    coord_r = xp.reshape(coord_r, (nf, nloc, nnei, 3))
+    coord_r = xp.reshape(coord_r, (-1, nloc, nnei, 3))
     # nf x nloc x 1 x 3
-    coord_l = xp.reshape(xp_take_first_n(coord, 1, nloc), (nf, -1, 1, 3))
+    coord_l = xp.reshape(xp_take_first_n(coord, 1, nloc), (-1, nloc, 1, 3))
     # nf x nloc x nnei x 3
     diff = coord_r - coord_l
     # nf x nloc x nnei