[LLM] Support master grad on dp. (#6650)

* support master grad on dp.
PaddlePaddle · liuzhenhai93 · May 29, 2023 · Jun 22, 2023 · Jun 22, 2023 · Jun 23, 2023
commit ed8ca957999248177dec02b7f2c529c2a72d072d
diff --git a/paddlenlp/trainer/trainer.py b/paddlenlp/trainer/trainer.py
@@ -897,6 +897,11 @@ def train(
                 # sharding
                 # stage1. the same as ddp
                 # stage2. manualy collect gradient on dp group
+
+                hack_dp_master_grad = self.args.amp_master_grad and not self.args.use_hybrid_parallel
+                if hack_dp_master_grad:
+                    is_no_sync = False
+
                 if is_no_sync:
                     # Avoid unnecessary DDP synchronization since there will be no backward pass on this example.
                     with model.no_sync():
@@ -949,6 +954,10 @@ def train(
                     self.timers and self.timers("all-reduce").stop()
                     self.timers and self.timers("optimizer-step").start()
 
+                    # Case 3: hack dp with master_grad
+                    if hack_dp_master_grad and not (args.recompute and availiable_no_sync):
+                        fused_allreduce_gradients(list(model.parameters()), None)
+
                     # pipeline parallel mode,  handle gradient merge here
                     if args.pipeline_parallel_degree > 1 and enable_delay_scale_loss:
                         for p in model._layers.parameters():
@@ -1518,7 +1527,12 @@ def _wrap_model(self, model, training=True):
 
         # Multi-gpu training
         if self.args.world_size > 1 and not self.args.use_hybrid_parallel:
-            model = paddle.DataParallel(model)
+            if self.args.amp_master_grad:
+                mix_precision_utils.MixPrecisionLayer(model, dtype=self.amp_dtype)  # return value has no use
+                logger.warning("Note amp_master_grad using in dp is an experimental support!")
+                self.optimizer = mix_precision_utils.MixPrecisionOptimizer(self.optimizer)
+            else:
+                model = paddle.DataParallel(model)
             # Distributed training (should be after fp16 initialization)
 
         in_pipeline_parallel_mode = self.args.pipeline_parallel_degree > 1

diff --git a/paddlenlp/trainer/training_args.py b/paddlenlp/trainer/training_args.py
@@ -714,13 +714,18 @@ def __post_init__(self):
             self.use_hybrid_parallel = True
 
         if self.amp_master_grad:
-            if self.pipeline_parallel_degree <= 1 and self.tensor_parallel_degree <= 1:
-                raise ValueError(
-                    "Temporarily amp master grad only suport for tensor/pipeline parallel. please set amp_master_grad to False."
-                )
-            # if not (self.bf16 or self.fp16):
+            # if (
+            #     self.pipeline_parallel_degree <= 1 and self.tensor_parallel_degree <= 1
+            # ) or self.fp16_opt_level != "O2":
+            #     raise ValueError(
+            #         "Temporarily amp master grad only suport for tensor/pipeline parallel with fp16_opt_level O2. please set amp_master_grad to False."
+            #     )
+            # if not (self.bf16 or self.fp16) or self.fp16_opt_level != "O2":
             #     logger.warning("set amp_master_grad to false since amp is disabled.")
             #     self.amp_master_grad = False
+            if self.pipeline_parallel_degree <= 1 and self.tensor_parallel_degree <= 1 and len(self.sharding) > 1:
+                logger.warning("set amp_master_grad to false, not support pure sharding yet.")
+                self.amp_master_grad = False
 
         if self.use_hybrid_parallel:
             world_size = paddle.distributed.get_world_size()