[AMP]register bf16 for communication ops (#52555)

* register bf16 for communication ops

* fix bfloat16 type finding compile error in c_allreduce_max_op

paddle/fluid/operators/collective/c_allreduce_max_op.cu.cc

@@ -28,7 +28,11 @@ PD_REGISTER_STRUCT_KERNEL(c_allreduce_max,
                           ALL_LAYOUT,
                           ops::CAllReduceMaxCUDAKernel,
                           float,
+#if NCCL_VERSION_CODE >= 21000
+                          plat::bfloat16,
+#endif
                           double,
                           int,
                           int64_t,
-                          plat::float16) {}
+                          plat::float16) {
+}

paddle/fluid/operators/collective/c_split_op.cu

@@ -123,4 +123,8 @@ PD_REGISTER_STRUCT_KERNEL(c_split,
                           double,
                           int,
                           int64_t,
-                          plat::float16) {}
+#if NCCL_VERSION_CODE >= 21000
+                          plat::bfloat16,
+#endif
+                          plat::float16) {
+}

paddle/fluid/operators/collective/partial_allgather_op.cu.cc

@@ -104,6 +104,9 @@ namespace plat = paddle::platform;
 
 REGISTER_OP_CUDA_KERNEL(partial_allgather,
                         ops::PartialAllGatherOpCUDAKernel<float>,
+#if NCCL_VERSION_CODE >= 21000
+                        ops::PartialAllGatherOpCUDAKernel<plat::bfloat16>,
+#endif
                         ops::PartialAllGatherOpCUDAKernel<double>,
                         ops::PartialAllGatherOpCUDAKernel<int>,
                         ops::PartialAllGatherOpCUDAKernel<int64_t>,

paddle/fluid/operators/collective/partial_recv_op.cu.cc

@@ -120,6 +120,9 @@ namespace plat = paddle::platform;
 
 REGISTER_OP_CUDA_KERNEL(partial_recv,
                         ops::PartialRecvOpCUDAKernel<float>,
+#if NCCL_VERSION_CODE >= 21000
+                        ops::PartialRecvOpCUDAKernel<plat::bfloat16>,
+#endif
                         ops::PartialRecvOpCUDAKernel<double>,
                         ops::PartialRecvOpCUDAKernel<int>,
                         ops::PartialRecvOpCUDAKernel<int64_t>,

paddle/fluid/operators/collective/partial_send_op.cu.cc

@@ -120,6 +120,9 @@ namespace plat = paddle::platform;
 REGISTER_OP_CUDA_KERNEL(partial_send,
                         ops::PartialSendCUDAKernel<float>,
                         ops::PartialSendCUDAKernel<double>,
+#if NCCL_VERSION_CODE >= 21000
+                        ops::PartialSendCUDAKernel<plat::bfloat16>,
+#endif
                         ops::PartialSendCUDAKernel<int>,
                         ops::PartialSendCUDAKernel<int64_t>,
                         ops::PartialSendCUDAKernel<plat::float16>);