diff --git a/mindspore/ccsrc/backend/common/pass/communication_op_fusion.cc b/mindspore/ccsrc/backend/common/pass/communication_op_fusion.cc index 4162114a7e7..216c3496fa5 100644 --- a/mindspore/ccsrc/backend/common/pass/communication_op_fusion.cc +++ b/mindspore/ccsrc/backend/common/pass/communication_op_fusion.cc @@ -35,7 +35,7 @@ namespace { constexpr auto kAttrDefaultGroup = "default_group"; constexpr auto kAttrDefaultOp = "default_op"; constexpr size_t kAlignSize = 2 << 9; -constexpr int64_t DEFAULT_THRESHOLD_MB_TO_BYTE = 262144; +constexpr int64_t kDefaultThresholdMb2Byte = 262144; kernel::KernelBuildInfoPtr GenerateKernelBuildInfo(const CommunicationOpInfo &communication_op_info, size_t start_index, size_t end_index) { @@ -120,13 +120,9 @@ void CheckInputs(const std::vector &fusion_inputs) { } } -bool CheckSegments(size_t segments, size_t communication_op_node_size, const std::vector *segment_index) { +bool CheckSegments(size_t communication_op_node_size, const std::vector *segment_index) { MS_EXCEPTION_IF_NULL(segment_index); - if (segments >= communication_op_node_size) { - MS_LOG(INFO) << "fusion not changed: segment_num=" << segments - << ", communication_op_node_size=" << communication_op_node_size; - return false; - } + auto segments = segment_index->size(); if (segment_index->at(segments - 1) != communication_op_node_size - 1) { MS_LOG(EXCEPTION) << "the last segment index is invalid."; } @@ -140,15 +136,13 @@ bool CheckSegments(size_t segments, size_t communication_op_node_size, const std } } // namespace -bool CommunicationOpFusion::GetSplitSegments(const CommunicationOpInfo &communication_op_info, size_t *segment_num, +bool CommunicationOpFusion::GetSplitSegments(const CommunicationOpInfo &communication_op_info, std::vector *segment_index, const std::string &group) const { - MS_EXCEPTION_IF_NULL(segment_num); MS_EXCEPTION_IF_NULL(segment_index); size_t communication_op_node_size = communication_op_info.communication_op_nodes.size(); MS_LOG(INFO) << "graph " << op_name_ << " node size " << communication_op_node_size; if (op_name_ == kHcomSendOpName || op_name_ == kReceiveOpName) { - *segment_num = 1; if (communication_op_node_size == 0) { return false; } @@ -163,7 +157,6 @@ bool CommunicationOpFusion::GetSplitSegments(const CommunicationOpInfo &communic split_indices = parallel_context->GetAllReduceFusionSplitIndices(group); } - size_t segments = 0; if (!split_indices.empty()) { uint32_t last_index = 0; for (size_t i = 0; i < split_indices.size(); ++i) { @@ -178,52 +171,44 @@ bool CommunicationOpFusion::GetSplitSegments(const CommunicationOpInfo &communic } segment_index->push_back(index); last_index = index; - segments++; } if (last_index != communication_op_node_size - 1) { segment_index->push_back(communication_op_node_size - 1); - segments++; } } else { - segments = groups_; - for (size_t i = 0; i < segments - 1; ++i) { - segment_index->push_back((i + 1) * (communication_op_node_size / segments) - 1); + for (size_t i = 0; i < groups_ - 1; ++i) { + segment_index->push_back((i + 1) * (communication_op_node_size / groups_) - 1); } segment_index->push_back(communication_op_node_size - 1); } auto parallel_mode = parallel_context->parallel_mode(); if (parallel_mode == parallel::kDataParallel && op_name_ == kAllReduceOpName) { - auto threshold = parallel_context->fusion_threshold_mb(); - GetAllReduceSplitSegment(communication_op_info.communication_op_nodes, threshold, &segments, segment_index); + auto threshold = parallel_context->dp_fusion_threshold_mb(); + GetAllReduceSplitSegment(communication_op_info.communication_op_nodes, threshold, segment_index); } - *segment_num = segments; - return CheckSegments(segments, communication_op_node_size, segment_index); + return CheckSegments(communication_op_node_size, segment_index); } void CommunicationOpFusion::GetAllReduceSplitSegment(const std::vector &nodes, int64_t threshold, - size_t *segment, std::vector *segment_index) const { - MS_EXCEPTION_IF_NULL(segment); + std::vector *segment_index) const { MS_EXCEPTION_IF_NULL(segment_index); if (threshold <= 0) { - MS_LOG(WARNING) << "Split threshold must be larger than 0, but got " << threshold; + MS_LOG(WARNING) << "Split threshold is " << threshold << ". AllReduce nodes will take default fusion strategy."; return; } - threshold *= DEFAULT_THRESHOLD_MB_TO_BYTE; + threshold *= kDefaultThresholdMb2Byte; std::vector real_segment_index; size_t start_index = 0; - size_t segment_num = 0; - for (size_t i = 0; i < segment_index->size(); i++) { - auto index = segment_index->at(i); + for (auto index : *segment_index) { if (index >= nodes.size()) { MS_LOG(WARNING) << "split index is greater than or equal to total gradient's number " << nodes.size(); continue; } size_t accumulate = 0; - for (size_t j = start_index; j <= index; j++) { + for (size_t j = start_index; j <= index; ++j) { auto tensor_size = AnfAlgo::GetOutputTensorMemSize(nodes[j], 0); if (accumulate + tensor_size > LongToSize(threshold)) { real_segment_index.push_back(j); - segment_num++; accumulate = 0; } else { accumulate += tensor_size; @@ -231,12 +216,10 @@ void CommunicationOpFusion::GetAllReduceSplitSegment(const std::vector } if (accumulate != 0) { real_segment_index.push_back(index); - segment_num++; } start_index = index + 1; } *segment_index = std::move(real_segment_index); - *segment = segment_num; } // Hard coded Load(%paraxxx, cnode()) to Load(%paraxxx, U) to prevent @@ -446,13 +429,13 @@ AnfNodePtr CommunicationOpFusion::CreateFusedCommunicationOp(const FuncGraphPtr } bool CommunicationOpFusion::DoFusion(const FuncGraphPtr &func_graph, const CommunicationOpInfo &communication_op_info, - size_t segment_num, const std::vector &segment_index) const { + const std::vector &segment_index) const { MS_EXCEPTION_IF_NULL(func_graph); auto manager = func_graph->manager(); MS_EXCEPTION_IF_NULL(manager); bool changed = false; size_t start_index = 0; - for (size_t segment_idx = 0; segment_idx < segment_num; ++segment_idx) { + for (size_t segment_idx = 0; segment_idx < segment_index.size(); ++segment_idx) { size_t end_index = segment_index.at(segment_idx); if (end_index - start_index < 1) { start_index = end_index + 1; @@ -534,10 +517,9 @@ bool CommunicationOpFusion::Run(const FuncGraphPtr &func_graph) { common::AnfAlgo::GetNodeAttr(b, kAttrIndex); }); } - size_t segment_num = 0; std::vector segment_index; - if (GetSplitSegments(it.second, &segment_num, &segment_index, it.first)) { - if (DoFusion(func_graph, it.second, segment_num, segment_index)) { + if (GetSplitSegments(it.second, &segment_index, it.first)) { + if (DoFusion(func_graph, it.second, segment_index)) { changed = true; } } diff --git a/mindspore/ccsrc/backend/common/pass/communication_op_fusion.h b/mindspore/ccsrc/backend/common/pass/communication_op_fusion.h index f6448de635f..589c5f6857f 100644 --- a/mindspore/ccsrc/backend/common/pass/communication_op_fusion.h +++ b/mindspore/ccsrc/backend/common/pass/communication_op_fusion.h @@ -40,15 +40,15 @@ class CommunicationOpFusion : public Pass { bool Run(const FuncGraphPtr &graph) override; private: - bool DoFusion(const FuncGraphPtr &func_graph, const CommunicationOpInfo &communication_op_info, size_t segment_num, + bool DoFusion(const FuncGraphPtr &func_graph, const CommunicationOpInfo &communication_op_info, const std::vector &segment_index) const; - void GetAllReduceSplitSegment(const std::vector &nodes, int64_t threshold, size_t *segment, + void GetAllReduceSplitSegment(const std::vector &nodes, int64_t threshold, std::vector *segment_index) const; AnfNodePtr CreateFusedCommunicationOp(const FuncGraphPtr &func_graph, const CommunicationOpInfo &communication_op_info, size_t start_index, size_t end_index) const; - bool GetSplitSegments(const CommunicationOpInfo &communication_op_info, size_t *segment_num, - std::vector *segment_index, const std::string &group) const; + bool GetSplitSegments(const CommunicationOpInfo &communication_op_info, std::vector *segment_index, + const std::string &group) const; std::string op_name_; size_t groups_ = 1; }; diff --git a/mindspore/ccsrc/include/common/utils/parallel_context.h b/mindspore/ccsrc/include/common/utils/parallel_context.h index 77a75620939..4f1a0390d41 100644 --- a/mindspore/ccsrc/include/common/utils/parallel_context.h +++ b/mindspore/ccsrc/include/common/utils/parallel_context.h @@ -54,6 +54,7 @@ constexpr char kFusionAuto[] = "auto"; constexpr char kFusionSize[] = "size"; constexpr char kFusionIndex[] = "index"; constexpr int64_t kFusionThreshold = 64; +constexpr int64_t kDataParallelFusionThreshold = 0; class COMMON_EXPORT ParallelContext { public: @@ -83,6 +84,8 @@ class COMMON_EXPORT ParallelContext { void set_fusion_threshold_mb(int64_t fusion_threshold); int64_t fusion_threshold_mb() const { return fusion_threshold_mb_; } + int64_t dp_fusion_threshold_mb() const { return dp_fusion_threshold_mb_; } + void set_allgather_fusion_threshold_mb(int64_t allgather_fusion_threshold); int64_t allgather_fusion_threshold_mb() const { return allgather_fusion_threshold_mb_; } @@ -185,6 +188,7 @@ class COMMON_EXPORT ParallelContext { bool gradient_fp32_sync_; bool loss_repeated_mean_; int64_t device_num_; + int64_t dp_fusion_threshold_mb_; int64_t fusion_threshold_mb_; int64_t allgather_fusion_threshold_mb_; int64_t reducescatter_fusion_threshold_mb_; // reducescatter diff --git a/mindspore/ccsrc/utils/parallel_context.cc b/mindspore/ccsrc/utils/parallel_context.cc index 4feba671e32..45279b2e763 100644 --- a/mindspore/ccsrc/utils/parallel_context.cc +++ b/mindspore/ccsrc/utils/parallel_context.cc @@ -74,6 +74,7 @@ void ParallelContext::Reset() { parallel_optimizer_threshold_ = -1; sharding_propagation_ = false; dataset_strategy_.clear(); + dp_fusion_threshold_mb_ = kDataParallelFusionThreshold; fusion_threshold_mb_ = kFusionThreshold; allgather_fusion_threshold_mb_ = kFusionThreshold; reducescatter_fusion_threshold_mb_ = kFusionThreshold; @@ -88,6 +89,7 @@ void ParallelContext::set_device_num(int64_t device_num) { void ParallelContext::set_fusion_threshold_mb(int64_t fusion_threshold) { fusion_threshold_mb_ = fusion_threshold; + dp_fusion_threshold_mb_ = fusion_threshold; fusion_threshold_is_set_ = true; enable_all_reduce_fusion_ = true; }