Refactor add straighten algo interface

docs/source/graph.rst

@@ -27,6 +27,7 @@ Base class for running neural networks in Static Graph Mode.
             set_zero_redundancy_optimizer_mode,
             set_zero_redundancy_optimizer_min_size_after_split,
             enable_cudnn_conv_heuristic_search_algo,
+            enable_random_straighten_nodes,
     :member-order: bysource
 
 

oneflow/core/graph/task_graph.cpp

@@ -420,7 +420,7 @@ void ForEachOpGraphNecessaryCtrlEdge(
 
 }  // namespace
 
-TaskGraph::TaskGraph() {
+TaskGraph::TaskGraph(bool random_straighten_nodes) {
   OpGraph* op_graph = Global<OpGraph>::Get();
   sub_tsk_gph_builder_ctx_.reset(new SubTskGphBuilderCtx(this));
   boxing_logger_ = CreateBoxingLogger();
@@ -451,7 +451,7 @@ TaskGraph::TaskGraph() {
     }
   });
 
-  if (ParseBooleanFromEnv("ONEFLOW_RANDOM_STRAIGHTEN_NODES", false)) {
+  if (random_straighten_nodes) {
     SetOrderInGraphForEachNode();
   } else {
     StraightenNodes(this, &ordered_task_nodes_);

oneflow/core/graph/task_graph.h

@@ -43,7 +43,7 @@ class TaskGraph final : public Graph<TaskNode, TaskEdge> {
   OF_DISALLOW_COPY_AND_MOVE(TaskGraph);
   ~TaskGraph() override;
 
-  explicit TaskGraph();
+  explicit TaskGraph(bool random_straighten_nodes);
 
   const char* TypeName() const override { return "TaskGraph"; }
   void RemoveEmptyRegsts();

oneflow/core/job/compiler.cpp

@@ -61,7 +61,8 @@ void Compiler::Compile(Job* job, Plan* plan, bool need_job_complete) const {
 
   // Step3: build task_gph.
   // TODO(levi): we can rewrite this part of code in visitor pattern.
-  auto task_gph = std::make_unique<TaskGraph>();
+  auto task_gph =
+      std::make_unique<TaskGraph>(job->job_conf().random_straighten_nodes_in_task_graph());
   using std::placeholders::_1;
   task_gph->ForEachNode(std::bind(&TaskNode::ProduceAllRegstsAndBindEdges, _1));
   task_gph->ForEachNode(std::bind(&TaskNode::ConsumeAllRegsts, _1));

oneflow/core/job/job_conf.proto

@@ -239,6 +239,8 @@ message JobConfigProto {
   optional bool cudnn_conv_enable_pseudo_half = 600 [default = true];
   optional bool enable_auto_mixed_precision = 602 [default = false];
   optional bool enable_quantization_aware_training = 603 [default = false];
+
+  optional bool random_straighten_nodes_in_task_graph = 700 [default = false];
 
   optional int64 concurrency_width = 1000 [default = 128];
 

python/oneflow/nn/graph/graph_config.py

@@ -269,6 +269,16 @@ def build(self, x):
         """
         self.proto.cudnn_conv_heuristic_search_algo = mode
 
+    def enable_random_straighten_nodes(self, mode: bool = False):
+        r""" Whether turn off the straighten algorithm.
+
+        If using nccl compute stream, turning it on will not speed up the training.
+        If not using nccl compute stream, turning it on will slow down data parallelism by 0.6% and slow down model parallelism by 6%.
+
+        The switch is off by default (i.e. use the straighten algorithm by default).
+        """
+        self.proto.random_straighten_nodes_in_task_graph = mode
+
     def _generate_optimizer_and_variable_configs(
         self, opt_dict: OptDict = None, variables_conf: OrderedDict = None,
     ):