Skip to content

Remove fp8_out from the LN API #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ksivaman merged 3 commits into from
Oct 12, 2022
Merged

Remove fp8_out from the LN API #8

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3f326b0
Remove fp8_out from LN API
ptrendx Sep 2, 2022
13f80c0
fix LN test
ksivaman Sep 7, 2022
f122b0d
Fixes
ptrendx Oct 10, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 9 additions & 6 deletions tests/cpp/operator/test_layernorm.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand All @@ -195,6 +195,7 @@ void performTest(const size_t N, const size_t H) {
float ref_amax;
compute_ref_stats(input.cpu_dptr<InputType>(), ref_mu.get(),
ref_rsigma.get(), N, H, epsilon);
float ref_scale = isFp8Type(otype) ? *(scale.cpu_dptr<float>()) : 1.f;
compute_ref_output(input.cpu_dptr<InputType>(),
gamma.cpu_dptr<WeightType>(),
beta.cpu_dptr<WeightType>(),
Expand All @@ -203,7 +204,7 @@ void performTest(const size_t N, const size_t H) {
rsigma.cpu_dptr<float>(),
N, H,
&ref_amax,
*(scale.cpu_dptr<float>()));
ref_scale);
compute_ref_backward(dz.cpu_dptr<WeightType>(), input.cpu_dptr<InputType>(),
mu.cpu_dptr<float>(), rsigma.cpu_dptr<float>(),
gamma.cpu_dptr<WeightType>(),
Expand All @@ -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<float>());
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<float>());
compareResults("scale_inv", scale_inv, &ref_scale_inv, atol_amax, rtol_amax);
}

auto [atol_stats, rtol_stats] = getTolerances(DType::kFloat32);
rtol_stats = 5e-5;
Expand Down
4 changes: 4 additions & 0 deletions tests/cpp/test_common.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
2 changes: 2 additions & 0 deletions tests/cpp/test_common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,8 @@ const std::string &typeName(DType type);

extern std::vector<DType> all_fp_types;

bool isFp8Type(DType type);

} // namespace test

#define TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(dtype, type, ...) \
Expand Down
4 changes: 1 addition & 3 deletions transformer_engine/common/include/transformer_engine/layer_norm.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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.
Expand Down
11 changes: 5 additions & 6 deletions transformer_engine/common/layer_norm/ln_api.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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<const Tensor*>(x),
*reinterpret_cast<const Tensor*>(gamma),
Expand All @@ -398,8 +398,7 @@ void nvte_layernorm_fwd(const NVTETensor x, // BxSxhidden_size
reinterpret_cast<Tensor*>(workspace),
reinterpret_cast<Tensor*>(barrier),
reinterpret_cast<Tensor*>(amax),
reinterpret_cast<Tensor*>(scale_inv),
fp8_out);
reinterpret_cast<Tensor*>(scale_inv));
}

void nvte_layernorm_bwd(const NVTETensor dz, // BxSxhidden_size
Expand Down
2 changes: 1 addition & 1 deletion transformer_engine/common/layer_norm/ln_fwd_kernels.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<compute_t*>(params.scale);
}
Expand Down
7 changes: 3 additions & 4 deletions transformer_engine/pytorch/csrc/common.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,8 +119,7 @@ void dispatch_layernorm(void* input, // i
void* scale_inv, // o
const std::vector<size_t>& 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);
Expand All @@ -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(),
Expand All @@ -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());
}


Expand Down
3 changes: 1 addition & 2 deletions transformer_engine/pytorch/csrc/common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,8 +165,7 @@ void dispatch_layernorm(void* input, // i
void* scale_inv, // o
const std::vector<size_t>& scale_inv_shape,
const transformer_engine::DType scale_inv_type,
const int multiProcessorCount,
const bool fp8_out
const int multiProcessorCount
);


Expand Down
6 changes: 2 additions & 4 deletions transformer_engine/pytorch/csrc/extensions.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -307,8 +307,7 @@ std::vector<at::Tensor> 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};
}
Expand Down Expand Up @@ -340,8 +339,7 @@ std::vector<at::Tensor> 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};
}
Expand Down