From 6e2222a30e96ada95c6547513dabcff256a26d7d Mon Sep 17 00:00:00 2001
From: yao_yf <yaoyifan1@huawei.com>
Date: Thu, 15 Apr 2021 20:49:45 +0800
Subject: [PATCH] fix node find param

---
 mindspore/ccsrc/frontend/parallel/step_parallel.cc | 7 ++++---
 mindspore/ccsrc/frontend/parallel/step_parallel.h  | 2 +-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/mindspore/ccsrc/frontend/parallel/step_parallel.cc b/mindspore/ccsrc/frontend/parallel/step_parallel.cc
index a1e89c75eb..79c17470be 100644
--- a/mindspore/ccsrc/frontend/parallel/step_parallel.cc
+++ b/mindspore/ccsrc/frontend/parallel/step_parallel.cc
@@ -2769,15 +2769,16 @@ bool IsCohesiveNode(const CNodePtr &cnode) {
          IsPrimitiveCNode(cnode, prim::kPrimAllGather) || IsPrimitiveCNode(cnode, prim::kPrimMiniStepAllGather);
 }
 
-std::vector<std::pair<std::string, int64_t>> NodeParameterName(const CNodePtr &node) {
+std::vector<std::pair<std::string, int64_t>> NodeParameterName(const CNodePtr &node, int64_t index) {
   std::vector<AnfNodePtr> node_inputs{node->inputs()};
   std::vector<std::pair<std::string, int64_t>> param_names;
   for (int64_t i = 0; i < UlongToLong(node_inputs.size()); ++i) {
+    int64_t idx = index > i ? index : i;
     auto input = node_inputs[i];
     if (input->isa<Parameter>()) {
       auto input_parameter = input->cast<ParameterPtr>();
       if (input_parameter->has_default() && ParameterRequireGrad(input_parameter)) {
-        param_names.push_back({input_parameter->name(), i});
+        param_names.push_back({input_parameter->name(), idx});
       }
     } else if (input->isa<CNode>()) {
       CNodePtr cnode = input->cast<CNodePtr>();
@@ -2785,7 +2786,7 @@ std::vector<std::pair<std::string, int64_t>> NodeParameterName(const CNodePtr &n
         continue;
       }
       if (IsCohesiveNode(cnode) && cnode->inputs().size() >= 1) {
-        auto input_param_names = NodeParameterName(cnode);
+        auto input_param_names = NodeParameterName(cnode, idx);
         param_names.insert(param_names.end(), input_param_names.begin(), input_param_names.end());
       }
     }
diff --git a/mindspore/ccsrc/frontend/parallel/step_parallel.h b/mindspore/ccsrc/frontend/parallel/step_parallel.h
index 4c4a71ed1b..3717b36575 100644
--- a/mindspore/ccsrc/frontend/parallel/step_parallel.h
+++ b/mindspore/ccsrc/frontend/parallel/step_parallel.h
@@ -134,7 +134,7 @@ bool IsLastStage();
 void ParallelCommunication(const FuncGraphPtr &root, const std::vector<AnfNodePtr> &all_nodes,
                            const FuncGraphManagerPtr &manager);
 
-std::vector<std::pair<std::string, int64_t>> NodeParameterName(const CNodePtr &node);
+std::vector<std::pair<std::string, int64_t>> NodeParameterName(const CNodePtr &node, int64_t index = -1);
 
 void CheckpointStrategy(const std::vector<AnfNodePtr> &all_nodes);