chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,72 @@
|
||||
// Copyright (c) Microsoft Corporation.
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// DeepSpeed Team
|
||||
|
||||
#include "dequantization_utils.h"
|
||||
#include "memory_access_utils.h"
|
||||
|
||||
namespace cg = cooperative_groups;
|
||||
|
||||
template <typename T, int numBits, dequantize::Type qType, int unroll, int threads>
|
||||
__global__ void dequantize_kernel(T* __restrict__ dequant_data,
|
||||
const int8_t* __restrict__ q_data,
|
||||
const float* __restrict__ q_params,
|
||||
int elems_per_group,
|
||||
int total_elems)
|
||||
{
|
||||
dequantize::to_global<T, numBits, qType, unroll, threads>(
|
||||
dequant_data, q_data, q_params, elems_per_group, total_elems);
|
||||
}
|
||||
|
||||
#define LAUNCH_DEQUANT_KERNEL(num_bits, q_type) \
|
||||
dequantize_kernel<T, num_bits, q_type, unroll, threads><<<grid, block, 0, stream>>>( \
|
||||
dequant_data, q_data, q_params, elems_per_group, total_elems);
|
||||
|
||||
template <typename T>
|
||||
void launch_dequantize_kernel(T* dequant_data,
|
||||
const int8_t* q_data,
|
||||
const float* q_params,
|
||||
quantize::Type q_type,
|
||||
int num_bits,
|
||||
int elems_per_group,
|
||||
int total_elems,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
constexpr int unroll = 8;
|
||||
constexpr int threads = 512;
|
||||
constexpr int elems_per_block = unroll * threads * dequantize::granularity / (sizeof(T));
|
||||
|
||||
const dim3 block(threads);
|
||||
const dim3 grid((total_elems + elems_per_block - 1) / elems_per_block);
|
||||
|
||||
// TODO(cmikeh2): It may make sense to tune unroll, there is perf benefit for large
|
||||
// problem sizes with this large unroll value.
|
||||
if (num_bits == 8 && q_type == quantize::Type::Symmetric) {
|
||||
LAUNCH_DEQUANT_KERNEL(8, quantize::Type::Symmetric);
|
||||
} else if (num_bits == 8 && q_type == quantize::Type::Asymmetric) {
|
||||
LAUNCH_DEQUANT_KERNEL(8, quantize::Type::Asymmetric);
|
||||
} else if (num_bits == 4 && q_type == quantize::Type::Symmetric) {
|
||||
LAUNCH_DEQUANT_KERNEL(4, quantize::Type::Symmetric);
|
||||
} else if (num_bits == 4 && q_type == quantize::Type::Asymmetric) {
|
||||
LAUNCH_DEQUANT_KERNEL(4, quantize::Type::Asymmetric);
|
||||
}
|
||||
}
|
||||
|
||||
template void launch_dequantize_kernel(__half* dequant_data,
|
||||
const int8_t* q_data,
|
||||
const float* q_params,
|
||||
quantize::Type q_type,
|
||||
int num_bits,
|
||||
int elems_per_group,
|
||||
int total_elems,
|
||||
cudaStream_t stream);
|
||||
|
||||
template void launch_dequantize_kernel(float* dequant_data,
|
||||
const int8_t* q_data,
|
||||
const float* q_params,
|
||||
quantize::Type q_type,
|
||||
int num_bits,
|
||||
int elems_per_group,
|
||||
int total_elems,
|
||||
cudaStream_t stream);
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,404 @@
|
||||
// Copyright (c) Microsoft Corporation.
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// DeepSpeed Team
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <torch/extension.h>
|
||||
#include <cassert>
|
||||
#include <vector>
|
||||
#include "quantization.h"
|
||||
|
||||
template <typename T>
|
||||
at::Tensor ds_quantize(at::Tensor& vals, int groups, int bits)
|
||||
{
|
||||
auto t_size = vals.sizes();
|
||||
int size = 1;
|
||||
for (auto dim : t_size) size *= dim;
|
||||
|
||||
if ((((size / groups) - 1) / 4096 + 1) <= 256) {
|
||||
launch_fake_quantize_kernel(
|
||||
(T*)vals.data_ptr(), size, groups, bits, at::cuda::getCurrentCUDAStream());
|
||||
}
|
||||
return vals;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
at::Tensor ds_sr_quantize(at::Tensor& vals, int groups, int bits)
|
||||
{
|
||||
auto t_size = vals.sizes();
|
||||
int size = 1;
|
||||
for (auto dim : t_size) size *= dim;
|
||||
|
||||
if (((size / groups) / 4 / 1024) <= 256) {
|
||||
launch_sr_fake_quantize_kernel(
|
||||
(T*)vals.data_ptr(), size, groups, bits, at::cuda::getCurrentCUDAStream());
|
||||
}
|
||||
return vals;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
at::Tensor ds_quantize_asym(at::Tensor& vals, int groups, int bits)
|
||||
{
|
||||
auto t_size = vals.sizes();
|
||||
int size = 1;
|
||||
for (auto dim : t_size) size *= dim;
|
||||
|
||||
if ((((size / groups) - 1) / 4096 + 1) <= 256) {
|
||||
launch_fake_quantize_kernel_asym(
|
||||
(T*)vals.data_ptr(), size, groups, bits, at::cuda::getCurrentCUDAStream());
|
||||
}
|
||||
return vals;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
at::Tensor ds_sr_quantize_asym(at::Tensor& vals, int groups, int bits)
|
||||
{
|
||||
auto t_size = vals.sizes();
|
||||
int size = 1;
|
||||
for (auto dim : t_size) size *= dim;
|
||||
|
||||
if (((size / groups) / 4 / 1024) <= 256) {
|
||||
launch_sr_fake_quantize_kernel_asym(
|
||||
(T*)vals.data_ptr(), size, groups, bits, at::cuda::getCurrentCUDAStream());
|
||||
}
|
||||
return vals;
|
||||
}
|
||||
|
||||
std::vector<at::Tensor> quantize_kernel(at::Tensor& input_vals,
|
||||
int groups,
|
||||
int numBits,
|
||||
quantize::Type quantType)
|
||||
{
|
||||
auto dtype = at::kFloat;
|
||||
auto params_options = at::TensorOptions()
|
||||
.dtype(dtype)
|
||||
.layout(at::kStrided)
|
||||
.device(at::kCUDA)
|
||||
.requires_grad(false);
|
||||
const int param_elems = (quantize::requires_offset(quantType)) ? 2 : 1;
|
||||
auto params = torch::empty({groups, param_elems}, params_options);
|
||||
|
||||
auto output_options = at::TensorOptions()
|
||||
.dtype(at::kChar)
|
||||
.layout(at::kStrided)
|
||||
.device(at::kCUDA)
|
||||
.requires_grad(false);
|
||||
|
||||
auto output_sizes = input_vals.sizes().vec();
|
||||
output_sizes[output_sizes.size() - 1] /= numBits == 8 ? 1 : 2;
|
||||
auto output = torch::empty(output_sizes, output_options);
|
||||
|
||||
const int elems_per_group = at::numel(input_vals) / groups;
|
||||
|
||||
launch_quant((int8_t*)output.data_ptr(),
|
||||
(float*)params.data_ptr(),
|
||||
(__half*)input_vals.data_ptr(),
|
||||
groups,
|
||||
elems_per_group,
|
||||
numBits,
|
||||
quantType,
|
||||
at::cuda::getCurrentCUDAStream());
|
||||
|
||||
return {output, params};
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
at::Tensor dequantize(at::Tensor& quantized_data,
|
||||
at::Tensor& params,
|
||||
int groups,
|
||||
int num_bits,
|
||||
quantize::Type quant_type)
|
||||
{
|
||||
auto dtype = (std::is_same<T, float>::value) ? torch::kFloat32 : torch::kFloat16;
|
||||
auto output_options = at::TensorOptions()
|
||||
.dtype(dtype)
|
||||
.layout(at::kStrided)
|
||||
.device(at::kCUDA)
|
||||
.requires_grad(false);
|
||||
|
||||
auto output_sizes = quantized_data.sizes().vec();
|
||||
output_sizes[output_sizes.size() - 1] *= num_bits == 8 ? 1 : 2;
|
||||
auto output = torch::empty(output_sizes, output_options);
|
||||
|
||||
const int total_elems = at::numel(output);
|
||||
const int elems_per_group = total_elems / groups;
|
||||
|
||||
launch_dequantize_kernel((T*)output.data_ptr(),
|
||||
(const int8_t*)quantized_data.data_ptr(),
|
||||
(const float*)params.data_ptr(),
|
||||
quant_type,
|
||||
num_bits,
|
||||
elems_per_group,
|
||||
total_elems,
|
||||
at::cuda::getCurrentCUDAStream());
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
at::Tensor dequantize_int4_to_half_experimental(at::Tensor& data_in,
|
||||
at::Tensor& scale_buffer,
|
||||
at::Tensor& min_val_buffer,
|
||||
int num_group,
|
||||
int group_size)
|
||||
{
|
||||
auto output_options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA);
|
||||
auto output = torch::empty({num_group, group_size}, output_options);
|
||||
|
||||
launch_dequantize_int4_to_half_experimental((uint8_t*)data_in.data_ptr(),
|
||||
(half*)output.data_ptr(),
|
||||
(half*)scale_buffer.data_ptr(),
|
||||
(half*)min_val_buffer.data_ptr(),
|
||||
num_group,
|
||||
group_size,
|
||||
at::cuda::getCurrentCUDAStream());
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
at::Tensor dequantize_int8_to_half_experimental(at::Tensor& data_in,
|
||||
at::Tensor& scale_buffer,
|
||||
at::Tensor& min_val_buffer,
|
||||
int num_group,
|
||||
int group_size)
|
||||
{
|
||||
auto output_options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA);
|
||||
auto output = torch::empty({num_group, group_size}, output_options);
|
||||
|
||||
launch_dequantize_int8_to_half_experimental((uint8_t*)data_in.data_ptr(),
|
||||
(half*)output.data_ptr(),
|
||||
(half*)scale_buffer.data_ptr(),
|
||||
(half*)min_val_buffer.data_ptr(),
|
||||
num_group,
|
||||
group_size,
|
||||
at::cuda::getCurrentCUDAStream());
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
std::vector<at::Tensor> ds_loco_swizzle_quant(at::Tensor& input_vals,
|
||||
at::Tensor& error_feedback,
|
||||
float err_beta,
|
||||
int groups,
|
||||
int num_bits,
|
||||
quantize::Type quant_type,
|
||||
int pipeline_size,
|
||||
int nodes,
|
||||
int devices_per_node)
|
||||
{
|
||||
auto scales_options = at::TensorOptions()
|
||||
.dtype(at::kFloat)
|
||||
.layout(at::kStrided)
|
||||
.device(at::kCUDA)
|
||||
.requires_grad(false);
|
||||
const int scales_elems = (quantize::requires_offset(quant_type)) ? 2 : 1;
|
||||
auto scales = torch::empty({groups, scales_elems}, scales_options);
|
||||
|
||||
auto output_options = at::TensorOptions()
|
||||
.dtype(at::kChar)
|
||||
.layout(at::kStrided)
|
||||
.device(at::kCUDA)
|
||||
.requires_grad(false);
|
||||
|
||||
const int quantization_scalar = 8 / num_bits;
|
||||
const int compressed_vals = at::numel(input_vals) / quantization_scalar;
|
||||
|
||||
auto output = torch::empty({compressed_vals}, output_options);
|
||||
const int elems_per_group = at::numel(input_vals) / groups;
|
||||
|
||||
launch_loco_swizzled_quant(reinterpret_cast<int8_t*>(output.data_ptr()),
|
||||
reinterpret_cast<float*>(scales.data_ptr()),
|
||||
reinterpret_cast<const __half*>(input_vals.data_ptr()),
|
||||
reinterpret_cast<__half*>(error_feedback.data_ptr()),
|
||||
err_beta,
|
||||
num_bits,
|
||||
quant_type,
|
||||
groups,
|
||||
elems_per_group,
|
||||
pipeline_size,
|
||||
nodes,
|
||||
devices_per_node,
|
||||
at::cuda::getCurrentCUDAStream());
|
||||
|
||||
return {output, scales};
|
||||
}
|
||||
|
||||
std::vector<at::Tensor> ds_swizzle_quant(at::Tensor& input_vals,
|
||||
int groups,
|
||||
int num_bits,
|
||||
quantize::Type quant_type,
|
||||
int pipeline_size,
|
||||
int nodes,
|
||||
int devices_per_node)
|
||||
{
|
||||
auto scales_options = at::TensorOptions()
|
||||
.dtype(at::kFloat)
|
||||
.layout(at::kStrided)
|
||||
.device(at::kCUDA)
|
||||
.requires_grad(false);
|
||||
const int scales_elems = (quantize::requires_offset(quant_type)) ? 2 : 1;
|
||||
auto scales = torch::empty({groups, scales_elems}, scales_options);
|
||||
|
||||
auto output_options = at::TensorOptions()
|
||||
.dtype(at::kChar)
|
||||
.layout(at::kStrided)
|
||||
.device(at::kCUDA)
|
||||
.requires_grad(false);
|
||||
|
||||
const int quantization_scalar = 8 / num_bits;
|
||||
const int compressed_vals = at::numel(input_vals) / quantization_scalar;
|
||||
|
||||
auto output = torch::empty({compressed_vals}, output_options);
|
||||
const int elems_per_group = at::numel(input_vals) / groups;
|
||||
|
||||
launch_swizzled_quant((int8_t*)output.data_ptr(),
|
||||
(float*)scales.data_ptr(),
|
||||
(__half*)input_vals.data_ptr(),
|
||||
num_bits,
|
||||
quant_type,
|
||||
groups,
|
||||
elems_per_group,
|
||||
pipeline_size,
|
||||
nodes,
|
||||
devices_per_node,
|
||||
at::cuda::getCurrentCUDAStream());
|
||||
|
||||
return {output, scales};
|
||||
}
|
||||
|
||||
std::vector<at::Tensor> quantized_reduction(at::Tensor& input_vals,
|
||||
at::Tensor& input_scales,
|
||||
int in_groups,
|
||||
int out_groups,
|
||||
int num_bits,
|
||||
quantize::Type quant_type,
|
||||
int devices_per_node)
|
||||
{
|
||||
auto scales_options = at::TensorOptions()
|
||||
.dtype(at::kFloat)
|
||||
.layout(at::kStrided)
|
||||
.device(at::kCUDA)
|
||||
.requires_grad(false);
|
||||
const int scales_elems = (quantize::requires_offset(quant_type)) ? 2 : 1;
|
||||
auto scales = torch::empty({out_groups, scales_elems}, scales_options);
|
||||
|
||||
auto output_options = at::TensorOptions()
|
||||
.dtype(at::kChar)
|
||||
.layout(at::kStrided)
|
||||
.device(at::kCUDA)
|
||||
.requires_grad(false);
|
||||
|
||||
std::vector<int64_t> sz(input_vals.sizes().begin(), input_vals.sizes().end());
|
||||
sz[sz.size() - 1] = sz.back() / devices_per_node; // num of GPU per nodes
|
||||
const int elems_per_in_tensor = at::numel(input_vals) / devices_per_node;
|
||||
auto output = torch::empty(sz, output_options);
|
||||
|
||||
const int elems_per_in_group = elems_per_in_tensor / (in_groups / devices_per_node);
|
||||
const int elems_per_out_group = elems_per_in_tensor / out_groups;
|
||||
|
||||
launch_dequant_reduce((int8_t*)output.data_ptr(),
|
||||
(float*)scales.data_ptr(),
|
||||
(const int8_t*)input_vals.data_ptr(),
|
||||
(const float*)input_scales.data_ptr(),
|
||||
devices_per_node,
|
||||
num_bits,
|
||||
quant_type,
|
||||
out_groups,
|
||||
elems_per_out_group,
|
||||
elems_per_in_tensor,
|
||||
in_groups / devices_per_node,
|
||||
elems_per_in_group,
|
||||
at::cuda::getCurrentCUDAStream());
|
||||
return {output, scales};
|
||||
}
|
||||
|
||||
std::vector<at::Tensor> loco_quantized_reduction(at::Tensor& input_vals,
|
||||
at::Tensor& input_scales,
|
||||
at::Tensor& error_feedback,
|
||||
float err_beta,
|
||||
int in_groups,
|
||||
int out_groups,
|
||||
int num_bits,
|
||||
quantize::Type quant_type,
|
||||
int devices_per_node)
|
||||
{
|
||||
auto scales_options = at::TensorOptions()
|
||||
.dtype(at::kFloat)
|
||||
.layout(at::kStrided)
|
||||
.device(at::kCUDA)
|
||||
.requires_grad(false);
|
||||
|
||||
const int scales_elems = (quantize::requires_offset(quant_type)) ? 2 : 1;
|
||||
|
||||
auto scales = torch::empty({out_groups, scales_elems}, scales_options);
|
||||
|
||||
auto output_options = at::TensorOptions()
|
||||
.dtype(at::kChar)
|
||||
.layout(at::kStrided)
|
||||
.device(at::kCUDA)
|
||||
.requires_grad(false);
|
||||
|
||||
std::vector<int64_t> sz(input_vals.sizes().begin(), input_vals.sizes().end());
|
||||
sz[sz.size() - 1] = sz.back() / devices_per_node;
|
||||
|
||||
const int elems_per_in_tensor = at::numel(input_vals) / devices_per_node;
|
||||
|
||||
auto output = torch::empty(sz, output_options);
|
||||
|
||||
const int elems_per_in_group = elems_per_in_tensor / (in_groups / devices_per_node);
|
||||
const int elems_per_out_group = elems_per_in_tensor / out_groups;
|
||||
|
||||
launch_loco_dequant_reduce((int8_t*)output.data_ptr(),
|
||||
(float*)scales.data_ptr(),
|
||||
(const int8_t*)input_vals.data_ptr(),
|
||||
(const float*)input_scales.data_ptr(),
|
||||
devices_per_node,
|
||||
num_bits,
|
||||
quant_type,
|
||||
out_groups,
|
||||
elems_per_out_group,
|
||||
elems_per_in_tensor,
|
||||
in_groups / devices_per_node,
|
||||
elems_per_in_group,
|
||||
(__half2*)error_feedback.data_ptr(),
|
||||
err_beta,
|
||||
at::cuda::getCurrentCUDAStream());
|
||||
|
||||
return {output, scales};
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
|
||||
{
|
||||
m.def("ds_quantize_fp32", &ds_quantize<float>, "DeepSpeed Quantize with fp32 (CUDA)");
|
||||
m.def("ds_quantize_fp16", &ds_quantize<__half>, "DeepSpeed Quantize with fp16 (CUDA)");
|
||||
m.def("ds_sr_quantize_fp32", &ds_sr_quantize<float>, "DeepSpeed Quantize with fp32 (CUDA)");
|
||||
m.def("ds_sr_quantize_fp16", &ds_sr_quantize<__half>, "DeepSpeed Quantize with fp16 (CUDA)");
|
||||
m.def("ds_quantize_asym_fp32", &ds_quantize_asym<float>, "DeepSpeed Quantize with fp32 (CUDA)");
|
||||
m.def(
|
||||
"ds_quantize_asym_fp16", &ds_quantize_asym<__half>, "DeepSpeed Quantize with fp16 (CUDA)");
|
||||
m.def("ds_sr_quantize_asym_fp32",
|
||||
&ds_sr_quantize_asym<float>,
|
||||
"DeepSpeed Quantize with fp32 (CUDA)");
|
||||
m.def("ds_sr_quantize_asym_fp16",
|
||||
&ds_sr_quantize_asym<__half>,
|
||||
"DeepSpeed Quantize with fp16 (CUDA)");
|
||||
pybind11::enum_<quantize::Type>(m, "QuantizationType")
|
||||
.value("Symmetric", quantize::Type::Symmetric)
|
||||
.value("Asymmetric", quantize::Type::Asymmetric)
|
||||
.export_values();
|
||||
m.def("quantize", &quantize_kernel);
|
||||
m.def("dequantize", &dequantize<__half>);
|
||||
m.def("dequantize_fp32", &dequantize<float>);
|
||||
m.def("dequantize_int4_to_half_experimental",
|
||||
&dequantize_int4_to_half_experimental,
|
||||
"Dequantize int4 to half (experimental)");
|
||||
m.def("dequantize_int8_to_half_experimental",
|
||||
&dequantize_int8_to_half_experimental,
|
||||
"Dequantize int8 to half (experimental)");
|
||||
m.def("swizzle_quant", &ds_swizzle_quant);
|
||||
m.def("quantized_reduction", &quantized_reduction);
|
||||
m.def("loco_swizzle_quant", &ds_loco_swizzle_quant, "LoCo Swizzled Quantization Kernel");
|
||||
m.def("loco_quantized_reduction",
|
||||
&loco_quantized_reduction,
|
||||
"LoCo Quantization and Reduction Kernel");
|
||||
}
|
||||
@@ -0,0 +1,557 @@
|
||||
// Copyright (c) Microsoft Corporation.
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// DeepSpeed Team
|
||||
|
||||
#include <cstdio>
|
||||
#include "dequantization_utils.h"
|
||||
#include "ds_kernel_utils.h"
|
||||
#include "memory_access_utils.h"
|
||||
#include "quantization_utils.h"
|
||||
#include "reduction_utils.h"
|
||||
|
||||
using rop = reduce::ROpType;
|
||||
|
||||
/*
|
||||
TODO(cmikeh2): Add implementation that better handles larger nodes. It would like make sense
|
||||
to leverage some parallel reductions here to improve performance.
|
||||
*/
|
||||
|
||||
template <int numBits, int numTensors, int totalChunks, quantize::Type quantType>
|
||||
__global__ void __launch_bounds__(1024) dequant_reduce(int8_t* reduced_data,
|
||||
float* reduced_scales,
|
||||
const int8_t* input_data,
|
||||
const float* input_scales,
|
||||
int elems_per_out_group,
|
||||
int elems_per_in_tensor,
|
||||
int groups_per_in_tensor,
|
||||
int elems_per_in_group,
|
||||
int num_tensors)
|
||||
{
|
||||
cg::thread_block tb = cg::this_thread_block();
|
||||
cg::thread_block_tile<hw_warp_size> warp = cg::tiled_partition<hw_warp_size>(tb);
|
||||
|
||||
// NOTE(cmikeh2): This probably could be hardcoded to a larger number,
|
||||
// but that means even stronger restrictions on the number of elements per group
|
||||
// A performance analysis here might be beneficial
|
||||
constexpr int mem_granularity = (numBits == 8) ? 8 : 4;
|
||||
constexpr int elems_per_load = mem_granularity / sizeof(int8_t); // div by 1
|
||||
constexpr int storage_values = 16 / sizeof(__half2);
|
||||
|
||||
const int block_offset = tb.group_index().x * elems_per_out_group;
|
||||
const int elem_offset = tb.thread_index().x * elems_per_load;
|
||||
const int base_offset = block_offset + elem_offset;
|
||||
const int stride = tb.group_dim().x * elems_per_load;
|
||||
|
||||
__half2 local_buffer[totalChunks * storage_values];
|
||||
|
||||
quantize::GroupStats<quantType> stats;
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < totalChunks; i++) {
|
||||
__half2* iteration_buffer = local_buffer + i * storage_values;
|
||||
|
||||
#pragma unroll
|
||||
for (int j = 0; j < storage_values; j++) {
|
||||
iteration_buffer[j] = reduce::init<rop::Add, __half2>();
|
||||
}
|
||||
|
||||
const int iter_offset = i * stride + base_offset;
|
||||
const int iter_scale_idx = iter_offset / elems_per_in_group;
|
||||
bool do_loads = i * stride + elem_offset < elems_per_out_group;
|
||||
|
||||
if (numTensors > 0) {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < numTensors; j++) {
|
||||
if (do_loads) {
|
||||
int8_t load_buffer[elems_per_load];
|
||||
|
||||
mem_access::load_global<mem_granularity>(
|
||||
load_buffer, input_data + j * elems_per_in_tensor + iter_offset);
|
||||
|
||||
quantize::Params<quantType, numBits> params(
|
||||
input_scales + j * groups_per_in_tensor, iter_scale_idx);
|
||||
|
||||
__half2 dequant_buffer[storage_values];
|
||||
dequantize::chunk<numBits, quantType>(dequant_buffer, load_buffer, params);
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < storage_values; k++) {
|
||||
iteration_buffer[k] =
|
||||
reduce::element<rop::Add>(iteration_buffer[k], dequant_buffer[k]);
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
#pragma unroll 4
|
||||
for (int j = 0; j < num_tensors; j++) {
|
||||
if (do_loads) {
|
||||
int8_t load_buffer[elems_per_load];
|
||||
|
||||
mem_access::load_global<mem_granularity>(
|
||||
load_buffer, input_data + j * elems_per_in_tensor + iter_offset);
|
||||
|
||||
quantize::Params<quantType, numBits> params(
|
||||
input_scales + j * groups_per_in_tensor, iter_scale_idx);
|
||||
|
||||
__half2 dequant_buffer[storage_values];
|
||||
dequantize::chunk<numBits, quantType>(dequant_buffer, load_buffer, params);
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < storage_values; k++) {
|
||||
iteration_buffer[k] =
|
||||
reduce::element<rop::Add>(iteration_buffer[k], dequant_buffer[k]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int j = 0; j < storage_values; j++) { stats.update(iteration_buffer[j]); }
|
||||
}
|
||||
|
||||
auto params = stats.template get_params<numBits, 1024>(tb, warp);
|
||||
|
||||
if (tb.thread_index().x == 0) { params.store(reduced_scales, tb.group_index().x); }
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < totalChunks; i++) {
|
||||
const int iter_offset = i * stride + base_offset;
|
||||
if (i * stride + elem_offset < elems_per_out_group) {
|
||||
int8_t local_output[elems_per_load];
|
||||
quantize::_chunk<numBits, quantType>(
|
||||
local_output, local_buffer + i * storage_values, params);
|
||||
mem_access::store_global<mem_granularity>(reduced_data + iter_offset, local_output);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int Power>
|
||||
int32_t pow2_round(int32_t raw_value)
|
||||
{
|
||||
return (((raw_value - 1) >> Power) + 1) << Power;
|
||||
}
|
||||
|
||||
#define LAUNCH_DEQUANT_REDUCE(num_chunks) \
|
||||
dequant_reduce<numBits, numTensors, num_chunks, quantType> \
|
||||
<<<grid, block, 0, stream>>>(reduced_data, \
|
||||
reduced_scales, \
|
||||
input_data, \
|
||||
input_scales, \
|
||||
elems_per_out_group, \
|
||||
elems_per_in_tensor, \
|
||||
groups_per_in_tensor, \
|
||||
elems_per_in_group, \
|
||||
num_tensors);
|
||||
|
||||
template <int numBits, int numTensors, quantize::Type quantType>
|
||||
void launch_dequant_reduce_impl(int8_t* reduced_data,
|
||||
float* reduced_scales,
|
||||
const int8_t* input_data,
|
||||
const float* input_scales,
|
||||
int out_groups,
|
||||
int elems_per_out_group,
|
||||
int elems_per_in_tensor,
|
||||
int groups_per_in_tensor,
|
||||
int elems_per_in_group,
|
||||
int num_tensors,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
// This is a coincidence. This is derived by 8 halves per 16 bytes with 2-way packing for int4
|
||||
constexpr int elems_per_thread = numBits;
|
||||
const int one_step_threads =
|
||||
next_pow2((elems_per_out_group + elems_per_thread - 1) / (elems_per_thread));
|
||||
// TODO(cmikeh2): Tune this
|
||||
const int threads = (one_step_threads < 1024) ? one_step_threads : 1024;
|
||||
|
||||
dim3 block(threads);
|
||||
dim3 grid(out_groups);
|
||||
|
||||
const int elems_per_step = threads * elems_per_thread;
|
||||
const int unroll_raw = (elems_per_out_group + elems_per_step - 1) / elems_per_step;
|
||||
|
||||
const int unroll = (unroll_raw >= 4) ? pow2_round<1>(unroll_raw) : unroll_raw;
|
||||
|
||||
if (unroll == 1) {
|
||||
// 0-4096 elems
|
||||
LAUNCH_DEQUANT_REDUCE(1);
|
||||
} else if (unroll == 2) {
|
||||
// 4097-8192 etc...
|
||||
LAUNCH_DEQUANT_REDUCE(2);
|
||||
} else if (unroll == 3) {
|
||||
LAUNCH_DEQUANT_REDUCE(3);
|
||||
} else if (unroll == 4) {
|
||||
LAUNCH_DEQUANT_REDUCE(4);
|
||||
} else if (unroll == 6) {
|
||||
LAUNCH_DEQUANT_REDUCE(6);
|
||||
} else if (unroll == 8) {
|
||||
LAUNCH_DEQUANT_REDUCE(8);
|
||||
} else if (unroll == 10) {
|
||||
LAUNCH_DEQUANT_REDUCE(10);
|
||||
} else if (unroll == 12) {
|
||||
// 48k limit
|
||||
LAUNCH_DEQUANT_REDUCE(12);
|
||||
} else {
|
||||
assert(false);
|
||||
}
|
||||
}
|
||||
|
||||
#define LAUNCH_DEQUANT_REDUCE_IMPL(NUM_BITS, NUM_GPUS, QUANT_TYPE) \
|
||||
launch_dequant_reduce_impl<NUM_BITS, NUM_GPUS, QUANT_TYPE>(reduced_data, \
|
||||
reduced_scales, \
|
||||
input_data, \
|
||||
input_scales, \
|
||||
out_groups, \
|
||||
elems_per_out_group, \
|
||||
elems_per_in_tensor, \
|
||||
groups_per_in_tensor, \
|
||||
elems_per_in_group, \
|
||||
num_gpus, \
|
||||
stream);
|
||||
|
||||
void launch_dequant_reduce(int8_t* reduced_data,
|
||||
float* reduced_scales,
|
||||
const int8_t* input_data,
|
||||
const float* input_scales,
|
||||
int num_gpus,
|
||||
int num_bits,
|
||||
quantize::Type quant_type,
|
||||
int out_groups,
|
||||
int elems_per_out_group,
|
||||
int elems_per_in_tensor,
|
||||
int groups_per_in_tensor,
|
||||
int elems_per_in_group,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
if (quant_type == quantize::Type::Symmetric) {
|
||||
if (num_bits == 4) {
|
||||
if (num_gpus == 8) {
|
||||
LAUNCH_DEQUANT_REDUCE_IMPL(4, 8, quantize::Type::Symmetric);
|
||||
} else if (num_gpus == 16) {
|
||||
LAUNCH_DEQUANT_REDUCE_IMPL(4, 16, quantize::Type::Symmetric);
|
||||
} else {
|
||||
LAUNCH_DEQUANT_REDUCE_IMPL(4, -1, quantize::Type::Symmetric);
|
||||
}
|
||||
} else if (num_bits == 8) {
|
||||
if (num_gpus == 8) {
|
||||
LAUNCH_DEQUANT_REDUCE_IMPL(8, 8, quantize::Type::Symmetric);
|
||||
} else if (num_gpus == 16) {
|
||||
LAUNCH_DEQUANT_REDUCE_IMPL(8, 16, quantize::Type::Symmetric);
|
||||
} else {
|
||||
LAUNCH_DEQUANT_REDUCE_IMPL(8, -1, quantize::Type::Symmetric);
|
||||
}
|
||||
}
|
||||
} else if (quant_type == quantize::Type::Asymmetric) {
|
||||
if (num_bits == 4) {
|
||||
if (num_gpus == 8) {
|
||||
LAUNCH_DEQUANT_REDUCE_IMPL(4, 8, quantize::Type::Asymmetric);
|
||||
} else if (num_gpus == 16) {
|
||||
LAUNCH_DEQUANT_REDUCE_IMPL(4, 16, quantize::Type::Asymmetric);
|
||||
} else {
|
||||
LAUNCH_DEQUANT_REDUCE_IMPL(4, -1, quantize::Type::Asymmetric);
|
||||
}
|
||||
} else if (num_bits == 8) {
|
||||
if (num_gpus == 8) {
|
||||
LAUNCH_DEQUANT_REDUCE_IMPL(8, 8, quantize::Type::Asymmetric);
|
||||
} else if (num_gpus == 16) {
|
||||
LAUNCH_DEQUANT_REDUCE_IMPL(8, 16, quantize::Type::Asymmetric);
|
||||
} else {
|
||||
LAUNCH_DEQUANT_REDUCE_IMPL(8, -1, quantize::Type::Asymmetric);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
Modified loco_dequant_reduce function that performs dequantization and reduction,
|
||||
and incorporates error-feedback by updating the error_feedback tensor in-place.
|
||||
*/
|
||||
|
||||
template <int numBits, int numTensors, int totalChunks, quantize::Type quantType>
|
||||
__global__ void __launch_bounds__(1024) loco_dequant_reduce(int8_t* reduced_data,
|
||||
float* reduced_scales,
|
||||
const int8_t* input_data,
|
||||
const float* input_scales,
|
||||
int elems_per_out_group,
|
||||
int elems_per_in_tensor,
|
||||
int groups_per_in_tensor,
|
||||
int elems_per_in_group,
|
||||
int num_tensors,
|
||||
__half2* error_feedback,
|
||||
const float err_beta)
|
||||
{
|
||||
cg::thread_block tb = cg::this_thread_block();
|
||||
cg::thread_block_tile<hw_warp_size> warp = cg::tiled_partition<hw_warp_size>(tb);
|
||||
|
||||
constexpr int mem_granularity = (numBits == 8) ? 8 : 4;
|
||||
constexpr int elems_per_load = mem_granularity / sizeof(int8_t);
|
||||
constexpr int storage_values = 16 / sizeof(__half2);
|
||||
|
||||
const int block_offset = tb.group_index().x * elems_per_out_group;
|
||||
const int elem_offset = tb.thread_index().x * elems_per_load;
|
||||
const int base_offset = block_offset + elem_offset;
|
||||
const int stride = tb.group_dim().x * elems_per_load;
|
||||
|
||||
constexpr int scaling_factor = elems_per_load / storage_values;
|
||||
const int block_offset_err = block_offset / scaling_factor;
|
||||
const int elem_offset_err = tb.thread_index().x * storage_values;
|
||||
const int base_offset_err = block_offset_err + elem_offset_err;
|
||||
const int stride_err = tb.group_dim().x * storage_values;
|
||||
|
||||
__half2 local_buffer[totalChunks * storage_values];
|
||||
__half2 err_buffer[totalChunks * storage_values];
|
||||
|
||||
quantize::GroupStats<quantType> stats;
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < totalChunks; i++) {
|
||||
__half2* iteration_buffer = local_buffer + i * storage_values;
|
||||
__half2* iter_err_buffer = err_buffer + i * storage_values;
|
||||
|
||||
#pragma unroll
|
||||
for (int j = 0; j < storage_values; j++) {
|
||||
iteration_buffer[j] = reduce::init<rop::Add, __half2>();
|
||||
}
|
||||
|
||||
const int iter_offset = i * stride + base_offset;
|
||||
const int iter_offset_err = i * stride_err + base_offset_err;
|
||||
const int iter_scale_idx = iter_offset / elems_per_in_group;
|
||||
bool do_loads = i * stride + elem_offset < elems_per_out_group;
|
||||
|
||||
if (numTensors > 0) {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < numTensors; j++) {
|
||||
if (do_loads) {
|
||||
int8_t load_buffer[elems_per_load];
|
||||
|
||||
mem_access::load_global<mem_granularity>(
|
||||
load_buffer, input_data + j * elems_per_in_tensor + iter_offset);
|
||||
|
||||
quantize::Params<quantType, numBits> params(
|
||||
input_scales + j * groups_per_in_tensor, iter_scale_idx);
|
||||
|
||||
__half2 dequant_buffer[storage_values];
|
||||
dequantize::chunk<numBits, quantType>(dequant_buffer, load_buffer, params);
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < storage_values; k++) {
|
||||
iteration_buffer[k] =
|
||||
reduce::element<rop::Add>(iteration_buffer[k], dequant_buffer[k]);
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
#pragma unroll 4
|
||||
for (int j = 0; j < num_tensors; j++) {
|
||||
if (do_loads) {
|
||||
int8_t load_buffer[elems_per_load];
|
||||
|
||||
mem_access::load_global<mem_granularity>(
|
||||
load_buffer, input_data + j * elems_per_in_tensor + iter_offset);
|
||||
|
||||
quantize::Params<quantType, numBits> params(
|
||||
input_scales + j * groups_per_in_tensor, iter_scale_idx);
|
||||
|
||||
__half2 dequant_buffer[storage_values];
|
||||
dequantize::chunk<numBits, quantType>(dequant_buffer, load_buffer, params);
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < storage_values; k++) {
|
||||
iteration_buffer[k] =
|
||||
reduce::element<rop::Add>(iteration_buffer[k], dequant_buffer[k]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
mem_access::load_global<quantize::granularity>(
|
||||
iter_err_buffer, error_feedback + iter_offset_err, do_loads);
|
||||
#pragma unroll
|
||||
for (int k = 0; k < storage_values; k++) {
|
||||
iteration_buffer[k] = __hadd2(iteration_buffer[k], iter_err_buffer[k]);
|
||||
stats.update(iteration_buffer[k]);
|
||||
}
|
||||
}
|
||||
|
||||
auto params = stats.template get_params<numBits, 1024>(tb, warp);
|
||||
|
||||
// Initialize dequantization parameters based on params
|
||||
auto de_params = params;
|
||||
de_params.scale = 1.0f / params.scale;
|
||||
if constexpr (quantType == quantize::Type::Asymmetric) { de_params.offset = params.offset; }
|
||||
|
||||
if (tb.thread_index().x == 0) { params.store(reduced_scales, tb.group_index().x); }
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < totalChunks; i++) {
|
||||
const int iter_offset = i * stride + base_offset;
|
||||
const int iter_offset_err = i * stride_err + base_offset_err;
|
||||
__half2* iteration_buffer = local_buffer + i * storage_values;
|
||||
__half2* iter_err_buffer = err_buffer + i * storage_values;
|
||||
|
||||
if (i * stride + elem_offset < elems_per_out_group) {
|
||||
// ----------- Begin Error-Feedback Modification -----------
|
||||
int8_t local_output[elems_per_load];
|
||||
quantize::_chunk<numBits, quantType>(local_output, iteration_buffer, params);
|
||||
mem_access::store_global<mem_granularity>(reduced_data + iter_offset, local_output);
|
||||
|
||||
// Dequantize the quantized output to compute the dequantized value
|
||||
__half2 dequant_buffer[storage_values];
|
||||
dequantize::chunk<numBits, quantType>(dequant_buffer, local_output, de_params);
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < storage_values; k++) {
|
||||
// __half2 to float2
|
||||
float2 iter_buf_f = __half22float2(iteration_buffer[k]);
|
||||
float2 dequant_buf_f = __half22float2(dequant_buffer[k]);
|
||||
|
||||
// Update within float precision
|
||||
float2 new_error_f;
|
||||
new_error_f.x = iter_buf_f.x - dequant_buf_f.x;
|
||||
new_error_f.y = iter_buf_f.y - dequant_buf_f.y;
|
||||
|
||||
float2 iter_err_buf_f = __half22float2(iter_err_buffer[k]);
|
||||
|
||||
iter_err_buf_f.x = err_beta * iter_err_buf_f.x + (1.0f - err_beta) * new_error_f.x;
|
||||
iter_err_buf_f.y = err_beta * iter_err_buf_f.y + (1.0f - err_beta) * new_error_f.y;
|
||||
|
||||
// float2 back to __half2
|
||||
iter_err_buffer[k] = __float22half2_rn(iter_err_buf_f);
|
||||
}
|
||||
mem_access::store_global<quantize::granularity>(error_feedback + iter_offset_err,
|
||||
iter_err_buffer);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define LAUNCH_LOCO_DEQUANT_REDUCE(num_chunks) \
|
||||
loco_dequant_reduce<numBits, numTensors, num_chunks, quantType> \
|
||||
<<<grid, block, 0, stream>>>(reduced_data, \
|
||||
reduced_scales, \
|
||||
input_data, \
|
||||
input_scales, \
|
||||
elems_per_out_group, \
|
||||
elems_per_in_tensor, \
|
||||
groups_per_in_tensor, \
|
||||
elems_per_in_group, \
|
||||
num_tensors, \
|
||||
error_feedback, \
|
||||
err_beta);
|
||||
|
||||
template <int numBits, int numTensors, quantize::Type quantType>
|
||||
void launch_loco_dequant_reduce_impl(int8_t* reduced_data,
|
||||
float* reduced_scales,
|
||||
const int8_t* input_data,
|
||||
const float* input_scales,
|
||||
int out_groups,
|
||||
int elems_per_out_group,
|
||||
int elems_per_in_tensor,
|
||||
int groups_per_in_tensor,
|
||||
int elems_per_in_group,
|
||||
int num_tensors,
|
||||
__half2* error_feedback,
|
||||
const float err_beta,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
constexpr int elems_per_thread = numBits;
|
||||
const int one_step_threads =
|
||||
next_pow2((elems_per_out_group + elems_per_thread - 1) / (elems_per_thread));
|
||||
const int threads = (one_step_threads < 1024) ? one_step_threads : 1024;
|
||||
|
||||
dim3 block(threads);
|
||||
dim3 grid(out_groups);
|
||||
|
||||
const int elems_per_step = threads * elems_per_thread;
|
||||
const int unroll_raw = (elems_per_out_group + elems_per_step - 1) / elems_per_step;
|
||||
|
||||
const int unroll = (unroll_raw >= 4) ? pow2_round<1>(unroll_raw) : unroll_raw;
|
||||
|
||||
if (unroll == 1) {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE(1);
|
||||
} else if (unroll == 2) {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE(2);
|
||||
} else if (unroll == 3) {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE(3);
|
||||
} else if (unroll == 4) {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE(4);
|
||||
} else if (unroll == 6) {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE(6);
|
||||
} else if (unroll == 8) {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE(8);
|
||||
} else if (unroll == 10) {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE(10);
|
||||
} else if (unroll == 12) {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE(12);
|
||||
} else {
|
||||
assert(false);
|
||||
}
|
||||
}
|
||||
|
||||
#define LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(NUM_BITS, NUM_GPUS, QUANT_TYPE) \
|
||||
launch_loco_dequant_reduce_impl<NUM_BITS, NUM_GPUS, QUANT_TYPE>(reduced_data, \
|
||||
reduced_scales, \
|
||||
input_data, \
|
||||
input_scales, \
|
||||
out_groups, \
|
||||
elems_per_out_group, \
|
||||
elems_per_in_tensor, \
|
||||
groups_per_in_tensor, \
|
||||
elems_per_in_group, \
|
||||
num_gpus, \
|
||||
error_feedback, \
|
||||
err_beta, \
|
||||
stream);
|
||||
|
||||
void launch_loco_dequant_reduce(int8_t* reduced_data,
|
||||
float* reduced_scales,
|
||||
const int8_t* input_data,
|
||||
const float* input_scales,
|
||||
int num_gpus,
|
||||
int num_bits,
|
||||
quantize::Type quant_type,
|
||||
int out_groups,
|
||||
int elems_per_out_group,
|
||||
int elems_per_in_tensor,
|
||||
int groups_per_in_tensor,
|
||||
int elems_per_in_group,
|
||||
__half2* error_feedback,
|
||||
const float err_beta,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
if (quant_type == quantize::Type::Symmetric) {
|
||||
if (num_bits == 4) {
|
||||
if (num_gpus == 8) {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, 8, quantize::Type::Symmetric);
|
||||
} else if (num_gpus == 16) {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, 16, quantize::Type::Symmetric);
|
||||
} else {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, -1, quantize::Type::Symmetric);
|
||||
}
|
||||
} else if (num_bits == 8) {
|
||||
if (num_gpus == 8) {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, 8, quantize::Type::Symmetric);
|
||||
} else if (num_gpus == 16) {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, 16, quantize::Type::Symmetric);
|
||||
} else {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, -1, quantize::Type::Symmetric);
|
||||
}
|
||||
}
|
||||
} else if (quant_type == quantize::Type::Asymmetric) {
|
||||
if (num_bits == 4) {
|
||||
if (num_gpus == 8) {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, 8, quantize::Type::Asymmetric);
|
||||
} else if (num_gpus == 16) {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, 16, quantize::Type::Asymmetric);
|
||||
} else {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, -1, quantize::Type::Asymmetric);
|
||||
}
|
||||
} else if (num_bits == 8) {
|
||||
if (num_gpus == 8) {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, 8, quantize::Type::Asymmetric);
|
||||
} else if (num_gpus == 16) {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, 16, quantize::Type::Asymmetric);
|
||||
} else {
|
||||
LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, -1, quantize::Type::Asymmetric);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
// Copyright (c) Microsoft Corporation.
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// DeepSpeed Team
|
||||
|
||||
#include "ds_kernel_utils.h"
|
||||
#include "memory_access_utils.h"
|
||||
#include "quantization.h"
|
||||
#include "quantization_utils.h"
|
||||
#include "reduction_utils.h"
|
||||
|
||||
namespace cg = cooperative_groups;
|
||||
|
||||
/*
|
||||
Pure quantization kernel with no fusion.
|
||||
*/
|
||||
template <int q_bits,
|
||||
quantize::Type quant_type,
|
||||
int UNROLL,
|
||||
int internal_unroll,
|
||||
int threads_per_group,
|
||||
int max_threads>
|
||||
__global__ void cached_quantization(int8_t* __restrict__ output_data,
|
||||
float* __restrict__ params,
|
||||
const __half* __restrict__ input_data,
|
||||
int groups,
|
||||
int elems_per_group)
|
||||
{
|
||||
cg::thread_block tb = cg::this_thread_block();
|
||||
cg::thread_block_tile<hw_warp_size> warp = cg::tiled_partition<hw_warp_size>(tb);
|
||||
|
||||
// Indexing offsets
|
||||
const int block_offset =
|
||||
(tb.group_index().x * (max_threads / threads_per_group) * elems_per_group) +
|
||||
(tb.thread_index().y * elems_per_group);
|
||||
const int elem_offset = tb.thread_index().x * quantize::h_per_load;
|
||||
const int base_offset = block_offset + elem_offset;
|
||||
const int stride = tb.size() * quantize::h_per_load;
|
||||
|
||||
const __half* input_base = input_data + base_offset; //..
|
||||
|
||||
__half2 local_buffer[UNROLL * internal_unroll * quantize::h2_per_load];
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < UNROLL; i++) {
|
||||
// Convenience helper, should resolve to register indices and not realize.
|
||||
__half2* iteration_buffer = local_buffer + i * internal_unroll * quantize::h2_per_load;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < internal_unroll; j++) {
|
||||
const int iteration = i * internal_unroll + j;
|
||||
mem_access::load_global<quantize::granularity>(
|
||||
iteration_buffer + j * quantize::h2_per_load,
|
||||
input_base + iteration * stride,
|
||||
elem_offset + iteration * stride < elems_per_group);
|
||||
}
|
||||
}
|
||||
|
||||
quantize::
|
||||
local_array<quant_type, q_bits, UNROLL * internal_unroll, threads_per_group, max_threads>(
|
||||
local_buffer, params, output_data, elems_per_group, groups);
|
||||
}
|
||||
|
||||
/********* Launcher methods ***********/
|
||||
#define LAUNCH_CACHED_QUANT_CALL(q_bits, quant_type) \
|
||||
cached_quantization<q_bits, \
|
||||
quant_type, \
|
||||
unroll_factor, \
|
||||
internal_unroll_l, \
|
||||
threads_per_group, \
|
||||
max_threads> \
|
||||
<<<grid, block, 0, stream>>>(output_data, params, input_data, groups, elems_per_group);
|
||||
|
||||
#define LAUNCH_CACHED_QUANT( \
|
||||
q_bits, quant_type, unroll_factor_in, internal_unroll_in, threads_per_group_in) \
|
||||
const int unroll_factor = unroll_factor_in; \
|
||||
const int internal_unroll_l = internal_unroll_in; \
|
||||
const int threads_per_group = threads_per_group_in; \
|
||||
if (q_bits == 4) { \
|
||||
if (quant_type == quantize::Type::Asymmetric) { \
|
||||
LAUNCH_CACHED_QUANT_CALL(4, quantize::Type::Asymmetric) \
|
||||
} else { \
|
||||
LAUNCH_CACHED_QUANT_CALL(4, quantize::Type::Symmetric) \
|
||||
} \
|
||||
} else { \
|
||||
if (quant_type == quantize::Type::Asymmetric) { \
|
||||
LAUNCH_CACHED_QUANT_CALL(8, quantize::Type::Asymmetric) \
|
||||
} else { \
|
||||
LAUNCH_CACHED_QUANT_CALL(8, quantize::Type::Symmetric) \
|
||||
} \
|
||||
}
|
||||
|
||||
void launch_quant(int8_t* output_data,
|
||||
float* params,
|
||||
const __half* input_data,
|
||||
const int groups,
|
||||
const int elems_per_group,
|
||||
const int num_bits,
|
||||
const quantize::Type quant_type,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
constexpr int max_threads = 256;
|
||||
|
||||
constexpr int internal_unroll = 2;
|
||||
|
||||
const bool is_subblock_schedule = (elems_per_group <= 128) ? true : false;
|
||||
const int h_per_step = is_subblock_schedule ? quantize::h_per_load
|
||||
: quantize::h_per_load * internal_unroll;
|
||||
|
||||
// Scheduling concern: may be slightly faster for some inputs to assign multiple stages of
|
||||
// warp-sized blocks rather than stepping up to 64/96 threads
|
||||
const int one_step_threads = next_pow2((elems_per_group + h_per_step - 1) / h_per_step);
|
||||
const int threads_per_group = (one_step_threads < max_threads) ? one_step_threads : max_threads;
|
||||
|
||||
const int groups_per_block =
|
||||
is_subblock_schedule ? (max_threads + threads_per_group - 1) / threads_per_group : 1;
|
||||
const int groups_launch = (groups_per_block + groups - 1) / groups_per_block;
|
||||
|
||||
dim3 block(threads_per_group, groups_per_block);
|
||||
dim3 grid(groups_launch);
|
||||
|
||||
const int elems_per_step = threads_per_group * h_per_step;
|
||||
const int external_unroll = (elems_per_group + elems_per_step - 1) / elems_per_step;
|
||||
|
||||
if (is_subblock_schedule) {
|
||||
// <=128
|
||||
if (threads_per_group == 1) {
|
||||
LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 1);
|
||||
} else if (threads_per_group == 2) {
|
||||
LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 2);
|
||||
} else if (threads_per_group == 4) {
|
||||
LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 4);
|
||||
} else if (threads_per_group == 8) {
|
||||
LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 8);
|
||||
} else if (threads_per_group == 16) {
|
||||
LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 16);
|
||||
}
|
||||
} else if (external_unroll == 1) {
|
||||
// 129 - 4096 elems
|
||||
// (this can launch with 1-7 warps as well)
|
||||
LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, internal_unroll, max_threads);
|
||||
} else if (external_unroll == 2) {
|
||||
// 4097 - 8192 elems
|
||||
LAUNCH_CACHED_QUANT(num_bits, quant_type, 2, internal_unroll, max_threads);
|
||||
} else if (external_unroll == 3) {
|
||||
// 8193 - 12288 elems
|
||||
LAUNCH_CACHED_QUANT(num_bits, quant_type, 3, internal_unroll, max_threads);
|
||||
} else if (external_unroll == 4) {
|
||||
// 12289 - 16384 elems
|
||||
LAUNCH_CACHED_QUANT(num_bits, quant_type, 4, internal_unroll, max_threads);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,281 @@
|
||||
// Copyright (c) Microsoft Corporation.
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// DeepSpeed Team
|
||||
|
||||
#include <assert.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include "memory_access_utils.h"
|
||||
|
||||
template <typename T, int N>
|
||||
struct alignas(sizeof(T) * N) AlignedArray {
|
||||
using Element = T;
|
||||
static const int kElements = N;
|
||||
|
||||
__device__ __host__ AlignedArray() {}
|
||||
|
||||
__device__ __host__ AlignedArray(const T& rhs)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < kElements; ++idx) { this->at(idx) = rhs; }
|
||||
}
|
||||
|
||||
__device__ __host__ T& operator[](int offset)
|
||||
{
|
||||
return reinterpret_cast<T&>(this->buffer[offset]);
|
||||
}
|
||||
|
||||
__device__ __host__ const T& operator[](int offset) const
|
||||
{
|
||||
return reinterpret_cast<const T&>(this->buffer[offset]);
|
||||
}
|
||||
|
||||
__device__ __host__ T& at(int offset) { return reinterpret_cast<T&>(this->buffer[offset]); }
|
||||
|
||||
__device__ __host__ const T& at(int offset) const
|
||||
{
|
||||
return reinterpret_cast<const T&>(this->buffer[offset]);
|
||||
}
|
||||
|
||||
__device__ __host__ AlignedArray<T, N> operator+(const AlignedArray<T, N>& rhs) const
|
||||
{
|
||||
AlignedArray<T, N> ret;
|
||||
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < kElements; ++idx) { ret[idx] = this->at(idx) + rhs.at(idx); }
|
||||
|
||||
return ret;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void clear()
|
||||
{
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < kElements; ++idx) { this->at(idx) = Element(0); }
|
||||
}
|
||||
|
||||
Element buffer[N];
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct reduce_max {
|
||||
__device__ __forceinline__ T operator()(const T& lhs, const T& rhs)
|
||||
{
|
||||
return lhs > rhs ? lhs : rhs;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct reduce_min {
|
||||
__device__ __forceinline__ T operator()(const T& lhs, const T& rhs)
|
||||
{
|
||||
return lhs < rhs ? lhs : rhs;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, int N>
|
||||
struct subtract {
|
||||
__device__ __forceinline__ AlignedArray<T, N> operator()(const AlignedArray<T, N>& lhs,
|
||||
const T& rhs)
|
||||
{
|
||||
AlignedArray<T, N> ret;
|
||||
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < N; ++idx) { ret[idx] = lhs[idx] - rhs; }
|
||||
|
||||
return ret;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, int N>
|
||||
struct plus {
|
||||
__device__ __forceinline__ AlignedArray<T, N> operator()(const AlignedArray<T, N>& lhs,
|
||||
const T& rhs)
|
||||
{
|
||||
AlignedArray<T, N> ret;
|
||||
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < N; ++idx) { ret[idx] = lhs[idx] + rhs; }
|
||||
|
||||
return ret;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, int N>
|
||||
struct multiply {
|
||||
__device__ __forceinline__ AlignedArray<T, N> operator()(const AlignedArray<T, N>& lhs,
|
||||
const T& rhs)
|
||||
{
|
||||
AlignedArray<T, N> ret;
|
||||
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < N; ++idx) { ret[idx] = lhs[idx] * rhs; }
|
||||
|
||||
return ret;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, int N>
|
||||
struct clamp {
|
||||
__device__ __forceinline__ AlignedArray<T, N> operator()(const AlignedArray<T, N>& lhs,
|
||||
const T& min_val,
|
||||
const T& max_val)
|
||||
{
|
||||
AlignedArray<T, N> ret;
|
||||
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < N; ++idx) {
|
||||
ret[idx] = reduce_max<T>()(reduce_min<T>()(lhs[idx], max_val), min_val);
|
||||
}
|
||||
|
||||
return ret;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, int N>
|
||||
struct round_int;
|
||||
|
||||
template <int N>
|
||||
struct round_int<half, N> {
|
||||
__device__ __forceinline__ AlignedArray<half, N> operator()(const AlignedArray<half, N>& lhs)
|
||||
{
|
||||
AlignedArray<half, N> ret;
|
||||
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < N; ++idx) { ret[idx] = hrint(lhs[idx]); }
|
||||
|
||||
return ret;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, int N>
|
||||
struct divide {
|
||||
__device__ __forceinline__ AlignedArray<T, N> operator()(const AlignedArray<T, N>& lhs,
|
||||
const T& rhs)
|
||||
{
|
||||
AlignedArray<T, N> ret;
|
||||
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < N; ++idx) { ret[idx] = lhs[idx] / rhs; }
|
||||
|
||||
return ret;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, int N, typename Reducer>
|
||||
__device__ __forceinline__ T to_scalar(const AlignedArray<T, N>& data)
|
||||
{
|
||||
Reducer re;
|
||||
T res = data[0];
|
||||
|
||||
#pragma unroll
|
||||
for (int idx = 1; idx < N; ++idx) { res = re(res, data[idx]); }
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
template <int N>
|
||||
__device__ __forceinline__ AlignedArray<half, N * 2> int4_to_half(
|
||||
const AlignedArray<uint8_t, N>& data)
|
||||
{
|
||||
AlignedArray<half, N * 2> ret;
|
||||
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < N * 2; idx += 2) {
|
||||
ret[idx] = half(int(data[idx / 2] >> 4));
|
||||
ret[idx + 1] = half(int(data[idx / 2] & 0xf));
|
||||
}
|
||||
|
||||
return ret;
|
||||
}
|
||||
|
||||
__global__ void dequantize_int4_to_half(uint8_t* data_in,
|
||||
half* data_out,
|
||||
half* scale_buffer,
|
||||
half* min_val_buffer,
|
||||
int num_group,
|
||||
int group_size)
|
||||
{
|
||||
using AccessType = AlignedArray<uint8_t, 4>;
|
||||
using AccessTypeOut = AlignedArray<half, 8>;
|
||||
|
||||
for (int idx = threadIdx.x + blockIdx.x * blockDim.x; idx < num_group * group_size / 8;
|
||||
idx += blockDim.x * gridDim.x) {
|
||||
int id_group = idx / (group_size / 8);
|
||||
AccessType value = reinterpret_cast<AccessType*>(data_in)[idx];
|
||||
half scale = scale_buffer[id_group];
|
||||
half min_value = min_val_buffer[id_group];
|
||||
|
||||
AccessTypeOut output = int4_to_half(value);
|
||||
output = divide<half, 8>()(output, scale);
|
||||
output = plus<half, 8>()(output, min_value);
|
||||
|
||||
reinterpret_cast<AccessTypeOut*>(data_out)[idx] = output;
|
||||
}
|
||||
}
|
||||
|
||||
void launch_dequantize_int4_to_half_experimental(uint8_t* data_in,
|
||||
half* data_out,
|
||||
half* scale_buffer,
|
||||
half* min_val_buffer,
|
||||
int num_group,
|
||||
int group_size,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
int num_warp = num_group / 4;
|
||||
int num_block = num_warp / 8; // 256 trd / block
|
||||
|
||||
dequantize_int4_to_half<<<num_block, 256, 0, stream>>>(
|
||||
data_in, data_out, scale_buffer, min_val_buffer, num_group, group_size);
|
||||
}
|
||||
|
||||
template <int N>
|
||||
__device__ __forceinline__ AlignedArray<half, N> int8_to_half(const AlignedArray<uint8_t, N>& data)
|
||||
{
|
||||
AlignedArray<half, N> ret;
|
||||
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < N; idx += 1) { ret[idx] = half(int(data[idx])); }
|
||||
|
||||
return ret;
|
||||
}
|
||||
|
||||
__global__ void dequantize_int8_to_half(uint8_t* data_in,
|
||||
half* data_out,
|
||||
half* scale_buffer,
|
||||
half* min_val_buffer,
|
||||
int num_group,
|
||||
int group_size)
|
||||
{
|
||||
using AccessType = AlignedArray<uint8_t, 8>;
|
||||
using AccessTypeOut = AlignedArray<half, 8>;
|
||||
|
||||
for (int idx = threadIdx.x + blockIdx.x * blockDim.x; idx < num_group * group_size / 8;
|
||||
idx += blockDim.x * gridDim.x) {
|
||||
int id_group = idx / (group_size / 8);
|
||||
AccessType value = reinterpret_cast<AccessType*>(data_in)[idx];
|
||||
half scale = scale_buffer[id_group];
|
||||
half min_value = min_val_buffer[id_group];
|
||||
|
||||
AccessTypeOut output = int8_to_half(value);
|
||||
output = divide<half, 8>()(output, scale);
|
||||
output = plus<half, 8>()(output, min_value);
|
||||
|
||||
reinterpret_cast<AccessTypeOut*>(data_out)[idx] = output;
|
||||
}
|
||||
}
|
||||
|
||||
void launch_dequantize_int8_to_half_experimental(uint8_t* data_in,
|
||||
half* data_out,
|
||||
half* scale_buffer,
|
||||
half* min_val_buffer,
|
||||
int num_group,
|
||||
int group_size,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
int num_warp = num_group / 4;
|
||||
int num_block = num_warp / 8; // 256 trd / block
|
||||
|
||||
dequantize_int8_to_half<<<num_block, 256, 0, stream>>>(
|
||||
data_in, data_out, scale_buffer, min_val_buffer, num_group, group_size);
|
||||
}
|
||||
@@ -0,0 +1,427 @@
|
||||
// Copyright (c) Microsoft Corporation.
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// DeepSpeed Team
|
||||
|
||||
#include "dequantization_utils.h"
|
||||
#include "memory_access_utils.h"
|
||||
#include "quantization_utils.h"
|
||||
#include "reduction_utils.h"
|
||||
|
||||
using rop = reduce::ROpType;
|
||||
|
||||
namespace swiz_quant {
|
||||
constexpr int max_threads = 512;
|
||||
constexpr int min_threads = 32;
|
||||
|
||||
constexpr int step_granularity = 2;
|
||||
constexpr int h_per_step = step_granularity * quantize::h_per_load;
|
||||
} // namespace swiz_quant
|
||||
|
||||
template <int numBits, int totalChunks, int threads, quantize::Type quantType>
|
||||
__global__ void swizzled_quant_kernel(int8_t* quantized_data,
|
||||
float* quantized_scales,
|
||||
const __half* uncompressed_data,
|
||||
int elems_per_group,
|
||||
int nodes,
|
||||
int devices_per_node)
|
||||
{
|
||||
cg::thread_block tb = cg::this_thread_block();
|
||||
cg::thread_block_tile<hw_warp_size> warp = cg::tiled_partition<hw_warp_size>(tb);
|
||||
|
||||
// Indexing offsets, same as normal quantization for in-case
|
||||
const int block_rank = blockIdx.x + blockIdx.y * gridDim.x + blockIdx.z * gridDim.x * gridDim.y;
|
||||
const int block_offset = block_rank * elems_per_group;
|
||||
const int elem_offset = tb.thread_index().x * quantize::h_per_load;
|
||||
const int base_offset = block_offset + elem_offset;
|
||||
const int stride = tb.size() * quantize::h_per_load;
|
||||
const __half* input_base = uncompressed_data + base_offset;
|
||||
|
||||
// Local buffer
|
||||
__half2 local_buffer[totalChunks * quantize::h2_per_load];
|
||||
|
||||
quantize::GroupStats<quantType> stats;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < totalChunks; i++) {
|
||||
__half2* iteration_buffer = local_buffer + i * quantize::h2_per_load;
|
||||
|
||||
mem_access::load_global<quantize::granularity>(
|
||||
iteration_buffer, input_base + i * stride, elem_offset + i * stride < elems_per_group);
|
||||
|
||||
#pragma unroll
|
||||
for (int j = 0; j < quantize::h2_per_load; j++) { stats.update(iteration_buffer[j]); }
|
||||
}
|
||||
|
||||
auto params = stats.template get_params<numBits, threads>(tb, warp);
|
||||
|
||||
const int partition_id = blockIdx.z;
|
||||
const int partition_offset = partition_id / devices_per_node;
|
||||
const int partition_base = (partition_id % devices_per_node) * nodes;
|
||||
const int pipelining_offset = blockIdx.y * (devices_per_node * nodes);
|
||||
const int output_partition = (pipelining_offset + partition_base + partition_offset);
|
||||
|
||||
constexpr int out_scalar_effect = 8 / numBits;
|
||||
const int out_block_rank = output_partition * gridDim.x + blockIdx.x;
|
||||
const int out_block_offset = out_block_rank * elems_per_group / out_scalar_effect;
|
||||
const int out_base_offset = out_block_offset + elem_offset / out_scalar_effect;
|
||||
int8_t* out_base = quantized_data + out_base_offset;
|
||||
|
||||
const int out_stride = stride / out_scalar_effect;
|
||||
constexpr int num_int8_out = quantize::h_per_load / out_scalar_effect;
|
||||
|
||||
if (tb.thread_index().x == 0) { params.store(quantized_scales, out_block_rank); }
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < totalChunks; i++) {
|
||||
if (i * stride + elem_offset < elems_per_group) {
|
||||
int8_t local_output[quantize::h_per_load / out_scalar_effect];
|
||||
quantize::_chunk<numBits, quantType>(
|
||||
local_output, local_buffer + i * quantize::h2_per_load, params);
|
||||
mem_access::store_global<num_int8_out>(out_base + i * out_stride, local_output);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define LAUNCH_SWIZZLE_QUANT(total_chunks, threads) \
|
||||
swizzled_quant_kernel<numBits, total_chunks, threads, qType><<<grid, block, 0, stream>>>( \
|
||||
q_data, q_scales, input_data, elems_per_group, nodes, devices_per_node);
|
||||
|
||||
/*
|
||||
Swizzled quantization reorganizes the quantized groups in order to better facilitate
|
||||
communication. As an example of the partitioning scheme we have the following example
|
||||
of 2 node, 4 device swizzling:
|
||||
|
||||
--- --- --- --- --- --- --- ---
|
||||
| 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 |
|
||||
--- --- --- --- --- --- --- ---
|
||||
becomes
|
||||
--- --- --- --- --- --- --- ---
|
||||
| 0 | 4 | 1 | 5 | 2 | 6 | 3 | 7 |
|
||||
--- --- --- --- --- --- --- ---
|
||||
|
||||
Multiple quantization groups may be mapped into a single partition. In order to better support
|
||||
later pipelining, we may also perform an additional slicing. In two-way slicing, for instance,
|
||||
the first halves of each partition are concatenated.
|
||||
*/
|
||||
|
||||
template <int numBits, quantize::Type qType>
|
||||
void launch_swizzled_quant_impl(int8_t* q_data,
|
||||
float* q_scales,
|
||||
const __half* input_data,
|
||||
int groups,
|
||||
int elems_per_group,
|
||||
int pipelining,
|
||||
int nodes,
|
||||
int devices_per_node,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
const int one_step_threads =
|
||||
next_pow2((elems_per_group + swiz_quant::h_per_step - 1) / swiz_quant::h_per_step);
|
||||
const int max_threads = (one_step_threads < swiz_quant::max_threads) ? one_step_threads
|
||||
: swiz_quant::max_threads;
|
||||
const int threads = (max_threads < swiz_quant::min_threads) ? swiz_quant::min_threads
|
||||
: max_threads;
|
||||
|
||||
dim3 block(threads);
|
||||
const int groups_per_partition = groups / (nodes * devices_per_node);
|
||||
assert(groups_per_partition % pipelining == 0);
|
||||
const int contiguous_groups = groups_per_partition / pipelining;
|
||||
const int partitions = nodes * devices_per_node;
|
||||
dim3 grid(contiguous_groups, pipelining, partitions);
|
||||
|
||||
const int elems_per_step = threads * swiz_quant::h_per_step;
|
||||
const int external_unroll = ((elems_per_group + elems_per_step - 1) / elems_per_step);
|
||||
const int total_unroll = external_unroll * swiz_quant::step_granularity;
|
||||
|
||||
assert(total_unroll % 2 == 0);
|
||||
|
||||
if (threads == 32) {
|
||||
LAUNCH_SWIZZLE_QUANT(2, 32);
|
||||
} else if (threads == 64) {
|
||||
LAUNCH_SWIZZLE_QUANT(2, 64);
|
||||
} else if (threads == 128) {
|
||||
LAUNCH_SWIZZLE_QUANT(2, 128);
|
||||
} else if (threads == 256) {
|
||||
LAUNCH_SWIZZLE_QUANT(2, 256);
|
||||
} else if (threads == 512) {
|
||||
if (total_unroll == 2) {
|
||||
LAUNCH_SWIZZLE_QUANT(2, 512);
|
||||
} else if (total_unroll == 4) {
|
||||
LAUNCH_SWIZZLE_QUANT(4, 512);
|
||||
} else if (total_unroll == 6) {
|
||||
LAUNCH_SWIZZLE_QUANT(6, 512);
|
||||
} else if (total_unroll == 8) {
|
||||
LAUNCH_SWIZZLE_QUANT(8, 512);
|
||||
} else if (total_unroll == 10) {
|
||||
LAUNCH_SWIZZLE_QUANT(10, 512);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define DISPATCH_SWIZZLE_QUANT(num_bits, qtype) \
|
||||
launch_swizzled_quant_impl<num_bits, qtype>(q_data, \
|
||||
q_scales, \
|
||||
input_data, \
|
||||
groups, \
|
||||
elems_per_group, \
|
||||
pipelining, \
|
||||
nodes, \
|
||||
devices_per_node, \
|
||||
stream);
|
||||
|
||||
void launch_swizzled_quant(int8_t* q_data,
|
||||
float* q_scales,
|
||||
const __half* input_data,
|
||||
int num_bits,
|
||||
quantize::Type q_type,
|
||||
int groups,
|
||||
int elems_per_group,
|
||||
int pipelining,
|
||||
int nodes,
|
||||
int devices_per_node,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
if (num_bits == 4) {
|
||||
if (q_type == quantize::Type::Asymmetric) {
|
||||
DISPATCH_SWIZZLE_QUANT(4, quantize::Type::Asymmetric);
|
||||
} else if (q_type == quantize::Type::Symmetric) {
|
||||
DISPATCH_SWIZZLE_QUANT(4, quantize::Type::Symmetric);
|
||||
}
|
||||
} else if (num_bits == 8) {
|
||||
if (q_type == quantize::Type::Asymmetric) {
|
||||
DISPATCH_SWIZZLE_QUANT(8, quantize::Type::Asymmetric);
|
||||
} else if (q_type == quantize::Type::Symmetric) {
|
||||
DISPATCH_SWIZZLE_QUANT(8, quantize::Type::Symmetric);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int numBits, int totalChunks, int threads, quantize::Type quantType>
|
||||
__global__ void loco_swizzled_quant_kernel(int8_t* quantized_data,
|
||||
float* quantized_scales,
|
||||
const __half* uncompressed_data,
|
||||
__half* error_feedback,
|
||||
const float err_beta,
|
||||
int groups,
|
||||
int elems_per_group,
|
||||
int pipelining,
|
||||
int nodes,
|
||||
int devices_per_node)
|
||||
{
|
||||
cg::thread_block tb = cg::this_thread_block();
|
||||
cg::thread_block_tile<hw_warp_size> warp = cg::tiled_partition<hw_warp_size>(tb);
|
||||
|
||||
// Indexing offsets, same as normal quantization for in-case
|
||||
const int block_rank_data =
|
||||
blockIdx.x + blockIdx.y * gridDim.x + blockIdx.z * gridDim.x * gridDim.y;
|
||||
const int block_offset_data = block_rank_data * elems_per_group;
|
||||
const int elem_offset = tb.thread_index().x * quantize::h_per_load;
|
||||
const int base_offset_data = block_offset_data + elem_offset;
|
||||
const int stride = tb.size() * quantize::h_per_load;
|
||||
const __half* uncompressed_data_base = uncompressed_data + base_offset_data;
|
||||
|
||||
const int partition_id = blockIdx.z;
|
||||
const int partition_offset = partition_id / devices_per_node;
|
||||
const int partition_base = (partition_id % devices_per_node) * nodes;
|
||||
const int pipelining_offset = blockIdx.y * (devices_per_node * nodes);
|
||||
const int output_partition = (pipelining_offset + partition_base + partition_offset);
|
||||
const int block_rank_err = output_partition * gridDim.x + blockIdx.x;
|
||||
|
||||
const int block_offset_err = block_rank_err * elems_per_group;
|
||||
const int base_offset_err = block_offset_err + elem_offset;
|
||||
__half* error_feedback_base = error_feedback + base_offset_err;
|
||||
|
||||
__half2 local_buffer[totalChunks * quantize::h2_per_load];
|
||||
__half2 err_buffer[totalChunks * quantize::h2_per_load];
|
||||
|
||||
quantize::GroupStats<quantType> stats;
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < totalChunks; i++) {
|
||||
__half2* iteration_buffer = local_buffer + i * quantize::h2_per_load;
|
||||
__half2* iter_err_buffer = err_buffer + i * quantize::h2_per_load;
|
||||
const int i_stride = i * stride;
|
||||
bool do_loads = (elem_offset + i_stride) < elems_per_group;
|
||||
|
||||
mem_access::load_global<quantize::granularity>(
|
||||
iteration_buffer, uncompressed_data_base + i_stride, do_loads);
|
||||
|
||||
mem_access::load_global<quantize::granularity>(
|
||||
iter_err_buffer, error_feedback_base + i_stride, do_loads);
|
||||
|
||||
#pragma unroll
|
||||
for (int j = 0; j < quantize::h2_per_load; j++) {
|
||||
iteration_buffer[j] = __hadd2(iteration_buffer[j], iter_err_buffer[j]);
|
||||
stats.update(iteration_buffer[j]);
|
||||
}
|
||||
}
|
||||
|
||||
auto params = stats.template get_params<numBits, threads>(tb, warp);
|
||||
|
||||
// Initialize dequantization parameters based on params
|
||||
auto de_params = params;
|
||||
de_params.scale = 1.0f / params.scale;
|
||||
if constexpr (quantType == quantize::Type::Asymmetric) { de_params.offset = params.offset; }
|
||||
|
||||
if (threadIdx.x == 0) { params.store(quantized_scales, block_rank_err); }
|
||||
|
||||
constexpr int out_scalar_effect = 8 / numBits;
|
||||
const int out_block_offset = block_rank_err * elems_per_group / out_scalar_effect;
|
||||
const int out_base_offset = out_block_offset + elem_offset / out_scalar_effect;
|
||||
int8_t* out_base = quantized_data + out_base_offset;
|
||||
|
||||
const int out_stride = stride / out_scalar_effect;
|
||||
constexpr int num_int8_out = quantize::h_per_load / out_scalar_effect;
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < totalChunks; i++) {
|
||||
const int i_stride = i * stride;
|
||||
__half2* iteration_buffer = local_buffer + i * quantize::h2_per_load;
|
||||
__half2* iter_err_buffer = err_buffer + i * quantize::h2_per_load;
|
||||
|
||||
if (i_stride + elem_offset < elems_per_group) {
|
||||
int8_t local_output[quantize::h_per_load / out_scalar_effect];
|
||||
quantize::_chunk<numBits, quantType>(local_output, iteration_buffer, params);
|
||||
mem_access::store_global<num_int8_out>(out_base + i * out_stride, local_output);
|
||||
|
||||
// Dequantize the quantized output to compute the dequantized value
|
||||
__half2 dequant_buffer[quantize::h2_per_load];
|
||||
dequantize::chunk<numBits, quantType>(dequant_buffer, local_output, de_params);
|
||||
|
||||
// Compute new error: sum - dequant_buffer
|
||||
#pragma unroll
|
||||
for (int k = 0; k < quantize::h2_per_load; k++) {
|
||||
// __half2 to float2
|
||||
float2 iter_buf_f = __half22float2(iteration_buffer[k]);
|
||||
float2 dequant_buf_f = __half22float2(dequant_buffer[k]);
|
||||
|
||||
// Update within float precision
|
||||
float2 new_error_f;
|
||||
new_error_f.x = iter_buf_f.x - dequant_buf_f.x;
|
||||
new_error_f.y = iter_buf_f.y - dequant_buf_f.y;
|
||||
|
||||
float2 iter_err_buf_f = __half22float2(iter_err_buffer[k]);
|
||||
|
||||
iter_err_buf_f.x = err_beta * iter_err_buf_f.x + (1.0f - err_beta) * new_error_f.x;
|
||||
iter_err_buf_f.y = err_beta * iter_err_buf_f.y + (1.0f - err_beta) * new_error_f.y;
|
||||
|
||||
// float2 back to __half2
|
||||
iter_err_buffer[k] = __float22half2_rn(iter_err_buf_f);
|
||||
}
|
||||
__half2* error_feedback_base_h2 = reinterpret_cast<__half2*>(error_feedback_base);
|
||||
mem_access::store_global<quantize::granularity>(error_feedback_base_h2 + i_stride / 2,
|
||||
iter_err_buffer);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define LAUNCH_LOCO_SWIZZLE_QUANT(total_chunks, threads) \
|
||||
loco_swizzled_quant_kernel<numBits, total_chunks, threads, qType> \
|
||||
<<<grid, block, 0, stream>>>(output_data, \
|
||||
params, \
|
||||
input_data, \
|
||||
error_feedback, \
|
||||
err_beta, \
|
||||
groups, \
|
||||
elems_per_group, \
|
||||
pipelining, \
|
||||
nodes, \
|
||||
devices_per_node);
|
||||
|
||||
template <int numBits, quantize::Type qType>
|
||||
void launch_loco_swizzled_quant_impl(int8_t* output_data,
|
||||
float* params,
|
||||
const __half* input_data,
|
||||
__half* error_feedback,
|
||||
const float err_beta,
|
||||
int groups,
|
||||
int elems_per_group,
|
||||
int pipelining,
|
||||
int nodes,
|
||||
int devices_per_node,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
const int one_step_threads =
|
||||
next_pow2((elems_per_group + swiz_quant::h_per_step - 1) / swiz_quant::h_per_step);
|
||||
const int max_threads = (one_step_threads < swiz_quant::max_threads) ? one_step_threads
|
||||
: swiz_quant::max_threads;
|
||||
const int threads = (max_threads < swiz_quant::min_threads) ? swiz_quant::min_threads
|
||||
: max_threads;
|
||||
|
||||
dim3 block(threads);
|
||||
const int groups_per_partition = groups / (nodes * devices_per_node);
|
||||
assert(groups_per_partition % pipelining == 0);
|
||||
const int contiguous_groups = groups_per_partition / pipelining;
|
||||
const int partitions = nodes * devices_per_node;
|
||||
dim3 grid(contiguous_groups, pipelining, partitions);
|
||||
|
||||
const int elems_per_step = threads * swiz_quant::h_per_step;
|
||||
const int external_unroll = ((elems_per_group + elems_per_step - 1) / elems_per_step);
|
||||
const int total_unroll = external_unroll * swiz_quant::step_granularity;
|
||||
|
||||
assert(total_unroll % 2 == 0);
|
||||
|
||||
if (threads == 32) {
|
||||
LAUNCH_LOCO_SWIZZLE_QUANT(2, 32);
|
||||
} else if (threads == 64) {
|
||||
LAUNCH_LOCO_SWIZZLE_QUANT(2, 64);
|
||||
} else if (threads == 128) {
|
||||
LAUNCH_LOCO_SWIZZLE_QUANT(2, 128);
|
||||
} else if (threads == 256) {
|
||||
LAUNCH_LOCO_SWIZZLE_QUANT(2, 256);
|
||||
} else if (threads == 512) {
|
||||
if (total_unroll == 2) {
|
||||
LAUNCH_LOCO_SWIZZLE_QUANT(2, 512);
|
||||
} else if (total_unroll == 4) {
|
||||
LAUNCH_LOCO_SWIZZLE_QUANT(4, 512);
|
||||
} else if (total_unroll == 6) {
|
||||
LAUNCH_LOCO_SWIZZLE_QUANT(6, 512);
|
||||
} else if (total_unroll == 8) {
|
||||
LAUNCH_LOCO_SWIZZLE_QUANT(8, 512);
|
||||
} else if (total_unroll == 10) {
|
||||
LAUNCH_LOCO_SWIZZLE_QUANT(10, 512);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define DISPATCH_LOCO_SWIZZLE_QUANT(num_bits, qtype) \
|
||||
launch_loco_swizzled_quant_impl<num_bits, qtype>(output_data, \
|
||||
params, \
|
||||
input_data, \
|
||||
error_feedback, \
|
||||
err_beta, \
|
||||
groups, \
|
||||
elems_per_group, \
|
||||
pipelining, \
|
||||
nodes, \
|
||||
devices_per_node, \
|
||||
stream);
|
||||
|
||||
void launch_loco_swizzled_quant(int8_t* output_data,
|
||||
float* params,
|
||||
const __half* input_data,
|
||||
__half* error_feedback,
|
||||
const float err_beta,
|
||||
int num_bits,
|
||||
quantize::Type q_type,
|
||||
int groups,
|
||||
int elems_per_group,
|
||||
int pipelining,
|
||||
int nodes,
|
||||
int devices_per_node,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
if (num_bits == 4) {
|
||||
if (q_type == quantize::Type::Asymmetric) {
|
||||
DISPATCH_LOCO_SWIZZLE_QUANT(4, quantize::Type::Asymmetric);
|
||||
} else if (q_type == quantize::Type::Symmetric) {
|
||||
DISPATCH_LOCO_SWIZZLE_QUANT(4, quantize::Type::Symmetric);
|
||||
}
|
||||
} else if (num_bits == 8) {
|
||||
if (q_type == quantize::Type::Asymmetric) {
|
||||
DISPATCH_LOCO_SWIZZLE_QUANT(8, quantize::Type::Asymmetric);
|
||||
} else if (q_type == quantize::Type::Symmetric) {
|
||||
DISPATCH_LOCO_SWIZZLE_QUANT(8, quantize::Type::Symmetric);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user