sharding stage 1 check diff lr and use param decay fn (#59537)

Author: FeixLiu

python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/dygraph_sharding_optimizer.py

@@ -14,6 +14,7 @@
 
 ######
 import os
+import warnings
 from distutils.util import strtobool
 from functools import reduce
 
@@ -115,6 +116,23 @@ def __init__(self, optimizer, hcg):
             self._set_inner_opt_attr('_parameter_list', local_params)
             self._set_inner_opt_attr('_param_groups', local_params)
         else:
+            if self.fuse_optimizer:
+                lr = None
+                for param in self._origin_parameter_list:
+                    if hasattr(param, "optimize_attr"):
+                        param_lr = param.optimize_attr['learning_rate']
+                        if lr is None:
+                            lr = param_lr
+                        elif lr != param_lr:
+                            warnings.warn(
+                                "Parameters have different learning rate, "
+                                "won't do fusion on the optimizer."
+                            )
+                            self.fuse_optimizer = False
+                            break
+            self.origin_decay_param_fun = getattr(
+                self._inner_opt, '_apply_decay_param_fun', None
+            )
             self._tensor_fusion()
 
             decay_params = [
@@ -138,10 +156,7 @@ def __init__(self, optimizer, hcg):
                 # Without comm overlap, all grads will be communicated after check_finite,
                 # which means each sharding rank should do check_finite to all grads.
                 self._local_parameter_list = local_fused_params
-            origin_decay_param_fun = getattr(
-                self._inner_opt, '_apply_decay_param_fun', None
-            )
-            if origin_decay_param_fun is not None:
+            if self.origin_decay_param_fun is not None:
                 self._set_inner_opt_attr(
                     '_apply_decay_param_fun', apply_decay_param_fun
                 )
@@ -191,6 +206,7 @@ def _tensor_fusion(self):
                 dst=dst,
                 acc_step=self.accumulate_steps,
                 scale_after_comm=False,
+                apply_decay_param_fun=self.origin_decay_param_fun,
             )
             if self.comm_overlap:
                 self._comm_buffers += all_buffer

python/paddle/distributed/fleet/utils/tensor_fusion_helper.py

@@ -550,6 +550,7 @@ def _fused_parameters_impl(
     dst=-1,
     acc_step=1,
     scale_after_comm=False,
+    apply_decay_param_fun=None,
 ):
     param_groups = []
     attrs = []
@@ -579,7 +580,9 @@ def _fused_parameters_impl(
         other_params = []
 
         for param in params:
-            if not any(nd in param.name for nd in ["bias", "norm", "b_0"]):
+            if apply_decay_param_fun is not None and apply_decay_param_fun(
+                param.name
+            ):
                 decay_params.append(param)
             else:
                 other_params.append(param)
@@ -632,6 +635,7 @@ def fused_parameters(
     acc_step=1,
     scale_after_comm=False,
     group_params=False,
+    apply_decay_param_fun=None,
 ):
     """
     Fuse gradients. Fuse parameters if be enabled. Prepare for comm overlap if be enabled.
@@ -645,6 +649,7 @@ def fused_parameters(
     :param fuse_param: fuse param or not
     :param scale_after_comm: if enable comm overlap, specify the location of grad scale
     :param group_params: the format of the input parameters is param group
+    :param apply_decay_param_fun: the funtion to filter decay param
     :return: param storage if fused, comm buffers if comm overlap, param groups if use group params
     """
     if act is None:
@@ -690,6 +695,7 @@ def fused_parameters(
                 dst=dst,
                 acc_step=acc_step,
                 scale_after_comm=scale_after_comm,
+                apply_decay_param_fun=apply_decay_param_fun,
             )
             if comm_overlap:
                 comm_buffers.extend(group_all_buffers)
@@ -709,6 +715,7 @@ def fused_parameters(
             dst=dst,
             acc_step=acc_step,
             scale_after_comm=scale_after_comm,
+            apply_decay_param_fun=apply_decay_param_fun,
         )
 
         return decay_fused, all_fused, all_buffers

test/collective/fleet/hybrid_parallel_sharding_model_with_fusion_amp.py

@@ -85,12 +85,18 @@ def build_optimizer(self, model):
         )
         return optimizer
 
-    def build_model_optimizer(self):
+    def build_model_optimizer(self, diff_lr=False):
         model = SimpleDPNet(vocab_size, hidden_size, inner_size, output_size)
         optimizer = self.build_optimizer(model)
         model, optimizer = paddle.amp.decorate(
             model, optimizers=optimizer, level="O2", dtype="float16"
         )
+        if diff_lr:
+            for param in model.parameters():
+                if 'w' in param.name:
+                    param.optimize_attr = {"learning_rate": 1.0}
+                else:
+                    param.optimize_attr = {"learning_rate": 2.0}
         scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
         scaler = fleet.distributed_scaler(scaler)
         model = fleet.distributed_model(model)
@@ -109,8 +115,13 @@ def sharding_model(self):
             scaler.update()
             optimizer.clear_grad()
 
+    def sharding_different_lr(self):
+        model, optimizer, scaler = self.build_model_optimizer(diff_lr=True)
+        assert optimizer._inner_opt.fuse_optimizer is False
+
     def test_sharding_adam(self):
         self.sharding_model()
+        self.sharding_different_lr()
 
 
 if __name__ == "__main__":