feat: enable configurable bfloat16 autocasting in ModelWrapper

Author: caic99

deepmd/pt/train/wrapper.py

@@ -5,6 +5,7 @@
     Union,
 )
 
+from deepmd.pt.utils.env import BF16_AUTOCAST
 import torch
 
 if torch.__version__.startswith("2"):
@@ -136,7 +137,7 @@ def share_params(self, shared_links, resume=False) -> None:
                             f"Shared params of {model_key_base}.{class_type_base} and {model_key_link}.{class_type_link}!"
                         )
 
-    @torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True)
+    @torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=BF16_AUTOCAST)
     def forward(
         self,
         coord,

deepmd/pt/utils/env.py

@@ -35,6 +35,7 @@
 CACHE_PER_SYS = 5  # keep at most so many sets per sys in memory
 ENERGY_BIAS_TRAINABLE = True
 CUSTOM_OP_USE_JIT = False
+BF16_AUTOCAST = False
 
 PRECISION_DICT = {
     "float16": torch.float16,
@@ -76,6 +77,7 @@
     torch.set_num_threads(intra_nthreads)
 
 __all__ = [
+    "BF16_AUTOCAST",
     "CACHE_PER_SYS",
     "CUSTOM_OP_USE_JIT",
     "DEFAULT_PRECISION",