make AllReduce ctring fallback to ctdirect if message size is small

Summary: As title.

Reviewed By: function47

Differential Revision: D85377269

fbshipit-source-id: 1df7f76d9ee632ddd0b0d55e93ba00c10842af91

Author: arttianezhu

comms/ctran/algos/AllReduce/AllReduce.cc

@@ -39,6 +39,15 @@ commResult_t ctranAllReduce(
         return ctranAllReduceDirect(
             sendbuff, recvbuff, count, datatype, redOp, comm, stream, timeout);
       }
+      if (count < comm->statex_->nRanks()) {
+        CLOGF(
+            WARN,
+            "AllReduce ctring requires count {} > nRanks {}, fallback to ctdirect",
+            count,
+            comm->statex_->nRanks());
+        return ctranAllReduceDirect(
+            sendbuff, recvbuff, count, datatype, redOp, comm, stream, timeout);
+      }
       return ctranAllReduceRing(
           sendbuff, recvbuff, count, datatype, redOp, comm, stream, timeout);
     case NCCL_ALLREDUCE_ALGO::ctdirect:

comms/ctran/tests/CtranAllReduceTest.cc

@@ -149,6 +149,17 @@ TEST_P(CtranAllReduceTest, BasicRunAbortEnabled) {
   }
 }
 
+TEST_P(CtranAllReduceTest, SmallMessageSize) {
+  auto [algoName, algo] = GetParam();
+  NCCL_ALLREDUCE_ALGO = algo;
+
+  startWorkers(/*abortEnabled=*/true);
+  for (int rank = 0; rank < kNRanks; ++rank) {
+    run(rank,
+        [this](PerRankState& state) { runAllReduce(/*nElem=*/1, state); });
+  }
+}
+
 void CtranAllReduceTest::runTestRanksAbsent(
     std::vector<int> ranksToRunCollective,
     std::vector<int> ranksAbsent,
@@ -425,78 +436,89 @@ class CtranAllReduceRingOneRankTest : public CtranStandaloneMultiRankBaseTest {
 
     CtranStandaloneMultiRankBaseTest::SetUp();
   }
-};
 
-TEST_F(CtranAllReduceRingOneRankTest, Basic) {
-  ASSERT_EQ(NCCL_ALLREDUCE_ALGO, NCCL_ALLREDUCE_ALGO::ctring);
-
-  CtranStandaloneMultiRankBaseTest::startWorkers(
-      kNRanks, /*aborts=*/{ctran::utils::createAbort(/*enabled=*/true)});
-
-  run(/*rank=*/0, [this](PerRankState& state) {
-    // set up src buffer to hold magic values, and zero out dst buffers
-    int magic = 0xdeadbeef;
-    int srcHost[kBufferNElem];
-    int dstHost[kBufferNElem];
-    for (int i = 0; i < kBufferNElem; ++i) {
-      srcHost[i] = magic + i;
-    }
-    memset(dstHost, 0, kBufferSize);
-    ASSERT_EQ(
-        cudaSuccess,
-        cudaMemcpy(
-            state.srcBuffer, srcHost, kBufferSize, cudaMemcpyHostToDevice));
-    ASSERT_EQ(cudaSuccess, cudaMemset(state.dstBuffer, 0, kBufferSize));
+  void runAllReduce(size_t nElem) {
+    ASSERT_EQ(NCCL_ALLREDUCE_ALGO, NCCL_ALLREDUCE_ALGO::ctring);
 
-    // warmup
-    void* srcHandle;
-    void* dstHandle;
-    ASSERT_EQ(
-        commSuccess,
-        state.ctranComm->ctran_->commRegister(
-            state.srcBuffer, kBufferSize, &srcHandle));
-    ASSERT_EQ(
-        commSuccess,
-        state.ctranComm->ctran_->commRegister(
-            state.dstBuffer, kBufferSize, &dstHandle));
-    SCOPE_EXIT {
-      // deregistering will happen after streamSync below
-      state.ctranComm->ctran_->commDeregister(dstHandle);
-      state.ctranComm->ctran_->commDeregister(srcHandle);
-    };
+    CtranStandaloneMultiRankBaseTest::startWorkers(
+        kNRanks, /*aborts=*/{ctran::utils::createAbort(/*enabled=*/true)});
 
-    CLOGF(INFO, "rank {} allReduce completed registration", state.rank);
+    run(/*rank=*/0, [this, nElem](PerRankState& state) {
+      // set up src buffer to hold magic values, and zero out dst buffers
+      int magic = 0xdeadbeef;
+      int srcHost[kBufferNElem];
+      int dstHost[kBufferNElem];
+      for (int i = 0; i < kBufferNElem; ++i) {
+        srcHost[i] = magic + i;
+      }
+      memset(dstHost, 0, kBufferSize);
+      ASSERT_EQ(
+          cudaSuccess,
+          cudaMemcpy(
+              state.srcBuffer, srcHost, kBufferSize, cudaMemcpyHostToDevice));
+      ASSERT_EQ(cudaSuccess, cudaMemset(state.dstBuffer, 0, kBufferSize));
 
-    EXPECT_EQ(
-        commSuccess,
-        ctranAllReduce(
-            state.srcBuffer,
-            state.dstBuffer,
-            kBufferNElem,
-            kDataType,
-            kReduceOpType,
-            state.ctranComm.get(),
-            state.stream,
-            std::nullopt,
-            /*timeout=*/std::nullopt));
+      // warmup
+      void* srcHandle;
+      void* dstHandle;
+      ASSERT_EQ(
+          commSuccess,
+          state.ctranComm->ctran_->commRegister(
+              state.srcBuffer, kBufferSize, &srcHandle));
+      ASSERT_EQ(
+          commSuccess,
+          state.ctranComm->ctran_->commRegister(
+              state.dstBuffer, kBufferSize, &dstHandle));
+      SCOPE_EXIT {
+        // deregistering will happen after streamSync below
+        state.ctranComm->ctran_->commDeregister(dstHandle);
+        state.ctranComm->ctran_->commDeregister(srcHandle);
+      };
+
+      CLOGF(INFO, "rank {} allReduce completed registration", state.rank);
+
+      EXPECT_EQ(
+          commSuccess,
+          ctranAllReduce(
+              state.srcBuffer,
+              state.dstBuffer,
+              nElem,
+              kDataType,
+              kReduceOpType,
+              state.ctranComm.get(),
+              state.stream,
+              std::nullopt,
+              /*timeout=*/std::nullopt));
+
+      CLOGF(INFO, "rank {} allReduce scheduled", state.rank);
+
+      // ensure async execution completion and no error
+      EXPECT_EQ(cudaSuccess, cudaStreamSynchronize(state.stream));
+      EXPECT_EQ(commSuccess, state.ctranComm->getAsyncResult());
 
-    CLOGF(INFO, "rank {} allReduce scheduled", state.rank);
+      CLOGF(INFO, "rank {} allReduce task completed", state.rank);
 
-    // ensure async execution completion and no error
-    EXPECT_EQ(cudaSuccess, cudaStreamSynchronize(state.stream));
-    EXPECT_EQ(commSuccess, state.ctranComm->getAsyncResult());
+      // validate results
+      ASSERT_EQ(
+          cudaSuccess,
+          cudaMemcpy(
+              dstHost, state.dstBuffer, kBufferSize, cudaMemcpyDeviceToHost));
+      for (int i = 0; i < nElem; ++i) {
+        EXPECT_EQ(srcHost[i], dstHost[i]);
+      }
+      for (int i = nElem; i < kBufferNElem; ++i) {
+        EXPECT_EQ(dstHost[i], 0);
+      }
+    });
+  }
+};
 
-    CLOGF(INFO, "rank {} allReduce task completed", state.rank);
+TEST_F(CtranAllReduceRingOneRankTest, Basic) {
+  this->runAllReduce(/*nElem=*/kBufferNElem);
+}
 
-    // validate results
-    ASSERT_EQ(
-        cudaSuccess,
-        cudaMemcpy(
-            dstHost, state.dstBuffer, kBufferSize, cudaMemcpyDeviceToHost));
-    for (int i = 0; i < kBufferNElem; ++i) {
-      EXPECT_EQ(srcHost[i], dstHost[i]);
-    }
-  });
+TEST_F(CtranAllReduceRingOneRankTest, SmallMessageSize) {
+  this->runAllReduce(/*nElem=*/1);
 }
 
 } // namespace ctran::testing