fix bug of ZeRO MP in complex case (#8404)

oneflow/core/graph/op_graph.cpp

@@ -443,6 +443,9 @@ void OpGraph::ForEachDataAndCtrlInNode(OpNode* node,
                                        const std::function<void(OpNode*)>& Handler) const {
   node->ForEachNodeOnInEdge(Handler);
   for (const auto& ctrl_in_op_name : node->op().op_conf().ctrl_in_op_name()) {
+    CHECK(op_name2op_node_.find(ctrl_in_op_name) != op_name2op_node_.end())
+        << " cannot find ctrl_in_op_name: [" << ctrl_in_op_name << "] of op: ["
+        << node->op().op_name() << "] in OpGraph. ";
     Handler(op_name2op_node_.at(ctrl_in_op_name));
   }
 }
@@ -453,6 +456,9 @@ void OpGraph::ForEachDataAndCtrlOutNode(OpNode* node,
   const auto& op_name_it = producer_op_name2ctrl_consumer_op_names_.find(node->op().op_name());
   if (op_name_it == producer_op_name2ctrl_consumer_op_names_.end()) { return; }
   for (const std::string& ctrl_consumer_op_name : op_name_it->second) {
+    CHECK(op_name2op_node_.find(ctrl_consumer_op_name) != op_name2op_node_.end())
+        << " cannot find ctrl_consumer_op_name: [" << ctrl_consumer_op_name << "] of op: ["
+        << node->op().op_name() << "] in OpGraph.";
     Handler(op_name2op_node_.at(ctrl_consumer_op_name));
   }
 }

oneflow/core/job_rewriter/optimizer_placement_optimization_pass.cpp

@@ -448,7 +448,9 @@ void ShardSequence(JobBuilder* builder, const int64_t threshold, const ParallelD
               << " with op conf " << new_var_op_conf.variable_conf().DebugString();
       continue;
     }
-    if (i != 0) { new_var_op_conf.add_ctrl_in_op_name(prev_allowed_op_name); }
+    if (!prev_allowed_op_name.empty()) {
+      new_var_op_conf.add_ctrl_in_op_name(prev_allowed_op_name);
+    }
     builder->MutOpsOnlyOnce({new_var_op_conf});
     // Set consumers to consum this variable op's cast op's output as Broadcast.
     if (new_split_signature != "") {