diff --git a/paddle/fluid/framework/details/multi_devices_graph_pass.cc b/paddle/fluid/framework/details/multi_devices_graph_pass.cc index 24977aabdacb265694d0130ccc28b5f2e05ce4f5..e0246740dd755b22347e975c2d2598a2f4587baf 100644 --- a/paddle/fluid/framework/details/multi_devices_graph_pass.cc +++ b/paddle/fluid/framework/details/multi_devices_graph_pass.cc @@ -731,6 +731,7 @@ bool DistSSAGraphBuilder::DealWithSpecialOp(ir::Graph *result, } } insert_op = true; + need_broadcast_var_ = true; } else if (OpHaveRole(*node, OpRole::kDist)) { int op_dev_id = CreateDistTrainOp(result, node); if (node->Op()->Type() == "concat") { @@ -925,14 +926,17 @@ void DistSSAGraphBuilder::InsertCollectiveOp(ir::Graph *result, void DistSSAGraphBuilder::InsertPostprocessOps(ir::Graph *result) const { // only GPU reduce mode need to broadcast parameters to each device. - if (UseGPU() && strategy_.reduce_ == BuildStrategy::ReduceStrategy::kReduce) { - if (strategy_.fuse_broadcast_op_) { - CreateFusedBroadcastOp(result, bcast_var_name_set_); - } else { - for (size_t dev_id = 0; dev_id < bcast_var_name_set_.size(); ++dev_id) { - auto &to_bcast_set = bcast_var_name_set_[dev_id]; - for (auto &bcast_name : to_bcast_set) { - CreateBroadcastOp(result, bcast_name, dev_id); + if (UseGPU()) { + if (need_broadcast_var_ || + strategy_.reduce_ == BuildStrategy::ReduceStrategy::kReduce) { + if (strategy_.fuse_broadcast_op_) { + CreateFusedBroadcastOp(result, bcast_var_name_set_); + } else { + for (size_t dev_id = 0; dev_id < bcast_var_name_set_.size(); ++dev_id) { + auto &to_bcast_set = bcast_var_name_set_[dev_id]; + for (auto &bcast_name : to_bcast_set) { + CreateBroadcastOp(result, bcast_name, dev_id); + } } } } diff --git a/paddle/fluid/framework/details/multi_devices_graph_pass.h b/paddle/fluid/framework/details/multi_devices_graph_pass.h index 21f85dc828687ba6224e5f353133546bedaa4a0a..6d4386538ea7d0cc318647c92282af9d598fa699 100644 --- a/paddle/fluid/framework/details/multi_devices_graph_pass.h +++ b/paddle/fluid/framework/details/multi_devices_graph_pass.h @@ -174,6 +174,7 @@ class DistSSAGraphBuilder : public BalanceVarSSAGraphBuilder { int CreateDistTrainOp(ir::Graph *result, ir::Node *node) const; mutable std::vector> bcast_var_name_set_; + mutable bool need_broadcast_var_{false}; }; std::unordered_set &MultiDevSSAGraphBuilder();