add split_norm_comm to split global comm to pp_comm & mp_comm (#72035)

Wennie396 · web-flow · commit f1ebd963aa46 · 2025-04-09T15:29:54.000+08:00
diff --git a/paddle/fluid/framework/distributed_strategy.proto b/paddle/fluid/framework/distributed_strategy.proto
@@ -115,6 +115,7 @@ message HybridConfig {
   optional PpConfig pp_configs = 7;
   optional DygraphShardingConfig sharding_configs = 8;
   optional bool enable_optimizer_timer = 9 [ default = false ];
+  optional bool split_norm_comm = 10 [ default = false ];
 }
 
 message AMPConfig {
diff --git a/python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/hybrid_parallel_optimizer.py b/python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/hybrid_parallel_optimizer.py
@@ -46,12 +46,13 @@
 
 
 class HybridParallelClipGrad:
-    def __init__(self, clip, hcg, timers=None):
+    def __init__(self, clip, hcg, split_norm_comm=False, timers=None):
         self._clip = clip
         self._hcg = hcg
         self.not_sharding_stage1 = True
         self._timers = timers
         self.processed_steps = 0
+        self.split_norm_comm = split_norm_comm
 
     def _global_norm(self, global_norm_var_dist, global_norm_var_not_dist):
         if self.processed_steps < g_profile_optimizer_details_steps:
@@ -82,7 +83,7 @@ def _global_norm(self, global_norm_var_dist, global_norm_var_not_dist):
             # dist should reduce among sharding group、mp group、pp group
 
             # the else branch would suffice, but this branch remains here for number precision backward compatibility
-            if not (dp_flag and sharding_flag):
+            if not (dp_flag and sharding_flag) and not self.split_norm_comm:
                 paddle.distributed.all_reduce(
                     global_norm_var_dist,
                     group=self._hcg.get_check_parallel_group(sharding_flag),
@@ -316,6 +317,8 @@ def __init__(self, optimizer, hcg, strategy):
 
         self._sep_enable = self._hcg.get_sep_parallel_world_size() > 1
 
+        split_norm_comm = strategy.hybrid_configs["split_norm_comm"]
+
         if (
             isinstance(self._inner_opt._grad_clip, ClipGradByGlobalNorm)
             and not self._use_dp_mode
@@ -347,19 +350,22 @@ def __init__(self, optimizer, hcg, strategy):
                 > 0
             ):
                 inner_opt._grad_clip = HybridParallelClipGrad(
-                    inner_opt._grad_clip, hcg, self._timers
+                    inner_opt._grad_clip, hcg, split_norm_comm, self._timers
                 )
             else:
                 inner_opt._grad_clip = HybridParallelClipGrad(
-                    inner_opt._grad_clip, hcg, self._timers
+                    inner_opt._grad_clip, hcg, split_norm_comm, self._timers
                 )
                 if inner_opt._parameter_list and isinstance(
                     inner_opt._parameter_list[0], dict
                 ):
                     for item in inner_opt._param_groups:
                         if "grad_clip" in item.keys():
                             item["grad_clip"] = HybridParallelClipGrad(
-                                inner_opt._grad_clip, hcg, self._timers
+                                inner_opt._grad_clip,
+                                hcg,
+                                split_norm_comm,
+                                self._timers,
                             )
         self.processed_steps = 0
 
diff --git a/python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py b/python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py
@@ -2531,7 +2531,6 @@ def forward_backward_pipeline(
 class OffloadQueue(queue.Queue):
     def __init__(self, offload=False, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        logger.info(f"Using OffloadQueue: {offload}")
         self.offload = offload
 
     def put(self, tensor, *args, **kwargs):

Original file line number	Diff line number	Diff line change
`@@ -115,6 +115,7 @@ message HybridConfig {`
`115`	`115`	`optional PpConfig pp_configs = 7;`
`116`	`116`	`optional DygraphShardingConfig sharding_configs = 8;`
`117`	`117`	`optional bool enable_optimizer_timer = 9 [ default = false ];`
	`118`	`+ optional bool split_norm_comm = 10 [ default = false ];`
`118`	`119`	`}`
`119`	`120`
`120`	`121`	`message AMPConfig {`