save torch_npu modifications

Author: caic99

deepmd/env.py

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

deepmd/pt/cxx_op.py

@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import platform
 
-import torch
+import torch 
+import torch_npu 
+# from torch_npu.contrib import transfer_to_npu
 from packaging.version import (
     Version,
 )

deepmd/pt/entrypoints/main.py

@@ -13,7 +13,9 @@
 )
 
 import h5py
-import torch
+import torch 
+import torch_npu 
+# from torch_npu.contrib import transfer_to_npu
 import torch.distributed as dist
 import torch.version
 from torch.distributed.elastic.multiprocessing.errors import (
@@ -108,7 +110,12 @@ def get_trainer(
     local_rank = os.environ.get("LOCAL_RANK")
     if local_rank is not None:
         local_rank = int(local_rank)
+<<<<<<< Updated upstream
         dist.init_process_group(backend="cuda:nccl,cpu:gloo")
+=======
+        assert dist.is_hccl_available()
+        dist.init_process_group(backend="hccl")
+>>>>>>> Stashed changes
 
     def prepare_trainer_input_single(
         model_params_single, data_dict_single, rank=0, seed=None

deepmd/pt/infer/deep_eval.py

@@ -9,7 +9,9 @@
 )
 
 import numpy as np
-import torch
+import torch 
+import torch_npu 
+# from torch_npu.contrib import transfer_to_npu
 
 from deepmd.dpmodel.common import PRECISION_DICT as NP_PRECISION_DICT
 from deepmd.dpmodel.output_def import (

deepmd/pt/infer/inference.py

@@ -4,7 +4,9 @@
     deepcopy,
 )
 
-import torch
+import torch 
+import torch_npu 
+# from torch_npu.contrib import transfer_to_npu
 
 from deepmd.pt.model.model import (
     get_model,

deepmd/pt/loss/denoise.py

@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import torch
+import torch 
+import torch_npu 
+# from torch_npu.contrib import transfer_to_npu
 import torch.nn.functional as F
 
 from deepmd.pt.loss.loss import (

deepmd/pt/loss/dos.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
-import torch
+import torch 
+import torch_npu 
+# from torch_npu.contrib import transfer_to_npu
 
 from deepmd.pt.loss.loss import (
     TaskLoss,

deepmd/pt/loss/ener.py

@@ -3,7 +3,9 @@
     Optional,
 )
 
-import torch
+import torch 
+import torch_npu 
+# from torch_npu.contrib import transfer_to_npu
 import torch.nn.functional as F
 
 from deepmd.pt.loss.loss import (

deepmd/pt/loss/ener_spin.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
-import torch
+import torch 
+import torch_npu 
+# from torch_npu.contrib import transfer_to_npu
 import torch.nn.functional as F
 
 from deepmd.pt.loss.loss import (

deepmd/pt/loss/loss.py

@@ -7,7 +7,9 @@
     NoReturn,
 )
 
-import torch
+import torch 
+import torch_npu 
+# from torch_npu.contrib import transfer_to_npu
 
 from deepmd.utils.data import (
     DataRequirementItem,