[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

deepmd/pd/loss/ener.py

@@ -4,6 +4,7 @@
 )
 
 import paddle
+import paddle.distributed as dist
 import paddle.nn.functional as F
 
 from deepmd.pd.loss.loss import (
@@ -21,7 +22,6 @@
 from deepmd.utils.version import (
     check_version_compatibility,
 )
-import paddle.distributed as dist
 
 
 def custom_huber_loss(predictions, targets, delta=1.0):
@@ -206,10 +206,9 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
             find_energy = label.get("find_energy", 0.0)
             pref_e = pref_e * find_energy
             if not self.use_l1_all:
-
                 logit = energy_pred - energy_label
                 logit = dist.reshard(tmp, tmp.process_mesh, [dist.Replicate()])
-                
+
                 l2_ener_loss = paddle.mean(paddle.square(logit))
                 if not self.inference:
                     more_loss["l2_ener_loss"] = self.display_if_exist(
@@ -264,7 +263,7 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
             force_label = label["force"]
             diff_f = (force_label - force_pred).reshape([-1])
             diff_f = dist.reshard(diff_f, diff_f.process_mesh, [dist.Replicate()])
-            
+
             if self.relative_f is not None:
                 force_label_3 = force_label.reshape([-1, 3])
                 norm_f = force_label_3.norm(axis=1, keepdim=True) + self.relative_f

deepmd/pd/model/model/make_model.py

@@ -1,9 +1,14 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import functools
 from typing import (
     Optional,
 )
 
 import paddle
+import paddle.distributed as dist
+from paddle.distributed import (
+    fleet,
+)
 
 from deepmd.dpmodel import (
     ModelOutputDef,
@@ -37,9 +42,6 @@
 from deepmd.utils.path import (
     DPPath,
 )
-import paddle.distributed as dist
-from paddle.distributed import fleet
-import functools
 
 
 def make_model(T_AtomicModel: type[BaseAtomicModel]):
@@ -182,10 +184,12 @@ def forward_common(
             #     box=bb,
             # )
             wrapped_func_1 = dist.local_map(
-                func=lambda a,b,c: extend_input_and_build_neighbor_list(a,b,self.get_rcut(), self.get_sel(), True, c),
+                func=lambda a, b, c: extend_input_and_build_neighbor_list(
+                    a, b, self.get_rcut(), self.get_sel(), True, c
+                ),
                 in_placements=[ele.placements for ele in [cc, atype, bb]],
                 out_placements=[[dist.Shard(0)] for _ in range(4)],
-                process_mesh=fleet.auto.get_mesh()
+                process_mesh=fleet.auto.get_mesh(),
             )
 
             (
@@ -209,11 +213,19 @@ def forward_common(
             # )
 
             wrapped_func_2 = dist.local_map(
-                func=functools.partial(self.forward_common_lower, do_atomic_virial=do_atomic_virial, fparam=fp, aparam=ap),
-                in_placements=[ele.placements for ele in [extended_coord, extended_atype, nlist, mapping]],
+                func=functools.partial(
+                    self.forward_common_lower,
+                    do_atomic_virial=do_atomic_virial,
+                    fparam=fp,
+                    aparam=ap,
+                ),
+                in_placements=[
+                    ele.placements
+                    for ele in [extended_coord, extended_atype, nlist, mapping]
+                ],
                 out_placements=[[dist.Shard(0)] for _ in range(6)],
                 process_mesh=fleet.auto.get_mesh(),
-                reshard_inputs=True
+                reshard_inputs=True,
             )
             model_predict_lower = wrapped_func_2(
                 extended_coord,

deepmd/pd/train/training.py

@@ -19,16 +19,13 @@
 from paddle.distributed import (
     fleet,
 )
-from paddle.distributed.fleet.utils import hybrid_parallel_util as hpu
 from paddle.framework import (
     core,
 )
 from paddle.io import (
     DataLoader,
 )
-import paddle.distributed as dist
-from paddle.distributed import fleet
-import functools
+
 from deepmd.common import (
     symlink_prefix_files,
 )
@@ -103,7 +100,6 @@ def __init__(
         Args:
         - config: The Dict-like configuration with training options.
         """
-        from paddle.distributed import fleet
         mesh_dims = [("dp", 32)]
         fleet.auto.create_mesh(mesh_dims)
         fleet.init(is_collective=True)
@@ -753,7 +749,7 @@ def step(_step_id, task_key="Default") -> None:
                     if self.world_size > 1
                     else contextlib.nullcontext
                 )
-                
+
                 # with sync_context():
                 #     with nvprof_context(enable_profiling, "Forward pass"):
                 #         model_pred, loss, more_loss = self.wrapper(
@@ -772,11 +768,19 @@ def step(_step_id, task_key="Default") -> None:
                 #     hpu.fused_allreduce_gradients(list(self.wrapper.parameters()), None)
 
                 with nvprof_context(enable_profiling, "Forward pass"):
-                    for __key in ('coord', 'atype', 'box'):
-                        input_dict[__key] = dist.shard_tensor(input_dict[__key], mesh=dist.get_mesh(), placements=[dist.Shard(0)])
+                    for __key in ("coord", "atype", "box"):
+                        input_dict[__key] = dist.shard_tensor(
+                            input_dict[__key],
+                            mesh=dist.get_mesh(),
+                            placements=[dist.Shard(0)],
+                        )
                     for __key, _ in label_dict.items():
                         if isinstance(label_dict[__key], paddle.Tensor):
-                            label_dict[__key] = dist.shard_tensor(label_dict[__key], mesh=dist.get_mesh(), placements=[dist.Shard(0)])
+                            label_dict[__key] = dist.shard_tensor(
+                                label_dict[__key],
+                                mesh=dist.get_mesh(),
+                                placements=[dist.Shard(0)],
+                            )
                     model_pred, loss, more_loss = self.wrapper(
                         **input_dict,
                         cur_lr=paddle.full([], pref_lr, DEFAULT_PRECISION),

examples/water/dpa3/run.sh

@@ -14,4 +14,4 @@ HDFS_USE_FILE_LOCKING=0 python -m paddle.distributed.launch --gpus="0,1,2,3,4,5,
 # python -m paddle.distributed.launch \
 #    --gpus=0,1,2,3 \
 #    --ips=10.67.200.17,10.67.200.11,10.67.200.13,10.67.200.15 \
-#    dp --pd train input_torch.json -l dp_train.log
+#    dp --pd train input_torch.json -l dp_train.log