fix(pt): optimize graph memory usage

deepmd/pt/model/model/dipole_model.py

@@ -111,6 +111,7 @@ def forward_lower(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            inference=True,
         )
         if self.get_fitting_net() is not None:
             model_predict = {}

deepmd/pt/model/model/dos_model.py

@@ -101,6 +101,7 @@ def forward_lower(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            inference=True,
         )
         if self.get_fitting_net() is not None:
             model_predict = {}

deepmd/pt/model/model/dp_zbl_model.py

@@ -112,6 +112,7 @@ def forward_lower(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            inference=True,
         )
 
         model_predict = {}

deepmd/pt/model/model/ener_model.py

@@ -115,6 +115,7 @@ def forward_lower(
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
             comm_dict=comm_dict,
+            inference=True,
         )
         if self.get_fitting_net() is not None:
             model_predict = {}

deepmd/pt/model/model/make_model.py

@@ -215,6 +215,7 @@ def forward_common_lower(
             aparam: Optional[torch.Tensor] = None,
             do_atomic_virial: bool = False,
             comm_dict: Optional[Dict[str, torch.Tensor]] = None,
+            inference: bool = False,
         ):
             """Return model prediction. Lower interface that takes
             extended atomic coordinates and types, nlist, and mapping
@@ -239,6 +240,8 @@ def forward_common_lower(
                 whether calculate atomic virial.
             comm_dict
                 The data needed for communication for parallel inference.
+            inference
+                Whether only perform inference rather than undergoing training.
 
             Returns
             -------
@@ -267,6 +270,7 @@ def forward_common_lower(
                 self.atomic_output_def(),
                 cc_ext,
                 do_atomic_virial=do_atomic_virial,
+                inference=inference,
             )
             model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict

deepmd/pt/model/model/polar_model.py

@@ -95,6 +95,7 @@ def forward_lower(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            inference=True,
         )
         if self.get_fitting_net() is not None:
             model_predict = {}

deepmd/pt/model/model/spin_model.py

@@ -467,6 +467,7 @@ def forward_common_lower(
         fparam: Optional[torch.Tensor] = None,
         aparam: Optional[torch.Tensor] = None,
         do_atomic_virial: bool = False,
+        inference: bool = False,
     ):
         nframes, nloc = nlist.shape[:2]
         (
@@ -487,6 +488,7 @@ def forward_common_lower(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            inference=inference,
         )
         model_output_type = self.backbone_model.model_output_type()
         if "mask" in model_output_type:
@@ -611,6 +613,7 @@ def forward_lower(
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=do_atomic_virial,
+            inference=True,
         )
         model_predict = {}
         model_predict["atom_energy"] = model_ret["energy"]

deepmd/pt/model/model/transform_output.py

@@ -33,15 +33,27 @@ def atomic_virial_corr(
     faked_grad = torch.ones_like(sumce0)
     lst = torch.jit.annotate(List[Optional[torch.Tensor]], [faked_grad])
     extended_virial_corr0 = torch.autograd.grad(
-        [sumce0], [extended_coord], grad_outputs=lst, create_graph=True
+        [sumce0],
+        [extended_coord],
+        grad_outputs=lst,
+        create_graph=False,
+        retain_graph=True,
     )[0]
     assert extended_virial_corr0 is not None
     extended_virial_corr1 = torch.autograd.grad(
-        [sumce1], [extended_coord], grad_outputs=lst, create_graph=True
+        [sumce1],
+        [extended_coord],
+        grad_outputs=lst,
+        create_graph=False,
+        retain_graph=True,
     )[0]
     assert extended_virial_corr1 is not None
     extended_virial_corr2 = torch.autograd.grad(
-        [sumce2], [extended_coord], grad_outputs=lst, create_graph=True
+        [sumce2],
+        [extended_coord],
+        grad_outputs=lst,
+        create_graph=False,
+        retain_graph=True,
     )[0]
     assert extended_virial_corr2 is not None
     extended_virial_corr = torch.concat(
@@ -61,11 +73,16 @@ def task_deriv_one(
     extended_coord: torch.Tensor,
     do_virial: bool = True,
     do_atomic_virial: bool = False,
+    inference: bool = False,
 ):
     faked_grad = torch.ones_like(energy)
     lst = torch.jit.annotate(List[Optional[torch.Tensor]], [faked_grad])
     extended_force = torch.autograd.grad(
-        [energy], [extended_coord], grad_outputs=lst, create_graph=True
+        [energy],
+        [extended_coord],
+        grad_outputs=lst,
+        create_graph=not inference,
+        retain_graph=True,
     )[0]
     assert extended_force is not None
     extended_force = -extended_force
@@ -106,6 +123,7 @@ def take_deriv(
     coord_ext: torch.Tensor,
     do_virial: bool = False,
     do_atomic_virial: bool = False,
+    inference: bool = False,
 ):
     size = 1
     for ii in vdef.shape:
@@ -123,6 +141,7 @@ def take_deriv(
             coord_ext,
             do_virial=do_virial,
             do_atomic_virial=do_atomic_virial,
+            inference=inference,
         )
         # nf x nloc x 1 x 3, nf x nloc x 1 x 9
         ffi = ffi.unsqueeze(-2)
@@ -146,6 +165,7 @@ def fit_output_to_model_output(
     fit_output_def: FittingOutputDef,
     coord_ext: torch.Tensor,
     do_atomic_virial: bool = False,
+    inference: bool = False,
 ) -> Dict[str, torch.Tensor]:
     """Transform the output of the fitting network to
     the model output.
@@ -169,6 +189,7 @@ def fit_output_to_model_output(
                     coord_ext,
                     do_virial=vdef.c_differentiable,
                     do_atomic_virial=do_atomic_virial,
+                    inference=inference,
                 )
                 model_ret[kk_derv_r] = dr
                 if vdef.c_differentiable: