fix ddp

Author: iProzd

deepmd/pt/train/training.py

@@ -693,10 +693,11 @@ def warm_up_linear(step: int, warmup_steps: int) -> float:
         # TODO add optimizers for multitask
         # author: iProzd
         if self.opt_type in ["Adam", "AdamW"]:
+            base_model = self._get_wrapper_module()
             parameter_list = [
-                {"params": self.wrapper.model.parameters()},
-                {"params": self.wrapper.loss.parameters(), "lr": 1e-4},
-            ]
+                {"params": base_model.model.parameters()},
+                {"params": base_model.loss.parameters(), "lr": 1e-4},
+            ]  # set a smaller lr for loss parameters
             if self.opt_type == "Adam":
                 self.optimizer = torch.optim.Adam(
                     # self.wrapper.parameters(),
@@ -1373,6 +1374,12 @@ def print_on_training(
         fout.write(print_str)
         fout.flush()
 
+    def _get_wrapper_module(self) -> torch.nn.Module:
+        """Compatible with DDP and normal model."""
+        if isinstance(self.wrapper, DDP):
+            return self.wrapper.module
+        return self.wrapper
+
 
 def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]:
     additional_data_requirement = []