[DCU] LLM train surpport bf16 (#65857)

yuguo-Jack · web-flow · commit 805cc7ae5bf4 · 2024-07-18T09:40:34.000+08:00
* [DCU] New features for LLM

* fix WeightQuantizeInferMeta

* [DCU] fix flashattn cmake

* fix

* fix code style

* [DCU] add FLAGS_batch_norm_use_miopen

* fix

* fix setup.py.in

* [DCU] LLM bf16

* fix code style

* fix

* fix
diff --git a/paddle/fluid/operators/collective/alltoall_op.cu.cc b/paddle/fluid/operators/collective/alltoall_op.cu.cc
@@ -146,7 +146,8 @@ PD_REGISTER_STRUCT_KERNEL(alltoall,
                           ops::AllToAllOpCUDAKernel,
                           float,
                           double,
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
                           phi::dtype::bfloat16,
 #endif
                           int,
diff --git a/paddle/fluid/operators/collective/c_allgather_op.cu.cc b/paddle/fluid/operators/collective/c_allgather_op.cu.cc
@@ -129,7 +129,8 @@ PD_REGISTER_STRUCT_KERNEL(c_allgather,
                           ops::CAllGatherOpCUDAKernel,
                           float,
                           double,
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
                           phi::dtype::bfloat16,
 #endif
                           int,
diff --git a/paddle/fluid/operators/collective/c_allreduce_max_op.cu.cc b/paddle/fluid/operators/collective/c_allreduce_max_op.cu.cc
@@ -27,7 +27,8 @@ PD_REGISTER_STRUCT_KERNEL(c_allreduce_max,
                           ALL_LAYOUT,
                           ops::CAllReduceMaxCUDAKernel,
                           float,
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
                           phi::dtype::bfloat16,
 #endif
                           double,
diff --git a/paddle/fluid/operators/collective/c_allreduce_op.h b/paddle/fluid/operators/collective/c_allreduce_op.h
@@ -411,7 +411,8 @@ class CAllReduceOpCUDAKernel : public framework::OpKernel<T> {
         nccl_red_type = ncclProd;
         break;
 
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
       case kRedAvg:
         nccl_red_type = ncclAvg;
         break;
diff --git a/paddle/fluid/operators/collective/c_allreduce_sum_op.cu.cc b/paddle/fluid/operators/collective/c_allreduce_sum_op.cu.cc
@@ -27,7 +27,8 @@ PD_REGISTER_STRUCT_KERNEL(c_allreduce_sum,
                           ALL_LAYOUT,
                           ops::CAllReduceSumCUDAKernel,
                           float,
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
                           phi::dtype::bfloat16,
 #endif
                           double,
diff --git a/paddle/fluid/operators/collective/c_broadcast_op.cu.cc b/paddle/fluid/operators/collective/c_broadcast_op.cu.cc
@@ -110,7 +110,8 @@ PD_REGISTER_STRUCT_KERNEL(c_broadcast,
                           int64_t,
                           float,
                           double,
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
                           phi::dtype::bfloat16,
 #endif
                           phi::dtype::float16) {
diff --git a/paddle/fluid/operators/collective/c_concat_op.cu.cc b/paddle/fluid/operators/collective/c_concat_op.cu.cc
@@ -174,7 +174,8 @@ PD_REGISTER_STRUCT_KERNEL(c_concat,
                           double,
                           int,
                           int64_t,
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
                           phi::dtype::bfloat16,
 #endif
                           phi::dtype::float16) {
diff --git a/paddle/fluid/operators/collective/c_reduce_op.h b/paddle/fluid/operators/collective/c_reduce_op.h
@@ -302,7 +302,8 @@ class CReduceOpCUDAKernel : public framework::OpKernel<T> {
         nccl_red_type = ncclProd;
         break;
 
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
       case kRedAvg:
         nccl_red_type = ncclAvg;
         break;
diff --git a/paddle/fluid/operators/collective/c_reducescatter_op.cu.cc b/paddle/fluid/operators/collective/c_reducescatter_op.cu.cc
@@ -134,7 +134,8 @@ PD_REGISTER_STRUCT_KERNEL(c_reducescatter,
                           ops::CReduceScatterOpCUDAKernel,
                           float,
                           double,
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
                           phi::dtype::bfloat16,
 #endif
                           int,
diff --git a/paddle/fluid/operators/collective/mp_allreduce_sum_op.cu.cc b/paddle/fluid/operators/collective/mp_allreduce_sum_op.cu.cc
@@ -31,7 +31,8 @@ PD_REGISTER_STRUCT_KERNEL(mp_allreduce_sum,
                           double,
                           int,
                           int64_t,
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
                           phi::dtype::bfloat16,
 #endif
                           phi::dtype::float16) {
diff --git a/paddle/fluid/operators/collective/partial_allgather_op.cu.cc b/paddle/fluid/operators/collective/partial_allgather_op.cu.cc
@@ -154,7 +154,8 @@ PD_REGISTER_STRUCT_KERNEL(partial_allgather,
                           ops::PartialAllGatherOpCUDAKernel,
                           float,
                           double,
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
                           phi::dtype::bfloat16,
 #endif
                           int,
diff --git a/paddle/fluid/operators/collective/partial_recv_op.cu.cc b/paddle/fluid/operators/collective/partial_recv_op.cu.cc
@@ -179,7 +179,8 @@ PD_REGISTER_STRUCT_KERNEL(partial_recv,
                           ops::PartialRecvOpCUDAKernel,
                           float,
                           double,
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
                           phi::dtype::bfloat16,
 #endif
                           int,
diff --git a/paddle/fluid/operators/collective/partial_send_op.cu.cc b/paddle/fluid/operators/collective/partial_send_op.cu.cc
@@ -174,7 +174,8 @@ PD_REGISTER_STRUCT_KERNEL(partial_send,
                           ops::PartialSendCUDAKernel,
                           float,
                           double,
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
                           phi::dtype::bfloat16,
 #endif
                           int,
diff --git a/paddle/fluid/operators/collective/recv_v2_op.cu.cc b/paddle/fluid/operators/collective/recv_v2_op.cu.cc
@@ -297,7 +297,8 @@ PD_REGISTER_STRUCT_KERNEL(recv_v2,
                           ops::RecvOpV2CUDAKernel,
                           float,
                           double,
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
                           phi::dtype::bfloat16,
 #endif
                           int,
diff --git a/paddle/fluid/operators/collective/send_v2_op.cu.cc b/paddle/fluid/operators/collective/send_v2_op.cu.cc
@@ -273,7 +273,8 @@ PD_REGISTER_STRUCT_KERNEL(send_v2,
                           ops::SendOpV2CUDAKernel,
                           float,
                           double,
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
                           phi::dtype::bfloat16,
 #endif
                           int,
diff --git a/paddle/fluid/operators/fused/fused_gate_attention.h b/paddle/fluid/operators/fused/fused_gate_attention.h
@@ -14,7 +14,7 @@ limitations under the License. */
 
 #pragma once
 
-#ifdef PADDLE_WITH_FLASHATTN
+#if defined(PADDLE_WITH_FLASHATTN) && !defined(PADDLE_WITH_HIP)
 #include "paddle/phi/backends/dynload/flashattn.h"
 #endif
 
@@ -202,7 +202,7 @@ struct GateAttentionConfig {
   }
 
   bool CanUseFlashAttn() const {
-#ifdef PADDLE_WITH_FLASHATTN
+#if defined(PADDLE_WITH_FLASHATTN) && !defined(PADDLE_WITH_HIP)
     if (!std::is_same<T, phi::dtype::bfloat16>::value &&
         !std::is_same<T, phi::dtype::float16>::value) {
       return false;
@@ -867,7 +867,7 @@ class FlashAttnWithGating {
                       phi::DenseTensor* softmax_lse,
                       phi::DenseTensor* fmha_out,
                       GateAttentionConfig<T>* config) {
-#ifdef PADDLE_WITH_FLASHATTN
+#if defined(PADDLE_WITH_FLASHATTN) && !defined(PADDLE_WITH_HIP)
     bool is_bf16 =
         qkv_transpose_out->dtype() == DataType::BFLOAT16 ? true : false;
     TypeDebugInfo<T>();
@@ -982,7 +982,7 @@ class FlashAttnWithGating {
                        phi::DenseTensor* src_mask_grad,
                        phi::DenseTensor* nonbatched_bias_grad,
                        GateAttentionGradConfig<T>* config) {
-#ifdef PADDLE_WITH_FLASHATTN
+#if defined(PADDLE_WITH_FLASHATTN) && !defined(PADDLE_WITH_HIP)
     bool is_bf16 =
         qkv_transpose_out->dtype() == DataType::BFLOAT16 ? true : false;
     TypeDebugInfo<T>();
diff --git a/paddle/fluid/platform/device/gpu/nccl_helper.h b/paddle/fluid/platform/device/gpu/nccl_helper.h
@@ -61,7 +61,8 @@ inline ncclDataType_t ToNCCLDataType(framework::proto::VarType::Type type) {
     return ncclUint8;
   } else if (type == framework::proto::VarType::BOOL) {
     return ncclUint8;
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
   } else if (type == framework::proto::VarType::BF16) {
     return ncclBfloat16;
 #endif
@@ -88,7 +89,8 @@ inline ncclDataType_t ToNCCLDataType(phi::DataType type) {
     return ncclInt8;
   } else if (type == phi::DataType::BOOL) {
     return ncclUint8;
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
   } else if (type == phi::DataType::BFLOAT16) {
     return ncclBfloat16;
 #endif
diff --git a/paddle/fluid/pybind/place.cc b/paddle/fluid/pybind/place.cc
@@ -374,7 +374,7 @@ void BindPlace(pybind11::module &m) {  // NOLINT
   m.def("is_bfloat16_supported", [](const phi::GPUPlace &place) -> bool {
   // Only GPUs with Compute Capability >= 80 support bfloat16
 #ifdef PADDLE_WITH_HIP
-    return false;
+    return true;
 #else
     return platform::GetGPUComputeCapability(place.device) >= 80;
 #endif
diff --git a/paddle/phi/core/distributed/nccl_tools.cc b/paddle/phi/core/distributed/nccl_tools.cc
@@ -62,7 +62,8 @@ std::string NCCLDTypeToString(ncclDataType_t dtype) {
   PD_NCCL_DTYPE_TO_STR(ncclFloat32, "float32");
   PD_NCCL_DTYPE_TO_STR(ncclHalf, "float16");
   PD_NCCL_DTYPE_TO_STR(ncclFloat16, "float16");
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
   PD_NCCL_DTYPE_TO_STR(ncclBfloat16, "bfloat16");
 #endif
   PD_NCCL_DTYPE_TO_STR(ncclDouble, "float64");
diff --git a/paddle/phi/core/utils/data_type.h b/paddle/phi/core/utils/data_type.h
@@ -243,7 +243,8 @@ inline ncclDataType_t ToNCCLDataType(DataType type) {
     return ncclInt8;
   } else if (type == DataType::BOOL) {
     return ncclUint8;
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
   } else if (type == DataType::BFLOAT16) {
     return ncclBfloat16;
 #endif
diff --git a/paddle/phi/kernels/funcs/blas/blas_impl.hip.h b/paddle/phi/kernels/funcs/blas/blas_impl.hip.h
@@ -757,9 +757,9 @@ inline void Blas<phi::GPUContext>::GEMM(CBLAS_TRANSPOSE transA,
   // TODO(zhiqiu): 80 has the same meaning for rocm and cuda?
   PADDLE_ENFORCE_GE(
       context_.GetComputeCapability(),
-      80,
+      53,
       phi::errors::InvalidArgument(
-          "rocblas bf16 gemm requires GPU compute capability >= 80,"
+          "rocblas bf16 gemm requires GPU compute capability >= 53,"
           "but received %d",
           context_.GetComputeCapability()));
 
@@ -1013,9 +1013,9 @@ inline void Blas<phi::GPUContext>::GEMM(bool transA,
       transB ? rocblas_operation_none : rocblas_operation_transpose;
   PADDLE_ENFORCE_GE(
       context_.GetComputeCapability(),
-      80,
+      53,
       phi::errors::InvalidArgument(
-          "rocblas bf16 gemm requires GPU compute capability >= 80,"
+          "rocblas bf16 gemm requires GPU compute capability >= 53,"
           "but received %d",
           context_.GetComputeCapability()));
 
diff --git a/paddle/phi/kernels/gpu/all_to_all_kernel.cu b/paddle/phi/kernels/gpu/all_to_all_kernel.cu
@@ -84,7 +84,8 @@ void AllToAllKernel(const Context& dev_ctx,
 
 }  // namespace phi
 
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
 PD_REGISTER_KERNEL(all_to_all,
                    GPU,
                    ALL_LAYOUT,
diff --git a/paddle/phi/kernels/gpu/c_embedding_grad_kernel.cu b/paddle/phi/kernels/gpu/c_embedding_grad_kernel.cu
@@ -140,7 +140,8 @@ void CEmbeddingGradKernel(const Context& dev_ctx,
 
 }  // namespace phi
 
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
 PD_REGISTER_KERNEL(c_embedding_grad,
                    GPU,
                    ALL_LAYOUT,
diff --git a/paddle/phi/kernels/gpu/c_embedding_kernel.cu b/paddle/phi/kernels/gpu/c_embedding_kernel.cu
@@ -113,7 +113,8 @@ void CEmbeddingKernel(const Context& ctx,
 }
 }  // namespace phi
 
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
 PD_REGISTER_KERNEL(c_embedding,
                    GPU,
                    ALL_LAYOUT,
diff --git a/paddle/phi/kernels/gpu/c_identity_kernel.cu b/paddle/phi/kernels/gpu/c_identity_kernel.cu
@@ -18,7 +18,8 @@ limitations under the License. */
 #include "paddle/phi/backends/gpu/gpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
 
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
 PD_REGISTER_KERNEL(c_identity,
                    GPU,
                    ALL_LAYOUT,
diff --git a/paddle/phi/kernels/gpu/c_split_kernel.cu b/paddle/phi/kernels/gpu/c_split_kernel.cu
@@ -103,7 +103,8 @@ void CSplitKernel(const Context& ctx,
 
 }  // namespace phi
 
-#if NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
 PD_REGISTER_KERNEL(c_split,
                    GPU,
                    ALL_LAYOUT,
diff --git a/paddle/phi/kernels/gpu/calc_reduced_attn_kernel.cu b/paddle/phi/kernels/gpu/calc_reduced_attn_kernel.cu
@@ -20,7 +20,7 @@
 
 namespace phi {
 
-#ifdef PADDLE_WITH_FLASHATTN
+#if defined(PADDLE_WITH_FLASHATTN) && !defined(PADDLE_WITH_HIP)
 struct CalcReducedAttnScoresParams : public FlashAttnParamsBase {
   bool return_softmax;
   DenseTensor* softmax;
@@ -55,7 +55,7 @@ void CalcReducedAttnScoresKernel(const Context& ctx,
                                  const DenseTensor& k,
                                  const DenseTensor& softmax_lse,
                                  DenseTensor* reduced_scores) {
-#ifdef PADDLE_WITH_FLASHATTN
+#if defined(PADDLE_WITH_FLASHATTN) && !defined(PADDLE_WITH_HIP)
   PADDLE_ENFORCE_EQ(q.dims().size(),
                     4,
                     phi::errors::InvalidArgument(
diff --git a/python/paddle/amp/auto_cast.py b/python/paddle/amp/auto_cast.py
@@ -127,7 +127,7 @@ def _is_gpu_float16_supported():
     Judge whether current gpu support float16 amp.
     """
     prop = paddle.device.cuda.get_device_capability()
-    return prop[0] >= 7
+    return prop[0] >= 7 or paddle.is_compiled_with_rocm()
 
 
 def _is_gpu_bfloat16_supported():
@@ -140,7 +140,7 @@ def _is_gpu_bfloat16_supported():
         cuda_version_check = int(cuda_version.split('.')[0]) >= 11
     else:
         cuda_version_check = False
-    return prop[0] >= 8 and cuda_version_check
+    return prop[0] >= 8 and cuda_version_check or paddle.is_compiled_with_rocm()
 
 
 def _is_xpu_float16_supported():
