WIP

anthonix · Jun 25, 2024 · 018738b · 018738b
1 parent 3829039
commit 018738b
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/llmc/zero.cuh b/llmc/zero.cuh
@@ -17,14 +17,16 @@ Utilities for ZeRO sharding
 #include <stdio.h>
 #include <stddef.h>
 
-#ifndef BUILD_AMD
 #ifdef MULTI_GPU
+#ifdef BUILD_AMD
+#include <rccl/rccl.h>
+#else
 #include <nccl.h>
+#endif
 #ifdef USE_MPI
 #include <mpi.h>
 #endif
 #endif
-#endif
 
 // ----------------------------------------------------------------------------
 // Multi-GPU related