diff --git a/mindspore/ccsrc/frontend/parallel/step_parallel.cc b/mindspore/ccsrc/frontend/parallel/step_parallel.cc index 294819e48d5..e0110d61c08 100644 --- a/mindspore/ccsrc/frontend/parallel/step_parallel.cc +++ b/mindspore/ccsrc/frontend/parallel/step_parallel.cc @@ -59,14 +59,25 @@ using mindspore::tensor::Tensor; namespace mindspore { namespace parallel { +// Set of communication operations static const std::set COMMUNICATION_OPS = {ALL_REDUCE, ALL_GATHER, ALL_TO_ALL, REDUCE_SCATTER}; + +// Set of invalid loss operations static const std::set INVALID_LOSS_OPS = {GET_NEXT, VIRTUALLOSS, LOAD, UPDATESTATE}; + +// Set of operations with no input tensors static const std::set NO_INPUT_TENSOR_OPS = {UNIFORM_REAL}; -// g_RefMap, for CNode B input i is a RefKey[Parameter C], -// it will be one item in map with key: C, and value: (B, i) + +// g_RefMap is a map that stores information about CNode inputs that are references. +// For a CNode B, where input i is a reference to Parameter C, it will be one item in the map +// with key: C, and value: (B, i) std::map> g_RefMap; + +// Maximum breadth-first search depth const uint32_t MAX_BFS_DEPTH = 7; +// SetMiniStepOpDoMirrorLabel function +// This function sets the 'do_mirror' and 'accu_flag' labels in the attributes of the given operator node. void SetMiniStepOpDoMirrorLabel(std::vector new_node_input, bool do_mirror, bool accu_flag) { if (new_node_input.empty()) { return; @@ -81,6 +92,8 @@ void SetMiniStepOpDoMirrorLabel(std::vector new_node_input, bool do_ prim->SetAttrs(attrs); } +// SetAllReduceRecomputeFlag function +// This function sets the 'RECOMPUTE' attribute in the operator's attributes based on the node and its inputs. void SetAllReduceRecomputeFlag(const std::vector &new_node_input, const CNodePtr &node) { if (new_node_input.empty()) { return; @@ -102,6 +115,8 @@ void SetAllReduceRecomputeFlag(const std::vector &new_node_input, co } } +// CreateInput function +// This function creates a vector of input nodes for an operator based on the given arguments and instance name. std::vector CreateInput(const Operator &op, const AnfNodePtr &node, const std::string &instance_name) { MS_EXCEPTION_IF_NULL(node); OperatorArgs arg_forward = op.second; @@ -119,11 +134,13 @@ std::vector CreateInput(const Operator &op, const AnfNodePtr &node, } } - // if the op have 'group' attr, set the rank list name for the op + // If the op has 'group' attribute, set the rank list name for the op SetCommunicationOpGroupLabel(new_node_input); return new_node_input; } +// GetAccuGrad function +// This function finds and returns the accumulation gradient node among the given parameters. AnfNodePtr GetAccuGrad(const std::vector ¶meters, const std::string &weight_name) { for (auto ¶m : parameters) { if (!ParameterIsCloned(param)) { @@ -141,6 +158,9 @@ AnfNodePtr GetAccuGrad(const std::vector ¶meters, const std::str return nullptr; } +// CreateMirrorInput function +// This function creates a vector of input nodes for a mirror operator based on the given arguments, instance name, +// and weight name. std::vector CreateMirrorInput(const FuncGraphPtr &root, const Operator &op, const AnfNodePtr &node, const std::string &instance_name, const std::string &weight_name) { MS_EXCEPTION_IF_NULL(root); @@ -163,63 +183,84 @@ std::vector CreateMirrorInput(const FuncGraphPtr &root, const Operat arg_forward.first.pop_back(); } else if (op_name == MINI_STEP_ALL_GATHER || op_name == MIRROR_MICRO_STEP_OPERATOR || op_name == MICRO_STEP_ALL_GATHER) { - MS_LOG(EXCEPTION) << "You should define `accu_grads` when use " << op_name << " parameter:" << weight_name; + MS_LOG(EXCEPTION) << "You should define `accu_grads` when using " << op_name << " parameter: " << weight_name; } } } - ValuePtr pyop_instance = CreateOpInstance(arg_forward.first, op_name, instance_name); - MS_EXCEPTION_IF_NULL(pyop_instance); - OperatorParams params = arg_forward.second; +// Create an instance of the Python operator using the provided arguments +ValuePtr pyop_instance = CreateOpInstance(arg_forward.first, op_name, instance_name); +MS_EXCEPTION_IF_NULL(pyop_instance); - std::vector new_node_input; - if (op_name == MIRROR_MINI_STEP_OPERATOR || op_name == MINI_STEP_ALL_GATHER || - op_name == MIRROR_MICRO_STEP_OPERATOR || op_name == MICRO_STEP_ALL_GATHER) { - new_node_input = {NewValueNode(pyop_instance), node, grad_accu}; - MS_LOG(INFO) << "Insert the grad accumulation node as the mirror op's input"; - } else { - new_node_input = {NewValueNode(pyop_instance), node}; +// Retrieve the operator parameters associated with the Python operator +OperatorParams params = arg_forward.second; + +// Create a vector to store the new input nodes for the CNode +std::vector new_node_input; + +// Check the type of the operator and adjust the new_node_input accordingly +if (op_name == MIRROR_MINI_STEP_OPERATOR || op_name == MINI_STEP_ALL_GATHER || + op_name == MIRROR_MICRO_STEP_OPERATOR || op_name == MICRO_STEP_ALL_GATHER) { + // If it's a mirror or all-gather operator, include additional input nodes for grad accumulation + new_node_input = {NewValueNode(pyop_instance), node, grad_accu}; + MS_LOG(INFO) << "Insert the grad accumulation node as the mirror op's input"; +} else { + // For other operators, only include the Python operator instance and the original node + new_node_input = {NewValueNode(pyop_instance), node}; +} + +// Check if there are additional parameters to be inserted into the new_node_input +if (!params.empty()) { + for (auto ¶m : params) { + // Create a ValueNode for the parameter and insert it at the specified position + AnfNodePtr val = NewValueNode(param.first.second); + MS_EXCEPTION_IF_NULL(val); + int64_t position = param.second; + (void)new_node_input.insert(new_node_input.begin() + position, val); } +} - if (!params.empty()) { - for (auto ¶m : params) { - AnfNodePtr val = NewValueNode(param.first.second); - MS_EXCEPTION_IF_NULL(val); - int64_t position = param.second; - (void)new_node_input.insert(new_node_input.begin() + position, val); - } - } - // if the op have 'group' attr, set the rank list name for the op + // If the op has 'group' attribute, set the rank list name for the op SetCommunicationOpGroupLabel(new_node_input); - // gradient accumulation + // Gradient accumulation if (grad_accumulation_step > 1) { bool add_accu = root->has_flag(kAccumulation); - // MiniStep need to do mirror at each micro step as we use the gradient accumulation sharding, + // MiniStep needs to do a mirror at each micro step as we use the gradient accumulation sharding SetMiniStepOpDoMirrorLabel(new_node_input, !add_accu, !add_accu); } return new_node_input; } +// InsertNode function +// This function inserts a new CNode into the given func_graph before the specified node. void InsertNode(const Operator &op, const CNodePtr &node, size_t index, const AnfNodePtr &pre_node, const FuncGraphPtr &func_graph, const std::string &instance_name, const std::string ¶m_name = "", const FuncGraphPtr &root = nullptr) { - // insert new node before the node + // Step 1: Get the function graph manager and scope FuncGraphManagerPtr manager = func_graph->manager(); MS_EXCEPTION_IF_NULL(manager); ScopePtr scope = node->scope(); MS_EXCEPTION_IF_NULL(scope); + + // Step 2: Create input nodes for the new CNode based on the given parameters std::vector node_input; if (root && !param_name.empty()) { node_input = CreateMirrorInput(root, op, pre_node, instance_name, param_name); } else { node_input = CreateInput(op, pre_node, instance_name); } + + // Step 3: Create a new CNode with the input nodes CNodePtr new_node = func_graph->NewCNode(node_input); MS_EXCEPTION_IF_NULL(new_node); + + // Step 4: Mark the new CNode as in the forward pass if (instance_name.find(SPLIT_SENS) == std::string::npos) { - new_node->set_in_forward_flag(true); // mark forward flag + new_node->set_in_forward_flag(true); } + + // Step 5: Set attributes of the new CNode's Primitive auto new_node_value = node_input[0]->cast(); MS_EXCEPTION_IF_NULL(new_node_value); PrimitivePtr new_node_prim = new_node_value->value()->cast(); @@ -228,52 +269,75 @@ void InsertNode(const Operator &op, const CNodePtr &node, size_t index, const An if (instance_name.find(NOT_RECOMPUTE) != std::string::npos) { new_node_prim->set_attr("recompute", MakeValue(false)); } + + // Step 6: Set the scope for the new CNode and input nodes new_node->set_scope(scope); node_input[0]->set_scope(scope); + + // Step 7: Replace the original node with the new CNode manager->SetEdge(node, SizeToInt(index), new_node); MS_LOG(INFO) << "Insert " << instance_name << " success"; } -// Replace pre_node with pre_node->op +// ReplaceNode function +// This function replaces the pre_node with a new CNode based on the given parameters. static CNodePtr ReplaceNode(const Operator &op, const AnfNodePtr &pre_node, const FuncGraphPtr &func_graph, const std::string &instance_name, const std::string ¶m_name = "", const FuncGraphPtr &root = nullptr) { - // insert new node before the node + // Step 1: Get the function graph manager and scope FuncGraphManagerPtr manager = func_graph->manager(); MS_EXCEPTION_IF_NULL(manager); ScopePtr scope = pre_node->scope(); MS_EXCEPTION_IF_NULL(scope); + + // Step 2: Create input nodes for the new CNode based on the given parameters std::vector node_input; if (root && !param_name.empty()) { node_input = CreateMirrorInput(root, op, pre_node, instance_name, param_name); } else { node_input = CreateInput(op, pre_node, instance_name); } + + // Step 3: Create a new CNode with the input nodes CNodePtr new_node = func_graph->NewCNode(node_input); MS_EXCEPTION_IF_NULL(new_node); + + // Step 4: Mark the new CNode as in the forward pass if (instance_name.find(SPLIT_SENS) == std::string::npos) { - new_node->set_in_forward_flag(true); // mark forward flag + new_node->set_in_forward_flag(true); } + + // Step 5: Set attributes of the new CNode's Primitive auto new_node_prim = GetValueNode(node_input[0]); new_node_prim->set_instance_name(instance_name); new_node_prim->set_attr("keep_value_node_input", MakeValue(true)); if (instance_name.find(NOT_RECOMPUTE) != std::string::npos) { new_node_prim->set_attr("recompute", MakeValue(false)); } + + // Step 6: Set the scope for the new CNode and input nodes new_node->set_scope(scope); node_input[0]->set_scope(scope); + + // Step 7: Replace the original pre_node with the new CNode manager->Replace(pre_node, new_node); MS_LOG(INFO) << "Insert " << instance_name << " success"; + return new_node; } +// ForwardCommunication function +// This function performs forward communication by inserting nodes into the given func_graph. void ForwardCommunication(OperatorVector forward_op, const CNodePtr &node) { MS_EXCEPTION_IF_NULL(node); - // step1:get graph manager distribute_operator + + // Step 1: Get the function graph and manager FuncGraphPtr func_graph = node->func_graph(); MS_EXCEPTION_IF_NULL(func_graph); FuncGraphManagerPtr manager = func_graph->manager(); MS_EXCEPTION_IF_NULL(manager); + + // Step 2: Find the appropriate node to insert the forward nodes auto uses_set = manager->node_users()[node]; CNodePtr node_to_insert = node; for (auto &uses_pair : uses_set) { @@ -294,23 +358,25 @@ void ForwardCommunication(OperatorVector forward_op, const CNodePtr &node) { MS_EXCEPTION_IF_NULL(node_to_insert); std::reverse(forward_op.begin(), forward_op.end()); - // step2:traverse op_list and insert node + // Step 3: Traverse the forward_op list and insert nodes for (size_t index = 0; index < forward_op.size(); ++index) { std::string instance_name_base = FORWARD_OP; std::string instance_name = instance_name_base + "_" + CreateInstanceName(node, index); std::vector forward_input = CreateInput(forward_op[index], node_to_insert, instance_name); SetAllReduceRecomputeFlag(forward_input, node_to_insert); - CNodePtr forward_node = func_graph->NewCNode(forward_input); // using NewCNode to create anfnode + CNodePtr forward_node = func_graph->NewCNode(forward_input); MS_EXCEPTION_IF_NULL(forward_node); ScopePtr scope = node->scope(); MS_EXCEPTION_IF_NULL(scope); forward_node->set_scope(scope); forward_node->set_in_forward_flag(true); forward_input[0]->set_scope(scope); - (void)manager->Replace(node_to_insert, forward_node); // using Replace function to insert node + (void)manager->Replace(node_to_insert, forward_node); } } +// InsertMakeTuple function +// This function inserts a MakeTuple node into the func_graph with the given parameters. CNodePtr InsertMakeTuple(const AnfNodePtr &prev, uint64_t num, const FuncGraphPtr &func_graph) { MS_EXCEPTION_IF_NULL(prev); MS_EXCEPTION_IF_NULL(func_graph); @@ -331,33 +397,51 @@ CNodePtr InsertMakeTuple(const AnfNodePtr &prev, uint64_t num, const FuncGraphPt return make_tuple; } +// InsertRedistribution function +// This function inserts redistribution nodes into the graph. void InsertRedistribution(const RedistributionOpListPtr &redistribution_oplist_ptr, const CNodePtr &node, const FuncGraphPtr &func_graph, int64_t pos, const CNodePtr &pre_node) { MS_EXCEPTION_IF_NULL(node); MS_EXCEPTION_IF_NULL(pre_node); MS_EXCEPTION_IF_NULL(func_graph); + + // Obtain the graph manager FuncGraphManagerPtr manager = func_graph->manager(); MS_EXCEPTION_IF_NULL(manager); + + // Check if the sizes of OperatorVector and OutPutInfoVector are the same if ((redistribution_oplist_ptr->first).size() != (redistribution_oplist_ptr->second).size()) { - MS_LOG(EXCEPTION) << "size of OperatorVector and OutPutInfoVector must be the same!"; + MS_LOG(EXCEPTION) << "Size of OperatorVector and OutPutInfoVector must be the same!"; } + + // Iterate through the redistribution operators for (size_t index = 0; index < (redistribution_oplist_ptr->first).size(); ++index) { + // Check if the position (pos) is valid if (pos >= SizeToLong(node->inputs().size())) { - MS_LOG(EXCEPTION) << "InsertRedistribution:pos can't be larger than node's inputs'size"; + MS_LOG(EXCEPTION) << "InsertRedistribution: Position (pos) cannot be larger than the node's inputs' size."; } - // Create new node + + // Create a new node based on the target node at the specified position (pos) AnfNodePtr target_node = node->input(LongToSize(pos)); MS_EXCEPTION_IF_NULL(target_node); - // Create instance_name + + // Extract information about the redistribution operator auto op = (redistribution_oplist_ptr->first)[index]; std::string op_name = (redistribution_oplist_ptr->first)[index].first; + + // Create an instance name for the redistribution node std::string instance_name_base = REDISTRIBUTION_OP; std::string instance_name = instance_name_base + "_" + CreateInstanceName(pre_node, index) + op_name; + + // Check if the output and input nodes have RECOMPUTE_COMM_OP attributes auto prim_out = GetCNodePrimitive(node); auto prim_in = GetCNodePrimitive(pre_node); + if (prim_out != nullptr && prim_in != nullptr) { auto prim_out_attr = prim_out->attrs(); auto prim_in_attr = prim_in->attrs(); + + // Check if the redistribution node should not be recomputed if (((prim_out_attr.find(RECOMPUTE_COMM_OP) != prim_out_attr.end() && !GetValue(prim_out_attr[RECOMPUTE_COMM_OP])) || (prim_in_attr.find(RECOMPUTE_COMM_OP) != prim_in_attr.end() && @@ -367,7 +451,11 @@ void InsertRedistribution(const RedistributionOpListPtr &redistribution_oplist_p instance_name = instance_name + "_" + NOT_RECOMPUTE; } } + + // Insert the redistribution node into the graph InsertNode(op, node, LongToSize(pos), target_node, func_graph, instance_name); + + // Check if additional tuple handling is required if ((redistribution_oplist_ptr->second)[index].first) { target_node = node->input(LongToSize(pos)); MS_EXCEPTION_IF_NULL(target_node); @@ -376,121 +464,216 @@ void InsertRedistribution(const RedistributionOpListPtr &redistribution_oplist_p } } +// InsertGetTensorSliceOp function +// This function inserts GetTensorSlice operator nodes into the graph. void InsertGetTensorSliceOp(const Operator &op, const CNodePtr &node, const FuncGraphPtr &func_graph, int64_t pos, const std::string &instance_name) { + // Check if the graph is null if (func_graph == nullptr) { - MS_LOG(EXCEPTION) << "InsertGetTensorSliceOp: the graph is null, the instance name is " << instance_name; + MS_LOG(EXCEPTION) << "InsertGetTensorSliceOp: The graph is null, the instance name is " << instance_name; } + // Obtain the graph manager FuncGraphManagerPtr manager = func_graph->manager(); MS_EXCEPTION_IF_NULL(manager); + + // Check if the position (pos) is valid if (pos >= SizeToLong(node->inputs().size())) { - MS_LOG(EXCEPTION) << "InsertGetTensorSliceOp: pos can't be larger than node's inputs'size, the instance name is " + MS_LOG(EXCEPTION) << "InsertGetTensorSliceOp: Position (pos) cannot be larger than the node's inputs' size, the instance name is " << instance_name; } - // Create new node + + // Create a new node based on the pre_node at the specified position (pos) AnfNodePtr pre_node = node->input(LongToSize(pos)); MS_EXCEPTION_IF_NULL(pre_node); + + // Insert the GetTensorSlice operator node into the graph InsertNode(op, node, LongToSize(pos), pre_node, func_graph, instance_name); } +// GetTensorInLayout function +// This function retrieves the input tensor layout of a given middle_node based on its primitive type. TensorLayout GetTensorInLayout(const CNodePtr &middle_node, const PrimitivePtr &middle_prim, const OperatorInfoPtr &distribute_operator) { + // Initialize a TensorInfo variable to store the input tensor information. TensorInfo tensorinfo_in; + + // Check if the middle primitive is a TupleGetItem operation. if (middle_prim->name() == prim::kTupleGetItem) { + // Extract the index from the TupleGetItem operation's input. auto value_node = middle_node->input(2)->cast(); MS_EXCEPTION_IF_NULL(value_node); size_t index_s = LongToSize(GetValue(value_node->value())); + + // Check if the index is out of range. if (index_s >= distribute_operator->outputs_tensor_info().size()) { - MS_LOG(EXCEPTION) << "The index out of range, index: " << index_s + MS_LOG(EXCEPTION) << "The index is out of range, index: " << index_s << ", vector size: " << distribute_operator->outputs_tensor_info().size(); } + + // Retrieve the input tensor information based on the index. tensorinfo_in = distribute_operator->outputs_tensor_info()[index_s]; } else { + // If not a TupleGetItem operation, use the first output tensor information. if (distribute_operator->outputs_tensor_info().empty()) { MS_LOG(EXCEPTION) << "The outputs tensor info is empty"; } tensorinfo_in = distribute_operator->outputs_tensor_info()[0]; } + + // Return the tensor layout of the input tensor information. return tensorinfo_in.tensor_layout(); } +// GetDistributeOperator function +// This function retrieves the distribute operator associated with a CNode. OperatorInfoPtr GetDistributeOperator(const CNodePtr &node) { MS_EXCEPTION_IF_NULL(node); + + // Check if the node is a parallel care node. if (!IsParallelCareNode(node)) { return nullptr; } + + // Retrieve the distribute operator from user data. OperatorInfoPtr distribute_operator = node->user_data(); return distribute_operator; } +/** + * Redistribution function performs tensor redistribution between two operators within the computation graph. + * + * @param node_pair - A pair consisting of the next operator node and its index. + * @param distribute_operator - OperatorInfoPtr representing the distributing operator. + * @param middle_node - CNodePtr representing the middle node in the computation graph. + * @param index - The index indicating the connection between operators. + * @param tensor_redistribution - TensorRedistribution object for computing the redistribution. + * @param pre_node - CNodePtr representing the previous node in the computation graph. + * + * This function is responsible for redistributing tensors between two operators within the computation graph. + * It takes the next operator node, the distributing operator, the middle node, the index, a tensor redistribution object, + * and the previous node as input. + * + * The steps involved in redistribution include: + * 1. Checking if the function graph is valid. + * 2. Getting the next operator node and validating it. + * 3. Extracting relevant information from the middle node and next node, such as their primitives and operator info. + * 4. Extracting tensor layout information from the distributing operator's outputs. + * 5. Initializing the tensor redistribution object with input and output tensor layouts and device list. + * 6. Inferring the tensor redistribution operator list. + * 7. Inserting the redistribution nodes before the next node if necessary. + */ void Redistribution(const std::pair &node_pair, const OperatorInfoPtr &distribute_operator, const CNodePtr &middle_node, int64_t index, TensorRedistribution tensor_redistribution, const CNodePtr &pre_node) { + // Check if the function graph is valid FuncGraphPtr func_graph = middle_node->func_graph(); if (func_graph == nullptr) { - MS_LOG(EXCEPTION) << "Redistribution:get graph failed"; + MS_LOG(EXCEPTION) << "Redistribution: Get graph failed"; } + + // Get the next operator node and validate it CNodePtr next_node = node_pair.first->cast(); MS_EXCEPTION_IF_NULL(next_node); + + // Extract primitive and operator information from the middle node auto middle_value = middle_node->input(0)->cast(); MS_EXCEPTION_IF_NULL(middle_value); PrimitivePtr middle_prim = middle_value->value()->cast(); MS_EXCEPTION_IF_NULL(middle_prim); + + // Get the distribute operator for the next node OperatorInfoPtr next_distribute_operator = GetDistributeOperator(next_node); if (next_distribute_operator == nullptr) { MS_LOG(EXCEPTION) << "Failure: " << next_node->ToString() << " GetDistributeOperator failed"; } + + // Get the device list from the distributing operator RankList dev_list = distribute_operator->stage_device_list(); + + // Get the name of the primitive for the next node std::string next_prim_name = GetValueNode(next_node->input(0))->name(); MS_LOG(DEBUG) << "Redistribution: middle_prim " << middle_prim->name() << " next_prim " << next_prim_name; MS_LOG(DEBUG) << "Redistribution: middle_node " << middle_node->ToString() << " next_node " << next_node->ToString(); - // extract tensor layout in and out + + // Extract tensor layout information from the distributing operator's outputs if (distribute_operator->outputs_tensor_info().empty()) { MS_LOG(WARNING) << "pre_node's tensorinfo_in is empty, operator name is " << distribute_operator->name(); return; } + // Check if the index is out of range if (LongToSize(index - 1) >= next_distribute_operator->inputs_tensor_info().size()) { MS_LOG(WARNING) << "The index is out of range, the index is " << (index - 1) << ", the vector size is " - << next_distribute_operator->inputs_tensor_info().size() << "next operator name is " + << next_distribute_operator->inputs_tensor_info().size() << " next operator name is " << next_distribute_operator->name(); return; } + + // Get tensor layout information for the input and output tensors TensorInfo tensorinfo_out = next_distribute_operator->inputs_tensor_info()[LongToSize(index - 1)]; TensorLayout tensorlayout_out = tensorinfo_out.tensor_layout(); TensorLayout tensorlayout_in = GetTensorInLayout(middle_node, middle_prim, distribute_operator); + + // Handle special case for Receive primitive if (IsPrimitiveCNode(middle_node, prim::kPrimReceive)) { tensorlayout_in = *(middle_node->user_data()); } + + // Initialize the tensor redistribution object and handle errors if (tensor_redistribution.Init(tensorlayout_in, tensorlayout_out, dev_list) == FAILED) { MS_LOG(ERROR) << "Redistribution: middle_prim " << middle_prim->name() << " next_prim : " << next_prim_name; MS_LOG(ERROR) << "Redistribution: middle_node " << middle_node->ToString() << " next_node " << next_node->ToString(); DumpGraph(func_graph, "redistribution_error"); - MS_LOG(EXCEPTION) << "Failure:tensor_redistribution init failed"; + MS_LOG(EXCEPTION) << "Failure: tensor_redistribution init failed"; } + + // Infer the tensor redistribution operator list RedistributionOpListPtr redistribution_oplist_ptr = tensor_redistribution.InferTensorRedistributionOperatorList(); if (redistribution_oplist_ptr == nullptr) { - MS_LOG(EXCEPTION) << "Failure:InferTensorRedistribution failed"; + MS_LOG(EXCEPTION) << "Failure: InferTensorRedistribution failed"; } + MS_LOG(DEBUG) << "Redistribution size " << redistribution_oplist_ptr->first.size(); + + // Insert redistribution nodes before the next node if necessary if (!redistribution_oplist_ptr->first.empty()) { - // insert node before next node InsertRedistribution(redistribution_oplist_ptr, next_node, func_graph, node_pair.second, pre_node); } } +/** + * @brief Checks if the 'IN_STRATEGY' attribute is found and not of type 'NONE' in the given attributes. + * + * @param attrs A HashMap containing the attributes to check. + * @return true if the 'IN_STRATEGY' attribute is found and not of type 'NONE', false otherwise. + */ bool StrategyFound(const mindspore::HashMap &attrs) { auto iter = attrs.find(IN_STRATEGY); return !((iter == attrs.end()) || (iter->second->type_name() == NONE)); } +/** + * @brief Checks if the specified target attribute is found and not of type 'NONE' in the given attributes. + * + * @param attrs A HashMap containing the attributes to check. + * @param target The name of the target attribute to check for. + * @return true if the target attribute is found and not of type 'NONE', false otherwise. + */ bool AttrFound(const mindspore::HashMap &attrs, const std::string &target) { auto iter = attrs.find(target); return !((iter == attrs.end()) || (iter->second->type_name() == NONE)); } +/** + * @brief Checks if any node within the FuncGraph has the 'IN_STRATEGY' attribute. + * + * This function performs a deep search within the FuncGraph to find nodes with the 'IN_STRATEGY' attribute. + * + * @param root The root FuncGraph to search within. + * @return true if any node within the FuncGraph has the 'IN_STRATEGY' attribute, false otherwise. + */ bool HasStrategy(const FuncGraphPtr &root) { AnfNodePtr ret = root->get_return(); MS_EXCEPTION_IF_NULL(ret); @@ -513,11 +696,33 @@ bool HasStrategy(const FuncGraphPtr &root) { return false; } +/** + * @brief Checks if a given primitive operation is a communication operation. + * + * This function takes a PrimitivePtr as input and checks if it matches the names of communication operations + * listed in the COMMUNICATION_OPS set. If the name of the primitive operation matches any communication operation, + * it returns true; otherwise, it returns false. + * + * @param prim The PrimitivePtr to be checked. + * @return True if the primitive operation is a communication operation; false otherwise. + */ bool IsCommunicationOp(const PrimitivePtr &prim) { MS_EXCEPTION_IF_NULL(prim); return (COMMUNICATION_OPS.find(prim->name()) != COMMUNICATION_OPS.end()); } +/** + * @brief Searches for communication operations in a list of AnfNodes. + * + * This function iterates through a list of AnfNodes and checks if any CNode represents a communication operation. + * It first verifies if the AnfNode is a CNode and if its input(0) is a ValueNode containing a Primitive operation. + * If both conditions are met and the Primitive operation is identified as a communication operation using + * the IsCommunicationOp function, it logs the occurrence and returns true. If no communication operation is found, + * it returns false. + * + * @param all_nodes The list of AnfNodes to search for communication operations. + * @return True if a communication operation is found; false otherwise. + */ bool FindCommunicationOp(const std::vector &all_nodes) { for (auto &node : all_nodes) { MS_EXCEPTION_IF_NULL(node); @@ -543,6 +748,19 @@ bool FindCommunicationOp(const std::vector &all_nodes) { return false; } +/** + * StepRedistribution function + * + * This function performs tensor redistribution for distributed computing. It recursively traverses the graph + * starting from a given `node` to identify potential redistribution points and inserts redistribution operations + * as needed. + * + * @param node: The current node being analyzed. + * @param distribute_operator: The operator information for distributed computing. + * @param insert_node: The insertion point for redistribution operations. + * @param tensor_redistribution: The tensor redistribution configuration. + * @param pre_node: The previous node in the graph traversal (used for tracking). + */ void StepRedistribution(const CNodePtr &node, const OperatorInfoPtr &distribute_operator, const CNodePtr &insert_node, const TensorRedistribution &tensor_redistribution, const CNodePtr &pre_node) { MS_EXCEPTION_IF_NULL(node->func_graph()); @@ -551,13 +769,17 @@ void StepRedistribution(const CNodePtr &node, const OperatorInfoPtr &distribute_ AnfNodeIndexSet node_set = manager->node_users()[node]; CNodePtr insert_node_new; + // Skip Send nodes in the analysis if (IsPrimitiveCNode(node, prim::kPrimSend)) { return; } + + // Skip nodes between 'make_tuple' and the next node, as no redistribution is needed. if (AnfNodeIsPrimitive(node, MAKE_TUPLE) || AnfNodeIsPrimitive(node, MAKE_LIST)) { MS_LOG(INFO) << "No need to insert redistribution op between make_tuple node and the next node"; return; } + if (IsValueNode(node->input(0))) { auto current_value = node->input(0)->cast(); MS_EXCEPTION_IF_NULL(current_value); @@ -567,10 +789,14 @@ void StepRedistribution(const CNodePtr &node, const OperatorInfoPtr &distribute_ } else { insert_node_new = insert_node; } + MS_EXCEPTION_IF_NULL(insert_node_new); + for (auto &node_pair : node_set) { CNodePtr use_cnode = node_pair.first->cast(); MS_EXCEPTION_IF_NULL(use_cnode); + + // Recursively analyze non-primitive nodes if (!IsValueNode(use_cnode->input(0))) { StepRedistribution(use_cnode, distribute_operator, insert_node_new, tensor_redistribution, pre_node); } else { @@ -578,26 +804,33 @@ void StepRedistribution(const CNodePtr &node, const OperatorInfoPtr &distribute_ MS_EXCEPTION_IF_NULL(prim_anf_node); PrimitivePtr node_prim = prim_anf_node->value()->cast(); MS_EXCEPTION_IF_NULL(node_prim); + + // Skip DEPEND and UPDATESTATE primitives if ((node_prim->name() == DEPEND && node_pair.second != 1) || node_prim->name() == UPDATESTATE) { continue; } + + // Perform redistribution for parallel care nodes that have OperatorInfo data if (IsParallelCareNode(use_cnode) && use_cnode->has_user_data()) { Redistribution(node_pair, distribute_operator, insert_node_new, node_pair.second, tensor_redistribution, pre_node); } else { + // Recursively analyze other nodes StepRedistribution(use_cnode, distribute_operator, insert_node_new, tensor_redistribution, pre_node); } } } } +// SplitTensor function +// This function is responsible for splitting a tensor into smaller slices for a specific operation. void SplitTensor(const AnfNodePtr &node, const CNodePtr &next_node, int64_t index) { MS_EXCEPTION_IF_NULL(node); MS_EXCEPTION_IF_NULL(next_node); OperatorInfoPtr op_info = next_node->user_data(); MS_EXCEPTION_IF_NULL(op_info); - // If the shape of tensor is [] or [1], no need to split it. + // Step 1: Check if the shape of the tensor is empty or has only one element, in which case no split is needed. Shapes shapes = GetNodeShape(node); if (shapes.size() != 1) { MS_LOG(EXCEPTION) << "Split tensor for " << op_info->name() @@ -613,7 +846,7 @@ void SplitTensor(const AnfNodePtr &node, const CNodePtr &next_node, int64_t inde MS_LOG(INFO) << "Split tensor for " << op_info->name() << ": The shape of tensor is " << shape_str; - // extract tensor layout + // Step 2: Extract tensor layout information. if (LongToSize(index - 1) >= op_info->inputs_tensor_info().size()) { MS_LOG(EXCEPTION) << "The index is out of range, index is " << (index - 1) << ", vector size is " << op_info->inputs_tensor_info().size(); @@ -621,11 +854,13 @@ void SplitTensor(const AnfNodePtr &node, const CNodePtr &next_node, int64_t inde TensorInfo tensor_info = op_info->inputs_tensor_info()[LongToSize(index - 1)]; TensorLayout tensor_layout = tensor_info.tensor_layout(); - // Use _GetTensorSlice operator to split the tensor + // Step 3: Create and insert the _GetTensorSlice operator to split the tensor. FuncGraphPtr func_graph = next_node->func_graph(); // only cnode can get the graph MS_EXCEPTION_IF_NULL(func_graph); Operator op = CreateGetTensorSliceOp(tensor_layout); InsertGetTensorSliceOp(op, next_node, func_graph, index, SPLIT_TENSOR); + + // Step 4: If the operation has sub-operators, insert _GetTensorSlice operators for them as well. if (!op_info->sub_ops().empty()) { auto sub_ops = op_info->sub_ops(); for (size_t i = 0; i < sub_ops.size(); i++) { @@ -636,7 +871,14 @@ void SplitTensor(const AnfNodePtr &node, const CNodePtr &next_node, int64_t inde } } +// SplitTensorList function +// This function splits a tensor list into individual tensors and replaces the original ValueNode with a MakeTuple operation. +// Parameters: +// - node: The ValueNode containing the tensor list to be split. +// - next_node: The CNode that uses the tensor list. +// - index: The index of the input in the next_node that corresponds to the tensor list. void SplitTensorList(const AnfNodePtr &node, const CNodePtr &next_node, int index) { + // Check if the inputs and index meet the expected conditions. MS_EXCEPTION_IF_NULL(node); MS_EXCEPTION_IF_NULL(next_node); if (next_node->inputs().size() != 2 || index != 1) { @@ -647,16 +889,21 @@ void SplitTensorList(const AnfNodePtr &node, const CNodePtr &next_node, int inde OperatorInfoPtr op_info = next_node->user_data(); MS_EXCEPTION_IF_NULL(op_info); + // Get the values from the input tensor list. std::vector inputs_values; if (IsValueNode(node)) { inputs_values = node->cast()->value()->cast()->value(); } else { inputs_values = node->cast()->value()->cast()->value(); } + + // Check if the number of input values matches the expected size. if (inputs_values.size() != op_info->inputs_tensor_info().size()) { MS_LOG(EXCEPTION) << "The inputs size " << inputs_values.size() << ", is not equal to inputs shape size " << op_info->inputs_tensor_info().size(); } + + // Create a MakeTuple operation to replace the original ValueNode. std::vector make_tuple_inputs = {NewValueNode(prim::kPrimMakeTuple)}; FuncGraphPtr func_graph = next_node->func_graph(); MS_EXCEPTION_IF_NULL(func_graph); @@ -688,23 +935,40 @@ void SplitTensorList(const AnfNodePtr &node, const CNodePtr &next_node, int inde manager->Replace(node, make_tuple); } +// StepSplitTensor function +// This function iterates over nodes that use the given node and splits tensors when needed. +// Parameters: +// - node: The AnfNode to be checked for tensor splitting. +// - manager: The FuncGraphManager for managing the FuncGraph. void StepSplitTensor(const AnfNodePtr &node, const FuncGraphManagerPtr &manager) { + // Check if the node and manager are valid. MS_EXCEPTION_IF_NULL(node); MS_EXCEPTION_IF_NULL(manager); + + // Get the set of nodes that use the given node. AnfNodeIndexSet node_set = manager->node_users()[node]; + + // Iterate over the nodes in the set. for (auto &node_pair : node_set) { CNodePtr use_cnode = node_pair.first->cast(); + + // Check if the node is a ValueNode with a Primitive input (e.g., a primitive operator). if (use_cnode == nullptr || !IsValueNode(use_cnode->input(0))) { continue; } + ValueNodePtr prim_anf_node = use_cnode->input(0)->cast(); MS_EXCEPTION_IF_NULL(prim_anf_node); PrimitivePtr use_cnode_prim = prim_anf_node->value()->cast(); MS_EXCEPTION_IF_NULL(use_cnode_prim); + + // Check if the node is dependent on the current node or if it's in the set of ops that don't take input tensors. if ((use_cnode_prim->name() == DEPEND && node_pair.second != 1) || NO_INPUT_TENSOR_OPS.find(use_cnode_prim->name()) != NO_INPUT_TENSOR_OPS.end()) { continue; } + + // Check if the node is a parallel care node and split tensors accordingly. if (IsParallelCareNode(use_cnode)) { if (IsValueNode(node) || IsValueNode(node)) { SplitTensorList(node, use_cnode, node_pair.second); @@ -716,19 +980,21 @@ void StepSplitTensor(const AnfNodePtr &node, const FuncGraphManagerPtr &manager) } void StepReplaceOp(OperatorVector replace_op, const CNodePtr &node) { - // step1:get graph manager distribute_operator + // Step 1: Get the distribute_operator associated with the CNode OperatorInfoPtr distribute_operator = node->user_data(); if (distribute_operator == nullptr) { - MS_LOG(EXCEPTION) << "Failure:AddNode error since distribute_operator is nullptr"; + MS_LOG(EXCEPTION) << "Failure: AddNode error since distribute_operator is nullptr"; } + + // Step 2: Obtain the function graph and manager FuncGraphPtr func_graph = node->func_graph(); MS_EXCEPTION_IF_NULL(func_graph); FuncGraphManagerPtr manager = func_graph->manager(); if (manager == nullptr) { - MS_LOG(EXCEPTION) << "Failure:AddNode error since manager is nullptr"; + MS_LOG(EXCEPTION) << "Failure: AddNode error since manager is nullptr"; } - // When reshape(bool), insert cast in the begin and end of op_list to avoid AllGather(bool). + // Step 3: Handle special case for boolean reshape operations auto reshape_type_str = node->abstract()->BuildType()->ToString(); auto replace_op_info = distribute_operator->replace_op_info(); if (reshape_type_str.find(BOOL) != std::string::npos) { @@ -740,7 +1006,7 @@ void StepReplaceOp(OperatorVector replace_op, const CNodePtr &node) { (void)replace_op_info.insert(replace_op_info.end(), {false, 1}); } - // step2:traverse op_list and insert node + // Step 4: Traverse the replace_op vector and insert new CNodes std::reverse(replace_op.begin(), replace_op.end()); std::reverse(replace_op_info.begin(), replace_op_info.end()); if (!replace_op_info.empty() && replace_op_info.size() != replace_op.size()) { @@ -748,13 +1014,18 @@ void StepReplaceOp(OperatorVector replace_op, const CNodePtr &node) { } bool replace_op_info_flag = !replace_op_info.empty(); for (size_t index = 0; index < replace_op.size(); ++index) { + // Step 4.1: Create a unique instance name for the new CNode std::string instance_name = CreateInstanceName(node, index); + + // Step 4.2: Create a vector of AnfNodePtr representing the inputs for the new CNode std::vector replace_input; if (index != replace_op.size() - 1) { replace_input = CreateInput(replace_op[index], node, instance_name); } else { replace_input = ReplaceOpInput(replace_op[index], instance_name, node); } + + // Step 4.3: Create the new CNode and set its attributes CNodePtr replace_node = func_graph->NewCNode(replace_input); MS_EXCEPTION_IF_NULL(replace_node); ScopePtr scope = node->scope(); @@ -763,6 +1034,8 @@ void StepReplaceOp(OperatorVector replace_op, const CNodePtr &node) { PrimitivePtr prim = GetValueNode(replace_node->input(0)); PrimitivePtr origin_prim = GetValueNode(node->input(0)); SetUserAttrs(origin_prim->attrs(), prim); + + // Step 4.4: Set recompute attribute for communication ops auto origin_prim_attrs = origin_prim->attrs(); if (origin_prim_attrs.find(RECOMPUTE_COMM_OP) != origin_prim_attrs.end() && !GetValue(origin_prim_attrs[RECOMPUTE_COMM_OP]) && @@ -770,18 +1043,22 @@ void StepReplaceOp(OperatorVector replace_op, const CNodePtr &node) { MS_LOG(INFO) << "The redistribution node in reshape would not be recomputed."; prim->set_attr("recompute", MakeValue(false)); } + + // Step 4.5: Set additional attributes and flags for the last replace_node if (index == replace_op.size() - 1) { replace_node->set_user_data(node->user_data()); replace_node->set_primal_attrs(node->primal_attrs()); } replace_node->set_in_forward_flag(true); replace_input[0]->set_scope(scope); + + // Step 4.6: Insert the new CNode using the FuncGraphManager if (replace_op_info_flag && replace_op_info[index].first) { auto new_cnode = InsertMakeTuple(replace_node, replace_op_info[index].second, func_graph); new_cnode->set_primal_attrs(node->primal_attrs()); - (void)manager->Replace(node, new_cnode); // using Replace function to insert node + (void)manager->Replace(node, new_cnode); // Using Replace function to insert the node } else { - (void)manager->Replace(node, replace_node); // using Replace function to insert node + (void)manager->Replace(node, replace_node); // Using Replace function to insert the node } } MS_LOG(INFO) << "Insert ReplaceOp success for " << distribute_operator->name(); @@ -795,7 +1072,7 @@ void StepReplaceGraph(const ReplaceGraphPtr &replace_graph, const CNodePtr &node MS_EXCEPTION_IF_NULL(func_graph); FuncGraphManagerPtr manager = func_graph->manager(); if (manager == nullptr) { - MS_LOG(EXCEPTION) << "Failure:AddNode error since manager is nullptr"; + MS_LOG(EXCEPTION) << "Failure: AddNode error since manager is nullptr"; } // Solve the input order // For example input_node:{segment_sum:1, segment_sum:2, gahter:2} @@ -803,6 +1080,7 @@ void StepReplaceGraph(const ReplaceGraphPtr &replace_graph, const CNodePtr &node // However, the segment_sum operation needs two inputs, To solve this // We maintain a dict to count the times of the same operations, // and bind the inputs according to the times of the op appears. + // Step 1: Solve the input order for replace_graph mindspore::HashMap input_map = {}; static int appear_count = 0; for (auto &replace_input : replace_graph->first) { @@ -825,23 +1103,34 @@ void StepReplaceGraph(const ReplaceGraphPtr &replace_graph, const CNodePtr &node input_map[replace_input.first] = appear_count; manager->SetEdge(replace_input.first, appear_count, pre_node); } - // "(void)manager->Replace(replace_graph->first, pre_node);" can not be called + // Step 2: Replace the original node with the replace_graph's second node auto replace_output = replace_graph->second->cast(); MS_EXCEPTION_IF_NULL(replace_output); replace_output->set_primal_attrs(node->primal_attrs()); (void)manager->Replace(node, replace_output); } +// GetTupleGetItemIndex function +// This function extracts the index value from a TupleGetItem CNode. +// It checks the input CNode for validity and retrieves the index as an integer. +// Parameters: +// - cnode: The TupleGetItem CNode from which the index needs to be extracted. +// Returns: +// - int64_t: The extracted index value as an integer. int64_t GetTupleGetItemIndex(const CNodePtr &cnode) { MS_EXCEPTION_IF_NULL(cnode); + + // Check if the input CNode has exactly 3 inputs, as expected for TupleGetItem. if (cnode->inputs().size() != 3) { MS_LOG(EXCEPTION) << cnode->ToString() << " size( " << cnode->inputs().size() << " ) is not 3"; } + // Check if the index of TupleGetItem is a ValueNode. if (!cnode->input(TUPLE_GETITEM_INDEX_POS)->isa()) { MS_LOG(EXCEPTION) << "The index of tuple getitem is not a value node"; } + // Extract the index value as an Int64Imm and return it. ValuePtr tuple_index_value = GetValueNode(cnode->input(TUPLE_GETITEM_INDEX_POS)); MS_EXCEPTION_IF_NULL(tuple_index_value); if (!tuple_index_value->isa()) { @@ -850,6 +1139,12 @@ int64_t GetTupleGetItemIndex(const CNodePtr &cnode) { return tuple_index_value->cast()->value(); } +// InsertVirtualDivOp function +// This function inserts virtual division operations into the provided CNode. +// It iterates through the inputs of the CNode and inserts virtual division operations as needed. +// Parameters: +// - virtual_div_op: The list of virtual division operations to insert. +// - node: The CNode into which virtual division operations will be inserted. void InsertVirtualDivOp(const VirtualDivOp &virtual_div_op, const CNodePtr &node) { MS_EXCEPTION_IF_NULL(node); size_t node_size = node->inputs().size(); @@ -858,20 +1153,24 @@ void InsertVirtualDivOp(const VirtualDivOp &virtual_div_op, const CNodePtr &node FuncGraphManagerPtr manager = func_graph->manager(); MS_EXCEPTION_IF_NULL(manager); + // Special handling for dropout do mask, only insert virtual division into input[0]. if (IsSomePrimitive(node, DROPOUT_DO_MASK)) { MS_LOG(INFO) << "Handle dropout do mask, only insert the virtual div to input[0]"; node_size = 2; } + // Iterate through the inputs of the node. for (size_t index = 1; index < node_size; ++index) { AnfNodePtr input = node->input(index); MS_EXCEPTION_IF_NULL(input); - // if it is not a tensor, continue + + // Check if the input is not a tensor or has an abstract monad. if ((!input->isa() && !input->isa()) || HasAbstractMonad(input)) { - MS_LOG(INFO) << "insert div op: the index " << index << " is not tensor, skip"; + MS_LOG(INFO) << "insert div op: the index " << index << " is not a tensor, skip"; continue; } + // Iterate through the virtual_div_op list and insert virtual division operations. for (size_t pos = 0; pos < virtual_div_op.size(); ++pos) { std::string instance_name = CreateInstanceName(node, pos); InsertNode(virtual_div_op[pos], node, index, node->input(index), func_graph, instance_name); @@ -880,6 +1179,13 @@ void InsertVirtualDivOp(const VirtualDivOp &virtual_div_op, const CNodePtr &node } } +// InsertRealDivOpToNodeInput function +// This function inserts a RealDiv operator as the input to the given CNode. +// It is used for scaling operations within a group. +// Parameters: +// - node: The CNode to which the RealDiv operator will be inserted. +// - scale: The scaling factor for the RealDiv operator. +// - instance_name: A unique identifier for the instance of the operator. void InsertRealDivOpToNodeInput(const CNodePtr &node, int64_t scale, const string &instance_name) { MS_EXCEPTION_IF_NULL(node); if (scale == 0) { @@ -888,14 +1194,15 @@ void InsertRealDivOpToNodeInput(const CNodePtr &node, int64_t scale, const strin size_t node_size = node->inputs().size(); FuncGraphPtr func_graph = node->func_graph(); MS_EXCEPTION_IF_NULL(func_graph); - // instance the real div operator + + // Instantiate the RealDiv operator Operator div_op = CreateDivOp(scale); // Insert it as the input of the node for (size_t index = 1; index < node_size; ++index) { AnfNodePtr input = node->input(index); MS_EXCEPTION_IF_NULL(input); - // if it is not a tensor, continue + // If it is not a tensor, continue if ((!input->isa() && !input->isa()) || HasAbstractMonad(input)) { continue; } @@ -903,12 +1210,20 @@ void InsertRealDivOpToNodeInput(const CNodePtr &node, int64_t scale, const strin } } +// InsertAllReduceToNodeInput function +// This function inserts an AllReduce operator as the input to the given CNode. +// It is used for reducing values across a group. +// Parameters: +// - node: The CNode to which the AllReduce operator will be inserted. +// - group: The communication group for the AllReduce operation. +// - instance_name: A unique identifier for the instance of the operator. void InsertAllReduceToNodeInput(const CNodePtr &node, const std::string &group, const std::string &instance_name) { MS_EXCEPTION_IF_NULL(node); size_t node_size = node->inputs().size(); FuncGraphPtr func_graph = node->func_graph(); MS_EXCEPTION_IF_NULL(func_graph); - // instance the real div operator + + // Instantiate the AllReduce operator CheckGlobalDeviceManager(); Operator allreduce_op = CreateAllReduceOp(REDUCE_OP_SUM, group); @@ -916,7 +1231,7 @@ void InsertAllReduceToNodeInput(const CNodePtr &node, const std::string &group, for (size_t index = 1; index < node_size; ++index) { AnfNodePtr input = node->input(index); MS_EXCEPTION_IF_NULL(input); - // if it is not a tensor, continue + // If it is not a tensor, continue if ((!input->isa() && !input->isa()) || HasAbstractMonad(input)) { continue; } @@ -925,6 +1240,13 @@ void InsertAllReduceToNodeInput(const CNodePtr &node, const std::string &group, } } +// PynativeParallelGraph function +// This function extracts the real graph from a hierarchy of graphs used in PyNative parallel execution. +// Parameters: +// - root: The root FuncGraph. +// - all_nodes: A vector of AnfNodePtr representing all nodes in the hierarchy. +// Returns: +// - FuncGraphPtr: The extracted real graph from the hierarchy. FuncGraphPtr PynativeParallelGraph(const FuncGraphPtr &root, const std::vector &all_nodes) { FuncGraphPtr real_graph = root; for (auto &node : all_nodes) { @@ -944,49 +1266,77 @@ FuncGraphPtr PynativeParallelGraph(const FuncGraphPtr &root, const std::vector &all_nodes) { - std::vector last_forward_node_ids; - std::vector last_indexs; + std::vector last_forward_node_ids; // Stores unique IDs of last forward nodes + std::vector last_indexs; // Stores indices of last forward nodes auto real_graph = PynativeParallelGraph(root, all_nodes); + + // Find unique IDs and indices of last forward nodes FindLastNodesUniqueId(real_graph, &last_forward_node_ids, &last_indexs); - MS_LOG(INFO) << "there are " << last_forward_node_ids.size() << " output nodes in eval/predict"; + MS_LOG(INFO) << "There are " << last_forward_node_ids.size() << " output nodes in eval/predict"; + + // Iterate through all nodes in the computation graph for (auto &node : all_nodes) { - // here insert virtualoutput node auto cnode = node->cast(); + if (cnode == nullptr) { continue; } + + // Check if the current node is one of the last forward nodes auto last_node_iter = std::find(last_forward_node_ids.begin(), last_forward_node_ids.end(), cnode->UniqueId()); + if (last_node_iter == last_forward_node_ids.end()) { continue; } + + // Iterate through the last forward nodes for (size_t last_node_index = 0; last_node_index < last_forward_node_ids.size(); ++last_node_index) { if (last_forward_node_ids[last_node_index] != cnode->UniqueId()) { continue; } - MS_LOG(INFO) << "find last node: " << cnode->fullname_with_scope() << ", the parallel care node is: " - << cnode->input(last_indexs[last_node_index])->fullname_with_scope(); + + MS_LOG(INFO) << "Found last node: " << cnode->fullname_with_scope() + << ", the parallel care node is: " << cnode->input(last_indexs[last_node_index])->fullname_with_scope(); + + // Handle the special case of tuple_get_item if (IsPrimitiveCNode(cnode, prim::kPrimTupleGetItem)) { FuncGraphManagerPtr manager = cnode->func_graph()->manager(); MS_EXCEPTION_IF_NULL(manager); auto node_pair = manager->node_users()[cnode].front(); + if (!node_pair.first->isa()) { - MS_LOG(EXCEPTION) << "the output of tuple_get_item is not a cnode"; + MS_LOG(EXCEPTION) << "The output of tuple_get_item is not a cnode"; } + cnode = node_pair.first->cast(); last_indexs[last_node_index] = IntToSize(node_pair.second); } + auto pre_node = cnode->input(last_indexs[last_node_index]); Shapes shape_outputs = GetNodeShape(pre_node); + if (shape_outputs[0].empty()) { continue; } + FuncGraphPtr func_graph = node->func_graph(); MS_EXCEPTION_IF_NULL(func_graph); OperatorParams params; OperatorAttrs attrs; OperatorArgs args = std::make_pair(attrs, params); Operator op = std::make_pair(VIRTUAL_OUTPUT, args); + + // Insert virtual output node InsertNode(op, cnode, last_indexs[last_node_index], pre_node, func_graph, VIRTUAL_OUTPUT); auto virtual_output_node = cnode->input(last_indexs[last_node_index]); AbstractBasePtr virtual_output_abstract = pre_node->abstract()->Clone(); @@ -1006,64 +1356,118 @@ CNodePtr SkipTrivialNodesMoveDown(const FuncGraphManagerPtr &manager, CNodePtr n return node; } -std::pair FindCNode(const AnfNodePtr &anode, const std::string &name, const FuncGraphPtr &func_graph, - size_t max_depth) { +/** + * @brief Finds a CNode with a specific name in the call hierarchy of a given AnfNode. + * + * This function searches for a CNode with the specified name in the call hierarchy of the provided AnfNode. + * It traverses the call graph up to a maximum depth (controlled by `max_depth`) to find the CNode. + * + * @param anode The AnfNode to start the search from. + * @param name The name of the Primitive to be found. + * @param func_graph The target FuncGraph to which the found CNode should belong. + * @param max_depth The maximum depth to traverse while searching. + * + * @return A pair consisting of a boolean value indicating if the CNode was found (`true` if found, `false` otherwise) + * and a CNodePtr representing the found CNode (or `nullptr` if not found). + */ +std::pair FindCNode(const AnfNodePtr &anode, const std::string &name, + const FuncGraphPtr &func_graph, size_t max_depth) { MS_EXCEPTION_IF_NULL(anode); MS_EXCEPTION_IF_NULL(anode->func_graph()); FuncGraphManagerPtr manager = anode->func_graph()->manager(); MS_EXCEPTION_IF_NULL(manager); + + // Check if the recursive depth exceeds the maximum allowed depth if (max_depth > MAX_RECURSIVE_DEPTH) { - MS_LOG(EXCEPTION) << "Recursive call is larger than 100000."; + MS_LOG(EXCEPTION) << "Recursive call depth exceeds the limit (100000)."; } + + // Retrieve the set of nodes that use the input AnfNode AnfNodeIndexSet node_set = manager->node_users()[anode]; bool result = false; CNodePtr cnode_return = nullptr; + + // Iterate through the nodes using the input AnfNode for (auto &node_pair : node_set) { CNodePtr use_apply = node_pair.first->cast(); + + // Check if the node is a CNode and if its first input is a Primitive if (use_apply == nullptr || !IsValueNode(use_apply->input(0))) { continue; } + + // Skip trivial nodes by moving down the graph use_apply = SkipTrivialNodesMoveDown(manager, use_apply); + + // Check if the updated node is still a CNode and if its first input is a Primitive if (use_apply == nullptr || !IsValueNode(use_apply->input(0))) { continue; } + ValueNodePtr prim_anf_node = use_apply->input(0)->cast(); MS_EXCEPTION_IF_NULL(prim_anf_node); PrimitivePtr node_prim = prim_anf_node->value()->cast(); MS_EXCEPTION_IF_NULL(node_prim); + + // Check if the Primitive's name matches the specified name and it's used only once if (node_prim->name() == name && node_pair.second == 1) { if (use_apply->func_graph() == func_graph) { result = true; cnode_return = use_apply; - MS_LOG(INFO) << "Find Primitive " << name << " in the same func_graph"; + MS_LOG(INFO) << "Found Primitive " << name << " in the same func_graph"; continue; } - MS_LOG(INFO) << "Find Primitive " << name << " in different func_graph"; + MS_LOG(INFO) << "Found Primitive " << name << " in a different func_graph"; } + + // Check if ParallelOptimizer is enabled and the node is in the AllGather node list if (ParallelContext::GetInstance()->enable_parallel_optimizer() && IsInAllGatherNodeList(use_apply)) { return FindCNode(node_pair.first, name, func_graph, max_depth + 1); } } + return std::make_pair(result, cnode_return); } +/** + * @brief Inserts a mirror node before a Cast operation if certain conditions are met. + * + * This function checks if it should insert a mirror node before a Cast operation based on specific conditions. + * The conditions include checking if gradient_fp32_sync is enabled and the previous node is a Cast operation + * with a type other than float32. + * + * @param node The CNode representing the Cast operation. + * @param index The input index of the Cast operation. + * + * @return `true` if a mirror node should be inserted, `false` otherwise. + */ bool InsertMirrorBeforeCast(const CNodePtr &node, size_t index) { - // only if gradient_fp32_sync is true, pre node is cast and type is not float32 return true + // Check if gradient_fp32_sync is enabled if (!ParallelContext::GetInstance()->gradient_fp32_sync()) { return false; } + + // Get the previous node of the Cast operation auto pre_node = node->input(index); MS_EXCEPTION_IF_NULL(pre_node); + + // Check if the previous node is a CNode and if its first input is a Primitive auto cnode = pre_node->cast(); if (cnode == nullptr || !IsValueNode(cnode->input(0))) { return false; } + + // If ParallelOptimizer is enabled and the node is in the AllGather node list, update the previous node if (ParallelContext::GetInstance()->enable_parallel_optimizer() && IsInAllGatherNodeList(cnode)) { pre_node = cnode->input(1); } + + // Check if the previous node is a Cast operation and its type is not float32 if (!IsPrimitiveCNode(pre_node, prim::kPrimCast)) { return false; } + + // Get the type of the previous node and check if it's not float32 auto node_type = pre_node->Type(); MS_EXCEPTION_IF_NULL(node_type); if (!node_type->isa()) { @@ -1072,10 +1476,26 @@ bool InsertMirrorBeforeCast(const CNodePtr &node, size_t index) { auto input_element_type = node_type->cast()->element(); MS_EXCEPTION_IF_NULL(input_element_type); auto type_id = input_element_type->type_id(); - + return (type_id != kNumberTypeFloat32); } +/** + * @brief Checks whether to insert mirror operations based on specific conditions. + * + * This function examines the given `node` and determines whether it is necessary to insert mirror operations. + * The conditions for insertion are as follows: + * - If the `node` is a Send primitive, insertion is required. + * - If the `node` has exactly 2 inputs and the second input is a ValueNode of type ValueSequence, it is skipped. + * - If the `node` has exactly 2 inputs and the second input is a primitive node (MAKE_TUPLE or MAKE_LIST), it is skipped. + * - If the size of `mirror_ops` is not equal to `node_size - 1`, an exception is thrown. + * + * @param[in] mirror_ops The MirrorOps container holding mirror operations. + * @param[in] node The CNode to be checked for insertion. + * @param[in] node_size The size of the CNode's inputs. + * + * @return True if mirror operations should be inserted, false otherwise. + */ static bool CheckInsertMirrorOps(const MirrorOps &mirror_ops, const CNodePtr &node, size_t node_size) { if (IsPrimitiveCNode(node, prim::kPrimSend)) { return true; @@ -1088,18 +1508,28 @@ static bool CheckInsertMirrorOps(const MirrorOps &mirror_ops, const CNodePtr &no if ((node->inputs().size() == kSingleArgCNodeSize) && (AnfNodeIsPrimitive(node->input(1), MAKE_TUPLE) || AnfNodeIsPrimitive(node->input(1), MAKE_LIST))) { - MS_LOG(INFO) << "The mirror for " << GetPrimName(node) << " has handle by make_tuple node"; + MS_LOG(INFO) << "The mirror for " << GetPrimName(node) << " has been handled by make_tuple node"; return false; } if (mirror_ops.size() != node_size - 1) { - MS_LOG(EXCEPTION) << "Mirrorops's size is wrong! mirror_ops size is " << mirror_ops.size() << ", node_size is " + MS_LOG(EXCEPTION) << "Mirrorops' size is incorrect! mirror_ops size is " << mirror_ops.size() << ", node_size is " << (node_size - 1); } return true; } -// only used for InsertMirrorOps +/** + * @brief Moves up the CNode `node`, skipping trivial nodes. + * + * This function is intended for use with InsertMirrorOps. It takes a CNode `node` and iterates upward, + * skipping trivial nodes found in the TrivialNodeList or AllGatherNodeList. It returns the first non-trivial node + * encountered or nullptr if none is found. + * + * @param[in] node The CNode to be moved up from. + * + * @return The first non-trivial CNode encountered after skipping trivial nodes, or nullptr if none is found. + */ CNodePtr SkipTrivialNodesMoveUp(CNodePtr node) { MS_EXCEPTION_IF_NULL(node); while (!IsSomePrimitive(node, LOAD)) { @@ -1120,6 +1550,14 @@ CNodePtr SkipTrivialNodesMoveUp(CNodePtr node) { return node; } +/** + * @brief Generates the name for the mirror operator based on pipeline settings. + * + * This function constructs the name for the mirror operator based on the current pipeline configuration, + * including the number of gradient accumulation steps and the number of pipeline stage splits. + * + * @return The generated mirror operator name. + */ std::string MirrorOpName() { int64_t grad_accumulation_step = ParallelContext::GetInstance()->grad_accumulation_step(); int64_t split_stage_num = ParallelContext::GetInstance()->pipeline_stage_split_num(); @@ -1134,151 +1572,210 @@ std::string MirrorOpName() { return mirror_op_name; } +// DoInsertMirrorOps function +// This function inserts mirror operations into the computation graph based on the given mirror_ops and node information. static void DoInsertMirrorOps(const FuncGraphPtr &root, const MirrorOps &mirror_ops, const CNodePtr &node, size_t node_size) { + // Step 1: Retrieve necessary information about the node and its graph context FuncGraphPtr func_graph = node->func_graph(); MS_EXCEPTION_IF_NULL(func_graph); FuncGraphManagerPtr manager = func_graph->manager(); MS_EXCEPTION_IF_NULL(manager); + // Step 2: Iterate over the mirror operations and insert them into the graph for (size_t index = 1; index < node_size; ++index) { + // Retrieve the backward_op corresponding to the current index OperatorVector backward_op = mirror_ops[index - 1]; + + // Handle special case for primitive "Send" if (IsPrimitiveCNode(node, prim::kPrimSend)) { auto param_index = GetValue(node->GetPrimalAttr(PARAM_INDEX)); backward_op = mirror_ops[IntToSize(param_index)]; } + + // Continue if the backward_op is empty if (backward_op.empty()) { continue; } + + // Find the parameter node connected to the current input of the node std::pair param_node_pair = FindParameter(node->input(index), func_graph); + + // Continue if no parameter node is found if (!param_node_pair.first) { continue; } + // Retrieve parameter information, including its name and gradient requirement auto param_ptr = param_node_pair.first->cast(); std::string param_name; bool is_shared_param = false; if (param_ptr) { param_name = param_ptr->name(); if (!param_ptr->param_info() || !param_ptr->param_info()->requires_grad()) { - MS_LOG(INFO) << param_name << " do not need gradient. Skip inserting mirror."; + MS_LOG(INFO) << param_name << " does not need gradient. Skip inserting mirror."; continue; } + + // Check if the parameter has shard mirror group information std::string opt_shard_mirror_group; if (param_ptr->user_data()) { opt_shard_mirror_group = param_ptr->user_data()->opt_shard_mirror_group(); is_shared_param = param_ptr->user_data()->is_shared_param(); } + + // If the parameter has shard mirror group information, create mirror operations based on group size if (!opt_shard_mirror_group.empty()) { - // mirror ops is covered in not fully use opt shard case uint32_t group_rank_size = 0; if (!CommManager::GetInstance().GetRankSize(opt_shard_mirror_group, &group_rank_size)) { - MS_LOG(EXCEPTION) << "Got the group size from the group " << opt_shard_mirror_group << " failed"; + MS_LOG(EXCEPTION) << "Failed to get group size from group " << opt_shard_mirror_group; } backward_op = CreateMirrorOps(opt_shard_mirror_group, static_cast(group_rank_size)); } } - // not a RefKey + + // Determine if a RefKey is used std::string mirror_op_name = MirrorOpName(); AnfNodePtr pre_node = node->input(index); + + // If there is no RefKey and a MirrorOp is found in the same graph, use the existing MirrorOp CNode as input if (!param_node_pair.second) { auto next_cnode = FindCNode(param_node_pair.first, mirror_op_name, func_graph, 0); - // if there is already a MirrorOp in the same graph, use MirrorOp CNode as a input instead if (next_cnode.first) { MS_EXCEPTION_IF_NULL(next_cnode.second); - // assume Load is inserted next to parameter - // skip Load moving up and insert mirror next to the parameter + // Assuming Load is inserted next to the parameter, skip Load moving up and insert mirror next to the parameter if (pre_node->cast()) { CNodePtr load_node = SkipTrivialNodesMoveUp(node->input(index)->cast()); manager->SetEdge(load_node, 1, next_cnode.second); } else { manager->SetEdge(node, static_cast(index), next_cnode.second); } - MS_LOG(INFO) << "Find parameter " << param_name << " for node " << GetPrimName(node->cast()) - << " and share the mirror."; + MS_LOG(INFO) << "Found parameter " << param_name << " for node " << GetPrimName(node->cast()) + << " and sharing the mirror."; continue; } } - // if the parameter found is a RefKey, or no MirrorOp is found in the same graph, insert a new MirrorOp - // only one MirrorOp in backward_op + + // If the parameter is a RefKey or no MirrorOp is found in the same graph, insert a new MirrorOp if (backward_op.size() != 1) { - MS_LOG(EXCEPTION) << "backward_op size must be 1, real is " << backward_op.size(); + MS_LOG(EXCEPTION) << "backward_op size must be 1, actual size is " << backward_op.size(); } + auto op = backward_op[0]; + + // Handle insertion when the parameter is a cast node or is_shared_param is true if (pre_node->cast() && (InsertMirrorBeforeCast(node, index) || is_shared_param)) { - // assume Load is inserted next to parameter - // skip Load moving up and insert mirror next to the parameter + // Assuming Load is inserted next to the parameter, skip Load moving up and insert mirror next to the parameter CNodePtr load_node = SkipTrivialNodesMoveUp(pre_node->cast()); InsertNode(op, load_node, 1, load_node->input(1), func_graph, mirror_op_name, param_name, root); auto comm_op = load_node->input(1)->cast(); - // add fusion flag + // Add fusion flag AddCommOpFusionType(comm_op, param_node_pair.first); - MS_LOG(INFO) << "Find parameter " << param_name << " for node " << GetPrimName(node->cast()) - << " and insert mirror before Load"; + MS_LOG(INFO) << "Found parameter " << param_name << " for node " << GetPrimName(node->cast()) + << " and inserted mirror before Load."; AddCommOpParamFlag(comm_op); continue; } + + // Insert a new MirrorOp before the node InsertNode(op, node, index, pre_node, func_graph, mirror_op_name, param_name, root); - MS_LOG(INFO) << "Find parameter " << param_name << " for node " << GetPrimName(node->cast()) - << " and insert mirror before the node"; + MS_LOG(INFO) << "Found parameter " << param_name << " for node " << GetPrimName(node->cast()) + << " and inserted mirror before the node."; auto comm_op = node->input(index)->cast(); - // add fusion flag - // pipeline mirror would not be set, which should be supported later + // Add fusion flag + // Pipeline mirror would not be set, which should be supported later AddCommOpFusionType(comm_op, param_node_pair.first); AddCommOpParamFlag(comm_op); } } +// InsertMirrorOps function +// This function inserts mirror operations into the computational graph. +// It checks whether to insert mirror operations and then performs the insertion. void InsertMirrorOps(const FuncGraphPtr &root, const MirrorOps &mirror_ops, const CNodePtr &node) { MS_EXCEPTION_IF_NULL(node); size_t node_size = node->inputs().size(); + // Calculate the number of inputs that are not abstract monads for (auto input : node->inputs()) { if (HasAbstractMonad(input)) { node_size--; } } + // Check if mirror operations should be inserted based on certain conditions if (!CheckInsertMirrorOps(mirror_ops, node, node_size)) { return; } + // Perform the insertion of mirror operations DoInsertMirrorOps(root, mirror_ops, node, node_size); } -void BackwardCommunication(const FuncGraphPtr &root, const OperatorInfoPtr &distribute_operator, const CNodePtr &node, - const std::vector> &sens_loss_pairs) { +// BackwardCommunication function +// This function handles backward communication for distributed training. +// It inserts mirror and virtual div operations based on certain conditions. +void BackwardCommunication(const FuncGraphPtr &root, const OperatorInfoPtr &distribute_operator, + const CNodePtr &node, const std::vector> &sens_loss_pairs) { MS_EXCEPTION_IF_NULL(distribute_operator); MS_EXCEPTION_IF_NULL(node); + // Skip the operation if it is a primitive "Receive" operation if (IsPrimitiveCNode(node, prim::kPrimReceive)) { return; } + + // Check if the current node is a loss node bool is_loss_cnode = std::any_of(sens_loss_pairs.begin(), sens_loss_pairs.end(), [node](const std::pair &element) { return element.second.loss_node == node; }); + // Retrieve mirror and virtual div operations from the distribute operator MirrorOps mirror_ops = distribute_operator->mirror_ops(); VirtualDivOp virtual_div_op = distribute_operator->virtual_div_op(); - // insert mirror op + + // Insert mirror operations if they exist if (!mirror_ops.empty()) { - MS_LOG(INFO) << "insert mirror op for " << distribute_operator->name(); + MS_LOG(INFO) << "Inserting mirror ops for " << distribute_operator->name(); InsertMirrorOps(root, mirror_ops, node); } - // insert virtual div op + + // Insert virtual div operations if they exist and the node is a loss node in the last stage if (!virtual_div_op.empty() && is_loss_cnode && IsLastStage()) { - MS_LOG(INFO) << "insert virtual div op for " << distribute_operator->name(); + MS_LOG(INFO) << "Inserting virtual div ops for " << distribute_operator->name(); InsertVirtualDivOp(virtual_div_op, node); } } +// GetDisOpName function +// This function retrieves the operator name from the primitive name. std::string GetDisOpName(const std::string &prim_name) { std::string op_name = prim_name; + // Remove the leading underscore from the primitive name if it exists if (!prim_name.empty() && (prim_name[0] == '_')) { op_name = prim_name.substr(1); } + // Append "Info" to the operator name and return it return op_name + "Info"; } +/** + * @brief Create an operator instance based on the operator name, attributes, and shape list. + * + * This function creates an instance of an operator with the given name, attributes, and shape lists. + * It first checks the size of the shape list to ensure it is of size 2. + * If the shape list size is not 2, it logs an error and returns nullptr. + * Then, it retrieves the distributed operator name using `GetDisOpName` based on the input name. + * It attempts to create the operator using the `DynCreator` singleton instance. + * If the creation is successful, it modifies the operator's name by appending a unique identifier to it, + * increments the total number of created operators, and logs a success message. + * If the creation fails, it logs an error message and returns nullptr. + * + * @param name The name of the operator. + * @param attrs The primitive attributes associated with the operator. + * @param shape_list The list of input and output shapes. + * + * @return An instance of OperatorInfoPtr pointing to the created operator, or nullptr if creation failed. + */ OperatorInfoPtr OperatorInstanceByName(const std::string &name, const PrimitiveAttrs &attrs, const std::vector &shape_list) { if (shape_list.size() != 2) { @@ -1302,6 +1799,23 @@ OperatorInfoPtr OperatorInstanceByName(const std::string &name, const PrimitiveA return operator_; } +/** + * @brief Create an operator instance based on the given primitive, attributes, and shape list. + * + * This function creates an operator instance using the provided primitive, attributes, and shape lists. + * It first checks if the primitive pointer is not null, and then calls `OperatorInstanceByName` to create the operator. + * If the creation is successful, it returns the operator instance. + * If the creation fails and the primitive is blacklisted for batch parallel, it logs an exception. + * Otherwise, it logs an error message and attempts to create the operator with the name 'BATCH_PARALLEL'. + * If this fallback creation is also unsuccessful, it returns nullptr. + * + * @param prim The primitive associated with the operator. + * @param attrs The primitive attributes associated with the operator. + * @param shape_list The list of input and output shapes. + * + * @return An instance of OperatorInfoPtr pointing to the created operator. + * @throws An exception if the primitive is blacklisted for batch parallel and creation fails. + */ OperatorInfoPtr OperatorInstance(const PrimitivePtr &prim, const PrimitiveAttrs &attrs, const std::vector &shape_list) { MS_EXCEPTION_IF_NULL(prim); @@ -1317,6 +1831,19 @@ OperatorInfoPtr OperatorInstance(const PrimitivePtr &prim, const PrimitiveAttrs return operator_; } +/** + * @brief Create a new operator instance based on the given primitive, attributes, and shape list. + * + * This function creates a new operator instance using the provided primitive, attributes, and shape lists. + * It first calls `OperatorInstance` to create the operator. + * After creation, it logs the input shapes for debugging purposes and returns the operator instance. + * + * @param prim The primitive associated with the operator. + * @param attrs The primitive attributes associated with the operator. + * @param shape_list The list of input and output shapes. + * + * @return An instance of OperatorInfoPtr pointing to the created operator. + */ OperatorInfoPtr NewOperatorInstance(const PrimitivePtr &prim, const PrimitiveAttrs &attrs, std::vector shape_list) { OperatorInfoPtr operator_ = OperatorInstance(prim, attrs, shape_list); @@ -1326,6 +1853,10 @@ OperatorInfoPtr NewOperatorInstance(const PrimitivePtr &prim, const PrimitiveAtt return operator_; } +// ExtractStrategy function +// This function extracts a strategy from a given ValuePtr, which represents a strategy. +// It handles the case where the input ValuePtr is null, not a ValueTuple, or has the wrong format. +// The extracted strategy is used for parallel execution. StrategyPtr ExtractStrategy(const ValuePtr &stra) { if (stra == nullptr) { return nullptr; @@ -1337,7 +1868,7 @@ StrategyPtr ExtractStrategy(const ValuePtr &stra) { } StrategyPtr strategyPtr; - int64_t stage_id = g_device_manager->stage_id(); + int64_t stage_id = g_device_manager->stage_id(); // Get the current stage ID from the device manager MS_LOG(INFO) << "Extract information: strategy " << stra->ToString(); if (var->size() > 0) { @@ -1348,9 +1879,10 @@ StrategyPtr ExtractStrategy(const ValuePtr &stra) { if (elements[index]->isa()) { auto value_tuple = elements[index]->cast(); std::vector value_vector = value_tuple->value(); + // Extract dimensions from ValueSequence and convert them to int64_t (void)std::transform(value_vector.begin(), value_vector.end(), std::back_inserter(dim), [](const ValuePtr &value) { return static_cast(GetValue(value)); }); - strategy.push_back(dim); + strategy.push_back(dim); // Append the extracted dimension to the strategy } else { MS_LOG(EXCEPTION) << "Failure: Strategy's format is wrong! Need ValueSequence"; } @@ -1358,105 +1890,133 @@ StrategyPtr ExtractStrategy(const ValuePtr &stra) { if (strategy.empty()) { MS_LOG(EXCEPTION) << "ExtractStrategy: failed to extract strategy"; } - strategyPtr = NewStrategy(stage_id, strategy); + strategyPtr = NewStrategy(stage_id, strategy); // Create a new strategy } - return strategyPtr; + return strategyPtr; // Return the extracted strategy } +// GetRefKeyNodeShape function +// This function retrieves the shape of a parameter node referred to by a given RefKey node within a FuncGraph. +// It first finds the parameter node associated with the RefKey node and then retrieves its shape. Shapes GetRefKeyNodeShape(const AnfNodePtr &node, const FuncGraphPtr &func_graph) { MS_EXCEPTION_IF_NULL(node); MS_EXCEPTION_IF_NULL(func_graph); - std::vector parameters = FindParameterByRefKeyNode(node, func_graph); + std::vector parameters = FindParameterByRefKeyNode(node, func_graph); // Find the associated parameter nodes if (parameters.size() != 1) { MS_LOG(EXCEPTION) << "Find parameter by ref key node failed"; } Shapes input_shapes; - input_shapes = GetNodeShape(parameters[0]); + input_shapes = GetNodeShape(parameters[0]); // Retrieve the shape of the parameter node if (input_shapes.size() != 1) { MS_LOG(EXCEPTION) << "Get input shape failed"; } MS_LOG(INFO) << "The parameter shape is " << ShapeToString(input_shapes[0]); - return input_shapes; + return input_shapes; // Return the shape of the parameter node } +// ExtractShape function +// This function extracts input and output shapes for a given CNode and returns them in a vector of Shapes. +// It processes the inputs and outputs of the CNode, handling different cases such as RefKey nodes, ValueNodes, Parameters, etc. std::vector ExtractShape(const CNodePtr &node) { MS_EXCEPTION_IF_NULL(node); Shapes shape_inputs, shape_outputs; std::vector shape_all; - std::vector all_inputs = node->inputs(); + std::vector all_inputs = node->inputs(); // Get all inputs of the CNode size_t inputs_size = all_inputs.size(); - for (size_t i = 1; i < inputs_size; ++i) { + for (size_t i = 1; i < inputs_size; ++i) { // Iterate through the inputs starting from index 1 (skip the first input) Shapes input_shapes; AnfNodePtr input = all_inputs[i]; if (HasAbstractMonad(input)) { continue; } - if (IsValueNode(input)) { + if (IsValueNode(input)) { // Check if the input is a RefKey node auto func_graph = node->func_graph(); MS_EXCEPTION_IF_NULL(func_graph); - std::vector parameters = FindParameterByRefKeyNode(input, func_graph); + std::vector parameters = FindParameterByRefKeyNode(input, func_graph); // Find associated parameters if (parameters.size() != 1) { MS_LOG(EXCEPTION) << "Find parameter by ref key node failed"; } std::pair node_pair = std::make_pair(node, SizeToLong(i)); - g_RefMap[parameters[0]] = node_pair; - input_shapes = GetRefKeyNodeShape(input, func_graph); + g_RefMap[parameters[0]] = node_pair; // Map the parameter node to the CNode and input index + input_shapes = GetRefKeyNodeShape(input, func_graph); // Get the shape of the referred parameter node } else if (input->isa() || IsValueNode(input) || input->isa() || ((IsValueNode(input) || IsValueNode(input)) && (inputs_size == 2))) { - input_shapes = GetNodeShape(input); + input_shapes = GetNodeShape(input); // Get the shape of the input node } else { continue; } if (input_shapes.size() != 1) { - if (inputs_size == 2) { // like concat - shape_inputs = input_shapes; + if (inputs_size == 2) { // If there are only two inputs (e.g., for concat operation) + shape_inputs = input_shapes; // Set input shape directly break; } else { MS_LOG(EXCEPTION) << "ExtractShape: Get input shape failed"; } } - shape_inputs.push_back(input_shapes[0]); + shape_inputs.push_back(input_shapes[0]); // Append the input shape to the input shapes vector } - shape_all.push_back(shape_inputs); - // extract out shape - shape_outputs = GetNodeShape(node); - shape_all.push_back(shape_outputs); - return shape_all; + shape_all.push_back(shape_inputs); // Add input shapes to the result vector + // Extract output shape + shape_outputs = GetNodeShape(node); // Get the shape of the CNode's output + shape_all.push_back(shape_outputs); // Add output shape to the result vector + return shape_all; // Return the vector containing input and output shapes } +// FindParallelCareNode function +// This function recursively searches for a parallel care node (CNode) in the graph starting from the given 'node'. +// A parallel care node is a CNode that represents an operator relevant for parallel execution. std::pair FindParallelCareNode(const AnfNodePtr &node, int32_t recursion_num) { + // Check if the recursion limit has been reached to prevent infinite recursion. if (recursion_num >= RECURSION_LIMIT) { return std::make_pair(nullptr, 0); } + // Check if 'node' is null. MS_EXCEPTION_IF_NULL(node); + + // Get the function graph associated with 'node'. FuncGraphPtr func_graph = node->func_graph(); MS_EXCEPTION_IF_NULL(func_graph); + + // Get the function graph manager. FuncGraphManagerPtr manager = func_graph->manager(); MS_EXCEPTION_IF_NULL(manager); + + // Get the set of nodes that use 'node' as an input. AnfNodeIndexSet node_set = manager->node_users()[node]; + + // Iterate through the nodes that use 'node' as an input. for (auto &node_pair : node_set) { CNodePtr cnode = node_pair.first->cast(); MS_EXCEPTION_IF_NULL(cnode); + + // Check if the input is a ValueNode containing a Primitive. if (!IsValueNode(cnode->input(0))) { continue; } + + // Get the Primitive associated with the CNode. ValueNodePtr prim_node_anf = cnode->input(0)->cast(); MS_EXCEPTION_IF_NULL(prim_node_anf); PrimitivePtr node_prim = prim_node_anf->value()->cast(); MS_EXCEPTION_IF_NULL(node_prim); + + // Check if the Primitive is not 'DEPEND', 'Receive', or 'Send'. if ((node_prim->name() == DEPEND && node_pair.second != 1) || IsPrimitiveCNode(cnode, prim::kPrimReceive) || IsPrimitiveCNode(cnode, prim::kPrimSend)) { continue; } + + // Check if the CNode is a parallel care node and has user data of OperatorInfo. if (IsParallelCareNode(cnode) && cnode->has_user_data()) { return node_pair; } else { + // Recursively call the function to search in the next level of nodes. auto tmp_pair = FindParallelCareNode(node_pair.first, recursion_num + 1); if (tmp_pair.first != nullptr) { return tmp_pair; @@ -1466,33 +2026,54 @@ std::pair FindParallelCareNode(const AnfNodePtr &node, int3 return std::make_pair(nullptr, 0); } +// FindSubGraph function +// This function finds a subgraph containing a parallel care node (CNode) connected to the given 'parameter'. +// It searches through the users of 'parameter' in the given 'graph'. std::pair FindSubGraph(const FuncGraphPtr &graph, const AnfNodePtr ¶meter) { + // Check if 'graph' or 'parameter' is null. MS_EXCEPTION_IF_NULL(graph); MS_EXCEPTION_IF_NULL(parameter); + + // Get the function graph manager. FuncGraphManagerPtr manager = graph->manager(); MS_EXCEPTION_IF_NULL(manager); + + // Find the parallel care node connected to 'parameter'. std::pair prim_anf_node_pair = FindParallelCareNode(parameter, 0); if (prim_anf_node_pair.first != nullptr) { return prim_anf_node_pair; } else { + // Get the set of nodes that use 'parameter'. AnfNodeIndexSet param_sub_set = manager->node_users()[parameter]; + + // Iterate through the users of 'parameter'. for (auto ¶m_pair : param_sub_set) { CNodePtr param_cnode = param_pair.first->cast(); AnfNodePtr graph_value_node; + + // Check if 'parameter' is connected to a ValueNode containing a FuncGraph. if (param_cnode->input(0)->isa()) { graph_value_node = param_cnode->input(0)->cast()->input(1); } else { graph_value_node = param_cnode->input(0); } + + // Check if 'graph_value_node' is a ValueNode containing a FuncGraph. if (!IsValueNode(graph_value_node)) { continue; } + + // Get the FuncGraph associated with 'graph_value_node'. FuncGraphPtr graph_sub = GetValueNode(graph_value_node); auto parameters = graph_sub->parameters(); + + // Check if the index is within the range of 'parameters'. if (LongToSize(param_pair.second - 1) >= parameters.size()) { - MS_LOG(EXCEPTION) << "The index is out of range, index is: " << (param_pair.second - 1) << ", vector size is " - << parameters.size(); + MS_LOG(EXCEPTION) << "The index is out of range, index is: " << (param_pair.second - 1) + << ", vector size is " << parameters.size(); } + + // Recursively call the function to search in the subgraph. std::pair res = FindSubGraph(graph_sub, parameters[LongToSize(param_pair.second - 1)]); if (res.first != nullptr) { return res; @@ -1502,30 +2083,48 @@ std::pair FindSubGraph(const FuncGraphPtr &graph, const Anf return std::make_pair(nullptr, 0); } +// InsertAllGatherAfterCast function +// This function inserts an 'AllGather' operation after a 'Cast' operation if certain conditions are met. CNodePtr InsertAllGatherAfterCast(const CNodePtr &cnode) { + // Check if 'cnode' is null. MS_EXCEPTION_IF_NULL(cnode); + + // Get the function graph containing 'cnode'. auto graph = cnode->func_graph(); MS_EXCEPTION_IF_NULL(graph); + + // Get the function graph manager. auto manager = graph->manager(); MS_EXCEPTION_IF_NULL(manager); - // skip Load moving down and assume it only has one node user + + // Initialize the result node to 'cnode'. CNodePtr res = cnode; + + // Skip Load operations by moving down the graph and assuming it has only one node user. if (IsSomePrimitive(res, LOAD)) { res = manager->node_users()[cnode].begin()->first->cast(); } - // return true only if cnode is Cast from fp32 to fp16 + + // Check if 'res' is a 'Cast' operation. if (!IsSomePrimitive(res, CAST)) { return nullptr; } + + // Get the type of the input element of the 'Cast' operation. auto node_type = res->Type(); MS_EXCEPTION_IF_NULL(node_type); + + // Check if the type is a TensorType. if (!node_type->isa()) { MS_LOG(EXCEPTION) << "Unknown type."; } + + // Get the type ID of the input element. auto input_element_type = node_type->cast()->element(); MS_EXCEPTION_IF_NULL(input_element_type); auto type_id = input_element_type->type_id(); + // Check if the type ID is kNumberTypeFloat32. if (type_id != kNumberTypeFloat32) { return res; } else { @@ -1533,18 +2132,31 @@ CNodePtr InsertAllGatherAfterCast(const CNodePtr &cnode) { } } +// InsertAllGatherOp function +// This function inserts an AllGather operation into the computation graph. +// It is responsible for creating the AllGather operator, determining the appropriate position for insertion, +// and handling special cases like shared parameters, gradient accumulation, and pipeline parallelism. static void InsertAllGatherOp(const FuncGraphPtr &root, const std::string &group, const std::pair &res, const AnfNodePtr &node, const std::string &op_name, bool is_shared_param) { + // Check for null pointers MS_EXCEPTION_IF_NULL(res.first); MS_EXCEPTION_IF_NULL(node); + + // Check if gradient accumulation shard is enabled bool grad_accumulation_shard = ParallelContext::GetInstance()->grad_accumulation_shard(); + + // Get information about the CNode associated with the node auto cnode = res.first->cast(); auto graph = cnode->func_graph(); MS_EXCEPTION_IF_NULL(graph); auto manager = graph->manager(); MS_EXCEPTION_IF_NULL(manager); + + // Get the primitive associated with the CNode auto cnode_prim = GetValueNode(cnode->input(0)); MS_EXCEPTION_IF_NULL(cnode_prim); + + // Create an operator based on the specified operation name (op_name) Operator op; CNodePtr allgather; auto param_name = node->cast()->name(); @@ -1555,6 +2167,8 @@ static void InsertAllGatherOp(const FuncGraphPtr &root, const std::string &group } else { op = CreateAllGatherOp(group); } + + // Insert AllGather operation after Cast if applicable CNodePtr cast_node = InsertAllGatherAfterCast(cnode); std::string opt_shard_mirror_group; auto param_ptr = node->cast(); @@ -1566,6 +2180,7 @@ static void InsertAllGatherOp(const FuncGraphPtr &root, const std::string &group allgather = ReplaceNode(op, cast_node, graph, PARALLEL_OPTIMIZER_ALLGATHER_NOT_COMPUTE, param_name, root); MS_LOG(INFO) << "Parallel optimizer is applied before Cast for " << param_name; } else { + // Handle special cases and insert AllGather auto pre_node = node; AnfNodePtr pre_node_ = node; auto node_user_map = manager->node_users(); @@ -1580,10 +2195,14 @@ static void InsertAllGatherOp(const FuncGraphPtr &root, const std::string &group allgather = cnode->input(IntToSize(res.second))->cast(); MS_LOG(INFO) << "Parallel optimizer is applied before " << GetPrimName(cnode) << " for " << param_name; } - // add fusion flag + + // Add fusion flag to the AllGather operation AddCommOpFusionType(allgather, node); - // add gradients mean + + // Add gradients mean flag to the AllGather operation AddCommOpMeanFlag(allgather); + + // Set mirror flag for AllGather based on the operation name if (op_name == MICRO_STEP_ALL_GATHER) { // When grad_accumulation_shard is enabled, the ReduceScatter is inserted at each micro step // so no need to do backward for the micro_step_allgather @@ -1597,17 +2216,22 @@ static void InsertAllGatherOp(const FuncGraphPtr &root, const std::string &group } } +// ApplyParallelOptOnParam function +// This function applies parallel optimization on a parameter node. +// It inserts AllGather operations based on the specified shard group and handles various optimization scenarios. static void ApplyParallelOptOnParam(const FuncGraphPtr &root, const AnfNodePtr ¶meter, const std::string &opt_shard_group) { + // Check if the shard group is empty (no optimization needed) if (opt_shard_group.empty()) { return; } - // set all gather type - MS_EXCEPTION_IF_NULL(parameter); + // Get parallel context information int64_t grad_accumulation_step = ParallelContext::GetInstance()->grad_accumulation_step(); int32_t split_stage_num = ParallelContext::GetInstance()->pipeline_stage_split_num(); std::string op_name; + + // Determine the AllGather operation type based on optimization scenarios if (grad_accumulation_step > 1) { op_name = MINI_STEP_ALL_GATHER; } else if (split_stage_num > 1) { @@ -1616,17 +2240,28 @@ static void ApplyParallelOptOnParam(const FuncGraphPtr &root, const AnfNodePtr & op_name = ALL_GATHER; } - // insert all gather + // Access the function graph manager FuncGraphManagerPtr manager = root->manager(); MS_EXCEPTION_IF_NULL(manager); + + // Get the set of nodes that use the parameter auto param_sub_set = manager->node_users()[parameter]; + + // Initialize insert_flag to track if an AllGather operation has already been inserted bool insert_flag = false; + + // Iterate through the nodes using the parameter and insert AllGather operations for (auto ¶m_pair : param_sub_set) { auto cnode = param_pair.first->cast(); MS_EXCEPTION_IF_NULL(cnode); + + // Check if the CNode is eligible for AllGather insertion if (cnode->in_forward_flag() && !IsPrimitiveCNode(cnode, prim::kPrimReceive) && !IsPrimitiveCNode(cnode, prim::kPrimDepend)) { + // Get the operator info associated with the CNode OperatorInfoPtr distribute_operator = cnode->user_data(); + + // Handle cases where the operator info is not available if (distribute_operator == nullptr) { MS_LOG(DEBUG) << "Parallel optimizer: " << GetPrimName(cnode) << " 's OperatorInfoPtr is nullptr"; } else if (IntToSize(param_pair.second - 1) >= distribute_operator->inputs_tensor_info().size()) { @@ -1635,7 +2270,7 @@ static void ApplyParallelOptOnParam(const FuncGraphPtr &root, const AnfNodePtr & } if (insert_flag) { - // if there are multiple node users, they share one same allgather + // If there are multiple node users, they share one same AllGather operation auto next_cnode = FindCNode(parameter, op_name, cnode->func_graph(), 0); if (next_cnode.first) { manager->SetEdge(cnode, param_pair.second, next_cnode.second); @@ -1645,7 +2280,7 @@ static void ApplyParallelOptOnParam(const FuncGraphPtr &root, const AnfNodePtr & MS_LOG(ERROR) << "Can not find the shared AllGather with multiple node users."; } } else { - // insert allgather operator between shard parameter and cnode + // Insert AllGather operation for the parameter auto param_ptr = parameter->cast(); MS_EXCEPTION_IF_NULL(param_ptr); bool is_shared_param = param_ptr->user_data()->is_shared_param(); @@ -1656,23 +2291,33 @@ static void ApplyParallelOptOnParam(const FuncGraphPtr &root, const AnfNodePtr & } } +// SetSharedParameterFlag function +// This function checks if a parameter is shared among multiple users in a computation graph +// and marks it as a shared parameter if necessary. void SetSharedParameterFlag(const FuncGraphPtr &root, const AnfNodePtr ¶meter) { MS_EXCEPTION_IF_NULL(root); MS_EXCEPTION_IF_NULL(parameter); FuncGraphManagerPtr manager = root->manager(); MS_EXCEPTION_IF_NULL(manager); ParameterPtr parameter_ptr = parameter->cast(); + + // Check if the given node is a parameter, log a message if not. if (parameter_ptr == nullptr) { - MS_LOG(INFO) << parameter->ToString() << ": cast to ptr failed. it may not be a parameter"; + MS_LOG(INFO) << parameter->ToString() << ": cast to ptr failed. It may not be a parameter."; return; } + auto user_set = manager->node_users()[parameter]; int32_t user_count = 0; + + // Count users of the parameter that are marked for forward execution. for (auto ¶m_pair : user_set) { CNodePtr cnode = param_pair.first->cast(); MS_EXCEPTION_IF_NULL(cnode); if (cnode->in_forward_flag()) user_count++; } + + // If there are multiple users, mark the parameter as shared and log a warning. if (user_count > 1) { auto tensor_layout = parameter_ptr->user_data(); tensor_layout->set_is_shared_param(true); @@ -1681,55 +2326,60 @@ void SetSharedParameterFlag(const FuncGraphPtr &root, const AnfNodePtr ¶mete } } -// When this function returns non-empty string, that means parallel optimizer is applied on this parameter. +// SetParallelShape function +// This function sets the parallel shape for a parameter based on the distributed operator information. +// It also generates a shard group for parallel optimization. std::string SetParallelShape(const AnfNodePtr ¶meter, const std::pair &res, const FuncGraphPtr &root) { - // check null for param and cnode + // Check for null values in parameter and cnode. auto param_shape = parameter->Shape(); - MS_EXCEPTION_IF_NULL(parameter); MS_EXCEPTION_IF_NULL(param_shape); CNodePtr cnode = res.first->cast(); MS_EXCEPTION_IF_NULL(cnode); - // get slice_shape + // Get the slice shape from the distributed operator. OperatorInfoPtr distribute_operator = cnode->user_data(); if (distribute_operator == nullptr) { - MS_LOG(EXCEPTION) << "node " << cnode->ToString() << " 's distribute_operator is nullptr"; + MS_LOG(EXCEPTION) << "Node " << cnode->ToString() << "'s distribute_operator is nullptr."; } if (LongToSize(res.second - 1) >= distribute_operator->inputs_tensor_info().size()) { - MS_LOG(EXCEPTION) << "The parameter index is not in inputs_tensor_info. index = " << (res.second - 1) + MS_LOG(EXCEPTION) << "The parameter index is not in inputs_tensor_info. Index = " << (res.second - 1) << ", inputs_tensor_info size = " << distribute_operator->inputs_tensor_info().size(); } TensorInfo tensorinfo_in = distribute_operator->inputs_tensor_info()[LongToSize(res.second - 1)]; TensorLayout tensor_layout = tensorinfo_in.tensor_layout(); Shape slice_shape = tensor_layout.slice_shape().array(); - // generate shard group + // Generate a shard group for parallel optimization. std::string opt_shard_group; MS_EXCEPTION_IF_NULL(ParallelContext::GetInstance()); bool enable_parallel_optimizer = ParallelContext::GetInstance()->enable_parallel_optimizer(); + if (enable_parallel_optimizer) { std::unique_ptr apOptParamMgr = createOptParamMgr(root); opt_shard_group = apOptParamMgr->ShardOptGroup(parameter, &tensor_layout, distribute_operator); - // set the shape of parameter to sliced shape + + // Set the parameter's shape to the sliced shape. if (!opt_shard_group.empty()) { slice_shape = tensor_layout.opt_shard_slice_shape(); } - MS_LOG(INFO) << "the shape of " << parameter->ToString() << "(original: " << param_shape->ToString() << ")" + + MS_LOG(INFO) << "The shape of " << parameter->ToString() << " (original: " << param_shape->ToString() << ")" << " will be sliced into " << MakeValue(slice_shape)->ToString() << " in op " << distribute_operator->name(); } + // Update the parameter's abstract, shape, and user data. AbstractBasePtr abstract = parameter->abstract(); if (abstract == nullptr) { - MS_LOG(EXCEPTION) << "parameter " << parameter->ToString() << ": abstract is nullptr"; + MS_LOG(EXCEPTION) << "Parameter " << parameter->ToString() << ": abstract is nullptr."; } AbstractBasePtr cloned_abstract = abstract->Clone(); if (cloned_abstract == nullptr) { - MS_LOG(EXCEPTION) << "parameter " << parameter->ToString() << ": abstract clone failed"; + MS_LOG(EXCEPTION) << "Parameter " << parameter->ToString() << ": abstract clone failed."; } cloned_abstract->set_shape(std::make_shared(slice_shape)); @@ -1737,9 +2387,18 @@ std::string SetParallelShape(const AnfNodePtr ¶meter, const std::paircast(); MS_EXCEPTION_IF_NULL(parameter_ptr); parameter_ptr->set_user_data(std::make_shared(tensor_layout)); + + // Return the generated shard group for parallel optimization. return opt_shard_group; } +// CoverSliceShape function +// This function iterates through the parameters of the root graph and performs the following actions: +// 1. Checks if a parameter has a reference in the g_RefMap. +// 2. If yes, it sets a parallel shape for the parameter based on the reference information. +// 3. Searches for forward nodes that use the parameter in graphs and inserts an allgather if the group is not empty. +// 4. Sets the shared parameter flag for the parameter. +// 5. Applies parallel optimization on the parameter based on the group. void CoverSliceShape(const FuncGraphPtr &root) { MS_EXCEPTION_IF_NULL(root); auto parameters = root->parameters(); @@ -1749,7 +2408,6 @@ void CoverSliceShape(const FuncGraphPtr &root) { auto iter = g_RefMap.find(parameter); if (iter != g_RefMap.end()) { std::string group = SetParallelShape(parameter, g_RefMap[parameter], root); - // find all forward nodes that use parameter in graphs and insert allgather if group is not empty SetSharedParameterFlag(root, parameter); ApplyParallelOptOnParam(root, parameter, group); continue; @@ -1757,10 +2415,9 @@ void CoverSliceShape(const FuncGraphPtr &root) { std::pair res = FindSubGraph(root, parameter); if (res.first == nullptr) { - MS_LOG(INFO) << "Parameter " << parameter->ToString() << " is not in graph, thus no need to set parallel shape"; + MS_LOG(INFO) << "Parameter " << parameter->ToString() << " is not in the graph, thus no need to set parallel shape"; } else { std::string group = SetParallelShape(parameter, res, root); - // find all forward nodes that use parameter in graphs and insert allgather if group is not empty SetSharedParameterFlag(root, parameter); ApplyParallelOptOnParam(root, parameter, group); MS_LOG(DEBUG) << "Parameter " << parameter->ToString() << " shape " << parameter->Shape()->ToString(); @@ -1769,6 +2426,9 @@ void CoverSliceShape(const FuncGraphPtr &root) { g_RefMap.clear(); } +// SetVirtualDatasetStrategy function +// This function sets the strategy attributes for VirtualDataset and VirtualOutput primitives. +// It checks whether the full batch mode is enabled and sets the appropriate strategy. void SetVirtualDatasetStrategy(const CNodePtr &node) { MS_EXCEPTION_IF_NULL(node); MS_EXCEPTION_IF_NULL(ParallelContext::GetInstance()); @@ -1804,7 +2464,7 @@ void SetVirtualDatasetStrategy(const CNodePtr &node) { } std::vector shape_list = ExtractShape(node); if (shape_list.empty()) { - MS_LOG(EXCEPTION) << "Failure:node " << node->ToString() << " failed to extract shape"; + MS_LOG(EXCEPTION) << "Failure: node " << node->ToString() << " failed to extract shape"; } std::vector elements; for (size_t i = 0; i < shape_list[0].size(); i++) { @@ -1830,10 +2490,20 @@ void SetVirtualDatasetStrategy(const CNodePtr &node) { } // find previous parallel care node's next node. +// Function: FindPreNodes +// This function recursively searches for previous nodes (upstream nodes) of the given 'node' and collects unique_ids and indexes. +// It stops the recursion when MAX_RECURSIVE_DEPTH is exceeded or when certain conditions are met. +// Parameters: +// - node: The current node being examined. +// - unique_ids: A vector to store unique IDs of valid nodes found. +// - indexes: A vector to store indexes of valid nodes found. +// - curr_depth: The current recursion depth. +// Returns: +// - 'true' if valid nodes were found in the previous nodes of 'node', 'false' otherwise. bool FindPreNodes(const AnfNodePtr &node, std::vector *unique_ids, std::vector *indexes, size_t curr_depth) { if (curr_depth > MAX_RECURSIVE_DEPTH) { - MS_LOG(WARNING) << "When find the previous node, exceeded the maximum recursion depth: " << MAX_RECURSIVE_DEPTH; + MS_LOG(WARNING) << "When finding the previous node, exceeded the maximum recursion depth: " << MAX_RECURSIVE_DEPTH; return false; } MS_EXCEPTION_IF_NULL(unique_ids); @@ -1871,15 +2541,29 @@ bool FindPreNodes(const AnfNodePtr &node, std::vector *unique_ids, return find; } +// Function: FindLastNodesUniqueId +// This function initiates the search for the unique IDs of the last parallel care nodes in the evaluation graph. +// Parameters: +// - root: The root of the evaluation graph (FuncGraph). +// - unique_ids: A vector to store the unique IDs of the last parallel care nodes found. +// - indexes: A vector to store the indexes of the last parallel care nodes found. void FindLastNodesUniqueId(const FuncGraphPtr &root, std::vector *unique_ids, std::vector *indexes) { MS_EXCEPTION_IF_NULL(unique_ids); CNodePtr cnode = root->get_return(); if (!FindPreNodes(cnode, unique_ids, indexes, 0)) { - MS_LOG(WARNING) << "cannot find the last parallel care node in eval graph"; + MS_LOG(WARNING) << "Cannot find the last parallel care node in the eval graph."; } } +// Function: GenerateBatchParallelStrategy +// This function generates a batch parallel strategy based on the given operator information and primitive. +// It uses operator-specific strategies and attributes to create the strategy. +// Parameters: +// - operator_: The operator information. +// - prim: The primitive associated with the operator. +// Returns: +// - A StrategyPtr representing the generated batch parallel strategy. StrategyPtr GenerateBatchParallelStrategy(const OperatorInfoPtr operator_, const PrimitivePtr prim) { MS_EXCEPTION_IF_NULL(operator_); MS_EXCEPTION_IF_NULL(prim); @@ -1892,93 +2576,147 @@ StrategyPtr GenerateBatchParallelStrategy(const OperatorInfoPtr operator_, const elements.push_back(MakeValue((*strategy_v_ptr)[i])); } ValueTuplePtr strategy = std::make_shared(elements); - // display the strategy generated by batch parallel + // Display the strategy generated by batch parallel auto attrs = prim->attrs(); attrs[GEN_STRATEGY] = strategy; (void)prim->SetAttrs(attrs); - MS_LOG(INFO) << "prim " << prim->name() << " batch parallel strategy is " << attrs[GEN_STRATEGY]->ToString(); + MS_LOG(INFO) << "Primitive " << prim->name() << " batch parallel strategy is " << attrs[GEN_STRATEGY]->ToString(); return strategyPtr; } -static bool CheckExtractInfomation(const CNodePtr &cnode) { - if ((cnode == nullptr) || !IsValueNode(cnode->input(0))) { +// Function: FindPreNodes +// This function recursively searches for previous nodes (upstream nodes) of the given 'node' and collects unique_ids and indexes. +// It stops the recursion when MAX_RECURSIVE_DEPTH is exceeded or when certain conditions are met. +// Parameters: +// - node: The current node being examined. +// - unique_ids: A vector to store unique IDs of valid nodes found. +// - indexes: A vector to store indexes of valid nodes found. +// - curr_depth: The current recursion depth. +// Returns: +// - 'true' if valid nodes were found in the previous nodes of 'node', 'false' otherwise. +bool FindPreNodes(const AnfNodePtr &node, std::vector *unique_ids, std::vector *indexes, + size_t curr_depth) { + if (curr_depth > MAX_RECURSIVE_DEPTH) { + MS_LOG(WARNING) << "When finding the previous node, exceeded the maximum recursion depth: " << MAX_RECURSIVE_DEPTH; return false; } - - ValueNodePtr prim_anf_node = cnode->input(0)->cast(); - PrimitivePtr prim = GetValueNode(prim_anf_node); - if ((prim->name() == MAKE_TUPLE) || (prim->name() == MAKE_LIST) || (prim->name() == RECEIVE)) { + MS_EXCEPTION_IF_NULL(unique_ids); + MS_EXCEPTION_IF_NULL(indexes); + if (!node->isa()) { return false; } - - if (!IsParallelCareNode(cnode)) { + CNodePtr pre_cnode = node->cast(); + if (!IsValueNode(pre_cnode->input(0))) { return false; } - return true; -} - -static void ExtractStrategyAndInit(const CNodePtr &cnode, const PrimitivePtr &prim, const OperatorInfoPtr &op_info) { - StrategyPtr in_strategy = nullptr, out_strategy = nullptr; - auto attrs = prim->attrs(); - - // load strategy map from checkpoint - StrategyMap stra_map; - if (StrategyCheckpoint::GetInstance().LoadCheckPointOn() && - (StrategyCheckpoint::GetInstance().Load(&stra_map) != SUCCESS)) { - MS_LOG(EXCEPTION) << "Load strategy checkpoint failed"; - } - - std::string strategy_key_name = ""; - auto param_names = NodeParameterName(cnode, -1, 0); - if (!param_names.empty()) { - strategy_key_name = prim->name() + "_" + param_names[0].first; - } - bool load_strategy_from_ckpt = - StrategyCheckpoint::GetInstance().LoadCheckPointOn() && stra_map.find(strategy_key_name) != stra_map.end(); - if ((!StrategyFound(attrs) && !load_strategy_from_ckpt) && !cnode->HasPrimalAttr(IN_STRATEGY)) { - MS_LOG(INFO) << "ExtractInformation: the strategy of node " << cnode->ToString() << " prim " << prim->name() - << " is empty, using batch parallel"; - in_strategy = GenerateBatchParallelStrategy(op_info, prim); - } else if (cnode->HasPrimalAttr(IN_STRATEGY)) { - in_strategy = ExtractStrategy(cnode->GetPrimalAttr(IN_STRATEGY)); - out_strategy = ExtractStrategy(cnode->GetPrimalAttr(OUT_STRATEGY)); - } else if (StrategyFound(attrs)) { - in_strategy = ExtractStrategy(attrs[IN_STRATEGY]); - out_strategy = ExtractStrategy(attrs[OUT_STRATEGY]); - } else { - in_strategy = stra_map[strategy_key_name]; - } - - MS_EXCEPTION_IF_NULL(in_strategy); - if (op_info->Init(in_strategy, out_strategy) == FAILED) { - MS_LOG(EXCEPTION) << "Failure:operator " << prim->name() << " init failed" << trace::DumpSourceLines(cnode); - } -} - -void ExtractInformation(const std::vector &all_nodes) { - SetStridedSliceSplitStrategy(all_nodes); - for (auto &node : all_nodes) { - auto cnode = node->cast(); - if (!CheckExtractInfomation(cnode) || IsPrimitiveCNode(node, prim::kPrimSend)) { + bool find = false; + for (size_t index = 1; index < pre_cnode->inputs().size(); ++index) { + auto next_node = pre_cnode->inputs()[index]; + if (!next_node->isa() || next_node->isa()) { + return false; + } + CNodePtr cnode = next_node->cast(); + if (!IsValueNode(cnode->input(0))) { + return false; + } + ValueNodePtr prim_anf_node = cnode->input(0)->cast(); + PrimitivePtr prim = prim_anf_node->value()->cast(); + if (IsParallelCareNode(cnode) && prim->name() != MAKE_TUPLE && prim->name() != MAKE_LIST) { + unique_ids->push_back(pre_cnode->UniqueId()); + indexes->push_back(index); + find = true; continue; } + if (FindPreNodes(cnode, unique_ids, indexes, ++curr_depth)) { + find = true; + continue; + } + } + return find; +} +// Function: FindLastNodesUniqueId +// This function initiates the search for the unique IDs of the last parallel care nodes in the evaluation graph. +// Parameters: +// - root: The root of the evaluation graph (FuncGraph). +// - unique_ids: A vector to store the unique IDs of the last parallel care nodes found. +// - indexes: A vector to store the indexes of the last parallel care nodes found. +void FindLastNodesUniqueId(const FuncGraphPtr &root, std::vector *unique_ids, + std::vector *indexes) { + MS_EXCEPTION_IF_NULL(unique_ids); + CNodePtr cnode = root->get_return(); + if (!FindPreNodes(cnode, unique_ids, indexes, 0)) { + MS_LOG(WARNING) << "Cannot find the last parallel care node in the eval graph."; + } +} + +// Function: GenerateBatchParallelStrategy +// This function generates a batch parallel strategy based on the given operator information and primitive. +// It uses operator-specific strategies and attributes to create the strategy. +// Parameters: +// - operator_: The operator information. +// - prim: The primitive associated with the operator. +// Returns: +// - A StrategyPtr representing the generated batch parallel strategy. +StrategyPtr GenerateBatchParallelStrategy(const OperatorInfoPtr operator_, const PrimitivePtr prim) { + MS_EXCEPTION_IF_NULL(operator_); + MS_EXCEPTION_IF_NULL(prim); + StrategyPtr strategyPtr; + std::shared_ptr strategy_v_ptr = operator_->GenerateBatchStrategies(); + MS_EXCEPTION_IF_NULL(strategy_v_ptr); + strategyPtr = NewStrategy(0, *strategy_v_ptr); + std::vector elements; + for (size_t i = 0; i < strategy_v_ptr->size(); i++) { + elements.push_back(MakeValue((*strategy_v_ptr)[i])); + } + ValueTuplePtr strategy = std::make_shared(elements); + // Display the strategy generated by batch parallel + auto attrs = prim->attrs(); + attrs[GEN_STRATEGY] = strategy; + (void)prim->SetAttrs(attrs); + MS_LOG(INFO) << "Primitive " << prim->name() << " batch parallel strategy is " << attrs[GEN_STRATEGY]->ToString(); + return strategyPtr; +} + +// ExtractInformation function +// This function extracts information from a list of AnfNodes and sets up corresponding data structures. +void ExtractInformation(const std::vector &all_nodes) { + // Step 1: Set the strided slice split strategy + SetStridedSliceSplitStrategy(all_nodes); + + // Step 2: Iterate through all nodes in the list + for (auto &node : all_nodes) { + auto cnode = node->cast(); + + // Step 3: Check if extraction information is necessary for this node + if (!CheckExtractInfomation(cnode) || IsPrimitiveCNode(node, prim::kPrimSend)) { + continue; // Skip extraction for this node + } + + // Step 4: Set the virtual dataset strategy for the node SetVirtualDatasetStrategy(cnode); + + // Step 5: Extract the primitive operation information ValueNodePtr prim_anf_node = cnode->input(0)->cast(); PrimitivePtr prim = GetValueNode(prim_anf_node); auto attrs = prim->attrs(); - MS_LOG(INFO) << "extract information: node: " << node->ToString() << " prim " << prim->name(); + MS_LOG(INFO) << "Extracting information: node: " << node->ToString() << " prim " << prim->name(); + // Step 6: Extract shapes information for the node std::vector shape_list = ExtractShape(cnode); if (shape_list.empty()) { - MS_LOG(EXCEPTION) << "Failure:node " << node->ToString() << " failed to extract shape"; + MS_LOG(EXCEPTION) << "Failure: Node " << node->ToString() << " failed to extract shape"; } + + // Step 7: Create an OperatorInfo instance with extracted information OperatorInfoPtr operator_ = OperatorInstance(prim, attrs, shape_list); MS_EXCEPTION_IF_NULL(operator_); auto &inputs = cnode->inputs(); std::vector input_value; + + // Step 8: Extract input values for the node for (size_t index = 1; index < inputs.size(); ++index) { if (inputs[index]->isa()) { input_value.push_back(GetValueNode(inputs[index])); @@ -1987,58 +2725,69 @@ void ExtractInformation(const std::vector &all_nodes) { input_value.emplace_back(nullptr); } + // Step 9: Set input values, outputs dtype, and associate the CNode with the OperatorInfo (*operator_).set_input_value(input_value); (*operator_).set_outputs_dtype(cnode->Type()); (*operator_).set_cnode(cnode); + + // Step 10: If the primitive operation is RESHAPE, associate the CNode with the OperatorInfo and continue if (prim->name() == RESHAPE) { cnode->set_user_data(operator_); continue; } + // Step 11: Extract strategy and initialize the OperatorInfo ExtractStrategyAndInit(cnode, prim, operator_); cnode->set_user_data(operator_); } } +// GetInputLayoutFromCNode function +// This function retrieves the input tensor layout from a CNode based on a node-index pair. TensorLayout GetInputLayoutFromCNode(const std::pair &node_pair) { CNodePtr cnode = node_pair.first->cast(); MS_EXCEPTION_IF_NULL(cnode); OperatorInfoPtr distribute_operator = GetDistributeOperator(cnode); MS_EXCEPTION_IF_NULL(distribute_operator); int64_t index = node_pair.second; + + // Step 1: Check if the index is out of range if (index > SizeToLong(distribute_operator->inputs_tensor_info().size())) { - MS_LOG(EXCEPTION) << "The index is out of range, the node_pair.second is " << (index - 1) - << ", the vector size is " << distribute_operator->inputs_tensor_info().size(); + MS_LOG(EXCEPTION) << "The index is out of range, the node_pair.second is " << (index - 1) + << ", the vector size is " << distribute_operator->inputs_tensor_info().size(); } + + // Step 2: Retrieve the tensor layout from the input tensor info TensorInfo tensorinfo_in = distribute_operator->inputs_tensor_info()[LongToSize(index - 1)]; TensorLayout tensorlayout_in = tensorinfo_in.tensor_layout(); return tensorlayout_in; } -// if reshape's output connect to several primitive, return the first layout found +// FindNextLayout function +// This function finds the next tensor layout based on the CNode and whether the next node is a RESHAPE operation. std::shared_ptr FindNextLayout(const CNodePtr &cnode, bool *next_is_reshape) { MS_EXCEPTION_IF_NULL(cnode); MS_EXCEPTION_IF_NULL(cnode->func_graph()); FuncGraphManagerPtr manager = cnode->func_graph()->manager(); MS_EXCEPTION_IF_NULL(manager); AnfNodeIndexSet node_set = manager->node_users()[cnode]; + + // Step 1: Iterate through nodes using the CNode for (auto &node_pair : node_set) { CNodePtr use_apply = node_pair.first->cast(); - if (use_apply == nullptr || !IsValueNode(use_apply->input(0))) { - continue; - } + + // Step 2: Check if the node is a RESHAPE operation if (IsPrimitiveCNode(use_apply, prim::kPrimReshape)) { *next_is_reshape = true; continue; } + ValueNodePtr prim_anf_node = use_apply->input(0)->cast(); MS_EXCEPTION_IF_NULL(prim_anf_node); PrimitivePtr node_prim = prim_anf_node->value()->cast(); MS_EXCEPTION_IF_NULL(node_prim); - MS_LOG(INFO) << "FindNextLayout prim " << node_prim->name(); - if (node_prim->name() == DEPEND && node_pair.second != 1) { - continue; - } + + // Step 3: Check if the node is a parallel-care node and has associated OperatorInfo if (IsParallelCareNode(use_apply) && use_apply->has_user_data()) { MS_LOG(INFO) << "FindNextLayout success prim " << node_prim->name(); *next_is_reshape = false; @@ -2048,6 +2797,7 @@ std::shared_ptr FindNextLayout(const CNodePtr &cnode, bool *next_i MS_LOG(DEBUG) << "FindNextLayout failed prim " << node_prim->name() << " " << IsParallelCareNode(use_apply) << " " << use_apply->has_user_data(); + // Step 4: Recursively search for the next layout auto layout_ptr = FindNextLayout(use_apply, next_is_reshape); if (layout_ptr) { return layout_ptr; @@ -2057,19 +2807,27 @@ std::shared_ptr FindNextLayout(const CNodePtr &cnode, bool *next_i return nullptr; } +// GetOutputLayoutFromCNode function +// This function retrieves the output tensor layout from a CNode based on the output index. std::shared_ptr GetOutputLayoutFromCNode(const CNodePtr &cnode, size_t output_index) { MS_EXCEPTION_IF_NULL(cnode); OperatorInfoPtr distribute_operator = GetDistributeOperator(cnode); MS_EXCEPTION_IF_NULL(distribute_operator); + + // Step 1: Check if the output index is out of range if (distribute_operator->outputs_tensor_info().size() <= output_index) { - MS_LOG(EXCEPTION) << "outputs_tensor_info size is " << distribute_operator->inputs_tensor_info().size() - << ", must be greater than output_index " << output_index; + MS_LOG(EXCEPTION) << "outputs_tensor_info size is " << distribute_operator->inputs_tensor_info().size() + << ", must be greater than output_index " << output_index; } + + // Step 2: Retrieve the tensor layout from the output tensor info TensorInfo tensorinfo_out = distribute_operator->outputs_tensor_info()[output_index]; TensorLayout tensorlayout_out = tensorinfo_out.tensor_layout(); return std::make_shared(tensorlayout_out); } +// FindPrevParallelCareNodeLayout function +// This function finds the previous parallel-care node's layout based on the current node and output index. std::shared_ptr FindPrevParallelCareNodeLayout(const AnfNodePtr &node, size_t output_index) { if (!node->isa()) { return nullptr; @@ -2078,16 +2836,20 @@ std::shared_ptr FindPrevParallelCareNodeLayout(const AnfNodePtr &n if (!IsValueNode(cnode->input(0))) { return nullptr; } + + // Step 1: Check if the node is a parallel-care node and has associated OperatorInfo if (IsParallelCareNode(cnode) && cnode->has_user_data()) { auto layout_ptr = GetOutputLayoutFromCNode(cnode, output_index); if (!layout_ptr) { - MS_LOG(EXCEPTION) << "Failure:GetLayoutFromCNode failed"; + MS_LOG(EXCEPTION) << "Failure: GetLayoutFromCNode failed"; } return layout_ptr; } return nullptr; } +// FindParameterNextLayout function +// This function finds the next tensor layout for a parameter node. std::shared_ptr FindParameterNextLayout(const AnfNodePtr &node, size_t curr_depth) { if (curr_depth > MAX_RECURSIVE_DEPTH) { MS_LOG(WARNING) << "When finding the next tensor layout for the parameter, exceeded the maximum recursion depth: " @@ -2097,6 +2859,8 @@ std::shared_ptr FindParameterNextLayout(const AnfNodePtr &node, si FuncGraphManagerPtr manager = node->func_graph()->manager(); MS_EXCEPTION_IF_NULL(manager); AnfNodeIndexSet node_set = manager->node_users()[node]; + + // Step 1: Iterate through nodes using the parameter node for (auto &node_pair : node_set) { if (IsPrimitiveCNode(node_pair.first, prim::kPrimLoad)) { auto layout_param = FindParameterNextLayout(node_pair.first, ++curr_depth); @@ -2113,9 +2877,8 @@ std::shared_ptr FindParameterNextLayout(const AnfNodePtr &node, si MS_EXCEPTION_IF_NULL(prim_anf_node); PrimitivePtr node_prim = prim_anf_node->value()->cast(); MS_EXCEPTION_IF_NULL(node_prim); - if ((node_prim->name() == DEPEND && node_pair.second != 1) || node_prim->name() == RESHAPE) { - continue; - } + + // Step 2: Check if the node is a parallel-care node and has associated OperatorInfo if (IsParallelCareNode(use_apply) && use_apply->has_user_data()) { auto layout = GetInputLayoutFromCNode(node_pair); return std::make_shared(layout); @@ -2124,72 +2887,101 @@ std::shared_ptr FindParameterNextLayout(const AnfNodePtr &node, si return nullptr; } +// CreateParameterLayout function +// This function creates a DataParallel tensor layout for a parameter node. std::shared_ptr CreateParameterLayout(const AnfNodePtr &node) { - // Create DataParallel tensor layout for parameter(support WideDeep). + // Step 1: Find the next layout for the parameter auto next_layout = FindParameterNextLayout(node, 0); if (next_layout != nullptr) { return next_layout; } + + // Step 2: Check global device manager CheckGlobalDeviceManager(); int64_t dev_num = g_device_manager->stage_device_num(); + + // Step 3: Create input tensor layout TensorLayout input_tensor_layout; - // create input_shape Shapes inputs_shape = GetNodeShape(node); Shape input_shape_array = inputs_shape[0]; + + // Step 4: Handle scalar parameter case if (input_shape_array.empty()) { MS_LOG(EXCEPTION) << "Don't support reshape a scalar parameter."; } - // create tensor_map + + // Step 5: Create tensor_map size_t shape_size = input_shape_array.size(); TensorMap input_tensor_map_array(SizeToLong(shape_size) - 1, -1); input_tensor_map_array.insert(input_tensor_map_array.begin(), 0); - // create dev_matrix + + // Step 6: Create dev_matrix Shape dev_matrix_array = {dev_num}; if (input_tensor_layout.InitFromVector(dev_matrix_array, input_tensor_map_array, input_shape_array) != SUCCESS) { MS_LOG(EXCEPTION) << "Create tensor layout for parameter failed."; } + + // Step 7: Return the created tensor layout return std::make_shared(input_tensor_layout); } +// InferSensRedistribution function +// This function infers the redistribution for sens (sensitivity) tensors. RedistributionOpListPtr InferSensRedistribution(const AnfNodePtr &node, const TensorLayout &loss_layout) { + // Step 1: Check if the node is valid MS_EXCEPTION_IF_NULL(node); + + // Step 2: Initialize tensor_redistribution TensorRedistribution tensor_redistribution; - // create stand alone layout:TensorMap:[all -1],dev_matrix:[dev_num]. CheckGlobalDeviceManager(); int64_t dev_num = g_device_manager->stage_device_num(); + + // Step 3: Create a stand-alone layout TensorLayout stand_alone_layout; Shapes inputs_shape = GetNodeShape(node); if (inputs_shape.empty()) { MS_LOG(EXCEPTION) << "InferSensRedistribution failed cause inputs shape is empty."; } Shape input_shape_array = inputs_shape[0]; + + // Step 4: Handle the case of an empty input shape (no redistribution needed) if (input_shape_array.empty()) { MS_LOG(INFO) << "No need to redistribution for sens."; return nullptr; } - // TensorMap + + // Step 5: Create stand-alone tensor_map TensorMap stand_alone_tensor_map_array(SizeToLong(input_shape_array.size()), -1); - // Dev_matrix + + // Step 6: Create dev_matrix Shape dev_matrix_array = {dev_num}; if (stand_alone_layout.InitFromVector(dev_matrix_array, stand_alone_tensor_map_array, input_shape_array) == FAILED) { MS_LOG(EXCEPTION) << "Create tensor layout for Sens failed."; } - // Infer Redistribution op list for stand alone and loss layout. + // Step 7: Initialize tensor redistribution for stand-alone and loss layout RankList dev_list = g_device_manager->GetDeviceListInThisStage(); if (tensor_redistribution.Init(stand_alone_layout, loss_layout, dev_list) == FAILED) { MS_LOG(EXCEPTION) << "Redistribution for Sens init failed."; } + + // Step 8: Infer redistribution operator list for sens tensor RedistributionOpListPtr sens_redistribution_list = tensor_redistribution.InferTensorRedistributionOperatorList(); MS_EXCEPTION_IF_NULL(sens_redistribution_list); + // Step 9: Return the inferred redistribution list return sens_redistribution_list; } +// FindPrevLayout function +// This function finds the previous layout for a given node, which can be a parameter or another node. std::shared_ptr FindPrevLayout(const AnfNodePtr &node) { + // Step 1: Check if the node is a parameter and create its layout if (node->isa()) { return CreateParameterLayout(node); } + + // Step 2: Check if the node is a CNode and its input is a primitive value node if (!node->isa()) { return nullptr; } @@ -2197,11 +2989,12 @@ std::shared_ptr FindPrevLayout(const AnfNodePtr &node) { if (!IsValueNode(cnode->input(0))) { return nullptr; } + + // Step 3: Handle Depend and other cases if (IsPrimitiveCNode(node, prim::kPrimReceive)) { return cnode->user_data(); } - if (IsParallelCareNode(cnode) && cnode->has_user_data() && - !IsPrimitiveCNode(node, prim::kPrimReshape)) { + if (IsParallelCareNode(cnode) && cnode->has_user_data() && !IsPrimitiveCNode(node, prim::kPrimReshape)) { auto layout_ptr = GetOutputLayoutFromCNode(cnode, 0); if (!layout_ptr) { MS_LOG(EXCEPTION) << "Failure:GetLayoutFromCNode failed"; @@ -2210,6 +3003,8 @@ std::shared_ptr FindPrevLayout(const AnfNodePtr &node) { } ValueNodePtr prim_anf_node = cnode->input(0)->cast(); PrimitivePtr prim = prim_anf_node->value()->cast(); + + // Step 4: Handle TupleGetItem if (prim->name() == prim::kTupleGetItem) { auto tuple_index = GetTupleGetItemIndex(cnode); auto layout_ptr = FindPrevParallelCareNodeLayout(cnode->input(1), LongToSize(tuple_index)); @@ -2220,6 +3015,8 @@ std::shared_ptr FindPrevLayout(const AnfNodePtr &node) { } return layout_ptr; } + + // Step 5: Recursively search for previous layout for (size_t index = 0; index < cnode->inputs().size(); ++index) { if (prim->name() == DEPEND && index != 1) { continue; @@ -2234,35 +3031,46 @@ std::shared_ptr FindPrevLayout(const AnfNodePtr &node) { return nullptr; } +// ReshapeInit function +// This function initializes the reshape information for all nodes in the graph. void ReshapeInit(const std::vector &all_nodes) { + // Step 1: Iterate through all nodes in the graph for (auto &node : all_nodes) { auto cnode = node->cast(); + + // Step 2: Check if the node is a valid CNode with a primitive value node as its input if ((cnode == nullptr) || !IsValueNode(cnode->input(0))) { continue; } ValueNodePtr prim_anf_node = cnode->input(0)->cast(); + + // Step 3: Check if the node is a parallel care node with operator info if (!IsParallelCareNode(cnode) || !cnode->has_user_data()) { continue; } PrimitivePtr prim = GetValueNode(prim_anf_node); MS_EXCEPTION_IF_NULL(prim); OperatorInfoPtr operator_info = cnode->user_data(); - if (operator_info == nullptr) { - MS_LOG(EXCEPTION) << "Failure:Primitive " << prim->ToString() << " OperatorInstance is nullptr"; - } + + // Step 4: Handle Reshape primitive if (prim->name() != RESHAPE) { continue; } + + // Step 5: Check if the strategy is already set auto attrs = prim->attrs(); if (StrategyFound(attrs)) { MS_LOG(EXCEPTION) << "Setting strategy for Reshape goes for nothing!"; } - MS_ASSERT(cnode->inputs().size() == RESHAPE_INPUT_SIZE); + + // Step 6: Find previous layout for input auto prev_layout_ptr = FindPrevLayout(cnode->input(1)); if (prev_layout_ptr) { auto reshape_info_ptr = std::dynamic_pointer_cast(operator_info); reshape_info_ptr->SetInputLayout(*prev_layout_ptr); } + + // Step 7: Check if next layout exists or use the previous layout bool is_next_reshape = false; auto next_layout_ptr = FindNextLayout(cnode, &is_next_reshape); if (next_layout_ptr) { @@ -2272,139 +3080,158 @@ void ReshapeInit(const std::vector &all_nodes) { auto reshape_info_ptr = std::dynamic_pointer_cast(operator_info); reshape_info_ptr->SetOutputLayout(*prev_layout_ptr); } + + // Step 8: Initialize the operator info if (operator_info->Init(nullptr, nullptr) == FAILED) { MS_LOG(EXCEPTION) << "Failure:operator " << prim->ToString() << " init failed"; } } } +// HandleDependLoss function +// This function handles the Depend nodes in the graph for loss calculation. CNodePtr HandleDependLoss(const CNodePtr &cnode, size_t curr_depth) { + // Step 1: Check if the current depth exceeds the maximum recursive depth if (curr_depth > MAX_RECURSIVE_DEPTH) { MS_LOG(WARNING) << "When handling the loss node of Depend, exceeded the max recursive depth: " << MAX_RECURSIVE_DEPTH; return nullptr; } - // Handle return->depend->loss + + // Step 2: Handle return->depend->loss pattern if (IsPrimitiveCNode(cnode, prim::kPrimDepend) || (IsPrimitiveCNode(cnode, prim::kPrimCast) && !cnode->has_user_data())) { auto depend_before = cnode->input(1)->cast(); MS_EXCEPTION_IF_NULL(depend_before); return HandleDependLoss(depend_before, ++curr_depth); } + + // Step 3: Return the current node if it doesn't match the pattern return cnode; } +// FindLossCNode function +// This function finds the loss CNode in a given FuncGraph based on certain conditions. LossNodeInfo FindLossCNode(const FuncGraphPtr &func_graph, size_t max_depth) { + // Check if the maximum recursive depth is exceeded if (max_depth > MAX_RECURSIVE_DEPTH) { MS_LOG(EXCEPTION) << "Recursive call is larger than 100000."; } - LossNodeInfo loss_node_info; - MS_EXCEPTION_IF_NULL(func_graph); - CNodePtr return_node = func_graph->get_return(); - MS_EXCEPTION_IF_NULL(return_node); + LossNodeInfo loss_node_info; // Initialize the loss node information + MS_EXCEPTION_IF_NULL(func_graph); // Ensure the input FuncGraph is not null + CNodePtr return_node = func_graph->get_return(); // Get the return CNode + MS_EXCEPTION_IF_NULL(return_node); // Ensure the return CNode is not null if (return_node->size() < 2) { MS_LOG(EXCEPTION) << "Failure: " << return_node->DebugString() << " size is smaller than 2"; } - AnfNodePtr pre_node = return_node->input(1); - MS_EXCEPTION_IF_NULL(pre_node); - auto pre_cnode = pre_node->cast(); - pre_cnode = HandleDependLoss(pre_cnode, 0); + AnfNodePtr pre_node = return_node->input(1); // Get the input node of the return CNode + MS_EXCEPTION_IF_NULL(pre_node); // Ensure the input node is not null + auto pre_cnode = pre_node->cast(); // Try to cast the input node to a CNode + pre_cnode = HandleDependLoss(pre_cnode, 0); // Handle potential Depend operations if (pre_cnode->input(0)->isa()) { auto switch_cnode = pre_cnode->input(0)->cast(); if (IsPrimitiveCNode(switch_cnode, prim::kPrimSwitch)) { MS_EXCEPTION_IF_NULL(switch_cnode); auto switch_graph = GetValueNode(switch_cnode->input(2)); - return FindLossCNode(switch_graph, max_depth + 1); + return FindLossCNode(switch_graph, max_depth + 1); // Recursively search in the switch graph } } if (pre_cnode == nullptr || !IsValueNode(pre_cnode->input(0))) { - return loss_node_info; + return loss_node_info; // If the previous CNode is null or not a primitive, return an empty result } if (!IsValueNode(pre_cnode->input(0))) { MS_LOG(DEBUG) << "pre_cnode:" << pre_cnode->ToString(); return loss_node_info; } - auto current_prim = GetValueNode(pre_cnode->input(0)); - // notice: the GetNext op has not input + auto current_prim = GetValueNode(pre_cnode->input(0)); // Get the primitive of the CNode + + // Check if the current primitive is in the set of invalid loss operations if (INVALID_LOSS_OPS.find(current_prim->name()) != INVALID_LOSS_OPS.end()) { MS_LOG(INFO) << "The loss is: " << current_prim->name(); loss_node_info.loss_node = pre_cnode; return loss_node_info; } - // size of common cnode is larger than 1 + // Check if the size of the common CNode is smaller than 2 if (pre_cnode->size() < 2) { MS_LOG(EXCEPTION) << pre_cnode->ToString() << " size( " << pre_cnode->inputs().size() << " ) is smaller than 2"; } - // return -> tuple_getitem -> loss + // Handle cases where the loss is wrapped in TupleGetItem, MakeTuple, or other operations if (current_prim->name() == prim::kTupleGetItem) { - auto tuple_index = GetTupleGetItemIndex(pre_cnode); - AnfNodePtr pre_pre_node = pre_cnode->input(1); + auto tuple_index = GetTupleGetItemIndex(pre_cnode); // Get the index of the TupleGetItem + AnfNodePtr pre_pre_node = pre_cnode->input(1); // Get the input node of TupleGetItem MS_EXCEPTION_IF_NULL(pre_pre_node); - auto pre_pre_cnode = pre_pre_node->cast(); + auto pre_pre_cnode = pre_pre_node->cast(); // Try to cast the input node to a CNode loss_node_info.has_tuple_getitem = true; loss_node_info.dout_index = tuple_index; loss_node_info.loss_node = pre_pre_cnode; return loss_node_info; } - // return -> make_tuple + // Handle cases where the loss contains MakeTuple if (current_prim->name() == MAKE_TUPLE) { - MS_LOG(WARNING) << "The loss have make_tuple, it is not supported"; + MS_LOG(WARNING) << "The loss contains MakeTuple, which is not supported"; return loss_node_info; } - // return -> loss + // Return the found loss node loss_node_info.loss_node = pre_cnode; MS_LOG(DEBUG) << "The loss name is " << current_prim->name(); return loss_node_info; } +// GetLossNodeGradOutputLayout function +// This function returns the gradient output layout of the loss node. TensorLayouts GetLossNodeGradOutputLayout(const LossNodeInfo &node_info) { - TensorLayouts ret; - auto loss_cnode = node_info.loss_node; - MS_EXCEPTION_IF_NULL(loss_cnode); + TensorLayouts ret; // Initialize the tensor layouts to be returned + auto loss_cnode = node_info.loss_node; // Get the loss CNode + MS_EXCEPTION_IF_NULL(loss_cnode); // Ensure the loss CNode is not null - ValueNodePtr prim_anf_node = loss_cnode->input(0)->cast(); - MS_EXCEPTION_IF_NULL(prim_anf_node); - PrimitivePtr prim = prim_anf_node->value()->cast(); - MS_EXCEPTION_IF_NULL(prim); + ValueNodePtr prim_anf_node = loss_cnode->input(0)->cast(); // Get the primitive ValueNode + MS_EXCEPTION_IF_NULL(prim_anf_node); // Ensure the primitive ValueNode is not null + PrimitivePtr prim = prim_anf_node->value()->cast(); // Get the primitive + MS_EXCEPTION_IF_NULL(prim); // Ensure the primitive is not null if (INVALID_LOSS_OPS.find(prim->name()) != INVALID_LOSS_OPS.end()) { MS_LOG(WARNING) << "The loss name is: " << prim->name() << ", do nothing for split sens now"; - return ret; + return ret; // Return an empty result if the loss is in the set of invalid loss operations } - OperatorInfoPtr operator_info = loss_cnode->user_data(); - MS_EXCEPTION_IF_NULL(operator_info); + OperatorInfoPtr operator_info = loss_cnode->user_data(); // Get operator info from loss CNode + MS_EXCEPTION_IF_NULL(operator_info); // Ensure the operator info is not null TensorInfo loss_grad_tensor_info; size_t op_output_size = operator_info->outputs_tensor_info().size(); MS_LOG(INFO) << "The loss name is " << operator_info->name() << ", the has tuple item is " << node_info.has_tuple_getitem << ", the output size is " << op_output_size << ", the dout_index is " << node_info.dout_index; + // Check if the output size and dout_index are valid if ((op_output_size == 0) || (op_output_size <= LongToSize(node_info.dout_index))) { MS_LOG(EXCEPTION) << "The index is " << node_info.dout_index << ", but the size of outputs is " << op_output_size; } + // Check if the sens is a tuple (currently not supported) if (!node_info.has_tuple_getitem && (op_output_size > 1)) { MS_LOG(EXCEPTION) << "Currently, it is not supported that the sens is a tuple."; } + // Get the tensor layout of the loss gradient loss_grad_tensor_info = operator_info->outputs_tensor_info()[LongToSize(node_info.dout_index)]; ret.push_back(loss_grad_tensor_info.tensor_layout()); - return ret; + return ret; // Return the gradient output layout of the loss node } +// SplitSens function +// This function handles the splitting of the sens tensor based on the loss gradient layout. void SplitSens(const CNodePtr &grad_sens_node, const TensorLayout &loss_grad_layout) { - MS_EXCEPTION_IF_NULL(grad_sens_node); + MS_EXCEPTION_IF_NULL(grad_sens_node); // Ensure the grad_sens_node is not null if (grad_sens_node->size() <= 1) { MS_LOG(EXCEPTION) << "The size of grad sens node is smaller than 2"; } - AnfNodePtr sens_tensor_node = grad_sens_node->input(1); - MS_EXCEPTION_IF_NULL(sens_tensor_node); - Shapes sens_shapes = GetNodeShape(sens_tensor_node); + AnfNodePtr sens_tensor_node = grad_sens_node->input(1); // Get the sens tensor node + MS_EXCEPTION_IF_NULL(sens_tensor_node); // Ensure the sens tensor node is not null + Shapes sens_shapes = GetNodeShape(sens_tensor_node); // Get the shape of the sens tensor if (sens_shapes.size() != 1) { MS_LOG(EXCEPTION) << "GetNodeShape for sens_tensor_node, output size is not 1"; } @@ -2457,39 +3284,44 @@ void SplitSens(const CNodePtr &grad_sens_node, const TensorLayout &loss_grad_lay } // Use _GetTensorSlice operator to split the sens tensor - FuncGraphPtr func_graph = grad_sens_node->func_graph(); // only cnode can get the graph + FuncGraphPtr func_graph = grad_sens_node->func_graph(); // Get the function graph MS_EXCEPTION_IF_NULL(func_graph); - Operator op = CreateGetTensorSliceOp(loss_grad_layout); - InsertGetTensorSliceOp(op, grad_sens_node, func_graph, 1, SPLIT_SENS); + Operator op = CreateGetTensorSliceOp(loss_grad_layout); // Create the GetTensorSlice operator + InsertGetTensorSliceOp(op, grad_sens_node, func_graph, 1, SPLIT_SENS); // Insert the GetTensorSlice operation } +// InsertForwardOps function +// This function inserts forward operations for a distribute operator. void InsertForwardOps(const OperatorInfoPtr &distribute_operator, const CNodePtr &cnode) { - MS_EXCEPTION_IF_NULL(distribute_operator); - MS_EXCEPTION_IF_NULL(cnode); + MS_EXCEPTION_IF_NULL(distribute_operator); // Ensure the distribute operator is not null + MS_EXCEPTION_IF_NULL(cnode); // Ensure the CNode is not null if (IsPrimitiveCNode(cnode, prim::kPrimReceive)) { - return; + return; // If the CNode is a Receive operation, do nothing } - OperatorVector forward_op = distribute_operator->forward_op(); + OperatorVector forward_op = distribute_operator->forward_op(); // Get the forward operations if (!forward_op.empty()) { - MS_LOG(INFO) << "Insert forward op for " << distribute_operator->name(); - ForwardCommunication(forward_op, cnode); + MS_LOG(INFO) << "Insert forward op for " << distribute_operator->name(); // Log the insertion of forward ops + ForwardCommunication(forward_op, cnode); // Perform forward communication } } +// StepReplace function +// This function performs replacements in the computation graph based on the given distribute_operator and CNode. void StepReplace(const OperatorInfoPtr &distribute_operator, const CNodePtr &cnode) { MS_EXCEPTION_IF_NULL(distribute_operator); MS_EXCEPTION_IF_NULL(cnode); - // StepReplaceOp + + // Step 1: Replace the CNode with the specified operators (replace_op) if available. OperatorVector replace_op = distribute_operator->replace_op(); if (!replace_op.empty()) { MS_LOG(INFO) << "StepReplaceOp " << cnode->ToString(); StepReplaceOp(replace_op, cnode); } - // StepReplaceGraph: after calling StepReplaceGraph, cnode can not be used anymore. + // Step 2: Replace the CNode with a new graph (replace_graph) if available. ReplaceGraphPtr replace_graph = distribute_operator->replace_graph(cnode); if (!replace_op.empty() && replace_graph) { - MS_LOG(EXCEPTION) << "Only one of replace_op or replace_op can be used"; + MS_LOG(EXCEPTION) << "Only one of replace_op or replace_graph can be used"; } if (replace_graph) { MS_LOG(INFO) << "StepReplaceGraph " << cnode->ToString(); @@ -2497,27 +3329,41 @@ void StepReplace(const OperatorInfoPtr &distribute_operator, const CNodePtr &cno } } +// FindForwardGraphByRootNodes function +// This function finds forward computation graphs based on a set of root nodes. std::set FindForwardGraphByRootNodes(const AnfNodeSet &root_all_nodes) { - // J->CNode->Graph std::set graph_set; + + // Step 1: Iterate through the root nodes to find potential forward graphs. for (auto &node : root_all_nodes) { MS_EXCEPTION_IF_NULL(node); + + // Skip non-CNode nodes. if (!node->isa()) { continue; } auto cnode = node->cast(); + + // Skip CNodes with less than 2 inputs or non-Primitive first input. if ((cnode->size() < 2) || !IsValueNode(cnode->input(0))) { continue; } + auto expect_prim = GetValueNode(cnode->input(0)); + + // Skip CNodes where the first input is neither "J" nor "SHARD". if (expect_prim->name() != J && expect_prim->name() != SHARD) { continue; } + + // Check if the second input is a FuncGraph. if (IsValueNode(cnode->input(1))) { auto graph = GetValueNode(cnode->input(1)); MS_LOG(DEBUG) << "Find the forward graph success"; graph_set.insert(graph); + + // Include all subgraphs used by the found graph. auto manager = graph->manager(); MS_EXCEPTION_IF_NULL(manager); auto graph_used = manager->func_graphs_used_total(graph); @@ -2529,31 +3375,42 @@ std::set FindForwardGraphByRootNodes(const AnfNodeSet &root_all_no return graph_set; } +// StepSplitSens function +// This function splits the sensitivity tensor (sens_node) based on the provided loss gradient layout. void StepSplitSens(const std::pair &sens_loss_pair) { CNodePtr sens_node = sens_loss_pair.first; auto loss_node = sens_loss_pair.second; + + // Step 1: Get the layout information for the loss gradient tensor. auto loss_grad_layout = GetLossNodeGradOutputLayout(loss_node); + + // Step 2: Split the sensitivity tensor if the layout information is available. if (!loss_grad_layout.empty()) { SplitSens(sens_node, loss_grad_layout[0]); } } +// IsPynativeParallel function +// This function checks if the current execution mode is PynativeParallel. bool IsPynativeParallel() { auto parallel_mode = ParallelContext::GetInstance()->parallel_mode(); auto execution_mode = MsContext::GetInstance()->get_param(MS_CTX_EXECUTION_MODE); return (execution_mode == kPynativeMode) && (parallel_mode == kSemiAutoParallel || parallel_mode == kAutoParallel); } -// Sens node satisfies the following conditions: cnode(sens)-->cnode(tuple_getitem)-->cnode-->cnode(J) +// GetSensLossPairs function +// This function finds pairs of sensitivity tensors and corresponding loss nodes within the given computation graph. std::vector> GetSensLossPairs(const FuncGraphPtr &root) { MS_EXCEPTION_IF_NULL(root); std::vector> sens_loss_pairs; + + // Step 1: Iterate through nodes in the computation graph. for (auto &node : root->nodes()) { if (!node->isa()) { continue; } - // cnode(sens)-->cnode(tuple_getitem) + // Step 2: Check if the node structure corresponds to the expected pattern (sens -> tuple_getitem -> J). auto sens_cnode = node->cast(); AnfNodePtr expect_tuple_getitem = sens_cnode->input(0); MS_EXCEPTION_IF_NULL(expect_tuple_getitem); @@ -2566,30 +3423,35 @@ std::vector> GetSensLossPairs(const FuncGraphP continue; } - // cnode(sens)-->cnode(tuple_getitem)-->cnode AnfNodePtr expect_anonymous = expect_tuple_getitem_cnode->input(1); MS_EXCEPTION_IF_NULL(expect_anonymous); if (!expect_anonymous->isa()) { continue; } - // cnode(sens)-->cnode(tuple_getitem)-->cnode-->cnode(J) auto expect_anonymous_cnode = expect_anonymous->cast(); AnfNodePtr expect_j = expect_anonymous_cnode->input(0); MS_EXCEPTION_IF_NULL(expect_j); if (!expect_j->isa()) { continue; } + auto expect_j_cnode = expect_j->cast(); + + // Step 3: Check if the node corresponds to the "J" primitive. if (!IsSomePrimitive(expect_j_cnode, J)) { continue; } + // Step 4: Ensure that the second input of "J" is a FuncGraph. if (!IsValueNode(expect_j_cnode->input(1))) { MS_LOG(EXCEPTION) << "Sens can't find the corresponding graph."; } + auto func_graph = GetValueNode(expect_j_cnode->input(1)); auto loss_node_info = FindLossCNode(func_graph, 0); + + // Step 5: Check if a loss node is found and create a sens-loss pair. if (loss_node_info.loss_node == nullptr) { MS_LOG(WARNING) << "Can not find the loss cnode"; continue; @@ -2665,42 +3527,62 @@ void ParallelCommunication(const FuncGraphPtr &root, const std::vector MAX_RECURSIVE_DEPTH) { - MS_LOG(WARNING) << "When finding the parameters' name of a operator, exceeded the maximum depth: " + MS_LOG(WARNING) << "When finding the parameters' name of an operator, exceeded the maximum depth: " << MAX_RECURSIVE_DEPTH; return {}; } + + // Get the inputs of the CNode std::vector node_inputs{node->inputs()}; ParameterMap param_names; + + // Iterate over the inputs for (int64_t i = 0; i < UlongToLong(node_inputs.size()); ++i) { int64_t idx = index > i ? index : i; auto input = node_inputs[LongToSize(i)]; + + // Check if the input is a Parameter node if (input->isa()) { auto input_parameter = input->cast(); + + // Check if the Parameter has a default value and requires gradient if (input_parameter->has_default() && ParameterRequireGrad(input_parameter)) { (void)param_names.emplace_back(std::make_pair(input_parameter->name(), input_parameter)); } } else if (input->isa()) { CNodePtr cnode = input->cast(); + + // Check if the input CNode is a cohesive node and has more than one input if (!IsValueNode(cnode->input(0))) { continue; } + if (IsCohesiveNode(cnode) && cnode->inputs().size() >= 1) { + // Recursively call NodeParameterName for cohesive nodes auto input_param_names = NodeParameterName(cnode, idx, 0); param_names.insert(param_names.end(), input_param_names.begin(), input_param_names.end()); } } } + return param_names; } +// Function: IsGatherInfo +// This function checks if a given name contains one of the specified strings indicating gather information. bool IsGatherInfo(const std::string &name) { std::vector gather_info_names = {"GatherInfo", "SparseGatherV2Info", "EmbeddingLookupInfo"}; for (std::string info_name : gather_info_names) { @@ -2711,40 +3593,61 @@ bool IsGatherInfo(const std::string &name) { return false; } +// Function: CheckpointStrategy +// This function computes and saves the strategy checkpoint based on operator information, tensor layouts, and manual shape information. void CheckpointStrategy(const std::vector &all_nodes, const FuncGraphPtr &root) { StrategyMap stra_map; TensorInfoMap tensor_info_map; ManualShapeMap manual_shape_map; + + // Iterate over all nodes in the graph for (auto &node : all_nodes) { MS_EXCEPTION_IF_NULL(node); auto cnode = node->cast(); + + // Check if the node is a CNode and has a Primitive input if ((cnode == nullptr) || !IsValueNode(cnode->input(0))) { continue; } + + // Retrieve parameter names associated with the operator node auto param_names = NodeParameterName(cnode, -1, 0); + + // Check if any parameter names are found if (param_names.empty()) { continue; } + string param_name = param_names[0].first; PrimitivePtr prim = GetValueNode(cnode->input(0)); MS_EXCEPTION_IF_NULL(prim); OperatorInfoPtr operator_info = cnode->user_data(); + + // Check if operator info exists if (operator_info) { if (operator_info->name().find(RESHAPEINFO) != std::string::npos) { continue; } std::string stratey_key_name = prim->name() + "_" + param_name; stra_map[stratey_key_name] = operator_info->strategy(); + + // Store tensor layout information for parameters for (auto param_name_pair : param_names) { tensor_info_map[param_name_pair.first] = param_name_pair.second->user_data(); } + + // Check if operator info represents gather information if (IsGatherInfo(operator_info->name())) { auto gather_info = std::dynamic_pointer_cast(operator_info); auto param_split_shapes = gather_info->param_split_shapes(); auto index_offsets = gather_info->index_offsets(); + + // Check consistency between param_split_shapes and index_offsets if (param_split_shapes.size() != index_offsets.size()) { - MS_LOG(EXCEPTION) << "In manual split, the param_split_shapes and index_offsets length should be same."; + MS_LOG(EXCEPTION) << "In manual split, the param_split_shapes and index_offsets length should be the same."; } + + // Store manual shape information std::vector> manual_shape; for (int64_t i = 0; i < UlongToLong(param_split_shapes.size()); ++i) { (void)manual_shape.emplace_back( @@ -2754,43 +3657,62 @@ void CheckpointStrategy(const std::vector &all_nodes, const FuncGrap } } } + + // Iterate over cloned parameters in the root graph for (auto &cloned_parameter_node : root->parameters()) { MS_EXCEPTION_IF_NULL(cloned_parameter_node); auto cloned_parameter = cloned_parameter_node->cast(); MS_EXCEPTION_IF_NULL(cloned_parameter); + // Check if the parameter is cloned if (!ParameterIsCloned(cloned_parameter_node)) { continue; } + std::string cloned_param_name = cloned_parameter_node->cast()->name(); auto cloned_param_layout = cloned_parameter_node->user_data(); + + // Check if the parameter has a tensor layout if (cloned_param_layout == nullptr) { continue; } + + // Store tensor layout information for cloned parameters tensor_info_map[cloned_param_name] = cloned_param_layout; } + + // Save the strategy checkpoint if (StrategyCheckpoint::GetInstance().Save(stra_map, tensor_info_map, &manual_shape_map) != SUCCESS) { MS_LOG(EXCEPTION) << "Save strategy checkpoint failed"; } } +// Function: SetForwardFlag +// This function sets the in_forward_flag for CNodes to indicate that they are part of the forward pass. void SetForwardFlag(const std::vector &all_nodes) { for (auto &node : all_nodes) { MS_EXCEPTION_IF_NULL(node); + + // Check if the node is a CNode if (!node->isa()) { continue; } + auto cnode = node->cast(); + + // Check if the CNode has a Primitive input if (!IsValueNode(cnode->input(0))) { continue; } - // CNode is globally unique. + // Set the in_forward_flag for the CNode MS_LOG(DEBUG) << "Set forward flag " << cnode->DebugString() << "."; cnode->set_in_forward_flag(true); } } +// SetForwardFlag function +// This function sets the in_forward_flag for CNodes in the provided set of nodes. void SetForwardFlag(const AnfNodeSet &all_nodes) { for (auto &node : all_nodes) { MS_EXCEPTION_IF_NULL(node); @@ -2802,11 +3724,13 @@ void SetForwardFlag(const AnfNodeSet &all_nodes) { continue; } - // CNode is globally unique. + // Set the in_forward_flag for the CNode, indicating that it is part of the forward pass. cnode->set_in_forward_flag(true); } } +// ForwardGraph function +// This function finds and returns a set of FuncGraphs connected to the provided root FuncGraph. std::set ForwardGraph(const FuncGraphPtr &root) { MS_EXCEPTION_IF_NULL(root); const auto &all_nodes = root->nodes(); @@ -2814,6 +3738,8 @@ std::set ForwardGraph(const FuncGraphPtr &root) { return graph_set; } +// FindRootForwardCNode function +// This function finds the root forward CNodes in the provided graph and returns them as a vector. std::vector FindRootForwardCNode(const FuncGraphPtr &graph, const AnfNodeSet &all_nodes) { MS_EXCEPTION_IF_NULL(graph); std::vector root_forward_nodes; @@ -2839,21 +3765,33 @@ std::vector FindRootForwardCNode(const FuncGraphPtr &graph, const An return root_forward_nodes; } +// InsertShapeOp function +// This function inserts a "shape" operation into the provided CNode with a specified input node and FuncGraph. void InsertShapeOp(const CNodePtr &node, const AnfNodePtr &pre_node, const FuncGraphPtr &root) { - // shape op doesn't have params and attrs. + // Create an empty parameter list and attribute map for the shape operation. OperatorParams params; OperatorAttrs attrs; + + // Extract the shape value from the CNode's input and convert it to a ValueSequence. auto shape_value = GetValueNode(node->input(2))->cast(); MS_EXCEPTION_IF_NULL(shape_value); auto shape = shape_value->value(); if (shape.empty()) { return; } + + // Create OperatorArgs using the empty attributes and parameters. OperatorArgs args = std::make_pair(attrs, params); + + // Create the "shape" operation. Operator op = std::make_pair(SHAPE_OP, args); + + // Insert the "shape" operation into the CNode's inputs. InsertNode(op, node, 2, pre_node, root, "shape"); } +// FindGrad function +// This recursive function searches for a "Grad" node within a CNode's inputs up to a specified depth. static AnfNodePtr FindGrad(const CNodePtr &cnode, size_t curr_depth) { if (curr_depth > MAX_RECURSIVE_DEPTH) { MS_LOG(WARNING) << "When finding Grad nodes, exceeded the maximum recursion depth: " << MAX_RECURSIVE_DEPTH; @@ -2872,9 +3810,11 @@ static AnfNodePtr FindGrad(const CNodePtr &cnode, size_t curr_depth) { return nullptr; } +// HandleRootReshapeAndSaveStrategy function +// This function handles root graph reshaping and saves strategy information. void HandleRootReshapeAndSaveStrategy(const std::vector &all_nodes) { - // If root graph has reshape op. Find the corresponding parameter. - // Reshape's shape is the shape of the parameter. + // Check if root graph has reshape operations and find the corresponding parameter nodes. + // For reshaping operations, save the strategy information in the executor. auto executor = pipeline::GraphExecutorPy::GetInstance(); for (auto &node : all_nodes) { if (!node->isa()) { @@ -2885,13 +3825,13 @@ void HandleRootReshapeAndSaveStrategy(const std::vector &all_nodes) continue; } if (cnode->in_forward_flag()) { - // Save strategy in executor + // Save strategy in executor for nodes in the forward pass. OperatorInfoPtr op_info = cnode->user_data(); if (op_info) { auto stra_ptr = op_info->strategy(); if (stra_ptr) { auto strategy = stra_ptr->GetInputDim(); - // fullname with scope should be found in step parallel end ir + // Fullname with scope should be found in step parallel end IR. executor->SetCNodeStrategy(cnode->fullname_with_scope(), strategy); } } @@ -2910,36 +3850,50 @@ void HandleRootReshapeAndSaveStrategy(const std::vector &all_nodes) auto root = node->func_graph(); auto grad_node = FindGrad(cnode, 0); if (grad_node) { + // Insert a "shape" operation for reshaping. InsertShapeOp(cnode, grad_node, root); } } } +// MarkForwardCNode function +// This function marks forward nodes in the provided FuncGraph. void MarkForwardCNode(const FuncGraphPtr &root) { MS_EXCEPTION_IF_NULL(root); auto all_nodes = root->nodes(); + + // Step 1: Find forward graphs rooted at the provided FuncGraph's nodes auto graph_set = FindForwardGraphByRootNodes(all_nodes); if (graph_set.empty()) { + // If no forward graphs are found, mark the ops in the root graph as forward MS_LOG(INFO) << "Can not find the forward graph, so mark the ops in root graph"; SetForwardFlag(all_nodes); } else { for (auto &func_graph : graph_set) { + // Step 2: For each forward graph, find return node and nodes reachable from it MS_LOG(INFO) << "The sub graph size of root is " << root->func_graphs_used().size(); auto return_node = func_graph->get_return(); MS_EXCEPTION_IF_NULL(return_node); auto all_dfs_nodes = DeepLinkedGraphSearch(return_node); + + // Step 3: Mark forward flag for the nodes reachable from the return node SetForwardFlag(all_dfs_nodes); + + // Step 4: Find forward nodes in the root graph associated with the forward graph auto root_forward_nodes = FindRootForwardCNode(func_graph, all_nodes); if (root_forward_nodes.empty()) { continue; } - // Mark forward flag for the nodes in root graph. + + // Step 5: Mark forward flag for the nodes in the root graph that correspond to the forward nodes SetForwardFlag(root_forward_nodes); } } } +// GetCommInfo function +// This function retrieves communication-related information such as device number and global rank. CommInfo GetCommInfo() { int64_t device_num = ParallelContext::GetInstance()->device_num(); int64_t global_rank = ParallelContext::GetInstance()->global_rank(); @@ -2948,6 +3902,8 @@ CommInfo GetCommInfo() { std::string backend = ms_context->get_param(MS_CTX_DEVICE_TARGET); std::string world_group; std::string communication_backend; + + // Determine the communication backend based on the device target if (backend == kAscendDevice || backend == kDavinciDevice) { world_group = HCCL_WORLD_GROUP; communication_backend = HCCL_BACKEND; @@ -2955,27 +3911,36 @@ CommInfo GetCommInfo() { world_group = NCCL_WORLD_GROUP; communication_backend = NCCL_BACKEND; } else { + // Raise an exception for an invalid communication backend MS_LOG(EXCEPTION) << "Invalid communication backend: " << backend; } + uint32_t world_rank_size = 0; + + // Retrieve the rank size from the communication model if (!CommManager::GetInstance().GetRankSize(world_group, &world_rank_size)) { MS_LOG(EXCEPTION) << "Get rank size failed"; } + // Set the device number from the rank size if it is not already set if (!ParallelContext::GetInstance()->device_num_is_set()) { device_num = UintToInt(world_rank_size); MS_LOG(INFO) << "Get device num from communication model, the device num is " << device_num; } + #if ENABLE_D || ENABLE_GPU if (ParallelContext::GetInstance()->device_num_is_set() && world_rank_size != device_num && !ParallelContext::GetInstance()->hccl_test_available()) { - // hccl_test_available is used when we compile graphs in real ascend card environment, but with hccl_test. - MS_LOG(EXCEPTION) << "The device_num " << device_num << " set in the context is not consist with " + // Check device number consistency for Ascend and GPU devices + MS_LOG(EXCEPTION) << "The device_num " << device_num << " set in the context is not consistent with " << world_rank_size << " devices you have" << ". Please check your rank_table file(for Ascend) or host file(for GPU)."; } #endif + uint32_t rank_id = 0; + + // Retrieve the global rank from the communication model if (!ParallelContext::GetInstance()->global_rank_is_set()) { if (!CommManager::GetInstance().GetRankID(world_group, &rank_id)) { MS_LOG(EXCEPTION) << "Get rank id failed"; @@ -2983,22 +3948,35 @@ CommInfo GetCommInfo() { global_rank = UintToInt(rank_id); MS_LOG(INFO) << "Get global rank from communication model, the global rank is " << global_rank; } + + // Create a CommInfo struct with the obtained information and return it CommInfo comm_info{device_num, global_rank, world_group, communication_backend}; return comm_info; } +// ParallelInit function +// This function initializes parallel execution based on the configuration provided in ParallelContext. Status ParallelInit() { + // Check if ParallelContext instance exists MS_EXCEPTION_IF_NULL(ParallelContext::GetInstance()); + + // Get split_stage_num and parallel_mode from ParallelContext int32_t split_stage_num = ParallelContext::GetInstance()->pipeline_stage_split_num(); std::string parallel_mode = ParallelContext::GetInstance()->parallel_mode(); + + // Check if split_stage_num is a positive number if (split_stage_num <= 0) { MS_LOG(ERROR) << "The parameter 'split_stage_num' must be a positive number, but got the value : " << split_stage_num; return FAILED; } + + // Get communication information auto comm_info = GetCommInfo(); int64_t device_num = comm_info.device_num; int64_t global_rank = comm_info.global_rank; + + // Check if device_num is positive and within a valid range if ((device_num <= 0) || (device_num > MAX_DEVICE_NUM)) { MS_LOG(ERROR) << "The context configuration parameter 'device_num' must be positive, " "but got the value of device_num: " @@ -3006,30 +3984,34 @@ Status ParallelInit() { return FAILED; } - // the device_num maybe get from communication interface + // Check if device_num is divisible by split_stage_num if (device_num % split_stage_num != 0) { MS_LOG(ERROR) << "The parameter 'device_num' must be divided by 'split_stage_num', but got the device_num : " - << device_num << "and the split_stage_num : " << split_stage_num; + << device_num << " and the split_stage_num : " << split_stage_num; return FAILED; } + // Check if global_rank is within a valid range if ((global_rank < 0) || (global_rank >= device_num)) { - MS_LOG(ERROR) << "The parameter 'global_rank' must be greater than 0 and less equal 'device num', " + MS_LOG(ERROR) << "The parameter 'global_rank' must be greater than 0 and less than 'device num', " "but got the global_rank : " - << global_rank << "and the device_num : " << device_num; + << global_rank << " and the device_num : " << device_num; return FAILED; } + // Create stages vector based on split_stage_num std::vector stages; for (int i = 0; i < split_stage_num; i++) { stages.push_back(device_num / split_stage_num); } + // Check if pipeline parallel is enabled and set parallel_mode to kSemiAutoParallel if ((split_stage_num > 1) && (parallel_mode != kSemiAutoParallel)) { MS_LOG(ERROR) << "To enable the pipeline parallel, please set the parallel mode to " << kSemiAutoParallel; return FAILED; } + // Initialize devices and log the parallel context information if (!InitDevice(device_num, global_rank, comm_info.communication_backend, stages)) { MS_LOG(ERROR) << "Init device failed"; return FAILED; @@ -3043,8 +4025,11 @@ Status ParallelInit() { return SUCCESS; } +// HandleForwardMakeTupleAndMakeList function +// This function handles forward make_tuple and make_list nodes by setting user data for operator info. void HandleForwardMakeTupleAndMakeList(const std::vector &all_nodes) { for (auto &node : all_nodes) { + // Check if the node is a make_tuple or make_list operation if (!AnfNodeIsPrimitive(node, MAKE_TUPLE) && !AnfNodeIsPrimitive(node, MAKE_LIST)) { continue; } @@ -3055,10 +4040,11 @@ void HandleForwardMakeTupleAndMakeList(const std::vector &all_nodes) continue; } + // Get the manager for the current function graph FuncGraphManagerPtr manager = cnode->func_graph()->manager(); MS_EXCEPTION_IF_NULL(manager); - // MakeTuple has multiple users, each user's TensorInfo must be same. + // Check and set operator info for make_tuple nodes with multiple users auto make_tuple_list_next_node = CheckMakeTupleSplit(node, manager); if (make_tuple_list_next_node == nullptr) { continue; @@ -3071,6 +4057,8 @@ void HandleForwardMakeTupleAndMakeList(const std::vector &all_nodes) } } +// CreateGroupsByCkptFile function +// This function creates groups based on the information loaded from a checkpoint file. bool CreateGroupsByCkptFile(const std::string &file) { GroupInfoMap group_info_map; if (StrategyCheckpoint::GetInstance().LoadGroupInfo(file, &group_info_map) != SUCCESS) { @@ -3084,6 +4072,8 @@ bool CreateGroupsByCkptFile(const std::string &file) { return true; } +// ReorderForPipelineSplit function +// This function reorders the function graph for pipeline split based on the number of pipeline stages. void ReorderForPipelineSplit(const FuncGraphPtr &root, const FuncGraphManagerPtr &manager, int64_t pipeline_stages) { if (!root->has_flag(BACKWARD) && pipeline_stages > 1) { root->set_flag(BACKWARD, true); @@ -3095,6 +4085,8 @@ void ReorderForPipelineSplit(const FuncGraphPtr &root, const FuncGraphManagerPtr } } +// IsInsertVirtualOutput function +// This function checks whether to insert a virtual output based on pipeline stage and parallel mode. bool IsInsertVirtualOutput(const FuncGraphPtr &root) { MS_EXCEPTION_IF_NULL(ParallelContext::GetInstance()); auto comm_info = GetCommInfo(); @@ -3112,6 +4104,8 @@ bool IsInsertVirtualOutput(const FuncGraphPtr &root) { IsPynativeParallel()); } +// HandleGroupInfo function +// This function handles group information, including saving it to a file if required. static void HandleGroupInfo(const FuncGraphPtr &root) { auto group_info = g_device_manager->group_info(); auto group_info_save_path = common::GetEnv("GROUP_INFO_FILE"); @@ -3127,6 +4121,8 @@ static void HandleGroupInfo(const FuncGraphPtr &root) { } } +// HandleDataParallel function +// This function handles data parallelism by saving group information to a file. static void HandleDataParallel() { std::string parallel_mode = ParallelContext::GetInstance()->parallel_mode(); if (parallel_mode == kDataParallel) { @@ -3146,38 +4142,69 @@ static void HandleDataParallel() { } } +// PipelinePreProcess function +// This function prepares the graph for pipeline parallelism, such as micro-batch handling, parameter start nodes, +// and the last stage's end node. static void PipelinePreProcess(const FuncGraphPtr &root, const FuncGraphManagerPtr &manager, const std::vector &all_nodes) { + // Step 1: Get the number of pipeline stages from ParallelContext auto pipeline_stages = ParallelContext::GetInstance()->pipeline_stage_split_num(); + + // Step 2: Check if pipeline parallelism is enabled (pipeline_stages > 1) if (pipeline_stages > 1) { + // Step 3: Handle micro-batch for all nodes HandleMicroBatch(all_nodes, manager); + + // Step 4: Create parameter start nodes ParameterStartNode(all_nodes, manager); + + // Step 5: Add the last stage's end node LastStageEndNode(all_nodes, manager, root); } } +// PipelinePostProcess function +// This function performs post-processing after pipeline parallelism, such as adding virtual assign add nodes, +// handling parameter receives, and generating label masks for micro-batch. static void PipelinePostProcess(const FuncGraphPtr &root, const std::vector &all_nodes) { + // Step 1: Get the number of pipeline stages from ParallelContext auto pipeline_stages = ParallelContext::GetInstance()->pipeline_stage_split_num(); + + // Step 2: Check if pipeline parallelism is enabled (pipeline_stages > 1) if (pipeline_stages > 1) { + // Step 3: Add virtual assign add nodes AddVirtualAssignAdd(root); + + // Step 4: Handle parameter receives HandleReceiveParam(root, all_nodes); + + // Step 5: Generate label masks for micro-batch LabelGenMaskMicro(root); } } +// InsertAllReduceForNormValue function +// This function inserts AllReduce operations for global norm value calculations. static void InsertAllReduceForNormValue(const AnfNodePtr &res_node) { + // Step 1: Extract information from the given result node auto cnode = res_node->cast(); auto graphs = res_node->func_graph(); MS_EXCEPTION_IF_NULL(graphs); auto manager = graphs->manager(); MS_EXCEPTION_IF_NULL(manager); auto node_user_map = manager->node_users(); + + // Step 2: Check if the result node corresponds to the EXPAND_DIMS primitive if (!IsSomePrimitive(cnode, EXPAND_DIMS)) { - MS_LOG(ERROR) << "Expected the operator expand_dims, but found the " << GetPrimName(cnode) - << "This may cause the calculation of the global norm incorrect"; + MS_LOG(ERROR) << "Expected the operator expand_dims, but found " << GetPrimName(cnode) + << ". This may cause incorrect global norm calculations."; return; } + + // Step 3: Get the number of pipeline stages from ParallelContext auto pipeline_stages = ParallelContext::GetInstance()->pipeline_stage_split_num(); + + // Step 4: Traverse the graph to find the SQRT node within limits auto find_node = res_node; uint32_t limits = 0; while (!IsSomePrimitive(find_node->cast(), SQRT) && limits < MAX_BFS_DEPTH) { @@ -3186,51 +4213,61 @@ static void InsertAllReduceForNormValue(const AnfNodePtr &res_node) { find_node = users.front().first; ++limits; } + + // Step 5: Check if a SQRT node is found, and if it already has an associated ALL_REDUCE node if (!find_node || !IsSomePrimitive(find_node->cast(), SQRT)) { return; } - auto anf_node = find_node->cast(); - if (anf_node->inputs().size() > 1 && IsSomePrimitive(anf_node->input(1)->cast(), ALL_REDUCE)) { + auto sqrt_node = find_node; + if (sqrt_node->inputs().size() > 1 && IsSomePrimitive(sqrt_node->input(1)->cast(), ALL_REDUCE)) { return; } - auto sqrt_node = find_node; + + // Step 6: Prepare the communication group for AllReduce auto cur_stage_rank_list = g_device_manager->GetDeviceListInThisStage(); Group cur_stage_device_list; if (g_device_manager->CreateGroup(cur_stage_rank_list, &cur_stage_device_list) != SUCCESS) { - MS_LOG(EXCEPTION) << "Create the communication group for allreduce in calculating global norm failed, " - "the rank_list is: " - << cur_stage_rank_list; + MS_LOG(EXCEPTION) << "Create the communication group for AllReduce in calculating global norm failed, " + "the rank_list is: " << cur_stage_rank_list; } + + // Step 7: Insert AllReduce operations for global norm value within the current stage InsertAllReduceToNodeInput(sqrt_node->cast(), cur_stage_device_list.name(), PARALLEL_GLOBALNORM); - MS_LOG(INFO) << "Insert the AllReduce for global norm value in stages succeed."; + MS_LOG(INFO) << "Inserted AllReduce for global norm value within stages succeed."; + + // Step 8: Check if there are multiple pipeline stages and insert AllReduce operations between stages if (pipeline_stages > 1) { - MS_LOG(INFO) << "Insert the AllReduce for global norm value between stages succeed."; + MS_LOG(INFO) << "Inserting AllReduce for global norm value between stages succeed."; auto ranks_between_stages = g_device_manager->GetDeviceListBetweenStage(); Group group_between_stages; if (g_device_manager->CreateGroup(ranks_between_stages, &group_between_stages)) { - MS_LOG(EXCEPTION) << "Create the communication group for allreduce in calculating global norm " - "with pipeline parallel failed, the rank_list is: " - << cur_stage_rank_list; + MS_LOG(EXCEPTION) << "Create the communication group for AllReduce in calculating global norm " + "with pipeline parallel failed, the rank_list is: " << cur_stage_rank_list; } InsertAllReduceToNodeInput(sqrt_node->cast(), group_between_stages.name(), PARALLEL_GLOBALNORM_BETWEEN); } } +// FindExpanDimsWIthGradScale function +// This function searches for an EXPAND_DIMS node in the graph with a specific attribute (GRAD_SCALE). AnfNodePtr FindExpanDimsWIthGradScale(const AnfNodePtr &node_ptr, const NodeUsersMap &node_users_map, uint32_t limits) { std::queue visited; AnfNodePtr queue_node = nullptr; CNodePtr cnode = nullptr; AnfNodePtr last_node = nullptr; uint32_t depth = 0; + if (!node_ptr) { return nullptr; } + visited.push(node_ptr); while (!visited.empty()) { queue_node = visited.front(); visited.pop(); cnode = queue_node->cast(); - // MAKE_TUPLE will not appear after the load in the forward graph + + // Check if the node corresponds to EXPAND_DIMS and has the GRAD_SCALE attribute if (IsSomePrimitive(cnode, EXPAND_DIMS)) { auto value = GetAttrsFromAnfNode(queue_node, GRAD_SCALE); if (!value || !GetValue(value)) { @@ -3238,13 +4275,20 @@ AnfNodePtr FindExpanDimsWIthGradScale(const AnfNodePtr &node_ptr, const NodeUser } return queue_node; } + + // Check if the node belongs to a predefined list of primitives if (!IsSomePrimitiveList(cnode, {ENVIRONGET, MUL, SQUARE, REDUCE_SUM, EXPAND_DIMS, DEPEND, CAST, REF_TO_EMBED})) { continue; } + auto node_set = node_users_map.at(queue_node); + + // Add users of the current node to the visited queue for (auto &node_user : node_set) { visited.push(node_user.first); } + + // Check if the current node is part of a sequence if (!last_node || last_node == queue_node) { if (++depth == limits) { break; @@ -3252,45 +4296,72 @@ AnfNodePtr FindExpanDimsWIthGradScale(const AnfNodePtr &node_ptr, const NodeUser last_node = visited.back(); } } + return nullptr; } - +// InsertDivAndAllReduceForNorm function +// This function inserts division and all-reduce operations for normalization based on node user map and device count. static void InsertDivAndAllReduceForNorm(const NodeUsersMap &node_user_map, const AnfNodePtr ¶meter, uint32_t dev_num) { + // Step 1: Get the users of the parameter node auto params_user_set = node_user_map.at(parameter); + + // Step 2: Iterate through the users of the parameter for (auto ¶m_pair : params_user_set) { auto cnode = param_pair.first->cast(); MS_EXCEPTION_IF_NULL(cnode); + + // Skip nodes that are part of the forward pass if (cnode->in_forward_flag()) { continue; } + + // Step 3: Find the expand_dims operation with grad_scale attribute auto expand_dims_node = FindExpanDimsWIthGradScale(cnode, node_user_map, MAX_BFS_DEPTH); + + // Skip if no expand_dims node with grad_scale attribute is found if (!expand_dims_node) continue; + + // Step 4: Check if the expand_dims node has the GRAD_SCALE attribute set to true auto value = GetAttrsFromAnfNode(expand_dims_node, GRAD_SCALE); if (!value || !GetValue(value)) continue; + + // Step 5: Insert a realdiv operation to the input of the expand_dims node if (dev_num > 0) { InsertRealDivOpToNodeInput(expand_dims_node->cast(), dev_num, PARALLEL_GLOBALNORM_DIV); MS_LOG(INFO) << "Insert the realdiv with " << dev_num << " for the parameter " << parameter->fullname_with_scope() << "succeed!"; } - // If already inserted allreduce, the pattern will not be matched and thus no allreduce will be inserted. + + // Step 6: Insert an all-reduce operation for norm value InsertAllReduceForNormValue(expand_dims_node); } } +// GetMirrorOp function +// This function retrieves the mirror operation corresponding to a parameter node from the node user map. static AnfNodePtr GetMirrorOp(const NodeUsersMap &node_user_map, const AnfNodePtr ¶meter) { + // Step 1: Get the users of the parameter node auto params_user_set = node_user_map.at(parameter); + + // Step 2: Iterate through the users of the parameter for (auto ¶m_pair : params_user_set) { auto cnode = param_pair.first->cast(); std::vector candidate = {cnode}; + + // Skip nodes that are not part of the forward pass if (!cnode->in_forward_flag()) { continue; } + + // Include additional candidates if the current node is trivial or a load node if (IsInTrivialNodeList(cnode) || IsSomePrimitive(cnode, LOAD)) { auto load_users = node_user_map.at(param_pair.first); std::transform(load_users.begin(), load_users.end(), std::back_inserter(candidate), [](const auto &v) { return v.first; }); } + + // Step 3: Find the mirror operation node among the candidates for (auto &node : candidate) { auto local_cnode = node->cast(); if (!IsPrimitiveCNode(local_cnode, prim::kPrimMirror) && @@ -3301,33 +4372,54 @@ static AnfNodePtr GetMirrorOp(const NodeUsersMap &node_user_map, const AnfNodePt return node; } } + + // Return nullptr if no mirror operation is found return nullptr; } +// HandlGlobalNormScale function +// This function handles global norm scaling for parameters in the computation graph. static void HandlGlobalNormScale(const FuncGraphPtr &root, const std::vector &all_nodes, const FuncGraphManagerPtr &manager) { + // Step 1: Get the list of parameters from the root graph auto parameters = root->parameters(); + + // Step 2: Get the node user map from the manager auto node_user_map = manager->node_users(); MS_LOG(INFO) << "Start to process the global norm"; + // Step 3: Iterate through the parameters for (auto ¶meter : parameters) { int64_t dev_num = 0; + + // Skip parameters that do not require gradients if (!ParameterRequireGrad(parameter)) continue; + + // Step 4: Get the mirror operation node corresponding to the parameter auto mirror_node = GetMirrorOp(node_user_map, parameter); + + // Step 5: Get the device number from the mirror node's attributes auto device_num_ptr = GetAttrsFromAnfNode(mirror_node, DEV_NUM); if (device_num_ptr && device_num_ptr->isa()) { dev_num = GetValue(device_num_ptr); } + + // Step 6: Insert division and all-reduce operations for norm based on device number InsertDivAndAllReduceForNorm(node_user_map, parameter, dev_num); } } +// StepParallel function +// This function performs step parallel optimization on the computation graph. bool StepParallel(const FuncGraphPtr &root, const opt::OptimizerPtr &optimizer) { + // Check if running in a distributed environment #if ((defined ENABLE_CPU) && (!defined _WIN32) && !defined(__APPLE__)) - if (ps::PSContext::instance()->is_server() || ps::PSContext::instance()->is_scheduler()) { + if (ps::PSContext::instance()->is_server() || ps::PSContext()->is_scheduler()) { return false; } #endif + + // Step 1: Initialize required variables and check parallel context MS_EXCEPTION_IF_NULL(root); MS_EXCEPTION_IF_NULL(optimizer); MS_EXCEPTION_IF_NULL(ParallelContext::GetInstance()); @@ -3340,7 +4432,8 @@ bool StepParallel(const FuncGraphPtr &root, const opt::OptimizerPtr &optimizer) auto pipeline_stages = ParallelContext::GetInstance()->pipeline_stage_split_num(); // assume no change to graph bool changes = false; - // control whether use model_parallel mode + + // Step 2: Handle cases based on parallel mode if (!root->has_flag(kAutoParallel) || ((parallel_mode != kAutoParallel) && (parallel_mode != kSemiAutoParallel)) || (root->has_flag(SEMI_AUTO_PARALLEL_RUN_ONCE_ONLY))) { if (!root->has_flag(CHECK_SET_STRATEGY_VALID_ONCE_ONLY)) { @@ -3356,6 +4449,7 @@ bool StepParallel(const FuncGraphPtr &root, const opt::OptimizerPtr &optimizer) struct timeval start_time, end_time; (void)gettimeofday(&start_time, nullptr); + // Step 3: Perform step parallel optimization MS_LOG(INFO) << "Now entering step parallel"; DumpGraph(root, std::string(STEP_PARALLEL_BEGIN)); AnfNodePtr ret = root->get_return(); @@ -3450,7 +4544,8 @@ bool StepParallel(const FuncGraphPtr &root, const opt::OptimizerPtr &optimizer) return changes; } -// Needed by rec_parser +// ExtractInputsTensorName function +// This function extracts tensor names from the inputs of a CNode and returns them in a vector. std::vector ExtractInputsTensorName(const CNodePtr &node) { std::vector name_inputs; std::vector all_inputs = node->inputs(); diff --git a/mindspore/ccsrc/frontend/parallel/step_parallel.h b/mindspore/ccsrc/frontend/parallel/step_parallel.h index 44debd1686e..b3d416569a2 100644 --- a/mindspore/ccsrc/frontend/parallel/step_parallel.h +++ b/mindspore/ccsrc/frontend/parallel/step_parallel.h @@ -38,146 +38,198 @@ using OperatorInfoPtr = std::shared_ptr; namespace mindspore { namespace parallel { -const uint64_t kUSecondInSecond = 1000000; -const int32_t RECURSION_LIMIT = 3; +// Constants +const uint64_t kUSecondInSecond = 1000000; // Microseconds in a second +const int32_t RECURSION_LIMIT = 3; // Recursion limit for some functions +// Struct to store information about a loss node struct LossNodeInfo { - bool has_tuple_getitem = false; - int64_t dout_index = 0; // now don't support the sens is a tuple - CNodePtr loss_node = nullptr; + bool has_tuple_getitem = false; // Indicates if the node has a tuple_getitem operation + int64_t dout_index = 0; // Index of the "dout" tensor in the node (currently doesn't support tuple "sens") + CNodePtr loss_node = nullptr; // Pointer to the loss node }; +// Struct to store communication information struct CommInfo { - int64_t device_num = 1; - int64_t global_rank = 0; - std::string world_group; - std::string communication_backend; + int64_t device_num = 1; // Number of devices + int64_t global_rank = 0; // Global rank of the current device + std::string world_group; // World group name + std::string communication_backend; // Communication backend used }; +// Function to create input nodes for an operator std::vector CreateInput(const Operator &op, const AnfNodePtr &node, const std::string &instance_name); + +// Function to perform forward communication for a list of operators void ForwardCommunication(OperatorVector forward_op, const CNodePtr &node); +// Function to insert redistribution operations into the graph void InsertRedistribution(const RedistributionOpListPtr &redistribution_oplist_ptr, const CNodePtr &node, const FuncGraphPtr &func_graph, int64_t pos, const CNodePtr &pre_node); +// Function to get the input tensor layout for a given node TensorLayout GetTensorInLayout(const CNodePtr &pre_node, const PrimitivePtr &pre_prim, const OperatorInfoPtr &distribute_operator_pre); +// Function to get the distribute operator associated with a CNode OperatorInfoPtr GetDistributeOperator(const CNodePtr &node); +// Function to perform tensor redistribution void Redistribution(const std::pair &node_pair, const OperatorInfoPtr &distribute_operator, const CNodePtr &middle_node, int64_t index, TensorRedistribution tensor_redistribution, const CNodePtr &pre_node); +// Function to check if a strategy is found in attributes bool StrategyFound(const mindspore::HashMap &attrs); +// Function to check if a specific attribute is found in attributes bool AttrFound(const mindspore::HashMap &attrs, const std::string &target); +// Function to get the accumulated gradient node for a parameter AnfNodePtr GetAccuGrad(const std::vector ¶meters, const std::string &weight_name); +// Function to mark forward CNodes in the computation graph void MarkForwardCNode(const FuncGraphPtr &root); +// Function to check if there are communication ops in the graph bool FindCommunicationOp(const std::vector &all_nodes); +// Function to perform step redistribution void StepRedistribution(const CNodePtr &node, const OperatorInfoPtr &distribute_operator, const CNodePtr &insert_node, const TensorRedistribution &tensor_redistribution, const CNodePtr &pre_node); +// Function to replace operators in a step void StepReplaceOp(OperatorVector replace_op, const CNodePtr &node); +// Function to insert virtual div operators void InsertVirtualDivOp(const VirtualDivOp &virtual_div_op, const CNodePtr &node); +// Function to find a CNode in a function graph std::pair FindCNode(const AnfNodePtr &anode, const std::string &name, const FuncGraphPtr &func_graph, size_t max_depth); -// Generate and init parallel operator +// Function to generate and initialize a parallel operator OperatorInfoPtr OperatorInstance(const PrimitivePtr &prim, const PrimitiveAttrs &attrs, const std::vector &shape_list); -// Generate without initing parallel operator +// Function to generate a parallel operator without initialization OperatorInfoPtr NewOperatorInstance(const PrimitivePtr &prim, const PrimitiveAttrs &attrs, std::vector shape_list); -// Extract strategy from attr +// Function to extract strategy from an attribute StrategyPtr ExtractStrategy(const ValuePtr &strategy); -// Extract shape from anfnode +// Function to extract shapes from a CNode std::vector ExtractShape(const CNodePtr &node); -// Find finally sub graph +// Function to find a sub-graph within a function graph std::pair FindSubGraph(const FuncGraphPtr &func_graph, const AnfNodePtr ¶meter); -// Set distribute shape for parameters abstract +// Function to set parallel shapes for parameter abstracts std::string SetParallelShape(const AnfNodePtr ¶meter, const std::pair &res); -// change parameters'shape in resource +// Function to change parameters' shapes in resource void CoverSliceShape(const FuncGraphPtr &root); +// Function to label batch size split void LableBatchSizeSplit(const CNodePtr &node); +// Function to set virtual dataset strategy void SetVirtualDatasetStrategy(const CNodePtr &node); + +// Function to check if virtual output should be inserted bool IsInsertVirtualOutput(const FuncGraphPtr &root); +// Function to set strided slice split strategy void SetStridedSliceSplitStrategy(const std::vector &all_nodes); -// Create parallel operator for primitive node(has strategy) +// Function to create parallel operators and extract information from the computation graph void ExtractInformation(const std::vector &all_nodes); +// Function to get input layout from a CNode TensorLayout GetInputLayoutFromCNode(const std::pair &node_pair); +// Function to find the next layout from a CNode std::shared_ptr FindNextLayout(const CNodePtr &node); +// Function to get output layout from a CNode std::shared_ptr GetOutputLayoutFromCNode(const CNodePtr &cnode, size_t output_index); +// Function to find the previous layout from an AnfNode std::shared_ptr FindPrevParallelCareNodeLayout(const AnfNodePtr &node, size_t output_index); +// Function to find the previous layout from an AnfNode std::shared_ptr FindPrevLayout(const AnfNodePtr &node); +// Function to initialize reshaping void ReshapeInit(const std::vector &all_nodes); +// Function to generate batch parallel strategy for an operator StrategyPtr GenerateBatchParallelStrategy(const OperatorInfoPtr operator_, const PrimitivePtr prim); -// Add node for whole graph +// Function to add parallel communication operations to the computation graph void ParallelCommunication(const FuncGraphPtr &root, const std::vector &all_nodes, const FuncGraphManagerPtr &manager); +// Function to create a mapping of parameter names to parameter nodes ParameterMap NodeParameterName(const CNodePtr &node, int64_t index, size_t curr_depth); +// Function to checkpoint the strategy for multi-train void CheckpointStrategy(const std::vector &all_nodes, const FuncGraphPtr &root); -// main step of Parallel +// Main function for performing step parallel optimization bool StepParallel(const FuncGraphPtr &func_graph, const opt::OptimizerPtr &optimizer); +// Function to get the index of a TupleGetItem operation int64_t GetTupleGetItemIndex(const CNodePtr &cnode); +// Function to initialize parallel context Status ParallelInit(); +// Function to forward the graph and return a set of forward graphs std::set ForwardGraph(const FuncGraphPtr &root); +// Function to extract input tensor names from a CNode std::vector ExtractInputsTensorName(const CNodePtr &node); +// Function to find the next layout for a parameter node std::shared_ptr FindParameterNextLayout(const AnfNodePtr &node); +// Function to check if a parameter node is used in a function graph bool IsUsedParameter(const FuncGraphPtr &graph, const AnfNodePtr ¶meter); +// Function to apply parallel optimization on a parameter void ApplyParallelOptOnParam(TensorLayout *tensor_layout, const OperatorInfoPtr &distribute_operator, const CNodePtr &cnode, const AnfNodePtr ¶meter, size_t index); +// Function to set the strategy for the last node void SetLastNodeStrategy(const StrategyPtr strategyPtr); +// Function to create groups based on a checkpoint file bool CreateGroupsByCkptFile(const std::string &file); +// Function to find the unique IDs of the last nodes in a function graph void FindLastNodesUniqueId(const FuncGraphPtr &root, std::vector *unique_ids, std::vector *indexes); +// Function to insert virtual output for the computation graph void InsertVirtualOutput(const FuncGraphPtr &root, const std::vector &all_nodes); +// Function to get the name of the Mirror operator std::string MirrorOpName(); +// Function to get communication information CommInfo GetCommInfo(); +// Function to get the name of a primitive operator in a CNode std::string GetPrimName(const CNodePtr &node); +// Function to reorder nodes for pipeline splitting void ReorderForPipelineSplit(const FuncGraphPtr &root, const FuncGraphManagerPtr &manager, int64_t pipeline_stages); + +} // namespace parallel +} // namespace mindspore + } // namespace parallel } // namespace mindspore