fix bug in dist op

Azure · Mar 11, 2024 · 1efec4a · 1efec4a
1 parent 363cb01
commit 1efec4a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 5 deletions.
diff --git a/csrc/dist_op/FindRCCL.cmake b/csrc/dist_op/FindRCCL.cmake
@@ -29,10 +29,10 @@ else()
 endif()
 
 find_path(RCCL_INCLUDE_DIR NAMES rccl.h
-        PATHS ${ROCM_PATH}/include/rccl)
+        PATHS ${ROCM_PATH}/include/rccl /usr/local/include/rccl)
 
 find_library(RCCL_LIBRARY NAMES rccl
-        PATHS ${ROCM_PATH}/lib)
+        PATHS ${ROCM_PATH}/lib /usr/local/lib)
 
 include(FindPackageHandleStandardArgs)
 find_package_handle_standard_args(RCCL DEFAULT_MSG RCCL_INCLUDE_DIR RCCL_LIBRARY)

diff --git a/csrc/dist_op/dist.cpp b/csrc/dist_op/dist.cpp
@@ -8,9 +8,6 @@
 #include <hip/hip_runtime.h>
 #include <rccl.h>
 #define stream_t hipStream_t
-// TODO: Remove ncclFp8E4M3 and ncclFp8E5M2 after rccl supports FP8
-#define ncclFp8E4M3 ncclUint8
-#define ncclFp8E5M2 ncclInt8
 #else
 #include <nccl.h>
 #define stream_t cudaStream_t