[Cherry-pick] Add config for broadcast data in model parallel (#57843)

* add config for broadcast data for mp * fix mp bug (#58037) --------- Co-authored-by: Yuang Liu <liuyuang@baidu.com>
PaddlePaddle · Oct 17, 2023 · 1be4c8d · 1be4c8d
1 parent ff9d102
commit 1be4c8d
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/paddle/fluid/framework/distributed_strategy.proto b/paddle/fluid/framework/distributed_strategy.proto
@@ -56,6 +56,8 @@ message MpConfig {
     optional bool sync_grad= 2 [ default = false ];
     optional bool sync_moment= 3 [ default = false ];
     optional string sync_mode= 4 [ default = 'broadcast' ];
+    // Broadcast mp input data
+    optional bool need_broadcast_data=8 [default = true];
 }
 
 message PpConfig {

diff --git a/python/paddle/distributed/fleet/meta_parallel/tensor_parallel.py b/python/paddle/distributed/fleet/meta_parallel/tensor_parallel.py
@@ -42,5 +42,10 @@ def _prepare_for_model(self):
         logger.info("mp's parameters is ready")
 
     def _pre_forward(self, *inputs, **kwargs):
-        logger.debug("mp start broadcast input data")
-        return broadcast_input_data(self._hcg, *inputs, **kwargs)
+        need_broadcast_data = True
+        if self._strategy is not None:
+            mp_configs = self._strategy.hybrid_configs["mp_configs"]
+            need_broadcast_data = mp_configs.need_broadcast_data
+        if need_broadcast_data:
+            logger.debug("mp start broadcast input data")
+            return broadcast_input_data(self._hcg, *inputs, **kwargs)