Stage1 no fuse opt (#58790)

FeixLiu · web-flow · commit 420e12b82fa2 · 2023-11-09T07:23:57.000+08:00
diff --git a/paddle/fluid/framework/distributed_strategy.proto b/paddle/fluid/framework/distributed_strategy.proto
@@ -87,6 +87,7 @@ message DygraphShardingConfig {
   optional int32 accumulate_steps = 2 [ default = 1 ];
   optional bool comm_overlap = 3 [ default = false ];
   optional bool split_param = 4 [ default = false ];
+  optional bool fuse_optimizer = 5 [ default = true ];
 }
 
 message HybridConfig {
diff --git a/python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/dygraph_sharding_optimizer.py b/python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/dygraph_sharding_optimizer.py
@@ -74,6 +74,7 @@ def __init__(self, optimizer, hcg):
             )
         # the self._parameter_list holds the whole model paramters
         self._parameter_list = optimizer._parameter_list
+        self._origin_parameter_list = self._parameter_list
         self._inner_opt = optimizer
         self._hcg = hcg
         self._sharding_world_size = self._hcg.get_sharding_parallel_world_size()
@@ -90,6 +91,9 @@ def __init__(self, optimizer, hcg):
         self.comm_overlap = strategy.hybrid_configs[
             'sharding_configs'
         ].comm_overlap
+        self.fuse_optimizer = strategy.hybrid_configs[
+            'sharding_configs'
+        ].fuse_optimizer
         pp_overlap = strategy.hybrid_configs['pp_configs'].sharding_comm_overlap
         if self.tensor_fusion or self.comm_overlap:
             assert (
@@ -339,9 +343,14 @@ def step(self):
         # otherwise the self._inner_opt will only grad_clip the self._rank2params[self._sharding_rank] params
         # TODO(pangengzheng): remove the hacked grad_clip codes here when there is no diff in calculating global norm values in HybridParallelClipGrad compared to dp.
         origin_clip = self._inner_opt._grad_clip
-        if not isinstance(self._parameter_list[0], dict):
+        target_param_list = (
+            self._origin_parameter_list
+            if (not self.tensor_fusion or not self.fuse_optimizer)
+            else self._parameter_list
+        )
+        if not isinstance(target_param_list[0], dict):
             params_grads = []
-            for param in self._parameter_list:
+            for param in target_param_list:
                 if (
                     hasattr(param, "regularizer")
                     and param.regularizer is not None
@@ -362,7 +371,7 @@ def step(self):
                 self._set_inner_opt_attr('_grad_clip', None)
             rank_params = (
                 self._rank2params[self._sharding_rank]
-                if not self.tensor_fusion
+                if (not self.tensor_fusion or not self.fuse_optimizer)
                 else self._rank2fused[self._sharding_rank]
             )
             update_param_names = [p.name for p in rank_params]
diff --git a/test/collective/fleet/hybrid_parallel_sharding_model_with_fusion.py b/test/collective/fleet/hybrid_parallel_sharding_model_with_fusion.py
@@ -101,6 +101,7 @@ def setUp(self):
         self.strategy.hybrid_configs["sharding_configs"].tensor_fusion = True
         self.strategy.hybrid_configs["sharding_configs"].comm_overlap = True
         self.strategy.hybrid_configs["sharding_configs"].accumulate_steps = 1
+        self.strategy.hybrid_configs["sharding_configs"].fuse_optimizer = False
         fleet.init(is_collective=True, strategy=self.strategy)
         self.data = np.random.randint(
             0,

Original file line number	Diff line number	Diff line change
`@@ -87,6 +87,7 @@ message DygraphShardingConfig {`
`87`	`87`	`optional int32 accumulate_steps = 2 [ default = 1 ];`
`88`	`88`	`optional bool comm_overlap = 3 [ default = false ];`
`89`	`89`	`optional bool split_param = 4 [ default = false ];`
	`90`	`+ optional bool fuse_optimizer = 5 [ default = true ];`
`90`	`91`	`}`
`91`	`92`
`92`	`93`	`message HybridConfig {`