Feat general basic communication

oneflow/core/auto_parallel/boxing_collector.cpp

@@ -18,6 +18,7 @@ limitations under the License.
 #include <string>
 #include "oneflow/core/auto_parallel/boxing_collector.h"
 #include "oneflow/core/common/data_type.h"
+#include "oneflow/core/common/device_type.pb.h"
 #include "oneflow/core/common/maybe.h"
 #include "oneflow/core/framework/nd_sbp.h"
 #include "oneflow/core/job/global_for.h"
@@ -49,7 +50,7 @@ void DfsSetNdSbp(const std::vector<::oneflow::SbpParallel>& id2sbp_parallel, int
 }
 
 // Let a nd sbp be consistent with the given hierarchy number
-Maybe<NdSbp> SetNdSbpDim(NdSbp nd_sbp, int32_t hierarchy_num) {
+Maybe<NdSbp> SetNdSbpDim(const NdSbp& nd_sbp, int32_t hierarchy_num) {
   // Do not need to change
   if (nd_sbp.sbp_parallel_size() == hierarchy_num) { return nd_sbp; }
   // (S0, S0) -> S0
@@ -71,6 +72,60 @@ Maybe<NdSbp> SetNdSbpDim(NdSbp nd_sbp, int32_t hierarchy_num) {
   return new_sbp;
 }
 
+int32_t TotalNumSplit(const NdSbp& nd_sbp, const ParallelDesc& parallel_desc) {
+  int32_t total_num_split = 1;
+  for (int32_t i = 0; i < nd_sbp.sbp_parallel_size(); i++) {
+    if (nd_sbp.sbp_parallel(i).has_split_parallel()) {
+      total_num_split *= parallel_desc.hierarchy()->At(i);
+    }
+  }
+  return total_num_split;
+}
+
+// Dealing with 1D sbp to 1D sbp
+// Specifically, S -> P.
+Maybe<void> AskSbpCombinationFor1DSbp(const NdSbp& sbp_producer, const NdSbp& sbp_consumer,
+                                      const ParallelDesc& producer_parallel_desc,
+                                      const ParallelDesc& consumer_parallel_desc,
+                                      std::vector<NdSbp>& middle_sbps, int32_t* diag_node_pos) {
+  if (sbp_consumer.sbp_parallel(0).has_partial_sum_parallel()) {
+    // Support [4]: P <--> [2, 2]: (P, P)
+    // Support {0, 1, 2, 3}: P <--> {2, 0, 6, 7}: (P, P)
+    if (producer_parallel_desc.parallel_num() == consumer_parallel_desc.parallel_num()
+        && sbp_producer.sbp_parallel(0).has_partial_sum_parallel()) {
+      return Maybe<void>::Ok();
+    }
+
+    if (!sbp_producer.sbp_parallel(0).has_broadcast_parallel()) {
+      // S -> B -> P (Large cost!)
+      // TODO: Please implement S -> P directly.
+      // We do not support [3]: P <--> [2, 2]: (P, P) as well.
+
+      int32_t hierarchy_size = 0;
+      if (producer_parallel_desc.hierarchy()->elem_cnt()
+          < consumer_parallel_desc.hierarchy()->elem_cnt()) {
+        // The diagonal node uses the parallel description from producer
+        // (S, S) -> (B, B) -> P/(P, P) or S -> B -> P/(P, P)
+        *diag_node_pos = 1;
+        hierarchy_size = producer_parallel_desc.hierarchy()->NumAxes();
+      } else {
+        // The diagonal node uses the parallel description from consumer
+        // S/(S, S) -> B -> P or S/(S, S) -> (B, B) -> (P, P)
+        *diag_node_pos = 0;
+        hierarchy_size = consumer_parallel_desc.hierarchy()->NumAxes();
+      }
+
+      NdSbp broadcast_nd;
+      for (int32_t i = 0; i < hierarchy_size; i++) {
+        broadcast_nd.add_sbp_parallel();
+        broadcast_nd.mutable_sbp_parallel(i)->mutable_broadcast_parallel();
+      }
+      middle_sbps.emplace_back(broadcast_nd);
+    }
+  }
+  return Maybe<void>::Ok();
+}
+
 }  // namespace
 
 // A constructor with init, designed for uncustomized boxing collector
@@ -190,10 +245,13 @@ Maybe<void> BoxingCollector::GenerateCombination4SamePlacement(int32_t max_middl
   // NOTE: The performance of this function are all the same with different hierarchy
   int32_t world_size = GlobalProcessCtx::WorldSize();
   Shape hierarchy44({4 * world_size, 4 * world_size});
+  int32_t virtual_range_size = hierarchy44.elem_cnt();
   std::shared_ptr<Shape> virtual_hierarchy = std::make_shared<Shape>(hierarchy44);
   auto parallel_desc = JUST(ParallelDesc::New(
       "cpu", {"0:0-" + std::to_string(hierarchy44.elem_cnt() - 1)}, virtual_hierarchy));
-  BlobDesc blob_desc({16, 16, 16, 16}, DataType::kInt8, /*is_dynamic=*/false);
+  BlobDesc blob_desc({virtual_range_size, virtual_range_size, virtual_range_size,
+                      virtual_range_size, virtual_range_size, virtual_range_size},
+                     DataType::kInt8, /*is_dynamic=*/false);
   JUST(GenerateCombination4SamePlacement(max_middle_node_num, blob_desc, *parallel_desc));
   return Maybe<void>::Ok();
 }
@@ -309,7 +367,10 @@ Maybe<void> BoxingCollector::GenerateCombination4DiffPlacement(
     BoxingCollector* boxing_collector_producer, BoxingCollector* boxing_collector_consumer) {
   // Virtual parallel and blob description
   int32_t world_size = GlobalProcessCtx::WorldSize();
-  BlobDesc blob_desc({16, 16, 16, 16}, DataType::kInt8, /*is_dynamic=*/false);
+  int32_t virtual_range_size = 4 * world_size * (4 * world_size + 1);
+  BlobDesc blob_desc({virtual_range_size, virtual_range_size, virtual_range_size,
+                      virtual_range_size, virtual_range_size, virtual_range_size},
+                     DataType::kInt8, /*is_dynamic=*/false);
   // Virtual placements before transfer
   Shape in_hierarchy44({4 * world_size + 1, 4 * world_size});
   std::shared_ptr<Shape> in_hierarchy = std::make_shared<Shape>(in_hierarchy44);
@@ -496,66 +557,39 @@ Maybe<void> BoxingCollector::AskSbpCombination(const NdSbp& sbp_producer, const
   if (ParseBooleanFromEnv("ONEFLOW_BOXING_DISABLE_MIDDLE_NODE_AND_CHECK", false)) {
     return Maybe<void>::Ok();
   }
-  // If compute_cost==false + 2D sbp + same placment + nccl logical + not (p->b),
-  // Use nccl logical send recv instead of middle node.
-  // Note that in op sbp inference, cost of middle nodes is still used for the moment.
-#ifdef WITH_CUDA
-  if (compute_cost == false && producer_parallel_desc.hierarchy()->NumAxes() == 2
-      && producer_parallel_desc == consumer_parallel_desc
-      && !(NdSbpHasPartialParallel(sbp_consumer)) &&
-      // TODO(): When same dim 0 finished dealing with (*, P) -> (*, S) in nccl logical pass, open
-      // this condition. When dealing with (P, P) -> (B, S0), middle node will change it to (P, P)
-      // -> (P, S0) -> (B, S0), neither same dim 0 or send recv in nccl logical pass can deal with
-      // (P, P) -> (P, S0) at the moment.
-      // !(NdSbpHasPartialParallel(sbp_producer) && NdSbpHasBroadcastParallel(sbp_consumer)) &&
-      Singleton<ResourceDesc, ForSession>::Get()->nccl_use_compute_stream()) {
-    VLOG(3) << "Middle node insertion is skipped when src sbp is " << NdSbpToString(sbp_producer)
-            << " dst sbp is " << NdSbpToString(sbp_consumer)
-            << ", because nccl logical send/recv can handle this.";
+  if (producer_parallel_desc == consumer_parallel_desc && sbp_producer == sbp_consumer) {
     return Maybe<void>::Ok();
   }
-#endif  // WITH_CUDA
 
   // Dealing with 1D sbp to 1D sbp
-  // Specifically, S -> P.
   if (Is1dSbp(sbp_producer) && Is1dSbp(sbp_consumer)) {
-    if (sbp_consumer.sbp_parallel(0).has_partial_sum_parallel()) {
-      // Support [4]: P <--> [2, 2]: (P, P)
-      // Support {0, 1, 2, 3}: P <--> {2, 0, 6, 7}: (P, P)
-      if (producer_parallel_desc.parallel_num() == consumer_parallel_desc.parallel_num()
-          && sbp_producer.sbp_parallel(0).has_partial_sum_parallel()) {
-        return Maybe<void>::Ok();
-      }
-
-      if (!sbp_producer.sbp_parallel(0).has_broadcast_parallel()) {
-        // S -> B -> P (Large cost!)
-        // TODO: Please implement S -> P directly.
-        // We do not support [3]: P <--> [2, 2]: (P, P) as well.
-
-        int32_t hierarchy_size = 0;
-        if (producer_parallel_desc.hierarchy()->elem_cnt()
-            < consumer_parallel_desc.hierarchy()->elem_cnt()) {
-          // The diagonal node uses the parallel description from producer
-          // (S, S) -> (B, B) -> P/(P, P) or S -> B -> P/(P, P)
-          *diag_node_pos = 1;
-          hierarchy_size = producer_parallel_desc.hierarchy()->NumAxes();
-        } else {
-          // The diagonal node uses the parallel description from consumer
-          // S/(S, S) -> B -> P or S/(S, S) -> (B, B) -> (P, P)
-          *diag_node_pos = 0;
-          hierarchy_size = consumer_parallel_desc.hierarchy()->NumAxes();
-        }
+    JUST(AskSbpCombinationFor1DSbp(sbp_producer, sbp_consumer, producer_parallel_desc,
+                                   consumer_parallel_desc, middle_sbps, diag_node_pos));
+    // No middle nodes for the other 1d-sbp combinations
+    return Maybe<void>::Ok();
+  }
 
-        NdSbp broadcast_nd;
-        for (int32_t i = 0; i < hierarchy_size; i++) {
-          broadcast_nd.add_sbp_parallel();
-          broadcast_nd.mutable_sbp_parallel(i)->mutable_broadcast_parallel();
-        }
-        middle_sbps.emplace_back(broadcast_nd);
-      }
-      return Maybe<void>::Ok();
+#ifdef WITH_CUDA
+  static const bool enable_general_basic_communication =
+      ParseBooleanFromEnv("ONEFLOW_BOXING_ENABLE_GENERAL_BASIC_COMMUNICATION", false);
+  // Use a general basic communication if no P in the consumer
+  if ((Singleton<ResourceDesc, ForSession>::Get()->nccl_use_compute_stream()
+       || enable_general_basic_communication)
+      && (!NdSbpHasPartialParallel(sbp_consumer))
+      && producer_parallel_desc.device_type() == DeviceType::kCUDA
+      && consumer_parallel_desc.device_type() == DeviceType::kCUDA) {
+    if (NdSbpHasPartialParallel(sbp_producer) && NdSbpHasBroadcastParallel(sbp_consumer)) {
+      // (?, P, ?)->(Si, Sj)->(?, B, ?), two-step transfer
+      // Directly applying general basic communication would have O(n^2) time complexity for P->B
+      // Using two-step transfer would reduce it to a linear cost
+      JUST(AskSbpCombination4GeneralBasicCommunication(
+          sbp_producer, sbp_consumer, logical_blob_desc, producer_parallel_desc,
+          consumer_parallel_desc, middle_sbps, diag_node_pos));
     }
+    // Otherwise, one-step transfer
+    return Maybe<void>::Ok();
   }
+#endif  // WITH_CUDA
 
   // Middle nodes algorithm supports transfer for different machines or devices or hierarchies
   if (producer_parallel_desc != consumer_parallel_desc) {
@@ -568,6 +602,7 @@ Maybe<void> BoxingCollector::AskSbpCombination(const NdSbp& sbp_producer, const
   // Transfer for the same machines, devices and hierarchy.
   if (sbp_producer == sbp_consumer) { return Maybe<void>::Ok(); }
   const auto& parallel_hierarchy = producer_parallel_desc.hierarchy();
+
   *diag_node_pos = 0;
   // Dealing with nD sbp, n>2
   if (parallel_hierarchy->NumAxes() > 2) {
@@ -1007,4 +1042,105 @@ Maybe<void> BoxingCollector::FilterNdSbpList4LogicalShape(const BlobDesc& logica
   return Maybe<void>::Ok();
 }
 
+// Ask for sbp combination for general basic communication
+Maybe<void> BoxingCollector::AskSbpCombination4GeneralBasicCommunication(
+    const NdSbp& sbp_producer, const NdSbp& sbp_consumer, const BlobDesc& logical_blob_desc,
+    const ParallelDesc& producer_parallel_desc, const ParallelDesc& consumer_parallel_desc,
+    std::vector<NdSbp>& middle_sbps, int32_t* diag_node_pos) {
+  // (P, X) -> (B, X) || (X , P) -> (X, B), X is any SBP
+  // One step transfer, at most 50% reduction in the transfer cost, do not use middle nodes
+  if (producer_parallel_desc == consumer_parallel_desc
+      && producer_parallel_desc.hierarchy()->NumAxes() == 2
+      && (sbp_producer.sbp_parallel(0) == sbp_consumer.sbp_parallel(0)
+          || sbp_producer.sbp_parallel(1) == sbp_consumer.sbp_parallel(1))) {
+    return Maybe<void>::Ok();
+  }
+
+  // Not enough gain in transfer cost, do not use middle nodes
+  int32_t partial_ratio4producer = PartialRatio4Producer(sbp_producer, producer_parallel_desc);
+  int32_t broadcast_ratio4consumer = BroadcastRatio4Consumer(sbp_consumer, consumer_parallel_desc);
+  if (2 * (partial_ratio4producer + broadcast_ratio4consumer)
+      >= partial_ratio4producer * broadcast_ratio4consumer) {
+    return Maybe<void>::Ok();
+  }
+
+  bool close2producer = true;
+  if (producer_parallel_desc.parallel_num() == consumer_parallel_desc.parallel_num()) {
+    // Get close to the one with more splits
+    close2producer = TotalNumSplit(sbp_producer, producer_parallel_desc)
+                     > TotalNumSplit(sbp_consumer, consumer_parallel_desc);
+  } else {
+    // Get close to the one with more machines
+    close2producer = producer_parallel_desc.parallel_num() > consumer_parallel_desc.parallel_num();
+  }
+  // Get the contiguous sbp
+  if (close2producer) {
+    JUST(AskCloseAllSplitSbp(sbp_producer, producer_parallel_desc, logical_blob_desc, middle_sbps));
+    *diag_node_pos = 1;
+  } else {
+    JUST(AskCloseAllSplitSbp(sbp_consumer, consumer_parallel_desc, logical_blob_desc, middle_sbps));
+    *diag_node_pos = 0;
+  }
+  return Maybe<void>::Ok();
+}
+
+// Ask for a all-split sbp which is close to the original one
+Maybe<void> BoxingCollector::AskCloseAllSplitSbp(const NdSbp& nd_sbp,
+                                                 const ParallelDesc& parallel_desc,
+                                                 const BlobDesc& logical_blob_desc,
+                                                 std::vector<NdSbp>& middle_sbps) {
+  Shape remain_shape = logical_blob_desc.shape();
+  Shape rest_split_shape = logical_blob_desc.shape();
+  int32_t dim_shape = remain_shape.NumAxes();
+  // Initialize the remains and splitting
+  // logical_blob_desc.shape() == remain_shape .* rest_split_shape;
+  for (int32_t i = 0; i < dim_shape; i++) { rest_split_shape.Set(i, 1); }
+  for (int32_t sbp_id = 0; sbp_id < nd_sbp.sbp_parallel_size(); sbp_id++) {
+    const auto& sbp = nd_sbp.sbp_parallel(sbp_id);
+    if (sbp.has_split_parallel()) {
+      int32_t axis = sbp.split_parallel().axis();
+      int32_t split_num = parallel_desc.hierarchy()->At(sbp_id);
+      remain_shape.Set(axis, remain_shape.At(axis) / split_num);
+      rest_split_shape.Set(axis, rest_split_shape.At(axis) * split_num);
+    }
+  }
+  // Get the contiguous sbp
+  NdSbp new_sbp = nd_sbp;
+  for (int32_t sbp_id = 0; sbp_id < nd_sbp.sbp_parallel_size(); sbp_id++) {
+    const auto& sbp = nd_sbp.sbp_parallel(sbp_id);
+    int32_t split_num = parallel_desc.hierarchy()->At(sbp_id);
+    if (sbp.has_split_parallel()) {
+      int32_t axis = sbp.split_parallel().axis();
+      // split shape is the total splitting number starting from sbp_id to the end
+      rest_split_shape.Set(axis, rest_split_shape.At(axis) / split_num);
+    } else {
+      // change P or B to S(axis)
+      int32_t axis = -1;
+      // 4096 is large enough, we might not have that much devices
+      int32_t min_split_num = 4096;
+      // We need to pick a suitable axis
+      for (int32_t i = 0; i < remain_shape.NumAxes(); i++) {
+        if (remain_shape.At(i) % split_num == 0) {
+          if (rest_split_shape.At(i) < min_split_num) {
+            // Pick the axis with smallest splitting number among the rest of the sbp
+            min_split_num = rest_split_shape.At(i);
+            axis = i;
+          }
+        }
+      }
+      // P, B -> S(axis)
+      if (axis >= 0) {
+        new_sbp.mutable_sbp_parallel(sbp_id)->mutable_split_parallel()->set_axis(axis);
+        remain_shape.Set(axis, remain_shape.At(axis) / split_num);
+      } else {
+        // Can not find a suitable contiguous sbp
+        return Maybe<void>::Ok();
+      }
+    }
+  }
+  // Add the new sbp into the middle node lists
+  middle_sbps.emplace_back(new_sbp);
+  return Maybe<void>::Ok();
+}
+
 }  // namespace oneflow

oneflow/core/auto_parallel/boxing_collector.h

@@ -129,6 +129,15 @@ class BoxingCollector final {
                                             BoxingCollector* boxing_collector_producer,
                                             BoxingCollector* boxing_collector_consumer,
                                             const std::vector<std::vector<int32_t>>& diag_nodes);
+  // Ask for sbp combination for general basic communication
+  Maybe<void> AskSbpCombination4GeneralBasicCommunication(
+      const NdSbp& sbp_producer, const NdSbp& sbp_consumer, const BlobDesc& logical_blob_desc,
+      const ParallelDesc& producer_parallel_desc, const ParallelDesc& consumer_parallel_desc,
+      std::vector<NdSbp>& middle_sbps, int32_t* diag_node_pos);
+  // Ask for a all-split sbp which is closed to the original one
+  Maybe<void> AskCloseAllSplitSbp(const NdSbp& nd_sbp, const ParallelDesc& parallel_desc,
+                                  const BlobDesc& logical_blob_desc,
+                                  std::vector<NdSbp>& middle_sbps);
   // Stores all the possible SbpParallel.
   HashMap<::oneflow::SbpParallel, int32_t> sbp_parallel_universe_;
   // Relationship between id and Sbp Parallel