fp8 support in all_gather

pytorch · Oct 14, 2023 · dd4c868 · dd4c868
1 parent 6b0edba
commit dd4c868
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp b/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp
@@ -3342,11 +3342,20 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::_allgather_base(
           c10::cuda::CUDACachingAllocator::recordStream(
               output.storage().data_ptr(), stream);
         }
+
+        ncclDataType_t nccl_dtype;
+        if (input.scalar_type() == at::kFloat8_e5m2 ||
+            input.scalar_type() == at::kFloat8_e4m3fn) {
+          nccl_dtype = ncclInt8;
+        } else {
+          nccl_dtype = getNcclDataType(input.scalar_type());
+        }
+
         return ncclAllGather(
             input.data_ptr(),
             output.data_ptr(),
             input.numel(),
-            getNcclDataType(input.scalar_type()),
+            nccl_dtype,
             comm,
             stream.stream());
       },