From 51238243b40c55e7ea35f94bd7735f36cb453725 Mon Sep 17 00:00:00 2001 From: lichenever Date: Tue, 6 Jul 2021 15:08:30 +0800 Subject: [PATCH] fix_pipeline_split_bug --- .../rec_core/rec_generate_strategy.cc | 2 +- .../graph_util/pipeline_split_utils.cc | 12 +--------- .../pipeline_transformer.cc | 2 +- .../frontend/parallel/step_auto_parallel.cc | 6 ++--- .../ccsrc/frontend/parallel/step_parallel.cc | 12 +++++----- .../ccsrc/frontend/parallel/step_parallel.h | 2 +- .../ccsrc/pipeline/jit/pipeline_split.cc | 22 +++++++++++++++++-- mindspore/ccsrc/pipeline/jit/pipeline_split.h | 2 ++ mindspore/nn/wrap/cell_wrapper.py | 2 +- .../pangu_alpha/src/pangu_alpha_wrapcell.py | 4 ++-- tests/ut/cpp/parallel/step_parallel_test.cc | 2 +- 11 files changed, 40 insertions(+), 28 deletions(-) diff --git a/mindspore/ccsrc/frontend/parallel/auto_parallel/rec_core/rec_generate_strategy.cc b/mindspore/ccsrc/frontend/parallel/auto_parallel/rec_core/rec_generate_strategy.cc index c48f1d2d93..f57913c4b9 100644 --- a/mindspore/ccsrc/frontend/parallel/auto_parallel/rec_core/rec_generate_strategy.cc +++ b/mindspore/ccsrc/frontend/parallel/auto_parallel/rec_core/rec_generate_strategy.cc @@ -49,7 +49,7 @@ void GenerateStrategy(const std::shared_ptr &graph, const std::vectorattrs(); if (StrategyFound(attrs)) { - StrategyPtr user_defined_stra = parallel::ExtractStrategy(attrs); + StrategyPtr user_defined_stra = parallel::ExtractStrategy(attrs[STRATEGY]); op->SetSelectedStrategyAndCost(user_defined_stra, op->selected_cost()); } // Set back to raw strategy for special node in predict/eval diff --git a/mindspore/ccsrc/frontend/parallel/graph_util/pipeline_split_utils.cc b/mindspore/ccsrc/frontend/parallel/graph_util/pipeline_split_utils.cc index 1457aa9528..0c8d691e77 100755 --- a/mindspore/ccsrc/frontend/parallel/graph_util/pipeline_split_utils.cc +++ b/mindspore/ccsrc/frontend/parallel/graph_util/pipeline_split_utils.cc @@ -60,10 +60,7 @@ void SetStridedSliceStrategy(const AnfNodePtr &node) { } auto cnode = node->cast(); MS_EXCEPTION_IF_NULL(cnode); - PrimitivePtr prim = GetValueNode(cnode->input(0)); - MS_EXCEPTION_IF_NULL(prim); int64_t dev_num = 1; - auto attrs_temp = prim->attrs(); std::vector shape_list = ExtractShape(cnode); if (shape_list.empty()) { MS_LOG(EXCEPTION) << "Failure:node " << cnode->ToString() << " failed to extract shape"; @@ -80,8 +77,7 @@ void SetStridedSliceStrategy(const AnfNodePtr &node) { elements.push_back(MakeValue(input_strategy)); } ValueTuplePtr strategy = std::make_shared(elements); - attrs_temp[STRATEGY] = strategy; - (void)prim->SetAttrs(attrs_temp); + cnode->AddPrimalAttr(STRATEGY, strategy); } void InsertVirtualAssignAdd(const std::pair &node_user, const FuncGraphManagerPtr &manager, @@ -467,12 +463,6 @@ void ParameterStartNode(const std::vector &all_nodes, const FuncGrap auto prim = GetCNodePrimitive(node); if (prim && prim->HasAttr(PARAMETER_START)) { auto micro = Micro(cnode, &node_users_map); - OperatorAttrs attrs_; - auto op = CreatOpInstance(attrs_, prim->name(), ""); - auto value_node = NewValueNode(op); - auto new_prim = GetValueNode(value_node)->cast(); - new_prim->SetAttrs(prim->attrs()); - manager->SetEdge(cnode, 0, value_node); cnode->AddPrimalAttr(MICRO, micro); cnode->AddPrimalAttr(PARAMETER_START, micro); } diff --git a/mindspore/ccsrc/frontend/parallel/pipeline_transformer/pipeline_transformer.cc b/mindspore/ccsrc/frontend/parallel/pipeline_transformer/pipeline_transformer.cc index a54951dce4..f91e3cf55a 100644 --- a/mindspore/ccsrc/frontend/parallel/pipeline_transformer/pipeline_transformer.cc +++ b/mindspore/ccsrc/frontend/parallel/pipeline_transformer/pipeline_transformer.cc @@ -295,7 +295,7 @@ OperatorInfoPtr PipelineTransformer::CreateOpInfo(const CNodePtr &cnode, int tup if (!StrategyFound(attrs)) { strategy = GenerateBatchParallelStrategy(op_info, prim); } else { - strategy = ExtractStrategy(attrs); + strategy = ExtractStrategy(attrs[STRATEGY]); } MS_EXCEPTION_IF_NULL(strategy); if (op_info->Init(strategy) == FAILED) { diff --git a/mindspore/ccsrc/frontend/parallel/step_auto_parallel.cc b/mindspore/ccsrc/frontend/parallel/step_auto_parallel.cc index ebc7aae043..8c5bbcf18b 100644 --- a/mindspore/ccsrc/frontend/parallel/step_auto_parallel.cc +++ b/mindspore/ccsrc/frontend/parallel/step_auto_parallel.cc @@ -242,12 +242,12 @@ void InitCostGraph() { } void SetStrategyToOperator(const OperatorInfoPtr &operator_info, const PrimitivePtr &prim, - const std::unordered_map &attrs, bool is_last_nodes, - StrategyMap *stra_map, const std::string &strategy_key_name) { + std::unordered_map attrs, bool is_last_nodes, StrategyMap *stra_map, + const std::string &strategy_key_name) { // In this case, the configured strategy should be extracted to help setting cost StrategyPtr strategyPtr; if (StrategyFound(attrs)) { - strategyPtr = parallel::ExtractStrategy(attrs); + strategyPtr = parallel::ExtractStrategy(attrs[STRATEGY]); } else { strategyPtr = (*stra_map)[strategy_key_name]; } diff --git a/mindspore/ccsrc/frontend/parallel/step_parallel.cc b/mindspore/ccsrc/frontend/parallel/step_parallel.cc index 7a29f6e10a..b289755521 100644 --- a/mindspore/ccsrc/frontend/parallel/step_parallel.cc +++ b/mindspore/ccsrc/frontend/parallel/step_parallel.cc @@ -1435,12 +1435,12 @@ OperatorInfoPtr NewOperatorInstance(const PrimitivePtr &prim, const PrimitiveAtt return operator_; } -StrategyPtr ExtractStrategy(std::unordered_map attrs) { - ValueTuplePtr var = attrs[STRATEGY]->cast(); +StrategyPtr ExtractStrategy(const ValuePtr &stra) { + ValueTuplePtr var = stra->cast(); StrategyPtr strategyPtr; int64_t stage_id = g_device_manager->stage_id(); - MS_LOG(INFO) << "Extract information: strategy " << attrs[STRATEGY]->ToString(); + MS_LOG(INFO) << "Extract information: strategy " << stra->ToString(); if (var == nullptr) { MS_LOG(EXCEPTION) << "Strategy value is nullptr"; } @@ -2180,12 +2180,14 @@ void ExtractInformation(const std::vector &all_nodes, bool is_traini } bool load_strategy_from_ckpt = StrategyCheckpoint::GetInstance().LoadCheckPointOn() && stra_map.find(strategy_key_name) != stra_map.end(); - if ((!StrategyFound(attrs) && !load_strategy_from_ckpt)) { + if ((!StrategyFound(attrs) && !load_strategy_from_ckpt) && !cnode->HasPrimalAttr(STRATEGY)) { MS_LOG(INFO) << "ExtractInformation: the strategy of node " << node->ToString() << " prim " << prim->name() << " is empty, using batch parallel"; strategyPtr = GenerateBatchParallelStrategy(operator_, prim); + } else if (cnode->HasPrimalAttr(STRATEGY)) { + strategyPtr = ExtractStrategy(cnode->GetPrimalAttr(STRATEGY)); } else if (StrategyFound(attrs)) { - strategyPtr = ExtractStrategy(attrs); + strategyPtr = ExtractStrategy(attrs[STRATEGY]); } else { strategyPtr = stra_map[strategy_key_name]; } diff --git a/mindspore/ccsrc/frontend/parallel/step_parallel.h b/mindspore/ccsrc/frontend/parallel/step_parallel.h index 2e9d745b81..f9b924bbee 100644 --- a/mindspore/ccsrc/frontend/parallel/step_parallel.h +++ b/mindspore/ccsrc/frontend/parallel/step_parallel.h @@ -99,7 +99,7 @@ OperatorInfoPtr NewOperatorInstance(const PrimitivePtr &prim, const PrimitiveAtt std::vector shape_list); // Extract strategy from attr -StrategyPtr ExtractStrategy(std::unordered_map attrs); +StrategyPtr ExtractStrategy(const ValuePtr &strategy); Shapes GetNodeShape(const AnfNodePtr &node); diff --git a/mindspore/ccsrc/pipeline/jit/pipeline_split.cc b/mindspore/ccsrc/pipeline/jit/pipeline_split.cc index 9290b5b78f..b7edc365f5 100644 --- a/mindspore/ccsrc/pipeline/jit/pipeline_split.cc +++ b/mindspore/ccsrc/pipeline/jit/pipeline_split.cc @@ -25,7 +25,7 @@ namespace mindspore { namespace pipeline { -static int64_t GetRank() { +std::string GetWorldGroup() { auto ms_context = MsContext::GetInstance(); MS_EXCEPTION_IF_NULL(ms_context); std::string world_group; @@ -37,6 +37,13 @@ static int64_t GetRank() { } else { MS_LOG(EXCEPTION) << "Invalid backend: " << backend; } + return world_group; +} + +static int64_t GetRank() { + auto ms_context = MsContext::GetInstance(); + MS_EXCEPTION_IF_NULL(ms_context); + auto world_group = GetWorldGroup(); int64_t global_rank = parallel::ParallelContext::GetInstance()->global_rank(); uint32_t rank_id = 0; if (!parallel::ParallelContext::GetInstance()->global_rank_is_set()) { @@ -75,7 +82,18 @@ bool PipelineSplit(const ResourcePtr &res) { auto manager = res->manager(); auto root = res->func_graph(); auto global_rank = GetRank(); - auto device_num = parallel::ParallelContext::GetInstance()->device_num(); + auto world_group = GetWorldGroup(); + uint32_t world_rank_size = 0; + int64_t device_num = 0; + if (!parallel::ParallelContext::GetInstance()->device_num_is_set()) { + if (!CommManager::GetInstance().GetRankSize(world_group, &world_rank_size)) { + MS_LOG(EXCEPTION) << "Get rank size failed"; + } + device_num = UintToInt(world_rank_size); + MS_LOG(INFO) << "Get device num from communication model, the device num is " << device_num; + } else { + device_num = parallel::ParallelContext::GetInstance()->device_num(); + } if (device_num < 1) { MS_LOG(EXCEPTION) << "Invalid device num: " << device_num; } diff --git a/mindspore/ccsrc/pipeline/jit/pipeline_split.h b/mindspore/ccsrc/pipeline/jit/pipeline_split.h index d148415f47..ac1b8431ae 100644 --- a/mindspore/ccsrc/pipeline/jit/pipeline_split.h +++ b/mindspore/ccsrc/pipeline/jit/pipeline_split.h @@ -17,11 +17,13 @@ #ifndef MINDSPORE_CCSRC_PIPELINE_JIT_PIPELINE_SPLIT_H_ #define MINDSPORE_CCSRC_PIPELINE_JIT_PIPELINE_SPLIT_H_ +#include #include "pipeline/jit/resource.h" namespace mindspore { namespace pipeline { bool PipelineSplit(const ResourcePtr &res); +std::string GetWorldGroup(); } // namespace pipeline } // namespace mindspore diff --git a/mindspore/nn/wrap/cell_wrapper.py b/mindspore/nn/wrap/cell_wrapper.py index 5dd385e18c..3e23789896 100644 --- a/mindspore/nn/wrap/cell_wrapper.py +++ b/mindspore/nn/wrap/cell_wrapper.py @@ -487,7 +487,7 @@ class _MicroBatch(Cell): input_shape = self.shape(each_input) micro_batch_begin = i * input_shape[0] // self.micro_size micro_batch_end = (i + 1) * input_shape[0] // self.micro_size - micro_input = each_input[micro_batch_begin:micro_batch_end, :] + micro_input = each_input[micro_batch_begin:micro_batch_end] micro_inputs += (micro_input,) return micro_inputs diff --git a/model_zoo/official/nlp/pangu_alpha/src/pangu_alpha_wrapcell.py b/model_zoo/official/nlp/pangu_alpha/src/pangu_alpha_wrapcell.py index bdc957fab0..4ea05370aa 100644 --- a/model_zoo/official/nlp/pangu_alpha/src/pangu_alpha_wrapcell.py +++ b/model_zoo/official/nlp/pangu_alpha/src/pangu_alpha_wrapcell.py @@ -76,14 +76,14 @@ def tensor_grad_scale_pipeline(scale, grad, accu_grad): new_grad = accu_grad * reciprocal(scale) accu_grad = F.depend(accu_grad, new_grad) zeros = F.tensor_mul(accu_grad, 0.0) - _ = F.assign(accu_grad, zeros) + new_grad = F.depend(new_grad, F.assign(accu_grad, zeros)) return new_grad @shard_grad_scale.register("Tensor", "Tensor", "Tensor") def tensor_shard_grad_scale_pipeline(scale, grad, accu_grad): new_grad = grad * reciprocal(scale) accu_grad = F.depend(accu_grad, new_grad) - _ = F.assign(accu_grad, F.zeros_like(accu_grad)) + new_grad = F.depend(new_grad, F.assign(accu_grad, F.zeros_like(accu_grad))) return new_grad class PanguAlphaTrainOneStepWithLossScaleCell(TrainOneStepWithLossScaleCell): diff --git a/tests/ut/cpp/parallel/step_parallel_test.cc b/tests/ut/cpp/parallel/step_parallel_test.cc index 78274863d1..4a63d54156 100644 --- a/tests/ut/cpp/parallel/step_parallel_test.cc +++ b/tests/ut/cpp/parallel/step_parallel_test.cc @@ -228,7 +228,7 @@ TEST_F(TestStepParallel, ExtractStrategy) { ValueTuplePtr strategy_tuple = std::make_shared(elements); attrs["strategy"] = strategy_tuple; Strategys strategy_expect = {v1, v2}; - StrategyPtr strategy = ExtractStrategy(attrs); + StrategyPtr strategy = ExtractStrategy(attrs["strategy"]); Strategys strategy_test = strategy->GetInputDim(); ASSERT_EQ(strategy_expect, strategy_test);