[just for review]Release/mt5 opt

oneflow/core/common/env_var/debug_mode.h

@@ -25,6 +25,9 @@ DEFINE_ENV_BOOL(ONEFLOW_DEBUG, false);
 
 inline bool IsInDebugMode() { return EnvBool<ONEFLOW_DEBUG_MODE>() || EnvBool<ONEFLOW_DEBUG>(); }
 
+DEFINE_ENV_BOOL(ENABLE_LOGICAL_CHAIN, true);
+inline bool EnableLogicalChain() { return EnvBool<ENABLE_LOGICAL_CHAIN>(); }
+
 }  // namespace oneflow
 
 #endif  // ONEFLOW_CORE_COMMON_ENV_VAR_DEBUG_MODE_H_

oneflow/core/framework/nn_graph.cpp

@@ -339,7 +339,6 @@ Maybe<void> NNGraph::CompileAndInitRuntime() {
     // TODO(chengcheng): test collective boxing for multi-job.
     PlanUtil::GenCollectiveBoxingPlan(&job_, &plan_);
     sub_compile_tc->Count("[GraphCompile]" + name_ + " GenCollectiveBoxingPlan", 1);
-    // PlanUtil::SetForceInplaceMemBlock(&plan_); NOTE(chengcheng): only for ssp.
     PlanUtil::DumpCtrlRegstInfoToPlan(&plan_);
     sub_compile_tc->Count("[GraphCompile]" + name_ + " DumpCtrlRegstInfoToPlan", 1);
     PlanUtil::PlanMemoryLog(&plan_, name_);

oneflow/core/functional/impl/common.cpp

@@ -17,6 +17,8 @@ limitations under the License.
 #include "oneflow/core/functional/impl/common.h"
 #include "oneflow/core/autograd/autograd_mode.h"
 #include "oneflow/core/common/wrap_dim_utils.h"
+#include "oneflow/core/ccl/ccl.h"
+#include "oneflow/core/job/rank_group.h"
 
 namespace oneflow {
 namespace one {
@@ -220,6 +222,16 @@ Maybe<std::tuple<Shape, bool, bool>> InferUnifiedShapeForBroadcasting(const Shap
   return std::make_tuple(target, need_to_broadcast.first, need_to_broadcast.second);
 }
 
+Maybe<void> BroadcastSeedToAllRanks(uint64_t* seed, int64_t root) {
+  CHECK_NOTNULL_OR_RETURN(seed) << "seed is not allowed to be nullptr";
+  const auto& rank_group = JUST(RankGroup::DefaultRankGroup());
+  const auto& parallel_desc = JUST(RankGroup::GetDefaultParallelDesc(DeviceType::kCPU, rank_group));
+  const auto& meta_transport_token =
+      JUST(TransportToken::NewTransportToken(kTransportTokenTypeMeta));
+  JUST(ccl::CpuBroadcast(seed, seed, sizeof(*seed), root, parallel_desc, meta_transport_token));
+  return Maybe<void>::Ok();
+}
+
 }  // namespace functional
 }  // namespace one
 }  // namespace oneflow

oneflow/core/functional/impl/common.h

@@ -43,6 +43,8 @@ Maybe<Shape> InferShapeUnspecifiedDim(const int64_t& elem_count, const Shape& sh
 Maybe<std::tuple<Shape, bool, bool>> InferUnifiedShapeForBroadcasting(const Shape& input_shape,
                                                                       const Shape& other_shape);
 
+Maybe<void> BroadcastSeedToAllRanks(uint64_t* seed, int64_t root = 0);
+
 }  // namespace functional
 }  // namespace one
 }  // namespace oneflow

oneflow/core/functional/impl/nn_functor.cpp

@@ -26,6 +26,8 @@ limitations under the License.
 #include "oneflow/user/kernels/dropout_kernel.h"
 #include "oneflow/core/common/container_util.h"
 #include "oneflow/user/kernels/distributions/common.h"
+#include "oneflow/user/kernels/random_seed_util.h"
+#include "oneflow/core/rpc/include/global_process_ctx.h"
 
 namespace oneflow {
 namespace one {
@@ -2105,20 +2107,31 @@ class GlobalNormalFunctor {
       dtype = output_tensor_dtype;
     }
 
-    const auto gen = optional_generator.value_or(JUST(one::DefaultAutoGenerator()));
     auto& attrs = THREAD_CACHED_MUTABLE_ATTR_MAP("mean", "std", "shape", "dtype", "seed", "nd_sbp");
 
-    const auto& distribution_state = std::make_shared<DistributionKernelState>(gen);
     const auto& nd_sbp = JUST(GetNdSbp(sbp_tuple));
 
+    std::shared_ptr<Generator> gen = optional_generator.value_or(JUST(one::DefaultAutoGenerator()));
+    uint64_t init_seed = JUST(gen->Get<CPUGeneratorImpl>(0))->engine()();
+
     if (LazyMode::is_enabled()) {
       attrs.SetAllAttrs(static_cast<double>(mean), static_cast<double>(std), shape,
-                        dtype->data_type(), static_cast<int64_t>(gen->current_seed()),
+                        dtype->data_type(), static_cast<int64_t>(init_seed),
                         *JUST(GetNdSbpStrList(nd_sbp)));
     } else {
+      uint64_t rank_seed = 0;
+      {
+        JUST(BroadcastSeedToAllRanks(&init_seed, /*root=*/0));
+        rank_seed =
+            JUST(GetRandomSeedForRank(*placement, *nd_sbp, init_seed, GlobalProcessCtx::Rank()));
+      }
       attrs.SetAllAttrs(static_cast<double>(mean), static_cast<double>(std), shape,
-                        dtype->data_type(), static_cast<int64_t>(gen->current_seed()), NullOpt);
+                        dtype->data_type(), static_cast<int64_t>(rank_seed), NullOpt);
+      gen = JUST(MakeGenerator(placement->device_type()));
+      gen->set_current_seed(rank_seed);
     }
+    const auto& distribution_state = std::make_shared<DistributionKernelState>(gen);
+
     if (out.has_value()) {
       std::shared_ptr<TensorTuple> outputs = std::make_shared<TensorTuple>(1);
       (*outputs)[0] = JUST(out);

oneflow/core/graph/plan_task_graph.cpp

@@ -19,7 +19,6 @@ namespace oneflow {
 
 int64_t PlanTaskNode::chain_id() const {
   int64_t chain_id = task_proto_->task_set_info().chain_id();
-  CHECK_NE(chain_id, -1);
   return chain_id;
 }
 

oneflow/core/graph/straighten_nodes.cpp

@@ -126,6 +126,7 @@ bool ShouldRunASAP(TaskType task_type) {
     case TaskType::kAcc:                         // 0
     case TaskType::kSourceTick:                  // 0
     case TaskType::kAccTick:                     // 0
+    case TaskType::kAccCtrlTick:                 // ?
     case TaskType::kCase:                        // 0
     case TaskType::kEsac:                        // 0
     case TaskType::kReentrantLock: return true;  // 0