msccl: adjust msccl threshold for bf16 (#1736)

* msccl: adjust msccl threshold for bf16 * Update src/collectives.cc Co-authored-by: corey-derochie-amd <161367113+corey-derochie-amd@users.noreply.github.com> --------- Co-authored-by: corey-derochie-amd <161367113+corey-derochie-amd@users.noreply.github.com>
2025-06-11 09:09:57 -05:00
parent 600ace7f19
commit 75c3c8215c
1 changed files with 5 additions and 3 deletions
@@ -125,9 +125,11 @@ ncclResult_t ncclAllReduce_impl(const void* sendbuff, void* recvbuff, size_t cou
  }

  if (mscclAvailable(comm) && !mscclIsCaller()) {
-    return mscclEnqueueCheck(
-      sendbuff, nullptr, nullptr, recvbuff, nullptr, nullptr,
-      count, datatype, 0, 0, op, mscclFuncAllReduce, comm, stream);
+    if (datatype != ncclBfloat16 || (count * ncclTypeSize(datatype) <= 8388608)) {
+	return mscclEnqueueCheck(
+                sendbuff, nullptr, nullptr, recvbuff, nullptr, nullptr,
+                count, datatype, 0, 0, op, mscclFuncAllReduce, comm, stream);	
+    }
  }

  return ncclEnqueueCheck(&info);