diff --git a/tests/cpp/operator/test_layernorm.cu b/tests/cpp/operator/test_layernorm.cu index 248461fe22..8340da03d4 100644 --- a/tests/cpp/operator/test_layernorm.cu +++ b/tests/cpp/operator/test_layernorm.cu @@ -163,12 +163,12 @@ void performTest(const size_t N, const size_t H) { float epsilon = 1e-5; nvte_layernorm_fwd(input.data(), gamma.data(), beta.data(), scale.data(), epsilon, z.data(), mu.data(), rsigma.data(), 0, prop.multiProcessorCount, - workspace.data(), barrier.data(), amax.data(), scale_inv.data(), true); + workspace.data(), barrier.data(), amax.data(), scale_inv.data()); workspace = Tensor(workspace.shape(), workspace.dtype()); barrier = Tensor(barrier.shape(), barrier.dtype()); nvte_layernorm_fwd(input.data(), gamma.data(), beta.data(), scale.data(), epsilon, z.data(), mu.data(), rsigma.data(), 0, prop.multiProcessorCount, - workspace.data(), barrier.data(), amax.data(), scale_inv.data(), true); + workspace.data(), barrier.data(), amax.data(), scale_inv.data()); // Backward kernel nvte_layernorm_bwd(dz.data(), input.data(), @@ -195,6 +195,7 @@ void performTest(const size_t N, const size_t H) { float ref_amax; compute_ref_stats(input.cpu_dptr(), ref_mu.get(), ref_rsigma.get(), N, H, epsilon); + float ref_scale = isFp8Type(otype) ? *(scale.cpu_dptr()) : 1.f; compute_ref_output(input.cpu_dptr(), gamma.cpu_dptr(), beta.cpu_dptr(), @@ -203,7 +204,7 @@ void performTest(const size_t N, const size_t H) { rsigma.cpu_dptr(), N, H, &ref_amax, - *(scale.cpu_dptr())); + ref_scale); compute_ref_backward(dz.cpu_dptr(), input.cpu_dptr(), mu.cpu_dptr(), rsigma.cpu_dptr(), gamma.cpu_dptr(), @@ -215,9 +216,11 @@ void performTest(const size_t N, const size_t H) { ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); auto [atol_amax, rtol_amax] = getTolerances(DType::kFloat32); - compareResults("amax", amax, &ref_amax, atol_amax, rtol_amax); - float ref_scale_inv = 1.f / (*scale.cpu_dptr()); - compareResults("scale_inv", scale_inv, &ref_scale_inv, atol_amax, rtol_amax); + if (isFp8Type(otype)) { + compareResults("amax", amax, &ref_amax, atol_amax, rtol_amax); + float ref_scale_inv = 1.f / (*scale.cpu_dptr()); + compareResults("scale_inv", scale_inv, &ref_scale_inv, atol_amax, rtol_amax); + } auto [atol_stats, rtol_stats] = getTolerances(DType::kFloat32); rtol_stats = 5e-5; diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index 8e1db6544e..0bd45fcda5 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -171,4 +171,8 @@ void fillUniform(const Tensor &t) { t.from_cpu(); } +bool isFp8Type(DType type) { + return type == DType::kFloat8E4M3 || type == DType::kFloat8E5M2; +} + } // namespace test diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index 441034733c..3416aae3b3 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -156,6 +156,8 @@ const std::string &typeName(DType type); extern std::vector all_fp_types; +bool isFp8Type(DType type); + } // namespace test #define TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(dtype, type, ...) \ diff --git a/transformer_engine/common/include/transformer_engine/layer_norm.h b/transformer_engine/common/include/transformer_engine/layer_norm.h index 91a1bc687e..5d5a5eb119 100644 --- a/transformer_engine/common/include/transformer_engine/layer_norm.h +++ b/transformer_engine/common/include/transformer_engine/layer_norm.h @@ -39,7 +39,6 @@ extern "C" { * \param[out] barrier Barrier tensor. * \param[in,out] amax AMAX value of the output tensor. * \param[out] scale_inv Inverse of the output's scaling factor. - * \param[in] fp8_out Whether to output FP8. */ void nvte_layernorm_fwd(const NVTETensor x, const NVTETensor gamma, @@ -54,8 +53,7 @@ void nvte_layernorm_fwd(const NVTETensor x, NVTETensor workspace, NVTETensor barrier, NVTETensor amax, - NVTETensor scale_inv, - bool fp8_out); + NVTETensor scale_inv); /*! \brief Compute backward of LayerNorm. diff --git a/transformer_engine/common/layer_norm/ln_api.cpp b/transformer_engine/common/layer_norm/ln_api.cpp index f542401a39..8eca903977 100644 --- a/transformer_engine/common/layer_norm/ln_api.cpp +++ b/transformer_engine/common/layer_norm/ln_api.cpp @@ -154,12 +154,13 @@ void layernorm_fwd(const Tensor& x, // BxSxhidden_size Tensor* workspace, Tensor* barrier, Tensor* amax, - Tensor *scale_inv, - bool fp8_out + Tensor *scale_inv ) { auto itype = x.dtype; auto wtype = gamma.dtype; auto otype = z->dtype; + bool fp8_out = otype == DType::kFloat8E4M3 || + otype == DType::kFloat8E5M2; auto ctype = layer_norm::DType::kFloat32; NVTE_CHECK(x.shape.size() == 2); @@ -382,8 +383,7 @@ void nvte_layernorm_fwd(const NVTETensor x, // BxSxhidden_size NVTETensor workspace, NVTETensor barrier, NVTETensor amax, - NVTETensor scale_inv, - bool fp8_out) { + NVTETensor scale_inv) { using namespace transformer_engine; layernorm_fwd(*reinterpret_cast(x), *reinterpret_cast(gamma), @@ -398,8 +398,7 @@ void nvte_layernorm_fwd(const NVTETensor x, // BxSxhidden_size reinterpret_cast(workspace), reinterpret_cast(barrier), reinterpret_cast(amax), - reinterpret_cast(scale_inv), - fp8_out); + reinterpret_cast(scale_inv)); } void nvte_layernorm_bwd(const NVTETensor dz, // BxSxhidden_size diff --git a/transformer_engine/common/layer_norm/ln_fwd_kernels.cuh b/transformer_engine/common/layer_norm/ln_fwd_kernels.cuh index 17ba92a840..7e45d3a089 100644 --- a/transformer_engine/common/layer_norm/ln_fwd_kernels.cuh +++ b/transformer_engine/common/layer_norm/ln_fwd_kernels.cuh @@ -69,7 +69,7 @@ void ln_fwd_tuned_kernel(FwdParams params) { constexpr compute_t rn = 1.f / compute_t(Ktraits::COLS); - compute_t scale; + compute_t scale = 1.f; if (params.fp8_out) { scale = *reinterpret_cast(params.scale); } diff --git a/transformer_engine/pytorch/csrc/common.cu b/transformer_engine/pytorch/csrc/common.cu index 528bc39b43..34141024ec 100644 --- a/transformer_engine/pytorch/csrc/common.cu +++ b/transformer_engine/pytorch/csrc/common.cu @@ -119,8 +119,7 @@ void dispatch_layernorm(void* input, // i void* scale_inv, // o const std::vector& scale_inv_shape, const transformer_engine::DType scale_inv_type, - const int multiProcessorCount, - const bool fp8_out + const int multiProcessorCount ) { auto input_cu = makeTransformerEngineTensor(input, input_shape, input_type); auto gamma_cu = makeTransformerEngineTensor(gamma, gamma_shape, gamma_type); @@ -139,7 +138,7 @@ void dispatch_layernorm(void* input, // i z_cu.data(), mu_cu.data(), rsigma_cu.data(), at::cuda::getCurrentCUDAStream(), multiProcessorCount, workspace.data(), barrier.data(), amax_cu.data(), - scale_inv_cu.data(), fp8_out); + scale_inv_cu.data()); // Fill workspace and barrier auto workspace_data = allocateSpace(workspace.shape(), @@ -160,7 +159,7 @@ void dispatch_layernorm(void* input, // i z_cu.data(), mu_cu.data(), rsigma_cu.data(), at::cuda::getCurrentCUDAStream(), multiProcessorCount, workspace.data(), barrier.data(), amax_cu.data(), - scale_inv_cu.data(), fp8_out); + scale_inv_cu.data()); } diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 9926122f78..594132d457 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -165,8 +165,7 @@ void dispatch_layernorm(void* input, // i void* scale_inv, // o const std::vector& scale_inv_shape, const transformer_engine::DType scale_inv_type, - const int multiProcessorCount, - const bool fp8_out + const int multiProcessorCount ); diff --git a/transformer_engine/pytorch/csrc/extensions.cu b/transformer_engine/pytorch/csrc/extensions.cu index 6147323c32..08467e4972 100644 --- a/transformer_engine/pytorch/csrc/extensions.cu +++ b/transformer_engine/pytorch/csrc/extensions.cu @@ -307,8 +307,7 @@ std::vector layernorm_fwd_fp8(const at::Tensor &input, rsigma.data_ptr(), {N}, DType::kFloat32, amax.data_ptr(), {1}, DType::kFloat32, scale_inv.data_ptr(), {1}, DType::kFloat32, - at::cuda::getCurrentDeviceProperties()->multiProcessorCount, - true); + at::cuda::getCurrentDeviceProperties()->multiProcessorCount); return {ln_out, mu, rsigma}; } @@ -340,8 +339,7 @@ std::vector layernorm_fwd(const at::Tensor &input, rsigma.data_ptr(), {N}, DType::kFloat32, nullptr, {1}, DType::kFloat32, nullptr, {1}, DType::kFloat32, - at::cuda::getCurrentDeviceProperties()->multiProcessorCount, - false); + at::cuda::getCurrentDeviceProperties()->multiProcessorCount); return {ln_out, mu, rsigma}; }