feat(pt): add torch.compile support

Author: caic99

deepmd/pt/entrypoints/main.py

@@ -251,7 +251,8 @@ def train(
     output: str = "out.json",
 ) -> None:
     log.info("Configuration path: %s", input_file)
-    env.CUSTOM_OP_USE_JIT = True
+    if not env.COMPILE:
+        env.CUSTOM_OP_USE_JIT = True
     if LOCAL_RANK == 0:
         SummaryPrinter()()
     with open(input_file) as fin:

deepmd/pt/train/training.py

@@ -53,6 +53,7 @@
     get_sampler_from_params,
 )
 from deepmd.pt.utils.env import (
+    COMPILE,
     DEVICE,
     JIT,
     LOCAL_RANK,
@@ -413,9 +414,13 @@ def get_lr(lr_params):
             self.lr_exp = get_lr(config["learning_rate"])
 
         # JIT
-        if JIT:
+        if JIT or COMPILE:
             self.model = torch.jit.script(self.model)
-
+        if COMPILE:
+            self.model = torch.compile(
+                self.model,
+                dynamic=True,
+            )
         # Model Wrapper
         self.wrapper = ModelWrapper(self.model, self.loss, model_params=model_params)
         self.start_step = 0

deepmd/pt/utils/env.py

@@ -46,6 +46,7 @@
 CACHE_PER_SYS = 5  # keep at most so many sets per sys in memory
 ENERGY_BIAS_TRAINABLE = True
 CUSTOM_OP_USE_JIT = False
+COMPILE = True
 
 PRECISION_DICT = {
     "float16": torch.float16,