diff --git a/transformer_engine/common/amd_detail/hip_float8.h b/transformer_engine/common/amd_detail/hip_float8.h index ab149ba05..84efaca3a 100644 --- a/transformer_engine/common/amd_detail/hip_float8.h +++ b/transformer_engine/common/amd_detail/hip_float8.h @@ -2,6 +2,11 @@ // FP8 header version 0.3, 2021/05/11 #define HIP_HOST_DEVICE __host__ __device__ +#define HIP_DEVICE __device__ +#define HIP_HOST __host__ + +#define E5M2_AMAX 57344.0 +#define E4M3_AMAX 240.0 namespace hip_f8_impl { @@ -49,8 +54,8 @@ enum class hip_f8_rounding_mode { // => bias = 15 for 152, 7 for 143 // => NAN/INF are represented as per IEEE conventions -static __device__ bool hip_f8_bias_mode_bit_device; -static bool hip_f8_bias_mode_bit_host; +static __device__ bool hip_f8_bias_mode_bit_device = true; +static bool hip_f8_bias_mode_bit_host = true; static __global__ void set_hip_f8_bias_mode_bit(bool v) { hip_f8_bias_mode_bit_device = v; @@ -69,7 +74,6 @@ static void set_hip_f8_bias_mode_optimal() { static inline HIP_HOST_DEVICE bool get_hip_f8_bias_mode() { #if defined(__HIP_DEVICE_COMPILE__) return hip_f8_bias_mode_bit_device; - #else return hip_f8_bias_mode_bit_host; #endif @@ -89,7 +93,53 @@ struct hip_f8 { } // constructor from float - explicit HIP_HOST_DEVICE hip_f8(float v, hip_f8_rounding_mode rm=hip_f8_rounding_mode::standard, uint32_t rng=0) { +#if defined(__gfx940__) + explicit HIP_DEVICE hip_f8(float v, hip_f8_rounding_mode rm=hip_f8_rounding_mode::standard, uint32_t rng=0) { + union { + float fval; + uint32_t i32val; + uint8_t i8val[4]; + } val; + uint32_t ival = 0; + val.fval = v; + + if (T == hip_f8_type::bf8) { // bf8 + if ((val.i32val & 0x7F800000) != 0x7F800000) // propagate NAN/INF, no clipping + val.fval = __builtin_amdgcn_fmed3f(val.fval, E5M2_AMAX, -E5M2_AMAX); + if (rm == hip_f8_rounding_mode::standard) { // RNE rounding + ival = __builtin_amdgcn_cvt_pk_bf8_f32( + val.fval, val.fval, ival, false); // false -> WORD0 + val.i32val = ival; + data = val.i8val[0]; + } + else { //stochastic rounding + ival = __builtin_amdgcn_cvt_sr_bf8_f32(val.fval, rng, ival, 0); // 0 pos + val.i32val = ival; + data = val.i8val[0]; // little endian + } + } + else { // fp8 + if ((val.i32val & 0x7F800000) != 0x7F800000) /// propagate NAN/INF, no clipping + val.fval = __builtin_amdgcn_fmed3f(val.fval, E4M3_AMAX, -E4M3_AMAX); + if (rm == hip_f8_rounding_mode::standard) { // RNE rounding + ival = __builtin_amdgcn_cvt_pk_fp8_f32( + val.fval, val.fval, ival, false); // false -> WORD0 + val.i32val = ival; + data = val.i8val[0]; + } + else { //stochastic rounding + ival = __builtin_amdgcn_cvt_sr_fp8_f32(val.fval, rng, ival, 0); // 0 pos + val.i32val = ival; + data = val.i8val[0]; // little endian + } + } + } + + explicit HIP_HOST //Code host still uses SW simulated conversion on gfx940 +#else // #if defined(__gfx940__) + explicit HIP_HOST_DEVICE // On architectures other than gfx940, both host and device still use SW simulated conversion +#endif // #if defined(__gfx940__) + hip_f8(float v, hip_f8_rounding_mode rm=hip_f8_rounding_mode::standard, uint32_t rng=0) { if (T == hip_f8_type::bf8) { if (get_hip_f8_bias_mode()) { data = hip_f8_impl::cast_to_f8<2, 5, float, true/*negative_zero_nan*/, true/*clip*/>(v, (rm == hip_f8_rounding_mode::stochastic), rng); @@ -106,7 +156,16 @@ struct hip_f8 { } // constructor from half - explicit HIP_HOST_DEVICE hip_f8(half v, hip_f8_rounding_mode rm=hip_f8_rounding_mode::standard, uint32_t rng=0) { +#if defined(__gfx940__) + explicit HIP_DEVICE hip_f8(half v, hip_f8_rounding_mode rm=hip_f8_rounding_mode::standard, uint32_t rng=0) + : hip_f8((float)v, rm, rng) + { + } + explicit HIP_HOST //Code host still uses SW simulated conversion on gfx940 +#else // #if defined(__gfx940__) + explicit HIP_HOST_DEVICE // On architectures other than gfx940, both host and device still use SW simulated conversion +#endif // #if defined(__gfx940__) + hip_f8(half v, hip_f8_rounding_mode rm=hip_f8_rounding_mode::standard, uint32_t rng=0) { if (T == hip_f8_type::bf8) { if (get_hip_f8_bias_mode()) { data = hip_f8_impl::cast_to_f8<2, 5, half, true/*negative_zero_nan*/, true/*clip*/>(v, (rm == hip_f8_rounding_mode::stochastic), rng); @@ -126,7 +185,32 @@ struct hip_f8 { explicit HIP_HOST_DEVICE hip_f8(hip_bfloat16 v, hip_f8_rounding_mode r=hip_f8_rounding_mode::standard, uint32_t rng=0); // convert to float - explicit inline HIP_HOST_DEVICE operator float() const { +#if defined(__gfx940__) + HIP_DEVICE operator float() const { + union + { + float fval; + uint32_t i32val; + uint8_t i8val[4]; // dependent of endian + } val; + + // assign 8bit data in position [7:0] + val.i32val = 0; + val.i8val[0] = data; // little endian + + // upcast + if(T == hip_f8_type::bf8) + val.i32val = __builtin_amdgcn_cvt_f32_bf8(val.i32val, 0); // 0 pos + else // fp8 + val.i32val = __builtin_amdgcn_cvt_f32_fp8(val.i32val, 0); // 0 pos + + return val.fval; + } + explicit inline HIP_HOST //Code host still uses SW simulated conversion on gfx940 +#else // #if defined(__gfx940__) + explicit inline HIP_HOST_DEVICE // On architectures other than gfx940, both host and device still use SW simulated conversion +#endif // #if defined(__gfx940__) + operator float() const { if (T == hip_f8_type::bf8) { if (get_hip_f8_bias_mode()) { return hip_f8_impl::cast_from_f8<2, 5, float, true/*negative_zero_nan*/>(data); @@ -143,7 +227,15 @@ struct hip_f8 { } // convert to half - explicit inline HIP_HOST_DEVICE operator half() const { +#if defined(__gfx940__) + explicit HIP_DEVICE inline operator half() const { + return __half(float(*this)); + } + explicit inline HIP_HOST //Code host still uses SW simulated conversion on gfx940 +#else // #if defined(__gfx940__) + explicit inline HIP_HOST_DEVICE // On architectures other than gfx940, both host and device still use SW simulated conversion +#endif // #if defined(__gfx940__) + operator half() const { if (T == hip_f8_type::bf8) { if (get_hip_f8_bias_mode()) { return hip_f8_impl::cast_from_f8<2, 5, half, true/*negative_zero_nan*/>(data); diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index 290535be5..6c55a73c3 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -4,13 +4,17 @@ * See LICENSE for license information. ************************************************************************/ +#include #include #include #include #ifndef __HIP_PLATFORM_HCC__ #include -#endif #include +#else +#define ROCBLAS_BETA_FEATURES_API +#include +#endif #include "../common.h" #include "../util/vectorized_pointwise.h" #ifdef __HIP_PLATFORM_HCC__ @@ -43,10 +47,15 @@ namespace transformer_engine { } \ break; \ case DType::kBFloat16: \ + { \ + using type = bf16; \ + {__VA_ARGS__} \ + } \ + break; \ case DType::kFloat8E5M2: \ case DType::kFloat8E4M3: \ { \ - NVTE_ERROR("Bfloat16 and FP8 type not instantiated"); \ + NVTE_ERROR("FP8 type not instantiated"); \ } \ break; \ default: \ @@ -77,18 +86,20 @@ float gelu_forward(float x) return x * cdf; } + template __global__ void gelu_forward_kernel(const Tin* in, T* out, int m, int n) { for(int id = blockIdx.x * blockDim.x + threadIdx.x; id < m * n; id += blockDim.x * gridDim.x) { - Tin x = (Tin)(__ldg(&in[id])); + Tin x = in[id]; float y = gelu_forward((float)x); out[id] = (T)(y); } } + template void gelu_forward_kernelLauncher(const Tin* in, T* out, int m, int n, hipStream_t stream) { int blocks_per_row = ceil(float(n)/1024); @@ -124,8 +135,8 @@ __global__ void gelu_backward_kernel(const Tin* dy, T* out, const T* __restrict pre_gelu_out, int m, int n) { for(int id = blockIdx.x * blockDim.x + threadIdx.x; id < m * n; id += blockDim.x * gridDim.x) { - Tin x = (Tin)(__ldg(&pre_gelu_out[id])); - Tin dx = gelu_backward((float)x, (float)dy[id]); + Tin x = (Tin)pre_gelu_out[id]; + Tin dx = (Tin)gelu_backward((float)x, (float)dy[id]); out[id] = (T)(dx); } } @@ -138,47 +149,47 @@ void gelu_backward_kernelLauncher(const Tin* in, T* out, const T* pre_gelu_out, hipLaunchKernelGGL(( gelu_backward_kernel), dim3(grid), dim3(block), 0, stream, in, out, pre_gelu_out, m, n); } -template +template __global__ -void add_bias_kernel(const Tin* in, T* out, const T* __restrict bias, int m, int n) +void add_bias_kernel(const Tin* in, T* out, const Tb* __restrict bias, int m, int n) { for(int id = blockIdx.x * blockDim.x + threadIdx.x; id < m * n; id += blockDim.x * gridDim.x) { - Tin reg_bias = (Tin)(__ldg(&bias[id % n])); + Tin reg_bias = (Tin)bias[id % n]; Tin val = in[id] + reg_bias; out[id] = (T)(val); } } -template -void add_bias_kernelLauncher(const Tin* in, T* out, const T* __restrict bias, int m, int n, hipStream_t stream) { +template +void add_bias_kernelLauncher(const Tin* in, T* out, const Tb* __restrict bias, int m, int n, hipStream_t stream) { dim3 block, grid; block.x = 1024; grid.x = ceil(m * n / 1024.); - hipLaunchKernelGGL(( add_bias_kernel), dim3(grid), dim3(block), 0, stream, in, out, bias, m, n); + hipLaunchKernelGGL(( add_bias_kernel), dim3(grid), dim3(block), 0, stream, in, out, bias, m, n); } -template +template __global__ -void add_bias_gelu_kernel(const Tin* in, T* out, T* pre_gelu_out, const T* __restrict bias, int m, int n) +void add_bias_gelu_kernel(const Tin* in, T* out, T* pre_gelu_out, const Tb* __restrict bias, int m, int n) { for(int id = blockIdx.x * blockDim.x + threadIdx.x; id < m * n; id += blockDim.x * gridDim.x) { - Tin reg_bias = (Tin)(__ldg(&bias[id % n])); + Tin reg_bias = (Tin)bias[id % n]; Tin val = in[id] + reg_bias; pre_gelu_out[id] = (T)(val); out[id] = (T)(gelu_forward(val)); } } -template -void add_bias_gelu_kernelLauncher(const Tin* in, T* out, T* pre_gelu_out, const T* __restrict bias, int m, int n, hipStream_t stream) { +template +void add_bias_gelu_kernelLauncher(const Tin* in, T* out, T* pre_gelu_out, const Tb* __restrict bias, int m, int n, hipStream_t stream) { dim3 block, grid; block.x = 1024; grid.x = ceil(m * n / 1024.); - hipLaunchKernelGGL(( add_bias_gelu_kernel), dim3(grid), dim3(block), 0, stream, in, out, pre_gelu_out, bias, m, n ); + hipLaunchKernelGGL(( add_bias_gelu_kernel), dim3(grid), dim3(block), 0, stream, in, out, pre_gelu_out, bias, m, n ); } @@ -278,12 +289,10 @@ transformer_engine::DType get_transformer_engine_dtype(const rocblas_datatype t) return DType::kFloat32; case rocblas_datatype_bf16_r: return DType::kBFloat16; - /* HIP-TODO: Add back after installing rocblas with FP8 types - case DType::kFloat8E4M3: - return CUDA_R_8F_E4M3; - case DType::kFloat8E5M2: - return CUDA_R_8F_E5M2; - */ + case rocblas_datatype_f8_r: + return DType::kFloat8E4M3; + case rocblas_datatype_bf8_r: + return DType::kFloat8E5M2; default: NVTE_ERROR("Invalid type"); } @@ -326,24 +335,38 @@ void cublas_gemm(void* A, float zero = 0.0; float beta = (accumulate) ? one : zero; + float alpha = 1.0; + if (use_fp8) { + float A_scale_inv, B_scale_inv; + hipMemcpy(&A_scale_inv, A_scale_inverse, sizeof(float), hipMemcpyDeviceToHost); + hipMemcpy(&B_scale_inv, B_scale_inverse, sizeof(float), hipMemcpyDeviceToHost); + alpha = A_scale_inv * B_scale_inv; + } rocblas_handle handle; NVTE_CHECK_CUBLAS(rocblas_create_handle(&handle)); - rocblas_datatype computeType = rocblas_datatype_f32_r; int64_t ld_gelumat = (int64_t) ldd; - // We don't deal with fp8 for now - NVTE_CHECK(!use_fp8, "fp8 gemm is unavailable right now!"); - NVTE_CHECK((A_type==rocblas_datatype_f16_r && B_type==rocblas_datatype_f16_r && D_type==rocblas_datatype_f16_r) || - (A_type==rocblas_datatype_f32_r && B_type==rocblas_datatype_f32_r && D_type==rocblas_datatype_f32_r), - "Only fp32 and fp16 GEMMs are available now!"); - - + (A_type==rocblas_datatype_f32_r && B_type==rocblas_datatype_f32_r && D_type==rocblas_datatype_f32_r) || + (A_type==rocblas_datatype_f8_r && B_type==rocblas_datatype_f8_r && D_type==rocblas_datatype_f32_r) || + (A_type==rocblas_datatype_f8_r && B_type==rocblas_datatype_bf8_r && D_type==rocblas_datatype_f32_r) || + (A_type==rocblas_datatype_bf8_r && B_type==rocblas_datatype_f8_r && D_type==rocblas_datatype_f32_r), + /* + (A_type==rocblas_datatype_f8_r && B_type==rocblas_datatype_f8_r && D_type==rocblas_datatype_f8_r) || + (A_type==rocblas_datatype_f8_r && B_type==rocblas_datatype_bf8_r && D_type==rocblas_datatype_bf8_r) || + (A_type==rocblas_datatype_bf8_r && B_type==rocblas_datatype_f8_r && D_type==rocblas_datatype_bf8_r), + */ + //Currently does not support output of fp8 tensors + "Only the following combinations of data types are enabled now!\n 1. input: fp16, output: fp16.\n \ + 2. input: fp32, output: fp32.\n 3. input: fp8, output: fp32"); + + + //If D is not fp32, then we need a temp buffer for GEMM result before applying epilogues. Otherwise, we can apply epilogues in-place. void* D_temp; - if ((bias || gelu) && (A_type==rocblas_datatype_f16_r && B_type==rocblas_datatype_f16_r && D_type==rocblas_datatype_f16_r)) { + if ((bias || gelu) && (D_type==rocblas_datatype_f16_r || D_type==rocblas_datatype_f8_r || D_type==rocblas_datatype_bf8_r)) { NVTE_CHECK_CUDA( hipMalloc(&D_temp, sizeof(float)*m*n) ); } else { @@ -356,14 +379,28 @@ void cublas_gemm(void* A, D_temp_type = rocblas_datatype_f16_r; } + // D = alpha * (A * B) + beta * C // TODO: Can we search for rocblas_gemm_algo?? - NVTE_CHECK_CUBLAS(rocblas_gemm_ex(handle, transa, transb, m, n, k, &one, + if (use_fp8) { + rocblas_computetype computeType = rocblas_compute_type_f32; + NVTE_CHECK_CUBLAS(rocblas_gemm_ex3(handle, transa, transb, m, n, k, &alpha, A, A_type, lda, B, B_type, ldb, &beta, D_temp, D_temp_type, ldd, D_temp, D_temp_type, ldd, computeType, rocblas_gemm_algo::rocblas_gemm_algo_standard,0,0)); + } + else { + rocblas_datatype computeType = rocblas_datatype_f32_r; + NVTE_CHECK_CUBLAS(rocblas_gemm_ex(handle, transa, transb, m, n, k, &alpha, + A, A_type, lda, + B, B_type, ldb, + &beta, D_temp, D_temp_type, ldd, D_temp, D_temp_type, ldd, + computeType, rocblas_gemm_algo::rocblas_gemm_algo_standard,0,0)); + } + + NVTE_CHECK_CUBLAS(rocblas_destroy_handle(handle)); int batch_size, input_dim, output_dim; @@ -389,6 +426,7 @@ void cublas_gemm(void* A, output_dim = k; DType input_dtype = get_transformer_engine_dtype(rocblas_datatype_f32_r); DType output_dtype = get_transformer_engine_dtype(D_type); + DType bias_dtype = get_transformer_engine_dtype(bias_type); TRANSFORMER_ENGINE_TYPE_SWITCH_ROCM_SIM(input_dtype, IType, TRANSFORMER_ENGINE_TYPE_SWITCH_ROCM_SIM(output_dtype, OType, detail::gelu_backward_kernelLauncher(reinterpret_cast(D_temp), @@ -401,7 +439,8 @@ void cublas_gemm(void* A, ); void* bias_tmp; - if (D_type == rocblas_datatype_f16_r) { + //if (D_type == rocblas_datatype_f16_r) { + if (bias_type != rocblas_datatype_f32_r) { NVTE_CHECK_CUDA( hipMalloc(&bias_tmp, sizeof(float)*input_dim) ); // The bias gradient is for the first linear layer } else { @@ -416,8 +455,9 @@ void cublas_gemm(void* A, 0); ); - if (D_type == rocblas_datatype_f16_r) { - TRANSFORMER_ENGINE_TYPE_SWITCH_ROCM_SIM(output_dtype, OType, + //if (D_type == rocblas_datatype_f16_r) { + if (bias_type != rocblas_datatype_f32_r) { + TRANSFORMER_ENGINE_TYPE_SWITCH_ROCM_SIM(bias_dtype, OType, detail::identity_kernelLauncher(reinterpret_cast(bias_tmp), reinterpret_cast(bias_ptr), input_dim, @@ -437,15 +477,18 @@ void cublas_gemm(void* A, output_dim = m; DType input_dtype = get_transformer_engine_dtype(rocblas_datatype_f32_r); DType output_dtype = get_transformer_engine_dtype(D_type); + DType bias_dtype = get_transformer_engine_dtype(bias_type); TRANSFORMER_ENGINE_TYPE_SWITCH_ROCM_SIM(input_dtype, IType, TRANSFORMER_ENGINE_TYPE_SWITCH_ROCM_SIM(output_dtype, OType, + TRANSFORMER_ENGINE_TYPE_SWITCH_ROCM_SIM(bias_dtype, BType, detail::add_bias_gelu_kernelLauncher(reinterpret_cast(D_temp), reinterpret_cast(D), reinterpret_cast(pre_gelu_out), - reinterpret_cast(bias_ptr), + reinterpret_cast(bias_ptr), batch_size, output_dim, 0); + ); ); ); } @@ -462,7 +505,8 @@ void cublas_gemm(void* A, input_dim = m; output_dim = n; void * bias_tmp; - if (B_type == rocblas_datatype_f16_r) { + //if (B_type == rocblas_datatype_f16_r) { + if (bias_type != rocblas_datatype_f32_r) { NVTE_CHECK_CUDA( hipMalloc(&bias_tmp, sizeof(float)*output_dim) ); } else { @@ -471,6 +515,7 @@ void cublas_gemm(void* A, DType input_dtype = get_transformer_engine_dtype(B_type); DType output_dtype = get_transformer_engine_dtype(D_type); + DType bias_dtype = get_transformer_engine_dtype(bias_type); TRANSFORMER_ENGINE_TYPE_SWITCH_ROCM_SIM(input_dtype, IType, detail::bias_gradient_kernelLauncher(reinterpret_cast(B), reinterpret_cast(bias_tmp), @@ -478,8 +523,8 @@ void cublas_gemm(void* A, output_dim, 0); ); - if (B_type == rocblas_datatype_f16_r) { - TRANSFORMER_ENGINE_TYPE_SWITCH_ROCM_SIM(output_dtype, OType, + if (bias_type != rocblas_datatype_f32_r) { + TRANSFORMER_ENGINE_TYPE_SWITCH_ROCM_SIM(bias_dtype, OType, detail::identity_kernelLauncher(reinterpret_cast(bias_tmp), reinterpret_cast(bias_ptr), output_dim, @@ -506,14 +551,17 @@ void cublas_gemm(void* A, output_dim = m; DType input_dtype = get_transformer_engine_dtype(rocblas_datatype_f32_r); DType output_dtype = get_transformer_engine_dtype(D_type); + DType bias_dtype = get_transformer_engine_dtype(bias_type); TRANSFORMER_ENGINE_TYPE_SWITCH_ROCM_SIM(input_dtype, IType, TRANSFORMER_ENGINE_TYPE_SWITCH_ROCM_SIM(output_dtype, OType, - detail::add_bias_kernelLauncher(reinterpret_cast(D_temp), + TRANSFORMER_ENGINE_TYPE_SWITCH_ROCM_SIM(bias_dtype, BType, + detail::add_bias_kernelLauncher(reinterpret_cast(D_temp), reinterpret_cast(D), - reinterpret_cast(bias_ptr), + reinterpret_cast(bias_ptr), batch_size, output_dim, 0); + ); ); ); } @@ -567,16 +615,9 @@ void cublas_gemm(void* A, ); } } - if ((bias || gelu) && (A_type==rocblas_datatype_f16_r && B_type==rocblas_datatype_f16_r && D_type==rocblas_datatype_f16_r)) { + if ((bias || gelu) && (D_type==rocblas_datatype_f16_r || D_type==rocblas_datatype_f8_r || D_type==rocblas_datatype_bf8_r)) { NVTE_CHECK_CUDA( hipFree(D_temp) ); } -/* - NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceDestroy(preference)); - NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Ddesc)); - NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Bdesc)); - NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Adesc)); - NVTE_CHECK_CUBLAS(cublasLtMatmulDescDestroy(operationDesc)); - */ } #else void cublas_gemm(void* A, @@ -767,12 +808,10 @@ rocblas_datatype get_cuda_dtype(const transformer_engine::DType t) { return rocblas_datatype_f32_r; case DType::kBFloat16: return rocblas_datatype_bf16_r; - /* HIP-TODO: Add back after installing rocblas with FP8 types case DType::kFloat8E4M3: - return CUDA_R_8F_E4M3; + return rocblas_datatype_f8_r; case DType::kFloat8E5M2: - return CUDA_R_8F_E5M2; - */ + return rocblas_datatype_bf8_r; default: NVTE_ERROR("Invalid type"); } @@ -853,18 +892,26 @@ void nvte_cublas_gemm(const NVTETensor A, nvte_log_gemm_config = true; } - if (nvte_log_gemm_config) + if (nvte_log_gemm_config) { + float A_scale_inv, B_scale_inv; + hipMemcpy(&A_scale_inv, Ainvscale->dptr, sizeof(float), hipMemcpyDeviceToHost); + hipMemcpy(&B_scale_inv, Binvscale->dptr, sizeof(float), hipMemcpyDeviceToHost); std::cout << "m=" << m << " k=" << k << " n=" << n << " transa=" << (transa?"T":"N") << " transb=" << (transb?"T":"N") << " A_type=" << (int)inputA->dtype << " B_type=" << (int)inputB->dtype << " D_type=" << (int)outputD->dtype + << " bias_type=" << (int)biasTensor->dtype << " grad=" << grad << " bias=" << (biasTensor->dptr != nullptr) << " gelu=" << (outputGelu->dptr != nullptr) + << " use_fp8=" << ( is_fp8_dtype(inputA->dtype) || is_fp8_dtype(inputB->dtype) ) + << " A_scale_inverse = " << A_scale_inv + << " B_scale_inverse = " << B_scale_inv << " accumulate=" << accumulate << std::endl; + } cublas_gemm(inputA->dptr, Ainvscale->dptr, inputB->dptr, Binvscale->dptr, diff --git a/transformer_engine/common/include/transformer_engine/logging.h b/transformer_engine/common/include/transformer_engine/logging.h index 90745a1cf..ff9636598 100644 --- a/transformer_engine/common/include/transformer_engine/logging.h +++ b/transformer_engine/common/include/transformer_engine/logging.h @@ -8,7 +8,12 @@ #define TRANSFORMER_ENGINE_LOGGING_H_ #include +#ifdef __HIP_PLATFORM_HCC__ +#define ROCBLAS_BETA_FEATURES_API +#include +#else #include +#endif #include #include diff --git a/transformer_engine/common/recipe.py b/transformer_engine/common/recipe.py index a1cccbdf2..3ed29e185 100644 --- a/transformer_engine/common/recipe.py +++ b/transformer_engine/common/recipe.py @@ -35,7 +35,7 @@ class Format(Enum): FP8 tensors in the backward pass are in e5m2 format """ - E4M3 = _FormatHelper(max_fwd=448, max_bwd=448) + E4M3 = _FormatHelper(max_fwd=240, max_bwd=240) E5M2 = _FormatHelper(max_fwd=57344, max_bwd=57344) HYBRID = _FormatHelper(max_fwd=E4M3.max_fwd, max_bwd=E5M2.max_bwd)