From 017aa359a6961a8e502bdf661c77c4bd0f2c64bb Mon Sep 17 00:00:00 2001 From: lilei Date: Tue, 21 Dec 2021 16:32:56 +0800 Subject: [PATCH] insert VirtualDataset node for master --- .../ccsrc/pipeline/jit/pipeline_split.cc | 131 +++++++++++++++++- mindspore/ccsrc/pipeline/jit/pipeline_split.h | 1 + mindspore/python/mindspore/train/amp.py | 8 +- mindspore/python/mindspore/train/model.py | 6 - ...t_auto_parallel_for_loop_multi_subgraph.py | 2 +- .../test_auto_parallel_partial_strategy.py | 2 +- .../parallel/test_auto_parallel_reshape.py | 2 + .../parallel/test_comparison_function_info.py | 1 + tests/ut/python/parallel/test_gather_v2.py | 5 +- .../python/parallel/test_manual_gatherv2.py | 2 + tests/ut/python/parallel/test_pack.py | 1 + .../python/parallel/test_sparse_gather_v2.py | 1 + .../python/parallel/test_split_grad_sens.py | 6 +- .../ut/python/parallel/test_virtual_output.py | 4 +- 14 files changed, 150 insertions(+), 22 deletions(-) diff --git a/mindspore/ccsrc/pipeline/jit/pipeline_split.cc b/mindspore/ccsrc/pipeline/jit/pipeline_split.cc index d882a4c149c..ea3b39642eb 100644 --- a/mindspore/ccsrc/pipeline/jit/pipeline_split.cc +++ b/mindspore/ccsrc/pipeline/jit/pipeline_split.cc @@ -14,6 +14,9 @@ * limitations under the License. */ +#include +#include +#include #include #include #include "pipeline/jit/pipeline_split.h" @@ -71,6 +74,124 @@ static int64_t InferStage(int64_t rank_id, int64_t stage_num, int64_t device_num return rank_id / per_stage_rank_num; } +static bool HasVirtualDataset(const std::vector &all_nodes) { + for (auto &node : all_nodes) { + auto cnode = node->cast(); + if (IsPrimitiveCNode(cnode, prim::kPrimVirtualDataset)) { + return true; + } + } + return false; +} + +static CNodePtr CreateTupleGetItem(const AnfNodePtr &node, size_t index, const FuncGraphPtr &func_graph) { + MS_EXCEPTION_IF_NULL(node); + MS_EXCEPTION_IF_NULL(func_graph); + auto idx = NewValueNode(SizeToLong(index)); + MS_EXCEPTION_IF_NULL(idx); + auto imm = std::make_shared(SizeToLong(index)); + auto abstract_scalar = std::make_shared(imm); + idx->set_abstract(abstract_scalar); + CNodePtr tuple_get_item = func_graph->NewCNode({NewValueNode(prim::kPrimTupleGetItem), node, idx}); + MS_EXCEPTION_IF_NULL(tuple_get_item); + tuple_get_item->set_scope(node->scope()); + auto input_abstract_tuple = node->abstract()->cast(); + MS_EXCEPTION_IF_NULL(input_abstract_tuple); + auto tuple_get_item_abstract = input_abstract_tuple->elements()[index]; + MS_EXCEPTION_IF_NULL(tuple_get_item_abstract); + tuple_get_item->set_abstract(tuple_get_item_abstract); + return tuple_get_item; +} + +static CNodePtr CreateVirtualDataset(const FuncGraphPtr &func_graph) { + mindspore::parallel::OperatorAttrs attrs; + ValuePtr pyop_instance = mindspore::parallel::CreateOpInstance(attrs, mindspore::parallel::VIRTUAL_DATA_SET, + mindspore::parallel::VIRTUAL_DATA_SET); + auto value_node = NewValueNode(pyop_instance); + std::vector abstract_list; + std::vector virtual_dataset_node_inputs = {value_node}; + for (size_t index = 0; index < func_graph->get_inputs().size(); index++) { + if (!HasAbstractMonad(func_graph->get_inputs()[index])) { + auto graph_input_index = func_graph->get_inputs()[index]; + auto virtual_dataset_abstract = graph_input_index->abstract()->Clone(); + MS_EXCEPTION_IF_NULL(virtual_dataset_abstract); + abstract_list.emplace_back(virtual_dataset_abstract); + virtual_dataset_node_inputs.push_back(func_graph->get_inputs()[index]); + } + } + CNodePtr virtual_dataset_node = func_graph->NewCNode(virtual_dataset_node_inputs); + MS_EXCEPTION_IF_NULL(virtual_dataset_node); + virtual_dataset_node->set_in_forward_flag(true); + virtual_dataset_node->set_abstract(std::make_shared(abstract_list)); + return virtual_dataset_node; +} + +static std::set FindForwardGraph(const FuncGraphPtr &root, const std::vector &all_nodes) { + std::set graph_sets; + for (auto &node : all_nodes) { + MS_EXCEPTION_IF_NULL(node); + if (!node->isa()) { + continue; + } + auto cnode = node->cast(); + if ((cnode->size() < NODE_INPUT_NUM) || !IsValueNode(cnode->input(0))) { + continue; + } + auto expect_j_prim = GetValueNode(cnode->input(0)); + FuncGraphPtr fun_graph = nullptr; + if (!root->has_flag(mindspore::parallel::TRAINING)) { + graph_sets.insert(root); + } + if (expect_j_prim->name() == mindspore::parallel::J) { + if (IsValueNode(cnode->inputs()[1])) { + fun_graph = GetValueNode(cnode->inputs()[1]); + } else { + fun_graph = node->func_graph(); + } + graph_sets.insert(fun_graph); + } + } + return graph_sets; +} + +static void InsertVirtualDataset(const FuncGraphPtr &root, const std::vector &all_nodes) { + MS_EXCEPTION_IF_NULL(root); + std::set forward_graph_set = FindForwardGraph(root, all_nodes); + for (auto forward_graph : forward_graph_set) { + FuncGraphManagerPtr manager = forward_graph->manager(); + MS_EXCEPTION_IF_NULL(manager); + std::vector graph_inputs = forward_graph->get_inputs(); + auto node_user_map = manager->node_users(); + auto virtual_dataset_node = CreateVirtualDataset(forward_graph); + std::map parameter_index_map; + for (size_t index = 0; index < graph_inputs.size(); index++) { + if (HasAbstractMonad(graph_inputs[index])) { + continue; + } + auto node_users = node_user_map[graph_inputs[index]]; + for (auto node_user : node_users) { + auto cnode = node_user.first->cast(); + for (size_t input_index = 1; input_index < cnode->inputs().size(); input_index++) { + bool is_node_input_flag = !(IsValueNode(cnode->inputs()[input_index]) || + IsValueNode(cnode->inputs()[input_index]) || + IsValueNode(cnode->inputs()[input_index])); + if (find(graph_inputs.begin(), graph_inputs.end(), cnode->inputs()[input_index]) != graph_inputs.end() && + is_node_input_flag && !HasAbstractMonad(cnode->inputs()[input_index])) { + auto node_input_iter = find(graph_inputs.begin(), graph_inputs.end(), cnode->inputs()[input_index]); + size_t node_input_index = node_input_iter - graph_inputs.begin(); + if (parameter_index_map.empty() || parameter_index_map.count(node_input_index) == 0) { + parameter_index_map[node_input_index] = + CreateTupleGetItem(virtual_dataset_node, node_input_index, forward_graph); + } + manager->SetEdge(cnode, input_index, parameter_index_map[node_input_index]); + manager->SetEdge(parameter_index_map[node_input_index], 1, virtual_dataset_node); + } + } + } + } + } +} + // Only auto_parallel and semi_auto_parallel support PipelineSplit bool PipelineSplit(const ResourcePtr &res) { #if ((defined ENABLE_CPU) && (!defined _WIN32) && !defined(__APPLE__)) @@ -84,13 +205,19 @@ bool PipelineSplit(const ResourcePtr &res) { MS_LOG(INFO) << "Only auto_parallel and semi_auto_parallel support pipeline split."; return true; } + auto manager = res->manager(); + auto root = res->func_graph(); + AnfNodePtr ret = root->get_return(); + MS_EXCEPTION_IF_NULL(ret); + std::vector all_nodes = DeepScopedGraphSearch(ret); + if (!HasVirtualDataset(all_nodes)) { + InsertVirtualDataset(root, all_nodes); + } auto stage_num = parallel::ParallelContext::GetInstance()->pipeline_stage_split_num(); if (stage_num <= 1) { MS_LOG(INFO) << "The parameter 'stage_num' is: " << stage_num << ". No need Pipeline split."; return true; } - auto manager = res->manager(); - auto root = res->func_graph(); auto global_rank = GetRank(); auto world_group = GetWorldGroup(); uint32_t world_rank_size = 0; diff --git a/mindspore/ccsrc/pipeline/jit/pipeline_split.h b/mindspore/ccsrc/pipeline/jit/pipeline_split.h index ac1b8431aed..4e63a97fd0e 100644 --- a/mindspore/ccsrc/pipeline/jit/pipeline_split.h +++ b/mindspore/ccsrc/pipeline/jit/pipeline_split.h @@ -22,6 +22,7 @@ namespace mindspore { namespace pipeline { +constexpr size_t NODE_INPUT_NUM = 2; bool PipelineSplit(const ResourcePtr &res); std::string GetWorldGroup(); } // namespace pipeline diff --git a/mindspore/python/mindspore/train/amp.py b/mindspore/python/mindspore/train/amp.py index c52c9ff8eff..6361b158899 100644 --- a/mindspore/python/mindspore/train/amp.py +++ b/mindspore/python/mindspore/train/amp.py @@ -17,12 +17,11 @@ from .. import nn from .._checkparam import Validator as validator from .._checkparam import Rel from ..common import dtype as mstype -from ..nn.wrap.cell_wrapper import _VirtualDatasetCell, _TrainPipelineAccuStepCell +from ..nn.wrap.cell_wrapper import _TrainPipelineAccuStepCell from ..nn.wrap.loss_scale import _TrainPipelineWithLossScaleCell from ..ops import functional as F -from ..parallel._utils import _get_parallel_mode, _get_pipeline_stages +from ..parallel._utils import _get_pipeline_stages from .loss_scale_manager import DynamicLossScaleManager, LossScaleManager -from ..context import ParallelMode from .. import boost from .. import context @@ -198,9 +197,6 @@ def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_leve if loss_fn: network = _add_loss_network(network, loss_fn, config["cast_model_type"]) - if _get_parallel_mode() in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL): - network = _VirtualDatasetCell(network) - loss_scale = 1.0 if config["loss_scale_manager"] is not None: loss_scale_manager = config["loss_scale_manager"] diff --git a/mindspore/python/mindspore/train/model.py b/mindspore/python/mindspore/train/model.py index f85d0c3c6e1..0eb3b4bc4bb 100644 --- a/mindspore/python/mindspore/train/model.py +++ b/mindspore/python/mindspore/train/model.py @@ -34,7 +34,6 @@ from ..parallel._utils import _get_parallel_mode, _get_device_num, _get_global_r from ..parallel._ps_context import _is_role_pserver, _is_role_sched from ..nn.metrics import Loss from .. import nn -from ..nn.wrap.cell_wrapper import _VirtualDatasetCell from ..boost import AutoBoost from ..context import ParallelMode from ..parallel._cost_model_context import _set_multi_subgraphs @@ -307,8 +306,6 @@ class Model: boost_level=self._boost_level, keep_batchnorm_fp32=self._keep_bn_fp32) elif self._loss_fn: - if self._parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL): - network = _VirtualDatasetCell(network) network = nn.WithLossCell(network, self._loss_fn) # If need to check if loss_fn is not None, but optimizer is None @@ -347,8 +344,6 @@ class Model: self._eval_indexes = [0, 1, 2] if self._parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL): - if self._optimizer: - self._eval_network = _VirtualDatasetCell(self._eval_network) if self._optimizer is None: # In this case, multiple optimizer(s) is supposed to be included in 'self._network' _set_multi_subgraphs() @@ -358,7 +353,6 @@ class Model: """Build the network for prediction.""" self._predict_network = self._network if self._parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL): - self._predict_network = _VirtualDatasetCell(self._network) # Unlike the cases in build_train_network() and build_eval_network(), 'multi_subgraphs' is not set self._predict_network.set_auto_parallel() diff --git a/tests/ut/python/parallel/test_auto_parallel_for_loop_multi_subgraph.py b/tests/ut/python/parallel/test_auto_parallel_for_loop_multi_subgraph.py index be5a70f8681..264718bfd95 100644 --- a/tests/ut/python/parallel/test_auto_parallel_for_loop_multi_subgraph.py +++ b/tests/ut/python/parallel/test_auto_parallel_for_loop_multi_subgraph.py @@ -131,5 +131,5 @@ def test_double_subgraphs(): net.set_train() _cell_graph_executor.compile(net, x, phase='train') num_ops = _cell_graph_executor._get_num_parallel_ops(net) - expected_num = 7 + expected_num = 9 assert expected_num == num_ops diff --git a/tests/ut/python/parallel/test_auto_parallel_partial_strategy.py b/tests/ut/python/parallel/test_auto_parallel_partial_strategy.py index cd366117303..98a471aa388 100644 --- a/tests/ut/python/parallel/test_auto_parallel_partial_strategy.py +++ b/tests/ut/python/parallel/test_auto_parallel_partial_strategy.py @@ -66,7 +66,7 @@ def test_four_matmul_linear(): return out size = 64 - context.set_auto_parallel_context(device_num=size, global_rank=0) + context.set_auto_parallel_context(dataset_strategy="full_batch", device_num=size, global_rank=0) strategy1 = ((2, 4), (4, 8)) x = Tensor(np.ones([128, 32]), dtype=ms.float32) y = Tensor(np.ones([32, 64]), dtype=ms.float32) diff --git a/tests/ut/python/parallel/test_auto_parallel_reshape.py b/tests/ut/python/parallel/test_auto_parallel_reshape.py index 6321ab9c573..4624239d42a 100644 --- a/tests/ut/python/parallel/test_auto_parallel_reshape.py +++ b/tests/ut/python/parallel/test_auto_parallel_reshape.py @@ -271,6 +271,7 @@ def test_reshape_auto_5(): return out size = 8 + context.set_auto_parallel_context(dataset_strategy="full_batch") x = Tensor(np.ones([4, 1024 * size, 1]), dtype=ms.float32) y = Tensor(np.ones([4, 1024 * size,]), dtype=ms.float32) net = GradWrapTwoInput(NetWithLossTwoInput(Net())) @@ -301,6 +302,7 @@ def test_reshape_auto_6(): return out size = 8 + context.set_auto_parallel_context(dataset_strategy="full_batch") x = Tensor(np.ones([4, 1024, 1]), dtype=ms.float32) y = Tensor(np.ones([4, 1024,]), dtype=ms.float32) net = GradWrapTwoInput(NetWithLossTwoInput(Net())) diff --git a/tests/ut/python/parallel/test_comparison_function_info.py b/tests/ut/python/parallel/test_comparison_function_info.py index 15c3880cbaa..e41dc9b76e7 100644 --- a/tests/ut/python/parallel/test_comparison_function_info.py +++ b/tests/ut/python/parallel/test_comparison_function_info.py @@ -392,6 +392,7 @@ def test_matmul_minimum_auto_parallel(): out = self.minimum(out, b) return out + context.set_auto_parallel_context(dataset_strategy="full_batch") context.set_auto_parallel_context(device_num=8, global_rank=0, parallel_mode="auto_parallel") net = GradWrap(NetWithLoss(Net())) diff --git a/tests/ut/python/parallel/test_gather_v2.py b/tests/ut/python/parallel/test_gather_v2.py index 1c50b021ea8..44ff4be4958 100644 --- a/tests/ut/python/parallel/test_gather_v2.py +++ b/tests/ut/python/parallel/test_gather_v2.py @@ -196,6 +196,7 @@ def test_gatherv2_forward_all_reduce(): """ strategy1 = ((8, 1), (1, 1)) strategy2 = ((2, 4, 1), (2, 4, 1)) + context.set_auto_parallel_context(dataset_strategy="full_batch") net = GradWrap(NetWithLoss(Net(0, strategy1, strategy2, shape=[2, 64]))) x = Tensor(np.ones([64, 64]), dtype=ms.float32) y = Tensor(np.ones([2, 64, 64]), dtype=ms.float32) @@ -210,6 +211,7 @@ def test_gatherv2_shard_batch_and_axis(): """ strategy1 = ((4, 1), (2, 1)) strategy2 = ((2, 4, 1), (2, 4, 1)) + context.set_auto_parallel_context(dataset_strategy="full_batch") net = GradWrap(NetWithLoss(Net(0, strategy1, strategy2, shape=[2, 64]))) x = Tensor(np.ones([64, 64]), dtype=ms.float32) y = Tensor(np.ones([2, 64, 64]), dtype=ms.float32) @@ -224,6 +226,7 @@ def test_gatherv2_split_axis_0_repeat_calc(): """ strategy1 = ((4, 1), (1, 1)) strategy2 = ((2, 4, 1), (2, 4, 1)) + context.set_auto_parallel_context(dataset_strategy="full_batch") net = GradWrap(NetWithLoss(Net(0, strategy1, strategy2, shape=[2, 64]))) x = Tensor(np.ones([64, 64]), dtype=ms.float32) y = Tensor(np.ones([2, 64, 64]), dtype=ms.float32) @@ -350,7 +353,7 @@ def test_gatherv2_target_cpu_reducescatter(): Description: axis is 0, split axis and batch, out strategy use reducescatter. Expectation: compile done without error. """ - strategy1 = ((4, 1), (2, 1)) + strategy1 = ((8, 1), (1, 1)) out_strategy = ((8, 1, 1),) strategy2 = ((2, 4, 1), (2, 4, 1)) net = GradWrap(NetWithLoss(Net(0, strategy1, strategy2, target="CPU", gather_out_strategy=out_strategy))) diff --git a/tests/ut/python/parallel/test_manual_gatherv2.py b/tests/ut/python/parallel/test_manual_gatherv2.py index fe4036ba021..34fa35a7262 100644 --- a/tests/ut/python/parallel/test_manual_gatherv2.py +++ b/tests/ut/python/parallel/test_manual_gatherv2.py @@ -87,6 +87,7 @@ def test_normal_split2(): def test_normal_split3(): + context.set_auto_parallel_context(dataset_strategy="full_batch") context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=32, global_rank=17) strategy1 = ((4, 8), (1, 4)) strategy2 = ((1, 4, 8), (1, 4, 8)) @@ -96,6 +97,7 @@ def test_normal_split3(): def test_normal_split_with_offset(): + context.set_auto_parallel_context(dataset_strategy="full_batch") context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=2, global_rank=0) strategy1 = ((2, 1), (1, 2)) strategy2 = ((1, 2, 1), (1, 2, 1)) diff --git a/tests/ut/python/parallel/test_pack.py b/tests/ut/python/parallel/test_pack.py index d8ef71d9977..337994d0c16 100644 --- a/tests/ut/python/parallel/test_pack.py +++ b/tests/ut/python/parallel/test_pack.py @@ -251,6 +251,7 @@ def test_pack_auto_parallel_axis1(): def test_pack_auto_parallel_3_tensor(): + context.set_auto_parallel_context(dataset_strategy="full_batch") context.set_auto_parallel_context(parallel_mode="auto_parallel", device_num=8, global_rank=0) net = Net2(_w1, _w2, _w3) compile_net2(net) diff --git a/tests/ut/python/parallel/test_sparse_gather_v2.py b/tests/ut/python/parallel/test_sparse_gather_v2.py index 3287795c72c..39764218d48 100644 --- a/tests/ut/python/parallel/test_sparse_gather_v2.py +++ b/tests/ut/python/parallel/test_sparse_gather_v2.py @@ -147,6 +147,7 @@ def test_gatherv2_auto0(): Description: gather net with strategy in semi auto parallel, gather axis is 1. Expectation: compile done without error. """ + context.set_auto_parallel_context(dataset_strategy="full_batch") net = GradWrap(NetWithLoss(Net(0))) compile_net(net, [64, 32], [64, 64, 32], parallel_mode="auto_parallel") diff --git a/tests/ut/python/parallel/test_split_grad_sens.py b/tests/ut/python/parallel/test_split_grad_sens.py index 0bbc7968f81..fc51b17ca3b 100644 --- a/tests/ut/python/parallel/test_split_grad_sens.py +++ b/tests/ut/python/parallel/test_split_grad_sens.py @@ -124,9 +124,9 @@ def test_grad_sens_parameter_type(): net.set_auto_parallel() net.set_train() _cell_graph_executor.compile(net, x, y, b, sens, phase='train', auto_parallel_mode=True) - x_layout = ([8, 8], [1, -1], [16, 32], 0, True, '') - y_layout = ([8, 8], [-1, 0], [32, 8], 0, True, '') - b_layout = ([8, 8], [0, -1], [8, 64], 0, True, '') + x_layout = ([64], [0, -1], [2, 32], 0, True, '') + y_layout = ([64], [-1, -1], [32, 64], 0, True, '') + b_layout = ([64], [0, -1], [1, 64], 0, True, '') sens_layout = ([8, 8], [1, -1], [16, 64], 0, True, '') expect_dict = {'x': x_layout, 'y': y_layout, 'b': b_layout, 'sens': sens_layout} assert net.parameter_layout_dict == expect_dict diff --git a/tests/ut/python/parallel/test_virtual_output.py b/tests/ut/python/parallel/test_virtual_output.py index fc5031abdd0..485d8e0bc74 100644 --- a/tests/ut/python/parallel/test_virtual_output.py +++ b/tests/ut/python/parallel/test_virtual_output.py @@ -319,7 +319,7 @@ def test_scalar_output_semi_auto(): if re.search('VirtualOutput-op', k) is not None: assert v[0][0] == 8 count += 1 - assert count == 1 + assert count == 2 def test_scalar_output_auto(): context.reset_auto_parallel_context() @@ -336,4 +336,4 @@ def test_scalar_output_auto(): if re.search('VirtualOutput-op', k) is not None: assert v[0][0] == 8 count += 1 - assert count == 1 + assert count == 2