wehub-resource-sync
125 lines
5.7 KiB
C++
125 lines
5.7 KiB
C++
// Copyright (c) Microsoft Corporation.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
// DeepSpeed Team
|
|
|
|
#include "fp_quantize.h"
|
|
|
|
#include <c10/cuda/CUDAStream.h>
|
|
#include <torch/extension.h>
|
|
#include <vector>
|
|
|
|
#define DISPATCH_QUANTIZE(T_TYPE, C_TYPE, mantisa, exponent) \
|
|
if (val.options().dtype() == torch::T_TYPE) { \
|
|
launch_quantization<C_TYPE, mantisa, exponent>((C_TYPE*)val.data_ptr(), \
|
|
(uint8_t*)out.data_ptr(), \
|
|
num_groups, \
|
|
group_size, \
|
|
at::cuda::getCurrentCUDAStream(), \
|
|
q_range, \
|
|
q_bits, \
|
|
q_mantisa_bits, \
|
|
stochastic_rounding); \
|
|
}
|
|
|
|
at::Tensor quantize(torch::Tensor& out,
|
|
torch::Tensor& val,
|
|
int group_size,
|
|
int stochastic_rounding,
|
|
int q_bits,
|
|
int q_mantisa_bits)
|
|
{
|
|
int total_elems = at::numel(val);
|
|
float q_range = q_bits == 8 ? (q_mantisa_bits == 3 ? 480.0 : 114688.0) : // fp8 ranges
|
|
(q_bits == 12 ? 510.0 : // fp12 range
|
|
(q_bits == 6 ? 28.0 : // fp6 range
|
|
6.0)); // fp4 range (using power 2); TODO (Reza): add the power-4
|
|
// in case accuracy is not matching!
|
|
int num_groups = total_elems / group_size;
|
|
|
|
DISPATCH_QUANTIZE(kHalf, __half, 23, 8);
|
|
#ifdef BF16_AVAILABLE
|
|
DISPATCH_QUANTIZE(kBFloat16, __nv_bfloat16, 23, 8);
|
|
#endif
|
|
|
|
return out;
|
|
}
|
|
|
|
#define DISPATCH_DEQUANTIZE(T_TYPE, C_TYPE, mantisa) \
|
|
if (val.options().dtype() == torch::T_TYPE) { \
|
|
launch_dequantization<C_TYPE, mantisa>((uint8_t*)val_q.data_ptr(), \
|
|
(C_TYPE*)val.data_ptr(), \
|
|
num_groups, \
|
|
group_size, \
|
|
q_mantisa_bits, \
|
|
q_exponent_bits, \
|
|
at::cuda::getCurrentCUDAStream()); \
|
|
return; \
|
|
}
|
|
|
|
void dequantize(torch::Tensor& val,
|
|
torch::Tensor& val_q,
|
|
int group_size,
|
|
int q_mantisa_bits,
|
|
int q_exponent_bits)
|
|
{
|
|
int total_elems = at::numel(val);
|
|
|
|
int num_groups = total_elems / group_size;
|
|
|
|
DISPATCH_DEQUANTIZE(kHalf, __half, 10);
|
|
#ifdef BF16_AVAILABLE
|
|
DISPATCH_DEQUANTIZE(kBFloat16, __nv_bfloat16, 7);
|
|
#endif
|
|
}
|
|
|
|
#define DISPATCH_DEQUANTIZE_INDEX(T_TYPE, C_TYPE, mantisa) \
|
|
if (val.options().dtype() == torch::T_TYPE) { \
|
|
launch_selective_dequantization<C_TYPE, mantisa>((uint8_t*)val_q.data_ptr(), \
|
|
(C_TYPE*)val.data_ptr(), \
|
|
(int32_t*)indexes.data_ptr(), \
|
|
num_groups, \
|
|
group_size, \
|
|
num_indexes, \
|
|
q_mantisa_bits, \
|
|
q_exponent_bits, \
|
|
at::cuda::getCurrentCUDAStream()); \
|
|
return; \
|
|
}
|
|
void selective_dequantize(torch::Tensor& val,
|
|
torch::Tensor& val_q,
|
|
torch::Tensor& indexes,
|
|
int group_size,
|
|
int q_mantisa_bits,
|
|
int q_exponent_bits)
|
|
{
|
|
int total_elems = at::numel(val);
|
|
int num_indexes = indexes.size(0);
|
|
int num_groups = total_elems / group_size;
|
|
|
|
DISPATCH_DEQUANTIZE_INDEX(kHalf, __half, 10);
|
|
#ifdef BF16_AVAILABLE
|
|
DISPATCH_DEQUANTIZE_INDEX(kBFloat16, __nv_bfloat16, 7);
|
|
#endif
|
|
}
|
|
|
|
at::Tensor get_scales(torch::Tensor& out, int num_groups)
|
|
{
|
|
auto options = at::TensorOptions()
|
|
.dtype(torch::kFloat)
|
|
.layout(at::kStrided)
|
|
.device(at::kCUDA)
|
|
.requires_grad(false);
|
|
auto scales =
|
|
torch::from_blob(out.data_ptr(), {num_groups, 1}, {out.stride(0) / 4, 1}, options);
|
|
return scales;
|
|
}
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
|
|
{
|
|
m.def("quantize", &quantize, "quantize function");
|
|
m.def("dequantize", &dequantize, "dequantize function");
|
|
m.def("get_scales", &get_scales, "get scales function");
|
|
m.def("selective_dequantize", &selective_dequantize, "selective dequantize function");
|
|
}
|