!31164 Fix the global norm missing insert allreduce

Merge pull request !31164 from huangxinjing/fx_global_norm_error

mindspore/ccsrc/frontend/parallel/step_parallel.cc

@@ -62,8 +62,6 @@ namespace parallel {
 static const std::set<std::string> COMMUNICATION_OPS = {ALL_REDUCE, ALL_GATHER, ALL_TO_ALL, REDUCE_SCATTER};
 static const std::set<std::string> INVALID_LOSS_OPS = {GET_NEXT, VIRTUALLOSS, LOAD, UPDATESTATE};
 static const std::set<std::string> NO_INPUT_TENSOR_OPS = {UNIFORM_REAL};
-static const std::vector<std::pair<const std::string, int64_t>> REDUCE_SUM_MATCH_PATTERN = {
-  std::make_pair(MAKE_TUPLE, 1), std::make_pair(ADDN, 1), std::make_pair(SQRT, 1)};
 // g_RefMap, for CNode B input i is a RefKey[Parameter C],
 // it will be one item in map with key: C, and value: (B, i)
 std::map<AnfNodePtr, std::pair<AnfNodePtr, int64_t>> g_RefMap;
@@ -3179,9 +3177,22 @@ static void InsertAllReduceForNormValue(const AnfNodePtr &res_node) {
     return;
   }
   auto pipeline_stages = ParallelContext::GetInstance()->pipeline_stage_split_num();
-  auto expand_dims_node = node_user_map.at(res_node).front().first;
-  auto sqrt_node = MatchPattern(expand_dims_node, node_user_map, REDUCE_SUM_MATCH_PATTERN);
-  if (!sqrt_node) return;
+  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;
+  }
+  if (!find_node || !IsSomePrimitive(find_node->cast<CNodePtr>(), SQRT)) {
+    return;
+  }
+  auto anf_node = find_node->cast<CNodePtr>();
+  if (anf_node->inputs().size() > 1 && IsSomePrimitive(anf_node->input(1)->cast<CNodePtr>(), ALL_REDUCE)) {
+    return;
+  }
+  auto sqrt_node = find_node;
   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) {
@@ -3245,27 +3256,22 @@ AnfNodePtr FindExpanDimsWIthGradScale(const AnfNodePtr &node_ptr, const NodeUser
 
 static void InsertDivAndAllReduceForNorm(const NodeUsersMap &node_user_map, const AnfNodePtr &parameter,
                                          uint32_t dev_num) {
-  AnfNodePtr expand_dims_node = nullptr;
-  AnfNodePtr prefix_node = nullptr;
   auto params_user_set = node_user_map.at(parameter);
   for (auto &param_pair : params_user_set) {
-    expand_dims_node = nullptr;
     auto cnode = param_pair.first->cast<CNodePtr>();
     MS_EXCEPTION_IF_NULL(cnode);
     if (cnode->in_forward_flag()) {
       continue;
     }
-    expand_dims_node = FindExpanDimsWIthGradScale(cnode, node_user_map, MAX_BFS_DEPTH);
-    if (!expand_dims_node) {
-      continue;
-    }
+    auto expand_dims_node = FindExpanDimsWIthGradScale(cnode, node_user_map, MAX_BFS_DEPTH);
+    if (!expand_dims_node) continue;
     auto value = GetAttrsFromAnfNode(expand_dims_node, GRAD_SCALE);
-    if (!value || !GetValue<bool>(value)) {
-      continue;
+    if (!value || !GetValue<bool>(value)) continue;
+    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!";
     }
-    InsertRealDivOpToNodeInput(expand_dims_node->cast<CNodePtr>(), dev_num, PARALLEL_GLOBALNORM_DIV);
-    MS_LOG(INFO) << "Insert the realdiv with " << dev_num << " for the parameter " << parameter->DebugString()
-                 << "succeed!";
     // If already inserted allreduce, the pattern will not be matched and thus no allreduce will be inserted.
     InsertAllReduceForNormValue(expand_dims_node);
   }
@@ -3302,22 +3308,15 @@ static void HandlGlobalNormScale(const FuncGraphPtr &root, const std::vector<Anf
   auto parameters = root->parameters();
   auto node_user_map = manager->node_users();
   MS_LOG(INFO) << "Start to process the global norm";
+
   for (auto &parameter : parameters) {
+    int64_t dev_num = 0;
     if (!ParameterRequireGrad(parameter)) continue;
     auto mirror_node = GetMirrorOp(node_user_map, parameter);
-    if (!mirror_node) continue;
     auto device_num_ptr = GetAttrsFromAnfNode(mirror_node, DEV_NUM);
-    if (!device_num_ptr) {
-      MS_LOG(ERROR) << "The mirror operator is excepted to have device number attribute, but found none. This "
-                       "will cause the global norm calculation with wrong precision.";
-      continue;
-    }
-    if (!device_num_ptr->isa<Int64Imm>()) {
-      MS_LOG(ERROR) << "The type of device number attribute of mirror operator is not int64.";
-      continue;
+    if (device_num_ptr && device_num_ptr->isa<Int64Imm>()) {
+      dev_num = GetValue<int64_t>(device_num_ptr);
     }
-    auto dev_num = device_num_ptr->cast<Int64ImmPtr>()->value();
-    if (dev_num == 0) continue;
     InsertDivAndAllReduceForNorm(node_user_map, parameter, dev_num);
   }
 }

mindspore/ccsrc/frontend/parallel/step_parallel_utils.cc

@@ -521,34 +521,5 @@ std::shared_ptr<Value> GetAttrsFromAnfNode(const std::shared_ptr<AnfNode> &node,
   }
   return nullptr;
 }
-
-AnfNodePtr MatchPattern(const AnfNodePtr &node, const NodeUsersMap &user_map,
-                        const std::vector<std::pair<const std::string, int64_t>> &match_pattern) {
-  AnfNodePtr start_node = node;
-  bool find = false;
-  for (uint32_t i = 0; i < match_pattern.size(); ++i) {
-    find = false;
-    if (!IsSomePrimitive(start_node->cast<CNodePtr>(), {match_pattern[i].first})) {
-      break;
-    } else if (i == match_pattern.size() - 1) {
-      find = true;
-      break;
-    }
-
-    auto next_node_users = user_map.at(start_node);
-    for (auto &next_node : next_node_users) {
-      if (i + 1 < match_pattern.size() &&
-          IsSomePrimitive(next_node.first->cast<CNodePtr>(), {match_pattern[i + 1].first}) &&
-          next_node.second == match_pattern[i + 1].second) {
-        start_node = next_node.first;
-        break;
-      }
-    }
-  }
-  if (!find) {
-    start_node = nullptr;
-  }
-  return start_node;
-}
 }  // namespace parallel
 }  // namespace mindspore

mindspore/ccsrc/frontend/parallel/step_parallel_utils.h

@@ -42,8 +42,6 @@ std::vector<AnfNodePtr> ReplaceOpInput(const Operator &replace_op, const std::st
 std::string CreateInstanceName(const CNodePtr &node, size_t index);
 TensorInfo GetInputsTensorInfo(const std::pair<AnfNodePtr, int64_t> &param_info);
 AnfNodePtr CheckMakeTupleSplit(const AnfNodePtr &node, const FuncGraphManagerPtr &manager);
-AnfNodePtr MatchPattern(const AnfNodePtr &node, const NodeUsersMap &user_map,
-                        const std::vector<std::pair<const std::string, int64_t>> &match_pattern);
 
 // for specific scenarios
 RankList FindCommonMirrorGroup(const FuncGraphPtr &root);

tests/ut/python/parallel/test_global_norm.py

@@ -30,6 +30,21 @@ from mindspore.ops import operations as P
 from mindspore.ops import composite as C
 from mindspore import context
 
+
+class OneParameterNet(nn.Cell):
+    """Net definition"""
+    def __init__(self, param_type, strategy1, strategy2):
+        super(OneParameterNet, self).__init__()
+        self.fc1 = P.MatMul().shard(strategy1)
+        self.p1 = Parameter(Tensor(np.ones([48, 16]).astype(param_type)), name="weight1")
+        self.sub = P.Sub().shard(strategy2)
+
+    def construct(self, x, y):
+        x = P.Cast()(x, ms.float16)
+        p1 = P.Cast()(self.p1, ms.float16)
+        x = self.fc1(x, p1)
+        return self.sub(x, 0)
+
 class Net(nn.Cell):
     """Net definition"""
     def __init__(self, param_type, strategy1, strategy2):
@@ -137,6 +152,16 @@ class TestGlobalNormInserted:
                     appear_count += 1
         assert appear_count == target_count
 
+    def test_nonpipeline_global_norm_one_parameter(self):
+        """
+        Feature: Parallel ClipByGlobalNorm
+        Description: Test the global norm using one parameter, there should be only one allreduce
+        Expectation:When there is no PARALLEL_GLOBALNORM_IN_STAGES inserted
+        """
+        auto_parallel_compile_net("semi_auto_parallel", 8, OneParameterNet, ((1, 8), (8, 1)), ((8, 1), ()),
+                                  interleaved_batch=1, param_type=np.float32)
+        self.run_count_check(target_count=1, pattern=r"PARALLEL_GLOBALNORM_IN_STAGES")
+
     def test_nonpipeline_global_norm(self):
         """
         Feature: Parallel ClipByGlobalNorm