Fixed issues found on AMD (#271)

microsoft · Mar 8, 2024 · 96e31bf · 96e31bf
1 parent d72f09f
commit 96e31bf
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 5 deletions.
diff --git a/apps/nccl/src/nccl.cu b/apps/nccl/src/nccl.cu
@@ -147,10 +147,8 @@ static std::shared_ptr<mscclpp::DeviceHandle<mscclpp::SmChannel>> setupSmChannel
                  [](const mscclpp::SmChannel& smChannel) { return mscclpp::deviceHandle(smChannel); });
   std::shared_ptr<mscclpp::DeviceHandle<mscclpp::SmChannel>> ptr =
       mscclpp::allocSharedCuda<mscclpp::DeviceHandle<mscclpp::SmChannel>>(smChannelDeviceHandles.size());
-  mscclpp::AvoidCudaGraphCaptureGuard guard;
-  CUDACHECK(cudaMemcpy(ptr.get(), smChannelDeviceHandles.data(),
-                       sizeof(mscclpp::DeviceHandle<mscclpp::SmChannel>) * smChannelDeviceHandles.size(),
-                       cudaMemcpyHostToDevice));
+  mscclpp::memcpyCuda<mscclpp::DeviceHandle<mscclpp::SmChannel>>(ptr.get(), smChannelDeviceHandles.data(),
+                       smChannelDeviceHandles.size(), cudaMemcpyHostToDevice);
   return ptr;
 }
 

diff --git a/include/mscclpp/gpu.hpp b/include/mscclpp/gpu.hpp
@@ -6,7 +6,7 @@
 
 #if defined(__HIP_PLATFORM_AMD__)
 
-#include <hip/hip_bf16.h>
+// #include <hip/hip_bf16.h>
 #include <hip/hip_fp16.h>
 #include <hip/hip_runtime.h>
 

diff --git a/src/fifo.cc b/src/fifo.cc
@@ -56,6 +56,7 @@ MSCCLPP_API_CPP void Fifo::pop() {
 MSCCLPP_API_CPP void Fifo::flushTail(bool sync) {
   // Flush the tail to device memory. This is either triggered every ProxyFlushPeriod to make sure that the fifo can
   // make progress even if there is no request mscclppSync. However, mscclppSync type is for flush request.
+  AvoidCudaGraphCaptureGuard cgcGuard;
   MSCCLPP_CUDATHROW(cudaMemcpyAsync(pimpl->tailReplica.get(), &pimpl->hostTail, sizeof(uint64_t),
                                     cudaMemcpyHostToDevice, pimpl->stream));
   if (sync) {