feat: add basic support to torch_npu for dp training

Author: caic99

deepmd/env.py

@@ -30,15 +30,15 @@
 
 
 # FLOAT_PREC
-dp_float_prec = os.environ.get("DP_INTERFACE_PREC", "high").lower()
+dp_float_prec = os.environ.get("DP_INTERFACE_PREC", "low").lower()
 if dp_float_prec in ("high", ""):
     # default is high
     GLOBAL_NP_FLOAT_PRECISION = np.float64
     GLOBAL_ENER_FLOAT_PRECISION = np.float64
     global_float_prec = "double"
 elif dp_float_prec == "low":
     GLOBAL_NP_FLOAT_PRECISION = np.float32
-    GLOBAL_ENER_FLOAT_PRECISION = np.float64
+    GLOBAL_ENER_FLOAT_PRECISION = np.float32
     global_float_prec = "float"
 else:
     raise RuntimeError(

deepmd/pt/utils/env.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 import torch
+import torch_npu
 
 from deepmd.common import (
     VALID_PRECISION,
@@ -37,10 +38,10 @@
 LOCAL_RANK = os.environ.get("LOCAL_RANK")
 LOCAL_RANK = int(0 if LOCAL_RANK is None else LOCAL_RANK)
 
-if os.environ.get("DEVICE") == "cpu" or torch.cuda.is_available() is False:
+if os.environ.get("DEVICE") == "cpu" or torch_npu.npu.is_available() is False:
     DEVICE = torch.device("cpu")
 else:
-    DEVICE = torch.device(f"cuda:{LOCAL_RANK}")
+    DEVICE = torch.device(f"npu:{LOCAL_RANK}")
 
 JIT = False
 CACHE_PER_SYS = 5  # keep at most so many sets per sys in memory