huawei
/
mindspore2022
0
0
Fork 53
Code Issues Pull Requests 29 Projects Releases Wiki Activity
mindspore2022/mindspore/ccsrc/frontend/parallel/step_parallel.cc

4565 lines
189 KiB
C++
Raw Blame History

/**
* Copyright 2019-2022 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "frontend/parallel/step_parallel.h"
#include <inttypes.h>
#include <sys/time.h>
#include <algorithm>
#include <map>
#include <memory>
#include <set>
#include <string>
#include <utility>
#include <queue>
#include "utils/hash_map.h"
#include "base/core_ops.h"
#include "frontend/operator/ops.h"
#include "frontend/optimizer/optimizer.h"
#include "frontend/parallel/auto_parallel/graph_costmodel.h"
#include "include/common/utils/parallel_context.h"
#include "frontend/parallel/device_manager.h"
#include "frontend/parallel/dynamic_creator.h"
#include "frontend/parallel/graph_util/generate_graph.h"
#include "frontend/parallel/graph_util/graph_info.h"
#include "frontend/parallel/graph_util/node_info.h"
#include "frontend/parallel/graph_util/pipeline_split_utils.h"
#include "frontend/parallel/node_check.h"
#include "frontend/parallel/parameter_manager.h"
#include "frontend/parallel/ops_info/matmul_info.h"
#include "ir/param_info.h"
#include "ir/tensor.h"
#include "utils/trace_base.h"
#include "include/common/utils/comm_manager.h"
#include "utils/ms_context.h"
#include "utils/symbolic.h"
#include "mindspore/core/utils/parallel_node_check.h"
#include "frontend/parallel/parallel_optimizer/opt_param_mgr.h"
#if ((defined ENABLE_CPU) && (!defined _WIN32))
#include "ps/util.h"
#include "ps/ps_context.h"
#endif
using mindspore::tensor::Tensor;
namespace mindspore {
namespace parallel {
// Set of communication operations
static const std::set<std::string> COMMUNICATION_OPS = {ALL_REDUCE, ALL_GATHER, ALL_TO_ALL, REDUCE_SCATTER};
// Set of invalid loss operations
static const std::set<std::string> INVALID_LOSS_OPS = {GET_NEXT, VIRTUALLOSS, LOAD, UPDATESTATE};
// Set of operations with no input tensors
static const std::set<std::string> NO_INPUT_TENSOR_OPS = {UNIFORM_REAL};
// 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<AnfNodePtr, std::pair<AnfNodePtr, int64_t>> 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<AnfNodePtr> new_node_input, bool do_mirror, bool accu_flag) {
if (new_node_input.empty()) {
return;
}
auto prim_anf_node = new_node_input[0]->cast<ValueNodePtr>();
auto prim = GetValueNode<PrimitivePtr>(prim_anf_node);
MS_EXCEPTION_IF_NULL(prim);
auto attrs = prim->attrs();
attrs[DO_MIRROR] = MakeValue<bool>(do_mirror);
attrs[ADD_ACCU] = MakeValue<bool>(accu_flag);
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<AnfNodePtr> &new_node_input, const CNodePtr &node) {
if (new_node_input.empty()) {
return;
}
auto prim_anf_node = new_node_input[0]->cast<ValueNodePtr>();
auto prim = GetValueNode<PrimitivePtr>(prim_anf_node);
MS_EXCEPTION_IF_NULL(prim);
auto attrs = prim->attrs();
auto anf_node = node->input(0)->cast<ValueNodePtr>();
auto prim_node = GetValueNode<PrimitivePtr>(anf_node);
MS_EXCEPTION_IF_NULL(prim_node);
auto node_attrs = prim_node->attrs();
if (node_attrs.find(RECOMPUTE_COMM_OP) != node_attrs.end() && !GetValue<bool>(node_attrs[RECOMPUTE_COMM_OP])) {
attrs[RECOMPUTE] = MakeValue<bool>(false);
prim->SetAttrs(attrs);
MS_LOG(INFO) << "Do not recompute the forward communication operator of " << prim_node->ToString();
}
}
// CreateInput function
// This function creates a vector of input nodes for an operator based on the given arguments and instance name.
std::vector<AnfNodePtr> CreateInput(const Operator &op, const AnfNodePtr &node, const std::string &instance_name) {
MS_EXCEPTION_IF_NULL(node);
OperatorArgs arg_forward = op.second;
ValuePtr pyop_instance = CreateOpInstance(arg_forward.first, op.first, instance_name);
MS_EXCEPTION_IF_NULL(pyop_instance);
OperatorParams params = arg_forward.second;
std::vector<AnfNodePtr> new_node_input = {NewValueNode(pyop_instance), node};
if (!params.empty()) {
for (auto &param : 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 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<AnfNodePtr> &parameters, const std::string &weight_name) {
for (auto &param : parameters) {
if (!ParameterIsCloned(param)) {
continue;
}
auto param_ptr = param->cast<ParameterPtr>();
MS_EXCEPTION_IF_NULL(param_ptr);
if (param_ptr->name().find(weight_name) != std::string::npos &&
param_ptr->name().find(ACCU_GRADS) != std::string::npos) {
MS_LOG(INFO) << "Find the accumulation grad node: " << param_ptr->name();
return param;
}
}
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<AnfNodePtr> 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);
MS_EXCEPTION_IF_NULL(node);
MS_EXCEPTION_IF_NULL(root->manager());
std::string op_name = op.first;
OperatorArgs arg_forward = op.second;
AnfNodePtr grad_accu = nullptr;
int64_t grad_accumulation_step = ParallelContext::GetInstance()->grad_accumulation_step();
int64_t split_stage_num = ParallelContext::GetInstance()->pipeline_stage_split_num();
if (grad_accumulation_step > 1 || split_stage_num > 1) {
auto parameters = root->parameters();
grad_accu = GetAccuGrad(parameters, weight_name);
if (!grad_accu) {
if (op_name == MIRROR_MINI_STEP_OPERATOR) {
op_name = MIRROR_OPERATOR;
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 using " << op_name << " parameter: " << weight_name;
}
}
}
// 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);
// 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<AnfNodePtr> 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 &param : 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 the op has 'group' attribute, set the rank list name for the op
SetCommunicationOpGroupLabel(new_node_input);
// Gradient accumulation
if (grad_accumulation_step > 1) {
bool add_accu = root->has_flag(kAccumulation);
// 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 &param_name = "",
const FuncGraphPtr &root = nullptr) {
// 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<AnfNodePtr> 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);
}
// Step 5: Set attributes of the new CNode's Primitive
auto new_node_value = node_input[0]->cast<ValueNodePtr>();
MS_EXCEPTION_IF_NULL(new_node_value);
PrimitivePtr new_node_prim = new_node_value->value()->cast<PrimitivePtr>();
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 node with the new CNode
manager->SetEdge(node, SizeToInt(index), new_node);
MS_LOG(INFO) << "Insert " << instance_name << " success";
}
// 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 &param_name = "",
const FuncGraphPtr &root = nullptr) {
// 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<AnfNodePtr> 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);
}
// Step 5: Set attributes of the new CNode's Primitive
auto new_node_prim = GetValueNode<PrimitivePtr>(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);
// 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) {
auto uses_cnode = uses_pair.first->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(uses_cnode);
if (!IsValueNode<Primitive>(uses_cnode->input(0))) {
break;
}
PrimitivePtr value_node_prim = GetValueNode<PrimitivePtr>(uses_cnode->input(0));
MS_EXCEPTION_IF_NULL(value_node_prim);
if (value_node_prim->name() == prim::kTupleGetItem) {
if (uses_set.size() > 1) {
MS_LOG(EXCEPTION) << "Now only support one output, but got " << uses_set.size();
}
node_to_insert = uses_cnode;
}
}
MS_EXCEPTION_IF_NULL(node_to_insert);
std::reverse(forward_op.begin(), forward_op.end());
// 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<AnfNodePtr> 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);
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);
}
}
// 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);
std::vector<AnfNodePtr> make_tuple_inputs;
make_tuple_inputs.push_back(NewValueNode(prim::kPrimMakeTuple));
for (uint64_t i = 0; i < num; i++) {
std::vector<AnfNodePtr> tuple_get_item_inputs{NewValueNode(prim::kPrimTupleGetItem), prev,
CreatInt64Imm(UlongToLong(i))};
auto tuple_get_item = func_graph->NewCNode(tuple_get_item_inputs);
MS_EXCEPTION_IF_NULL(tuple_get_item);
make_tuple_inputs.push_back(tuple_get_item);
}
auto make_tuple = func_graph->NewCNode(make_tuple_inputs);
MS_EXCEPTION_IF_NULL(make_tuple);
FuncGraphManagerPtr manager = func_graph->manager();
MS_EXCEPTION_IF_NULL(manager);
(void)manager->Replace(prev, make_tuple);
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!";
}
// 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: Position (pos) cannot be larger than the node's inputs' size.";
}
// 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);
// 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<bool>(prim_out_attr[RECOMPUTE_COMM_OP])) ||
(prim_in_attr.find(RECOMPUTE_COMM_OP) != prim_in_attr.end() &&
!GetValue<bool>(prim_in_attr[RECOMPUTE_COMM_OP]))) &&
COMMUNICATION_OPS.find(op_name) != COMMUNICATION_OPS.end()) {
MS_LOG(INFO) << "The redistribution node would not be recomputed.";
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);
(void)InsertMakeTuple(target_node, (redistribution_oplist_ptr->second)[index].second, func_graph);
}
}
}
// 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;
}
// 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: Position (pos) cannot be larger than the node's inputs' size, the instance name is "
<< instance_name;
}
// 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<ValueNodePtr>();
MS_EXCEPTION_IF_NULL(value_node);
size_t index_s = LongToSize(GetValue<int64_t>(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 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<OperatorInfo>();
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<AnfNodePtr, int64_t> &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";
}
// Get the next operator node and validate it
CNodePtr next_node = node_pair.first->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(next_node);
// Extract primitive and operator information from the middle node
auto middle_value = middle_node->input(0)->cast<ValueNodePtr>();
MS_EXCEPTION_IF_NULL(middle_value);
PrimitivePtr middle_prim = middle_value->value()->cast<PrimitivePtr>();
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<PrimitivePtr>(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 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->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<TensorLayout>());
}
// 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";
}
// 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(DEBUG) << "Redistribution size " << redistribution_oplist_ptr->first.size();
// Insert redistribution nodes before the next node if necessary
if (!redistribution_oplist_ptr->first.empty()) {
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<std::string, ValuePtr> &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<std::string, ValuePtr> &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);
std::vector<AnfNodePtr> all_nodes = DeepScopedGraphSearch(ret);
for (auto &node : all_nodes) {
auto cnode = node->cast<CNodePtr>();
if ((cnode == nullptr) || !IsValueNode<Primitive>(cnode->input(0))) {
continue;
}
ValueNodePtr prim_anf_node = cnode->input(0)->cast<ValueNodePtr>();
PrimitivePtr prim = GetValueNode<PrimitivePtr>(prim_anf_node);
auto attrs = prim->attrs();
if (StrategyFound(attrs)) {
return true;
}
}
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<AnfNodePtr> &all_nodes) {
for (auto &node : all_nodes) {
MS_EXCEPTION_IF_NULL(node);
if (!node->isa<CNode>()) {
continue;
}
auto cnode = node->cast<CNodePtr>();
if (!IsValueNode<Primitive>(cnode->input(0))) {
continue;
}
ValueNodePtr prim_value_node = cnode->input(0)->cast<ValueNodePtr>();
MS_EXCEPTION_IF_NULL(prim_value_node);
PrimitivePtr prim = GetValueNode<PrimitivePtr>(prim_value_node);
MS_EXCEPTION_IF_NULL(prim);
if (IsCommunicationOp(prim) && cnode->in_forward_flag()) {
MS_EXCEPTION_IF_NULL(prim_value_node->scope());
MS_LOG(INFO) << "The graph contain communication op: " << prim->name() << ", scope name is "
<< prim_value_node->scope()->name();
return true;
}
}
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());
FuncGraphManagerPtr manager = node->func_graph()->manager();
MS_EXCEPTION_IF_NULL(manager);
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<Primitive>(node->input(0))) {
auto current_value = node->input(0)->cast<ValueNodePtr>();
MS_EXCEPTION_IF_NULL(current_value);
PrimitivePtr current_prim = current_value->value()->cast<PrimitivePtr>();
MS_EXCEPTION_IF_NULL(current_prim);
insert_node_new = ((current_prim->name() == prim::kTupleGetItem) ? node : insert_node);
} 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<CNodePtr>();
MS_EXCEPTION_IF_NULL(use_cnode);
// Recursively analyze non-primitive nodes
if (!IsValueNode<Primitive>(use_cnode->input(0))) {
StepRedistribution(use_cnode, distribute_operator, insert_node_new, tensor_redistribution, pre_node);
} else {
ValueNodePtr prim_anf_node = use_cnode->input(0)->cast<ValueNodePtr>();
MS_EXCEPTION_IF_NULL(prim_anf_node);
PrimitivePtr node_prim = prim_anf_node->value()->cast<PrimitivePtr>();
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<OperatorInfo>()) {
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<OperatorInfo>();
MS_EXCEPTION_IF_NULL(op_info);
// 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()
<< ": GetNodeShape for tensor_node, output size is not 1";
}
Shape shape = shapes[0];
std::string shape_str = ShapeToString(shape);
if (shape.empty() || ((shape.size() == 1) && (shape[0] == 1))) {
MS_LOG(INFO) << "Split tensor for " << op_info->name() << ": The shape is " << shape_str
<< ", no need to split it.";
return;
}
MS_LOG(INFO) << "Split tensor for " << op_info->name() << ": The shape of tensor is " << shape_str;
// 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();
}
TensorInfo tensor_info = op_info->inputs_tensor_info()[LongToSize(index - 1)];
TensorLayout tensor_layout = tensor_info.tensor_layout();
// 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++) {
if (!sub_ops.at(i).empty()) {
InsertGetTensorSliceOp(sub_ops.at(i).at(0), next_node, func_graph, index, SUB);
}
}
}
}
// 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) {
MS_LOG(INFO) << next_node->fullname_with_scope() << " Inputs must have only one input, get "
<< (next_node->inputs().size() - 1) << " index should be 1, get " << index;
return;
}
OperatorInfoPtr op_info = next_node->user_data<OperatorInfo>();
MS_EXCEPTION_IF_NULL(op_info);
// Get the values from the input tensor list.
std::vector<ValuePtr> inputs_values;
if (IsValueNode<ValueList>(node)) {
inputs_values = node->cast<ValueNodePtr>()->value()->cast<ValueListPtr>()->value();
} else {
inputs_values = node->cast<ValueNodePtr>()->value()->cast<ValueTuplePtr>()->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<AnfNodePtr> make_tuple_inputs = {NewValueNode(prim::kPrimMakeTuple)};
FuncGraphPtr func_graph = next_node->func_graph();
MS_EXCEPTION_IF_NULL(func_graph);
FuncGraphManagerPtr manager = func_graph->manager();
MS_EXCEPTION_IF_NULL(manager);
ScopePtr scope = next_node->scope();
MS_EXCEPTION_IF_NULL(scope);
for (size_t i = 0; i < inputs_values.size(); ++i) {
auto value_ptr = inputs_values[i];
auto tensor = value_ptr->cast<tensor::TensorPtr>();
MS_EXCEPTION_IF_NULL(tensor);
TensorInfo tensor_info = op_info->inputs_tensor_info()[i];
TensorLayout tensor_layout = tensor_info.tensor_layout();
auto value_node = NewValueNode(value_ptr)->cast<AnfNodePtr>();
Operator op = CreateGetTensorSliceOp(tensor_layout);
std::vector<AnfNodePtr> node_input = CreateInput(op, value_node, SPLIT_TENSOR);
CNodePtr new_node = func_graph->NewCNode(node_input);
new_node->set_in_forward_flag(true);
auto new_node_value = node_input[0]->cast<ValueNodePtr>();
MS_EXCEPTION_IF_NULL(new_node_value);
PrimitivePtr new_node_prim = new_node_value->value()->cast<PrimitivePtr>();
new_node_prim->set_instance_name(SPLIT_TENSOR);
new_node_prim->set_attr("keep_value_node_input", MakeValue(true));
new_node->set_scope(scope);
node_input[0]->set_scope(scope);
make_tuple_inputs.push_back(new_node);
}
CNodePtr make_tuple = func_graph->NewCNode(make_tuple_inputs);
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<CNodePtr>();
// Check if the node is a ValueNode with a Primitive input (e.g., a primitive operator).
if (use_cnode == nullptr || !IsValueNode<Primitive>(use_cnode->input(0))) {
continue;
}
ValueNodePtr prim_anf_node = use_cnode->input(0)->cast<ValueNodePtr>();
MS_EXCEPTION_IF_NULL(prim_anf_node);
PrimitivePtr use_cnode_prim = prim_anf_node->value()->cast<PrimitivePtr>();
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<ValueList>(node) || IsValueNode<ValueTuple>(node)) {
SplitTensorList(node, use_cnode, node_pair.second);
} else {
SplitTensor(node, use_cnode, node_pair.second);
}
}
}
}
void StepReplaceOp(OperatorVector replace_op, const CNodePtr &node) {
// Step 1: Get the distribute_operator associated with the CNode
OperatorInfoPtr distribute_operator = node->user_data<OperatorInfo>();
if (distribute_operator == 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";
}
// 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) {
auto cast_int = CreateCastOp(kInt32);
auto cast_bool = CreateCastOp(kBool);
(void)replace_op.insert(replace_op.begin(), cast_int);
(void)replace_op.insert(replace_op.end(), cast_bool);
(void)replace_op_info.insert(replace_op_info.begin(), {false, 1});
(void)replace_op_info.insert(replace_op_info.end(), {false, 1});
}
// 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()) {
MS_LOG(EXCEPTION) << "replace_op_info is not empty and size not equal to replace_op!";
}
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<AnfNodePtr> 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();
MS_EXCEPTION_IF_NULL(scope);
replace_node->set_scope(scope);
PrimitivePtr prim = GetValueNode<PrimitivePtr>(replace_node->input(0));
PrimitivePtr origin_prim = GetValueNode<PrimitivePtr>(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<bool>(origin_prim_attrs[RECOMPUTE_COMM_OP]) &&
COMMUNICATION_OPS.find(prim->name()) != COMMUNICATION_OPS.end()) {
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<OperatorInfo>(node->user_data<OperatorInfo>());
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 the node
} else {
(void)manager->Replace(node, replace_node); // Using Replace function to insert the node
}
}
MS_LOG(INFO) << "Insert ReplaceOp success for " << distribute_operator->name();
}
void StepReplaceGraph(const ReplaceGraphPtr &replace_graph, const CNodePtr &node) {
MS_EXCEPTION_IF_NULL(replace_graph);
MS_EXCEPTION_IF_NULL(node);
MS_EXCEPTION_IF_NULL(replace_graph->second);
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";
}
// Solve the input order
// For example input_node:{segment_sum:1, segment_sum:2, gahter:2}
// The Original code here will bind the all operations to the first inputs of these operatos
// 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<AnfNodePtr, int> input_map = {};
static int appear_count = 0;
for (auto &replace_input : replace_graph->first) {
auto pre_node = node->input(LongToSize(replace_input.second));
auto it = input_map.find(replace_input.first);
if (it != input_map.end()) {
appear_count = 1 + it->second;
} else {
appear_count = 1;
}
auto replace_input_cnode = replace_input.first->cast<CNodePtr>();
size_t inputs_size = replace_input_cnode->inputs().size();
while (IntToSize(appear_count) < inputs_size && replace_input_cnode->input(appear_count)->func_graph() != nullptr) {
++appear_count;
}
if (IntToSize(appear_count) >= inputs_size) {
MS_LOG(EXCEPTION) << "No replaceable virtual_input_node";
}
input_map[replace_input.first] = appear_count;
manager->SetEdge(replace_input.first, appear_count, pre_node);
}
// Step 2: Replace the original node with the replace_graph's second node
auto replace_output = replace_graph->second->cast<CNodePtr>();
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<ValueNode>()) {
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<Int64Imm>()) {
MS_LOG(EXCEPTION) << "The index of tuple getitem is not int32";
}
return tuple_index_value->cast<Int64ImmPtr>()->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();
FuncGraphPtr func_graph = node->func_graph();
MS_EXCEPTION_IF_NULL(func_graph);
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);
// Check if the input is not a tensor or has an abstract monad.
if ((!input->isa<CNode>() && !input->isa<Parameter>()) || HasAbstractMonad(input)) {
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);
}
MS_LOG(INFO) << "insert div op for input index " << index << " of 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) {
MS_LOG(EXCEPTION) << "Find the scale value is 0, you should check the mirror operators's group size.";
}
size_t node_size = node->inputs().size();
FuncGraphPtr func_graph = node->func_graph();
MS_EXCEPTION_IF_NULL(func_graph);
// 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 ((!input->isa<CNode>() && !input->isa<Parameter>()) || HasAbstractMonad(input)) {
continue;
}
InsertNode(div_op, node, index, node->input(index), func_graph, instance_name);
}
}
// 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);
// Instantiate the AllReduce operator
CheckGlobalDeviceManager();
Operator allreduce_op = CreateAllReduceOp(REDUCE_OP_SUM, group);
// 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 ((!input->isa<CNode>() && !input->isa<Parameter>()) || HasAbstractMonad(input)) {
continue;
}
InsertNode(allreduce_op, node, index, node->input(index), func_graph, instance_name);
}
}
// 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<AnfNodePtr> &all_nodes) {
FuncGraphPtr real_graph = root;
for (auto &node : all_nodes) {
if (!node->isa<CNode>()) {
continue;
}
auto cnode = node->cast<CNodePtr>();
if (!IsValueNode<Primitive>(cnode->input(0))) {
continue;
}
auto expect_shard_prim = GetValueNode<PrimitivePtr>(cnode->input(0));
if (expect_shard_prim->name() != SHARD) {
continue;
}
real_graph = GetValueNode<FuncGraphPtr>(cnode->input(1));
}
return real_graph;
}
/**
* @brief Inserts virtual output nodes into a computation graph for parallel execution.
*
* This function identifies the last forward nodes in the computation graph and inserts virtual output nodes
* accordingly. It ensures that these nodes are available as outputs during the parallel execution.
*
* @param root The root of the computation graph.
* @param all_nodes A vector containing all the nodes in the computation graph.
*/
void InsertVirtualOutput(const FuncGraphPtr &root, const std::vector<AnfNodePtr> &all_nodes) {
std::vector<std::string> last_forward_node_ids; // Stores unique IDs of last forward nodes
std::vector<size_t> 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";
// Iterate through all nodes in the computation graph
for (auto &node : all_nodes) {
auto cnode = node->cast<CNodePtr>();
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) << "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<CNode>()) {
MS_LOG(EXCEPTION) << "The output of tuple_get_item is not a cnode";
}
cnode = node_pair.first->cast<CNodePtr>();
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();
std::shared_ptr<abstract::BaseShape> virtual_output_shape = std::make_shared<abstract::Shape>(shape_outputs[0]);
virtual_output_abstract->set_shape(virtual_output_shape);
virtual_output_node->set_abstract(virtual_output_abstract);
}
}
}
// only used for FindCNode
CNodePtr SkipTrivialNodesMoveDown(const FuncGraphManagerPtr &manager, CNodePtr node) {
MS_EXCEPTION_IF_NULL(node);
while (IsInTrivialNodeList(node) || IsSomePrimitive(node, LOAD)) {
node = manager->node_users()[node].begin()->first->cast<CNodePtr>();
}
return node;
}
/**
* @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<bool, CNodePtr> 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 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<CNodePtr>();
// Check if the node is a CNode and if its first input is a Primitive
if (use_apply == nullptr || !IsValueNode<Primitive>(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<Primitive>(use_apply->input(0))) {
continue;
}
ValueNodePtr prim_anf_node = use_apply->input(0)->cast<ValueNodePtr>();
MS_EXCEPTION_IF_NULL(prim_anf_node);
PrimitivePtr node_prim = prim_anf_node->value()->cast<PrimitivePtr>();
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) << "Found Primitive " << name << " in the same func_graph";
continue;
}
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) {
// 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<CNodePtr>();
if (cnode == nullptr || !IsValueNode<Primitive>(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<mindspore::TensorType>()) {
MS_LOG(EXCEPTION) << "Unknown type.";
}
auto input_element_type = node_type->cast<mindspore::TensorTypePtr>()->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;
}
constexpr size_t kSingleArgCNodeSize = 2;
if ((node->inputs().size() == kSingleArgCNodeSize) && (IsValueNode<ValueSequence>(node->input(1)))) {
MS_LOG(INFO) << "Input is ValueList, skip it.";
return false;
}
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 been handled by make_tuple node";
return false;
}
if (mirror_ops.size() != node_size - 1) {
MS_LOG(EXCEPTION) << "Mirrorops' size is incorrect! mirror_ops size is " << mirror_ops.size() << ", node_size is "
<< (node_size - 1);
}
return true;
}
/**
* @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)) {
if (IsInTrivialNodeList(node) || IsInAllGatherNodeList(node)) {
node = node->input(1)->cast<CNodePtr>();
}
}
auto prev_node = node->input(1)->cast<CNodePtr>();
if (prev_node != nullptr) {
if (IsSomePrimitive(prev_node, DEPEND)) {
auto prev_prev_node = prev_node->input(1)->cast<CNodePtr>();
if (IsSomePrimitive(node, LOAD)) {
node = prev_prev_node;
MS_LOG(INFO) << "Moving to the Load node before Depend 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();
std::string mirror_op_name;
if (grad_accumulation_step > 1) {
mirror_op_name = MIRROR_MINI_STEP_OPERATOR;
} else if (split_stage_num > 1) {
mirror_op_name = MIRROR_MICRO_STEP_OPERATOR;
} else {
mirror_op_name = MIRROR_OPERATOR;
}
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<int>(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<AnfNodePtr, bool> 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<ParameterPtr>();
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 << " 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<TensorLayout>()) {
opt_shard_mirror_group = param_ptr->user_data<TensorLayout>()->opt_shard_mirror_group();
is_shared_param = param_ptr->user_data<TensorLayout>()->is_shared_param();
}
// If the parameter has shard mirror group information, create mirror operations based on group size
if (!opt_shard_mirror_group.empty()) {
uint32_t group_rank_size = 0;
if (!CommManager::GetInstance().GetRankSize(opt_shard_mirror_group, &group_rank_size)) {
MS_LOG(EXCEPTION) << "Failed to get group size from group " << opt_shard_mirror_group;
}
backward_op = CreateMirrorOps(opt_shard_mirror_group, static_cast<size_t>(group_rank_size));
}
}
// 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 (next_cnode.first) {
MS_EXCEPTION_IF_NULL(next_cnode.second);
// Assuming Load is inserted next to the parameter, skip Load moving up and insert mirror next to the parameter
if (pre_node->cast<CNodePtr>()) {
CNodePtr load_node = SkipTrivialNodesMoveUp(node->input(index)->cast<CNodePtr>());
manager->SetEdge(load_node, 1, next_cnode.second);
} else {
manager->SetEdge(node, static_cast<int>(index), next_cnode.second);
}
MS_LOG(INFO) << "Found parameter " << param_name << " for node " << GetPrimName(node->cast<CNodePtr>())
<< " and sharing the mirror.";
continue;
}
}
// 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, 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<CNodePtr>() && (InsertMirrorBeforeCast(node, index) || is_shared_param)) {
// 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<CNodePtr>());
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<CNodePtr>();
// Add fusion flag
AddCommOpFusionType(comm_op, param_node_pair.first);
MS_LOG(INFO) << "Found parameter " << param_name << " for node " << GetPrimName(node->cast<CNodePtr>())
<< " 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) << "Found parameter " << param_name << " for node " << GetPrimName(node->cast<CNodePtr>())
<< " and inserted mirror before the node.";
auto comm_op = node->input(index)->cast<CNodePtr>();
// 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);
}
// 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<std::pair<CNodePtr, LossNodeInfo>> &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<CNodePtr, LossNodeInfo> &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 operations if they exist
if (!mirror_ops.empty()) {
MS_LOG(INFO) << "Inserting mirror ops for " << distribute_operator->name();
InsertMirrorOps(root, mirror_ops, node);
}
// 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) << "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<Shapes> &shape_list) {
if (shape_list.size() != 2) {
MS_LOG(ERROR) << "The size of shape list is not 2";
return nullptr;
}
if (name.length() == 0) {
MS_LOG(EXCEPTION) << "Length of name is zero!";
}
std::string distribute_opname = GetDisOpName(name);
OperatorInfoPtr operator_ =
(OperatorInfoPtr)DynCreator::Instance().Create(distribute_opname, shape_list[0], shape_list[1], attrs, TOTAL_OPS);
if (operator_ == nullptr) {
MS_LOG(INFO) << "Create " << name << " failed";
return nullptr;
}
std::string origin_name = operator_->name();
operator_->set_name(origin_name + std::to_string(TOTAL_OPS));
MS_LOG(INFO) << "Successfully created operator " << origin_name;
++TOTAL_OPS;
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<Shapes> &shape_list) {
MS_EXCEPTION_IF_NULL(prim);
OperatorInfoPtr operator_ = OperatorInstanceByName(prim->name(), attrs, shape_list);
if (operator_ == nullptr) {
if (IsInBatchParallelBlackList(prim)) {
MS_LOG(EXCEPTION) << "Operator " << prim->name() << " is not supported yet in auto parallel mode.";
}
MS_LOG(INFO) << "Create " << prim->name() << " failed, use batch parallel";
operator_ = OperatorInstanceByName(BATCH_PARALLEL, attrs, shape_list);
MS_EXCEPTION_IF_NULL(operator_);
}
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<Shapes> shape_list) {
OperatorInfoPtr operator_ = OperatorInstance(prim, attrs, shape_list);
for (size_t i = 0; i < shape_list[0].size(); ++i) {
MS_LOG(INFO) << "No: " << i << " input's shape: " << ShapeToString(shape_list[0][i]);
}
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;
}
auto var = stra->cast<ValueTuplePtr>();
if (var == nullptr) {
return nullptr;
}
StrategyPtr strategyPtr;
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) {
std::vector<ValuePtr> elements = var->value();
Strategys strategy;
for (uint64_t index = 0; index < elements.size(); ++index) {
Dimensions dim;
if (elements[index]->isa<ValueSequence>()) {
auto value_tuple = elements[index]->cast<ValueTuplePtr>();
std::vector<ValuePtr> 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<int64_t>(GetValue<int64_t>(value)); });
strategy.push_back(dim); // Append the extracted dimension to the strategy
} else {
MS_LOG(EXCEPTION) << "Failure: Strategy's format is wrong! Need ValueSequence";
}
}
if (strategy.empty()) {
MS_LOG(EXCEPTION) << "ExtractStrategy: failed to extract strategy";
}
strategyPtr = NewStrategy(stage_id, strategy); // Create a new strategy
}
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<AnfNodePtr> 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]); // 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 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<Shapes> ExtractShape(const CNodePtr &node) {
MS_EXCEPTION_IF_NULL(node);
Shapes shape_inputs, shape_outputs;
std::vector<Shapes> shape_all;
std::vector<AnfNodePtr> 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) { // 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<RefKey>(input)) { // Check if the input is a RefKey node
auto func_graph = node->func_graph();
MS_EXCEPTION_IF_NULL(func_graph);
std::vector<AnfNodePtr> parameters = FindParameterByRefKeyNode(input, func_graph); // Find associated parameters
if (parameters.size() != 1) {
MS_LOG(EXCEPTION) << "Find parameter by ref key node failed";
}
std::pair<AnfNodePtr, int64_t> node_pair = std::make_pair(node, SizeToLong(i));
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<CNode>() || IsValueNode<Tensor>(input) || input->isa<Parameter>() ||
((IsValueNode<ValueList>(input) || IsValueNode<ValueTuple>(input)) && (inputs_size == 2))) {
input_shapes = GetNodeShape(input); // Get the shape of the input node
} else {
continue;
}
if (input_shapes.size() != 1) {
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]); // Append the input shape to the input shapes vector
}
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<AnfNodePtr, int64_t> 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<CNodePtr>();
MS_EXCEPTION_IF_NULL(cnode);
// Check if the input is a ValueNode containing a Primitive.
if (!IsValueNode<Primitive>(cnode->input(0))) {
continue;
}
// Get the Primitive associated with the CNode.
ValueNodePtr prim_node_anf = cnode->input(0)->cast<ValueNodePtr>();
MS_EXCEPTION_IF_NULL(prim_node_anf);
PrimitivePtr node_prim = prim_node_anf->value()->cast<PrimitivePtr>();
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<OperatorInfo>()) {
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;
}
}
}
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<AnfNodePtr, int64_t> FindSubGraph(const FuncGraphPtr &graph, const AnfNodePtr &parameter) {
// 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<AnfNodePtr, int64_t> 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 &param_pair : param_sub_set) {
CNodePtr param_cnode = param_pair.first->cast<CNodePtr>();
AnfNodePtr graph_value_node;
// Check if 'parameter' is connected to a ValueNode containing a FuncGraph.
if (param_cnode->input(0)->isa<CNode>()) {
graph_value_node = param_cnode->input(0)->cast<CNodePtr>()->input(1);
} else {
graph_value_node = param_cnode->input(0);
}
// Check if 'graph_value_node' is a ValueNode containing a FuncGraph.
if (!IsValueNode<FuncGraph>(graph_value_node)) {
continue;
}
// Get the FuncGraph associated with 'graph_value_node'.
FuncGraphPtr graph_sub = GetValueNode<FuncGraphPtr>(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();
}
// Recursively call the function to search in the subgraph.
std::pair<AnfNodePtr, int64_t> res = FindSubGraph(graph_sub, parameters[LongToSize(param_pair.second - 1)]);
if (res.first != nullptr) {
return res;
}
}
}
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);
// 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<CNodePtr>();
}
// 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<mindspore::TensorType>()) {
MS_LOG(EXCEPTION) << "Unknown type.";
}
// Get the type ID of the input element.
auto input_element_type = node_type->cast<mindspore::TensorTypePtr>()->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 {
return nullptr;
}
}
// 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<AnfNodePtr, int> &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<CNodePtr>();
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<PrimitivePtr>(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<ParameterPtr>()->name();
if (op_name == MINI_STEP_ALL_GATHER) {
op = CreateMiniStepAllGatherOp(group);
} else if (op_name == MICRO_STEP_ALL_GATHER) {
op = CreateMicroStepAllGatherOp(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<ParameterPtr>();
MS_EXCEPTION_IF_NULL(param_ptr);
if (param_ptr->user_data<TensorLayout>()) {
opt_shard_mirror_group = param_ptr->user_data<TensorLayout>()->opt_shard_mirror_group();
}
if (!is_shared_param && cast_node) {
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();
TypePtr next_node_dtype = FindChildCastWithFP32ToFP16(cnode, node_user_map);
if (next_node_dtype) {
MS_LOG(INFO) << "Inserting Cast from float32 to float16 for node " << node->fullname_with_scope() << " for saving"
<< " communication.";
pre_node_ = CreateFP16Cast(cnode, pre_node, next_node_dtype);
}
InsertNode(op, cnode, IntToSize(res.second), pre_node_, graph, PARALLEL_OPTIMIZER_ALLGATHER_NOT_COMPUTE, param_name,
root);
allgather = cnode->input(IntToSize(res.second))->cast<CNodePtr>();
MS_LOG(INFO) << "Parallel optimizer is applied before " << GetPrimName(cnode) << " for " << param_name;
}
// Add fusion flag to the AllGather operation
AddCommOpFusionType(allgather, node);
// 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
AddCommOpMirrorFlag(allgather, !grad_accumulation_shard);
} else if (op_name == MINI_STEP_ALL_GATHER) {
// We need to manually set the add_accu to be false if it's father node is MirrorMiniStep
bool add_accu = root->has_flag(kAccumulation);
bool is_with_mirror = opt_shard_mirror_group.size() > 1;
AddCommOpAddAccuFlag(allgather, !add_accu && !is_with_mirror);
AddCommOpMirrorFlag(allgather, grad_accumulation_shard || !add_accu);
}
}
// 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 &parameter,
const std::string &opt_shard_group) {
// Check if the shard group is empty (no optimization needed)
if (opt_shard_group.empty()) {
return;
}
// 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) {
op_name = MICRO_STEP_ALL_GATHER;
} else {
op_name = 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 &param_pair : param_sub_set) {
auto cnode = param_pair.first->cast<CNodePtr>();
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<OperatorInfo>();
// 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()) {
MS_LOG(EXCEPTION) << "The index is out of range, index is " << (param_pair.second - 1) << ", vector size is "
<< distribute_operator->inputs_tensor_info().size();
}
if (insert_flag) {
// 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);
MS_LOG(INFO) << "Parallel optimizer is shared between " << parameter->ToString() << " and "
<< GetPrimName(cnode);
} else {
MS_LOG(ERROR) << "Can not find the shared AllGather with multiple node users.";
}
} else {
// Insert AllGather operation for the parameter
auto param_ptr = parameter->cast<ParameterPtr>();
MS_EXCEPTION_IF_NULL(param_ptr);
bool is_shared_param = param_ptr->user_data<TensorLayout>()->is_shared_param();
InsertAllGatherOp(root, opt_shard_group, param_pair, parameter, op_name, is_shared_param);
insert_flag = true;
}
}
}
}
// 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 &parameter) {
MS_EXCEPTION_IF_NULL(root);
MS_EXCEPTION_IF_NULL(parameter);
FuncGraphManagerPtr manager = root->manager();
MS_EXCEPTION_IF_NULL(manager);
ParameterPtr parameter_ptr = parameter->cast<ParameterPtr>();
// 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.";
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 &param_pair : user_set) {
CNodePtr cnode = param_pair.first->cast<CNodePtr>();
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<TensorLayout>();
tensor_layout->set_is_shared_param(true);
MS_LOG(WARNING) << "There are multiple users for " << parameter->ToString()
<< ". Mixed precision optimization is not valid here.";
}
}
// 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 &parameter, const std::pair<AnfNodePtr, int64_t> &res,
const FuncGraphPtr &root) {
// 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<CNodePtr>();
MS_EXCEPTION_IF_NULL(cnode);
// Get the slice shape from the distributed operator.
OperatorInfoPtr distribute_operator = cnode->user_data<OperatorInfo>();
if (distribute_operator == 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)
<< ", 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 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<OptParamMgr> apOptParamMgr = createOptParamMgr(root);
opt_shard_group = apOptParamMgr->ShardOptGroup(parameter, &tensor_layout, distribute_operator);
// 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() << ")"
<< " 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.";
}
AbstractBasePtr cloned_abstract = abstract->Clone();
if (cloned_abstract == nullptr) {
MS_LOG(EXCEPTION) << "Parameter " << parameter->ToString() << ": abstract clone failed.";
}
cloned_abstract->set_shape(std::make_shared<abstract::Shape>(slice_shape));
parameter->set_abstract(cloned_abstract);
ParameterPtr parameter_ptr = parameter->cast<ParameterPtr>();
MS_EXCEPTION_IF_NULL(parameter_ptr);
parameter_ptr->set_user_data<TensorLayout>(std::make_shared<TensorLayout>(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();
for (auto &parameter : parameters) {
MS_EXCEPTION_IF_NULL(parameter->Shape());
auto iter = g_RefMap.find(parameter);
if (iter != g_RefMap.end()) {
std::string group = SetParallelShape(parameter, g_RefMap[parameter], root);
SetSharedParameterFlag(root, parameter);
ApplyParallelOptOnParam(root, parameter, group);
continue;
}
std::pair<AnfNodePtr, int64_t> res = FindSubGraph(root, parameter);
if (res.first == nullptr) {
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);
SetSharedParameterFlag(root, parameter);
ApplyParallelOptOnParam(root, parameter, group);
MS_LOG(DEBUG) << "Parameter " << parameter->ToString() << " shape " << parameter->Shape()->ToString();
}
}
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());
bool full_batch = ParallelContext::GetInstance()->full_batch();
PrimitivePtr prim = GetValueNode<PrimitivePtr>(node->input(0));
MS_EXCEPTION_IF_NULL(prim);
if (prim->name() == VIRTUAL_DATA_SET || prim->name() == VIRTUAL_OUTPUT) {
CheckGlobalDeviceManager();
auto attrs_temp = prim->attrs();
if (!ParallelContext::GetInstance()->dataset_strategy().empty() && prim->name() == VIRTUAL_DATA_SET) {
std::vector<ValuePtr> elements;
auto dataset_strategy = ParallelContext::GetInstance()->dataset_strategy();
(void)std::transform(dataset_strategy.begin(), dataset_strategy.end(), std::back_inserter(elements),
[](auto input_stra) { return MakeValue(input_stra); });
ValueTuplePtr strategy = std::make_shared<ValueTuple>(elements);
attrs_temp[IN_STRATEGY] = strategy;
(void)prim->SetAttrs(attrs_temp);
if (prim->HasAttr(REPEAT_DIM_DIRECT) && GetValue<std::string>(prim->GetAttr(REPEAT_DIM_DIRECT)) == RIGHT) {
ParallelContext::GetInstance()->set_dataset_repeat_dim_right(true);
MS_LOG(INFO) << "dataset repeat dim is right";
}
return;
}
int64_t dev_num;
if (full_batch) {
dev_num = 1;
} else {
dev_num = g_device_manager->stage_device_num();
}
if (dev_num == 0) {
MS_LOG(EXCEPTION) << "Device Num must be larger than 0, but got 0.";
}
std::vector<Shapes> shape_list = ExtractShape(node);
if (shape_list.empty()) {
MS_LOG(EXCEPTION) << "Failure: node " << node->ToString() << " failed to extract shape";
}
std::vector<ValuePtr> elements;
for (size_t i = 0; i < shape_list[0].size(); i++) {
if (shape_list[0][i].empty()) {
(void)elements.emplace_back(MakeValue(Dimensions()));
continue;
}
Dimensions input_strategy;
if (shape_list[0][i][0] % dev_num == 0) {
input_strategy.push_back(dev_num);
} else {
input_strategy.push_back(1);
}
for (size_t j = 1; j < shape_list[0][i].size(); j++) {
input_strategy.push_back(1);
}
(void)elements.emplace_back(MakeValue(input_strategy));
}
ValueTuplePtr strategy = std::make_shared<ValueTuple>(elements);
attrs_temp[IN_STRATEGY] = strategy;
(void)prim->SetAttrs(attrs_temp);
}
}
// 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<std::string> *unique_ids, std::vector<size_t> *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;
}
MS_EXCEPTION_IF_NULL(unique_ids);
MS_EXCEPTION_IF_NULL(indexes);
if (!node->isa<CNode>()) {
return false;
}
CNodePtr pre_cnode = node->cast<CNodePtr>();
if (!IsValueNode<Primitive>(pre_cnode->input(0))) {
return false;
}
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<CNode>() || next_node->isa<Parameter>()) {
return false;
}
CNodePtr cnode = next_node->cast<CNodePtr>();
if (!IsValueNode<Primitive>(cnode->input(0))) {
return false;
}
ValueNodePtr prim_anf_node = cnode->input(0)->cast<ValueNodePtr>();
PrimitivePtr prim = prim_anf_node->value()->cast<PrimitivePtr>();
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<std::string> *unique_ids,
std::vector<size_t> *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<Strategys> strategy_v_ptr = operator_->GenerateBatchStrategies();
MS_EXCEPTION_IF_NULL(strategy_v_ptr);
strategyPtr = NewStrategy(0, *strategy_v_ptr);
std::vector<ValuePtr> 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<ValueTuple>(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;
}
// 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<std::string> *unique_ids, std::vector<size_t> *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;
}
MS_EXCEPTION_IF_NULL(unique_ids);
MS_EXCEPTION_IF_NULL(indexes);
if (!node->isa<CNode>()) {
return false;
}
CNodePtr pre_cnode = node->cast<CNodePtr>();
if (!IsValueNode<Primitive>(pre_cnode->input(0))) {
return false;
}
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<CNode>() || next_node->isa<Parameter>()) {
return false;
}
CNodePtr cnode = next_node->cast<CNodePtr>();
if (!IsValueNode<Primitive>(cnode->input(0))) {
return false;
}
ValueNodePtr prim_anf_node = cnode->input(0)->cast<ValueNodePtr>();
PrimitivePtr prim = prim_anf_node->value()->cast<PrimitivePtr>();
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<std::string> *unique_ids,
std::vector<size_t> *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<Strategys> strategy_v_ptr = operator_->GenerateBatchStrategies();
MS_EXCEPTION_IF_NULL(strategy_v_ptr);
strategyPtr = NewStrategy(0, *strategy_v_ptr);
std::vector<ValuePtr> 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<ValueTuple>(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<AnfNodePtr> &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<CNodePtr>();
// 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<ValueNodePtr>();
PrimitivePtr prim = GetValueNode<PrimitivePtr>(prim_anf_node);
auto attrs = prim->attrs();
MS_LOG(INFO) << "Extracting information: node: " << node->ToString() << " prim " << prim->name();
// Step 6: Extract shapes information for the node
std::vector<Shapes> shape_list = ExtractShape(cnode);
if (shape_list.empty()) {
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<ValuePtr> input_value;
// Step 8: Extract input values for the node
for (size_t index = 1; index < inputs.size(); ++index) {
if (inputs[index]->isa<ValueNode>()) {
input_value.push_back(GetValueNode(inputs[index]));
continue;
}
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<OperatorInfo>(operator_);
continue;
}
// Step 11: Extract strategy and initialize the OperatorInfo
ExtractStrategyAndInit(cnode, prim, operator_);
cnode->set_user_data<OperatorInfo>(operator_);
}
}
// GetInputLayoutFromCNode function
// This function retrieves the input tensor layout from a CNode based on a node-index pair.
TensorLayout GetInputLayoutFromCNode(const std::pair<AnfNodePtr, int64_t> &node_pair) {
CNodePtr cnode = node_pair.first->cast<CNodePtr>();
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();
}
// 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;
}
// 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<TensorLayout> 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<CNodePtr>();
// 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<ValueNodePtr>();
MS_EXCEPTION_IF_NULL(prim_anf_node);
PrimitivePtr node_prim = prim_anf_node->value()->cast<PrimitivePtr>();
MS_EXCEPTION_IF_NULL(node_prim);
// Step 3: Check if the node is a parallel-care node and has associated OperatorInfo
if (IsParallelCareNode(use_apply) && use_apply->has_user_data<OperatorInfo>()) {
MS_LOG(INFO) << "FindNextLayout success prim " << node_prim->name();
*next_is_reshape = false;
auto layout = GetInputLayoutFromCNode(node_pair);
return std::make_shared<TensorLayout>(layout);
}
MS_LOG(DEBUG) << "FindNextLayout failed prim " << node_prim->name() << " " << IsParallelCareNode(use_apply)
<< " " << use_apply->has_user_data<OperatorInfo>();
// Step 4: Recursively search for the next layout
auto layout_ptr = FindNextLayout(use_apply, next_is_reshape);
if (layout_ptr) {
return layout_ptr;
}
}
MS_LOG(WARNING) << "FindNextLayout return nullptr, if reshape is not the last primitive, there must be some error";
return nullptr;
}
// GetOutputLayoutFromCNode function
// This function retrieves the output tensor layout from a CNode based on the output index.
std::shared_ptr<TensorLayout> 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;
}
// 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>(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<TensorLayout> FindPrevParallelCareNodeLayout(const AnfNodePtr &node, size_t output_index) {
if (!node->isa<CNode>()) {
return nullptr;
}
CNodePtr cnode = node->cast<CNodePtr>();
if (!IsValueNode<Primitive>(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<OperatorInfo>()) {
auto layout_ptr = GetOutputLayoutFromCNode(cnode, output_index);
if (!layout_ptr) {
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<TensorLayout> 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: "
<< MAX_RECURSIVE_DEPTH;
return nullptr;
}
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);
if (!layout_param) {
continue;
}
return layout_param;
}
CNodePtr use_apply = node_pair.first->cast<CNodePtr>();
if (use_apply == nullptr || !IsValueNode<Primitive>(use_apply->input(0))) {
continue;
}
ValueNodePtr prim_anf_node = use_apply->input(0)->cast<ValueNodePtr>();
MS_EXCEPTION_IF_NULL(prim_anf_node);
PrimitivePtr node_prim = prim_anf_node->value()->cast<PrimitivePtr>();
MS_EXCEPTION_IF_NULL(node_prim);
// Step 2: Check if the node is a parallel-care node and has associated OperatorInfo
if (IsParallelCareNode(use_apply) && use_apply->has_user_data<OperatorInfo>()) {
auto layout = GetInputLayoutFromCNode(node_pair);
return std::make_shared<TensorLayout>(layout);
}
}
return nullptr;
}
// CreateParameterLayout function
// This function creates a DataParallel tensor layout for a parameter node.
std::shared_ptr<TensorLayout> CreateParameterLayout(const AnfNodePtr &node) {
// 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;
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.";
}
// 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);
// 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<TensorLayout>(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;
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;
}
// Step 5: Create stand-alone tensor_map
TensorMap stand_alone_tensor_map_array(SizeToLong(input_shape_array.size()), -1);
// 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.";
}
// 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<TensorLayout> FindPrevLayout(const AnfNodePtr &node) {
// Step 1: Check if the node is a parameter and create its layout
if (node->isa<Parameter>()) {
return CreateParameterLayout(node);
}
// Step 2: Check if the node is a CNode and its input is a primitive value node
if (!node->isa<CNode>()) {
return nullptr;
}
CNodePtr cnode = node->cast<CNodePtr>();
if (!IsValueNode<Primitive>(cnode->input(0))) {
return nullptr;
}
// Step 3: Handle Depend and other cases
if (IsPrimitiveCNode(node, prim::kPrimReceive)) {
return cnode->user_data<TensorLayout>();
}
if (IsParallelCareNode(cnode) && cnode->has_user_data<OperatorInfo>() && !IsPrimitiveCNode(node, prim::kPrimReshape)) {
auto layout_ptr = GetOutputLayoutFromCNode(cnode, 0);
if (!layout_ptr) {
MS_LOG(EXCEPTION) << "Failure:GetLayoutFromCNode failed";
}
return layout_ptr;
}
ValueNodePtr prim_anf_node = cnode->input(0)->cast<ValueNodePtr>();
PrimitivePtr prim = prim_anf_node->value()->cast<PrimitivePtr>();
// Step 4: Handle TupleGetItem
if (prim->name() == prim::kTupleGetItem) {
auto tuple_index = GetTupleGetItemIndex(cnode);
auto layout_ptr = FindPrevParallelCareNodeLayout(cnode->input(1), LongToSize(tuple_index));
if (!layout_ptr) {
MS_LOG(EXCEPTION) << " Failure:FindPrevLayout failed, tuple_getitem before reshape, but there does not exit a "
"parallel care node "
"before tuple_getitem!";
}
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;
}
auto layout_ptr = FindPrevLayout(cnode->inputs()[index]);
if (!layout_ptr) {
continue;
}
return layout_ptr;
}
MS_LOG(WARNING) << "FindPrevLayout return nullptr, if reshape is not the first primitive, there must be some error";
return nullptr;
}
// ReshapeInit function
// This function initializes the reshape information for all nodes in the graph.
void ReshapeInit(const std::vector<AnfNodePtr> &all_nodes) {
// Step 1: Iterate through all nodes in the graph
for (auto &node : all_nodes) {
auto cnode = node->cast<CNodePtr>();
// Step 2: Check if the node is a valid CNode with a primitive value node as its input
if ((cnode == nullptr) || !IsValueNode<Primitive>(cnode->input(0))) {
continue;
}
ValueNodePtr prim_anf_node = cnode->input(0)->cast<ValueNodePtr>();
// Step 3: Check if the node is a parallel care node with operator info
if (!IsParallelCareNode(cnode) || !cnode->has_user_data<OperatorInfo>()) {
continue;
}
PrimitivePtr prim = GetValueNode<PrimitivePtr>(prim_anf_node);
MS_EXCEPTION_IF_NULL(prim);
OperatorInfoPtr operator_info = cnode->user_data<OperatorInfo>();
// 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!";
}
// 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<ReshapeInfo>(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) {
auto reshape_info_ptr = std::dynamic_pointer_cast<ReshapeInfo>(operator_info);
reshape_info_ptr->SetOutputLayout(*next_layout_ptr);
} else if (is_next_reshape && prev_layout_ptr != nullptr) {
auto reshape_info_ptr = std::dynamic_pointer_cast<ReshapeInfo>(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;
}
// Step 2: Handle return->depend->loss pattern
if (IsPrimitiveCNode(cnode, prim::kPrimDepend) ||
(IsPrimitiveCNode(cnode, prim::kPrimCast) && !cnode->has_user_data<OperatorInfo>())) {
auto depend_before = cnode->input(1)->cast<CNodePtr>();
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; // 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); // 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<CNodePtr>(); // 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<CNode>()) {
auto switch_cnode = pre_cnode->input(0)->cast<CNodePtr>();
if (IsPrimitiveCNode(switch_cnode, prim::kPrimSwitch)) {
MS_EXCEPTION_IF_NULL(switch_cnode);
auto switch_graph = GetValueNode<FuncGraphPtr>(switch_cnode->input(2));
return FindLossCNode(switch_graph, max_depth + 1); // Recursively search in the switch graph
}
}
if (pre_cnode == nullptr || !IsValueNode<Primitive>(pre_cnode->input(0))) {
return loss_node_info; // If the previous CNode is null or not a primitive, return an empty result
}
if (!IsValueNode<Primitive>(pre_cnode->input(0))) {
MS_LOG(DEBUG) << "pre_cnode:" << pre_cnode->ToString();
return loss_node_info;
}
auto current_prim = GetValueNode<PrimitivePtr>(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;
}
// 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";
}
// 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); // 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<CNodePtr>(); // 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;
}
// Handle cases where the loss contains MakeTuple
if (current_prim->name() == MAKE_TUPLE) {
MS_LOG(WARNING) << "The loss contains MakeTuple, which is not supported";
return loss_node_info;
}
// 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; // 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<ValueNodePtr>(); // 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<PrimitivePtr>(); // 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 an empty result if the loss is in the set of invalid loss operations
}
OperatorInfoPtr operator_info = loss_cnode->user_data<OperatorInfo>(); // 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 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); // 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); // 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";
}
// If the shape of sens tensor is [] or [1], no need to split it.
Shape sens_shape = sens_shapes[0];
if (sens_shape.empty() || ((sens_shape.size() == 1) && (sens_shape[0] == 1))) {
if (sens_tensor_node->isa<Parameter>()) {
auto sens_tensor_param = sens_tensor_node->cast<ParameterPtr>();
MS_LOG(DEBUG) << "loss layout " << loss_grad_layout.ToString();
sens_tensor_param->set_user_data<TensorLayout>(std::make_shared<TensorLayout>(loss_grad_layout));
}
MS_LOG(INFO) << "The shape of sens is " << ShapeToString(sens_shape) << ", no need to split sens";
return;
}
auto loss_shape = loss_grad_layout.tensor_shape().array();
if (loss_shape != sens_shape) {
MS_LOG(EXCEPTION) << "The shape of sens is not equal to loss output, it is unsupported now. Sens shape is "
<< ShapeToString(sens_shape) << ", loss shape is " << ShapeToString(loss_shape);
}
MS_LOG(INFO) << "The shape of sens is " << ShapeToString(sens_shape) << ", split it.";
if (!IsValueNode<Tensor>(sens_tensor_node)) {
if (sens_tensor_node->isa<Parameter>()) {
MS_LOG(DEBUG) << "loss layout " << loss_grad_layout.ToString();
AbstractBasePtr abstract = sens_tensor_node->abstract();
MS_EXCEPTION_IF_NULL(abstract);
auto slice_shape = loss_grad_layout.slice_shape().array();
std::shared_ptr<abstract::BaseShape> parallel_shape = std::make_shared<abstract::Shape>(slice_shape);
MS_EXCEPTION_IF_NULL(parallel_shape);
auto cloned_abstract = abstract->Clone();
MS_EXCEPTION_IF_NULL(cloned_abstract);
cloned_abstract->set_shape(parallel_shape);
sens_tensor_node->set_abstract(cloned_abstract);
auto sens_tensor_param = sens_tensor_node->cast<ParameterPtr>();
sens_tensor_param->set_user_data<TensorLayout>(std::make_shared<TensorLayout>(loss_grad_layout));
return;
}
if (sens_tensor_node->isa<CNode>()) {
auto op_list_ptr = InferSensRedistribution(sens_tensor_node, loss_grad_layout);
if (op_list_ptr == nullptr) {
return;
}
auto sens_tensor_cnode = sens_tensor_node->cast<CNodePtr>();
auto func_graph = grad_sens_node->func_graph();
MS_EXCEPTION_IF_NULL(func_graph);
InsertRedistribution(op_list_ptr, grad_sens_node, func_graph, 1, sens_tensor_cnode);
return;
}
MS_LOG(EXCEPTION) << "The type of sens node is not Tensor or Parameter or CNode, it is unsupported now.";
}
// Use _GetTensorSlice operator to split the sens tensor
FuncGraphPtr func_graph = grad_sens_node->func_graph(); // Get the function graph
MS_EXCEPTION_IF_NULL(func_graph);
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); // Ensure the distribute operator is not null
MS_EXCEPTION_IF_NULL(cnode); // Ensure the CNode is not null
if (IsPrimitiveCNode(cnode, prim::kPrimReceive)) {
return; // If the CNode is a Receive operation, do nothing
}
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(); // 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);
// 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);
}
// 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_graph can be used";
}
if (replace_graph) {
MS_LOG(INFO) << "StepReplaceGraph " << cnode->ToString();
StepReplaceGraph(replace_graph, cnode);
}
}
// FindForwardGraphByRootNodes function
// This function finds forward computation graphs based on a set of root nodes.
std::set<FuncGraphPtr> FindForwardGraphByRootNodes(const AnfNodeSet &root_all_nodes) {
std::set<FuncGraphPtr> 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<CNode>()) {
continue;
}
auto cnode = node->cast<CNodePtr>();
// Skip CNodes with less than 2 inputs or non-Primitive first input.
if ((cnode->size() < 2) || !IsValueNode<Primitive>(cnode->input(0))) {
continue;
}
auto expect_prim = GetValueNode<PrimitivePtr>(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<FuncGraph>(cnode->input(1))) {
auto graph = GetValueNode<FuncGraphPtr>(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);
for (auto &sub_graph : graph_used) {
graph_set.insert(sub_graph);
}
}
}
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<CNodePtr, LossNodeInfo> &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<int>(MS_CTX_EXECUTION_MODE);
return (execution_mode == kPynativeMode) && (parallel_mode == kSemiAutoParallel || parallel_mode == kAutoParallel);
}
// GetSensLossPairs function
// This function finds pairs of sensitivity tensors and corresponding loss nodes within the given computation graph.
std::vector<std::pair<CNodePtr, LossNodeInfo>> GetSensLossPairs(const FuncGraphPtr &root) {
MS_EXCEPTION_IF_NULL(root);
std::vector<std::pair<CNodePtr, LossNodeInfo>> sens_loss_pairs;
// Step 1: Iterate through nodes in the computation graph.
for (auto &node : root->nodes()) {
if (!node->isa<CNode>()) {
continue;
}
// Step 2: Check if the node structure corresponds to the expected pattern (sens -> tuple_getitem -> J).
auto sens_cnode = node->cast<CNodePtr>();
AnfNodePtr expect_tuple_getitem = sens_cnode->input(0);
MS_EXCEPTION_IF_NULL(expect_tuple_getitem);
if (!expect_tuple_getitem->isa<CNode>()) {
continue;
}
auto expect_tuple_getitem_cnode = expect_tuple_getitem->cast<CNodePtr>();
if (!IsSomePrimitive(expect_tuple_getitem_cnode, prim::kTupleGetItem)) {
continue;
}
AnfNodePtr expect_anonymous = expect_tuple_getitem_cnode->input(1);
MS_EXCEPTION_IF_NULL(expect_anonymous);
if (!expect_anonymous->isa<CNode>()) {
continue;
}
auto expect_anonymous_cnode = expect_anonymous->cast<CNodePtr>();
AnfNodePtr expect_j = expect_anonymous_cnode->input(0);
MS_EXCEPTION_IF_NULL(expect_j);
if (!expect_j->isa<CNode>()) {
continue;
}
auto expect_j_cnode = expect_j->cast<CNodePtr>();
// 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<FuncGraph>(expect_j_cnode->input(1))) {
MS_LOG(EXCEPTION) << "Sens can't find the corresponding graph.";
}
auto func_graph = GetValueNode<FuncGraphPtr>(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;
}
std::pair<CNodePtr, LossNodeInfo> sens_loss_pair = std::make_pair(sens_cnode, loss_node_info);
sens_loss_pairs.push_back(sens_loss_pair);
}
return sens_loss_pairs;
}
void ParallelCommunication(const FuncGraphPtr &root, const std::vector<AnfNodePtr> &all_nodes,
const FuncGraphManagerPtr &manager) {
MS_EXCEPTION_IF_NULL(root);
MS_EXCEPTION_IF_NULL(manager);
TensorRedistribution tensor_redistribution;
std::vector<std::pair<CNodePtr, LossNodeInfo>> sens_loss_pairs = GetSensLossPairs(root);
bool has_backward = !sens_loss_pairs.empty();
// split sens must before inserting the operators.
for (auto &pair : sens_loss_pairs) {
// If the shape of grad-sens tensor is not [] or [1], use get tensor slice to handle it.
// If the type of sens node is not Tensor, it is unsupported now, do nothing default.
if (IsLastStage()) {
StepSplitSens(pair);
}
}
for (auto &node : all_nodes) {
MS_EXCEPTION_IF_NULL(node);
if (node->isa<CNode>()) {
auto cnode = node->cast<CNodePtr>();
// the make_tuple is parallel care node, but it may have not operator info
if (!IsParallelCareNode(cnode) || !cnode->has_user_data<OperatorInfo>()) {
continue;
}
OperatorInfoPtr distribute_operator = GetDistributeOperator(cnode);
MS_EXCEPTION_IF_NULL(distribute_operator);
// skip Send Receive
if (!cnode->HasPrimalAttr(PIPELINE_PARAM)) {
// insert forward ops
InsertForwardOps(distribute_operator, cnode);
// insert redistribution ops
StepRedistribution(cnode, distribute_operator, cnode, tensor_redistribution, cnode);
}
// insert backward ops
if (has_backward || IsPynativeParallel()) {
BackwardCommunication(root, distribute_operator, cnode, sens_loss_pairs);
}
distribute_operator->ReplaceNodeInputOrAttrs();
} else if (IsValueNode<Tensor>(node) || IsValueNode<ValueList>(node) || IsValueNode<ValueTuple>(node)) {
StepSplitTensor(node, manager);
}
}
for (auto &node : all_nodes) {
MS_EXCEPTION_IF_NULL(node);
if (node->isa<CNode>()) {
auto cnode = node->cast<CNodePtr>();
if (!IsParallelCareNode(cnode) || !cnode->has_user_data<OperatorInfo>() || IsSomePrimitive(cnode, RECEIVE) ||
IsSomePrimitive(cnode, SEND)) {
continue;
}
OperatorInfoPtr distribute_operator = GetDistributeOperator(cnode);
MS_EXCEPTION_IF_NULL(distribute_operator);
// StepReplace
StepReplace(distribute_operator, cnode);
}
}
}
// Function: IsCohesiveNode
// This function checks if a CNode is a cohesive node, which includes operations like cast, load, all-gather, mini-step all-gather, and micro-step all-gather.
bool IsCohesiveNode(const CNodePtr &cnode) {
return IsPrimitiveCNode(cnode, prim::kPrimCast) || IsPrimitiveCNode(cnode, prim::kPrimLoad) ||
IsPrimitiveCNode(cnode, prim::kPrimAllGather) || IsPrimitiveCNode(cnode, prim::kPrimMiniStepAllGather) ||
IsPrimitiveCNode(cnode, prim::kPrimMicroStepAllGather);
}
// Function: NodeParameterName
// This recursive function retrieves parameter names associated with an operator node.
ParameterMap NodeParameterName(const CNodePtr &node, int64_t index, size_t curr_depth) {
// Check if recursion depth exceeds the maximum allowed depth
if (curr_depth > MAX_RECURSIVE_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<AnfNodePtr> 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<Parameter>()) {
auto input_parameter = input->cast<ParameterPtr>();
// 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<CNode>()) {
CNodePtr cnode = input->cast<CNodePtr>();
// Check if the input CNode is a cohesive node and has more than one input
if (!IsValueNode<Primitive>(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<std::string> gather_info_names = {"GatherInfo", "SparseGatherV2Info", "EmbeddingLookupInfo"};
for (std::string info_name : gather_info_names) {
if (name.find(info_name) != std::string::npos) {
return true;
}
}
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<AnfNodePtr> &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<CNodePtr>();
// Check if the node is a CNode and has a Primitive input
if ((cnode == nullptr) || !IsValueNode<Primitive>(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<PrimitivePtr>(cnode->input(0));
MS_EXCEPTION_IF_NULL(prim);
OperatorInfoPtr operator_info = cnode->user_data<OperatorInfo>();
// 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<TensorLayout>();
}
// Check if operator info represents gather information
if (IsGatherInfo(operator_info->name())) {
auto gather_info = std::dynamic_pointer_cast<GatherInfo>(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 the same.";
}
// Store manual shape information
std::vector<std::pair<int64_t, int64_t>> manual_shape;
for (int64_t i = 0; i < UlongToLong(param_split_shapes.size()); ++i) {
(void)manual_shape.emplace_back(
std::make_pair(param_split_shapes[LongToSize(i)], index_offsets[LongToSize(i)]));
}
manual_shape_map[param_name] = manual_shape;
}
}
}
// 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<ParameterPtr>();
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<ParameterPtr>()->name();
auto cloned_param_layout = cloned_parameter_node->user_data<TensorLayout>();
// 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<AnfNodePtr> &all_nodes) {
for (auto &node : all_nodes) {
MS_EXCEPTION_IF_NULL(node);
// Check if the node is a CNode
if (!node->isa<CNode>()) {
continue;
}
auto cnode = node->cast<CNodePtr>();
// Check if the CNode has a Primitive input
if (!IsValueNode<Primitive>(cnode->input(0))) {
continue;
}
// 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);
if (!node->isa<CNode>()) {
continue;
}
auto cnode = node->cast<CNodePtr>();
if (!IsValueNode<Primitive>(cnode->input(0))) {
continue;
}
// 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<FuncGraphPtr> ForwardGraph(const FuncGraphPtr &root) {
MS_EXCEPTION_IF_NULL(root);
const auto &all_nodes = root->nodes();
std::set<FuncGraphPtr> graph_set = FindForwardGraphByRootNodes(all_nodes);
return graph_set;
}
// FindRootForwardCNode function
// This function finds the root forward CNodes in the provided graph and returns them as a vector.
std::vector<AnfNodePtr> FindRootForwardCNode(const FuncGraphPtr &graph, const AnfNodeSet &all_nodes) {
MS_EXCEPTION_IF_NULL(graph);
std::vector<AnfNodePtr> root_forward_nodes;
auto loss_cnode = FindLossCNode(graph, 0).loss_node;
if (loss_cnode == nullptr) {
MS_LOG(WARNING) << "Can not find the loss cnode";
return root_forward_nodes;
}
auto loss_cnode_id = loss_cnode->UniqueIdThroughCopy();
for (auto &node : all_nodes) {
MS_EXCEPTION_IF_NULL(node);
if (!node->isa<CNode>()) {
continue;
}
auto cnode = node->cast<CNodePtr>();
auto root_node_id = node->UniqueIdThroughCopy();
if (loss_cnode_id == root_node_id) {
root_forward_nodes = DeepLinkedGraphSearch(cnode);
break;
}
}
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) {
// 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<ValueSequencePtr>();
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;
return nullptr;
}
for (auto &node : cnode->inputs()) {
if (!node->isa<CNode>()) {
continue;
}
if (!IsPrimitiveCNode(node, prim::kPrimEnvironGet)) {
return FindGrad(node->cast<CNodePtr>(), ++curr_depth);
} else {
return node;
}
}
return nullptr;
}
// HandleRootReshapeAndSaveStrategy function
// This function handles root graph reshaping and saves strategy information.
void HandleRootReshapeAndSaveStrategy(const std::vector<AnfNodePtr> &all_nodes) {
// 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<CNode>()) {
continue;
}
auto cnode = node->cast<CNodePtr>();
if (!IsValueNode<Primitive>(cnode->input(0)) || cnode == nullptr) {
continue;
}
if (cnode->in_forward_flag()) {
// Save strategy in executor for nodes in the forward pass.
OperatorInfoPtr op_info = cnode->user_data<OperatorInfo>();
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.
executor->SetCNodeStrategy(cnode->fullname_with_scope(), strategy);
}
}
continue;
}
auto prim = GetValueNode<PrimitivePtr>(cnode->input(0));
if (prim->name() != RESHAPE) {
continue;
}
Shape origin_dst_shape = GetValue<std::vector<int64_t>>(cnode->input(2)->cast<ValueNodePtr>()->value());
if (origin_dst_shape.size() == 1 && origin_dst_shape[0] == -1) {
continue;
}
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;
}
// 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();
auto ms_context = MsContext::GetInstance();
MS_EXCEPTION_IF_NULL(ms_context);
std::string backend = ms_context->get_param<std::string>(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;
} else if (backend == kGPUDevice) {
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()) {
// 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";
}
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: "
<< device_num;
return FAILED;
}
// 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;
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 than 'device num', "
"but got the global_rank : "
<< global_rank << " and the device_num : " << device_num;
return FAILED;
}
// Create stages vector based on split_stage_num
std::vector<int64_t> 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;
}
MS_LOG(INFO) << "The parallel context: device_num: " << device_num << ", global_rank: " << global_rank
<< ", communication_backend: " << comm_info.communication_backend
<< ", gradients_mean: " << ParallelContext::GetInstance()->gradients_mean()
<< ", gradient_fp32_sync: " << ParallelContext::GetInstance()->gradient_fp32_sync();
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<AnfNodePtr> &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;
}
auto cnode = node->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(cnode);
if (!cnode->in_forward_flag()) {
continue;
}
// Get the manager for the current function graph
FuncGraphManagerPtr manager = cnode->func_graph()->manager();
MS_EXCEPTION_IF_NULL(manager);
// 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;
}
auto make_tuple_list_next_cnode = make_tuple_list_next_node->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(make_tuple_list_next_cnode);
OperatorInfoPtr op_info = GetDistributeOperator(make_tuple_list_next_cnode);
MS_EXCEPTION_IF_NULL(op_info);
cnode->set_user_data<OperatorInfo>(op_info);
}
}
// 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) {
return false;
}
if (CreateGroups(group_info_map) != SUCCESS) {
return false;
}
MS_LOG(INFO) << "Create groups by checkpoint file success";
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);
if (root->has_flag(kTraining)) {
Reorder(root);
} else {
ReorderForPredict(root, manager);
}
}
}
// 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();
int64_t split_stage_num = ParallelContext::GetInstance()->pipeline_stage_split_num();
int64_t per_stage_device_num = comm_info.device_num / split_stage_num;
int64_t current_stage = comm_info.global_rank / per_stage_device_num;
MS_LOG(INFO) << "The current stage is: " << current_stage;
if (!root->has_flag(kTraining) && !ParallelContext::GetInstance()->dataset_strategy().empty()) {
MS_LOG(WARNING) << "In eval/predict net, the output parallel strategy would not follow "
"the input parallel strategy when using context.set_auto_parallel_context(dataset_strategy)"
" to configure the input strategy.";
}
return ((!root->has_flag(kTraining) && ParallelContext::GetInstance()->dataset_strategy().empty() &&
current_stage == split_stage_num - 1) ||
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");
if (!group_info_save_path.empty()) {
ParallelContext::GetInstance()->set_group_ckpt_save_file(group_info_save_path);
}
if (StrategyCheckpoint::GetInstance().group_info_save_on()) {
RankList comm_group = FindCommonMirrorGroup(root);
if (StrategyCheckpoint::GetInstance().SaveGroupInfo(group_info, comm_group) != SUCCESS) {
MS_LOG(EXCEPTION) << "Save group info failed";
}
}
}
// 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) {
auto group_info_save_path = common::GetEnv("GROUP_INFO_FILE");
if (!group_info_save_path.empty()) {
std::vector<std::pair<std::string, std::vector<uint32_t>>> group_info;
int64_t device_num = GetCommInfo().device_num;
RankList comm_group;
for (size_t i = 0; i < size_t(device_num); ++i) {
comm_group.push_back(i);
}
ParallelContext::GetInstance()->set_group_ckpt_save_file(group_info_save_path);
if (StrategyCheckpoint::GetInstance().SaveGroupInfo(group_info, comm_group) != SUCCESS) {
MS_LOG(EXCEPTION) << "Save group info failed";
}
}
}
}
// 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<AnfNodePtr> &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<AnfNodePtr> &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<CNodePtr>();
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 " << 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<CNodePtr>(), SQRT) && limits < MAX_BFS_DEPTH) {
auto users = node_user_map.at(find_node);
if (users.empty()) return;
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<CNodePtr>(), SQRT)) {
return;
}
auto sqrt_node = find_node;
if (sqrt_node->inputs().size() > 1 && IsSomePrimitive(sqrt_node->input(1)->cast<CNodePtr>(), ALL_REDUCE)) {
return;
}
// 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;
}
// Step 7: Insert AllReduce operations for global norm value within the current stage
InsertAllReduceToNodeInput(sqrt_node->cast<CNodePtr>(), cur_stage_device_list.name(), PARALLEL_GLOBALNORM);
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) << "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;
}
InsertAllReduceToNodeInput(sqrt_node->cast<CNodePtr>(), 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<AnfNodePtr> 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<CNodePtr>();
// 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<bool>(value)) {
continue;
}
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;
}
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 &parameter,
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 &param_pair : params_user_set) {
auto cnode = param_pair.first->cast<CNodePtr>();
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<bool>(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<CNodePtr>(), dev_num, PARALLEL_GLOBALNORM_DIV);
MS_LOG(INFO) << "Insert the realdiv with " << dev_num << " for the parameter " << parameter->fullname_with_scope()
<< "succeed!";
}
// 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 &parameter) {
// 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 &param_pair : params_user_set) {
auto cnode = param_pair.first->cast<CNodePtr>();
std::vector<AnfNodePtr> 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<CNodePtr>();
if (!IsPrimitiveCNode(local_cnode, prim::kPrimMirror) &&
!IsPrimitiveCNode(local_cnode, prim::kPrimMirrorMicroStep) &&
!IsPrimitiveCNode(local_cnode, prim::kPrimMirrorMiniStep)) {
continue;
}
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<AnfNodePtr> &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 &parameter : 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<Int64Imm>()) {
dev_num = GetValue<int64_t>(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()->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());
std::string parallel_mode = ParallelContext::GetInstance()->parallel_mode();
HandleDataParallel();
pipeline::ResourceBasePtr res = optimizer->resource();
MS_EXCEPTION_IF_NULL(res);
FuncGraphManagerPtr manager = res->manager();
MS_EXCEPTION_IF_NULL(manager);
auto pipeline_stages = ParallelContext::GetInstance()->pipeline_stage_split_num();
// assume no change to graph
bool changes = false;
// 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)) {
MS_LOG(INFO) << "Strategies would be ignored in " << parallel_mode
<< ", shard() only valid in [semi_]auto_parallel.";
root->set_flag(CHECK_SET_STRATEGY_VALID_ONCE_ONLY, true);
}
ReorderForPipelineSplit(root, manager, pipeline_stages);
return changes;
}
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();
MS_EXCEPTION_IF_NULL(ret);
std::vector<AnfNodePtr> all_nodes = DeepScopedGraphSearch(ret);
std::reverse(all_nodes.begin(), all_nodes.end());
if (parallel_mode != kAutoParallel) {
TOTAL_OPS = 0;
if (pipeline_stages <= 1 && ParallelInit() != SUCCESS) {
MS_LOG(EXCEPTION) << "Parallel init failed";
}
PipelinePreProcess(root, manager, all_nodes);
// mark the forward cnodes, parallel only care these nodes
MarkForwardCNode(root);
if (FindCommunicationOp(all_nodes)) {
MS_LOG(EXCEPTION) << "The graph contain communication op";
}
if (IsInsertVirtualOutput(root)) {
InsertVirtualOutput(root, all_nodes);
AnfNodePtr ret_after = root->get_return();
MS_EXCEPTION_IF_NULL(ret_after);
all_nodes = DeepScopedGraphSearch(ret_after);
std::reverse(all_nodes.begin(), all_nodes.end());
}
// extract shape and strategy, set operator_info
ExtractInformation(all_nodes);
ReshapeInit(all_nodes);
}
SetCastForParamNotRecompute(all_nodes);
HandleRootReshapeAndSaveStrategy(all_nodes);
HandleForwardMakeTupleAndMakeList(all_nodes);
// if the input or parameter has multiple users, check whether its split strategies are consistent.
CheckParameterSplit(all_nodes);
HandleSymbolicKeyInstance(root, all_nodes);
// cover Parallel shape
CoverSliceShape(root);
// handle input is not used
HandleNoUsedParameter(root);
// set the shape for optimizer's clone tensor
SetClonedTensorShapeForOptimizer(root);
HandleAdaFactorOpt(root);
auto adasum_param_tensor_layout_map = AdaSumParamTensorLayout(root);
bool is_apply_adasum = HandleAdaSum(root, all_nodes, &adasum_param_tensor_layout_map);
// save strategy as checkpoint for multi-train
if (StrategyCheckpoint::GetInstance().SaveCheckPointOn()) {
CheckpointStrategy(all_nodes, root);
}
// ForwardCommunication BackwardCommunication TensorRedistribution
ParallelCommunication(root, all_nodes, manager);
if (is_apply_adasum) {
HandleMirrorInAdaSum(root, &adasum_param_tensor_layout_map);
}
PipelinePostProcess(root, all_nodes);
HandleGroupInfo(root);
// handle full split parammeters in grad accumulation, do not contain optimizer-sharding's parameter
HandleFullySplitParameters(root);
HandlGlobalNormScale(root, all_nodes, manager);
DumpGraph(root, std::string(STEP_PARALLEL_END));
// step parallel only run once
root->set_flag(SEMI_AUTO_PARALLEL_RUN_ONCE_ONLY, true);
res->SetResult(pipeline::kStepParallelGraph, root);
// in auto parallel mode, no need to check if stategies set
root->set_flag(CHECK_SET_STRATEGY_VALID_ONCE_ONLY, true);
(void)gettimeofday(&end_time, nullptr);
uint64_t time = kUSecondInSecond * static_cast<uint64_t>(end_time.tv_sec - start_time.tv_sec);
time += static_cast<uint64_t>(end_time.tv_usec - start_time.tv_usec);
MS_LOG(INFO) << "Now leaving step parallel, used time: " << time << " us";
return changes;
}
// ExtractInputsTensorName function
// This function extracts tensor names from the inputs of a CNode and returns them in a vector.
std::vector<std::string> ExtractInputsTensorName(const CNodePtr &node) {
std::vector<std::string> name_inputs;
std::vector<AnfNodePtr> all_inputs = node->inputs();
std::vector<AnfNodePtr> node_inputs{all_inputs.begin() + 1, all_inputs.end()};
std::string node_id = node->UniqueId();
name_inputs.push_back(node_id);
for (auto &input : node_inputs) {
std::string name = input->UniqueId();
name_inputs.push_back(name);
}
return name_inputs;
}
} // namespace parallel
} // namespace mindspore
Reference in New Issue View Git Blame Copy Permalink