From 4114beac5bc3239771377140ed880d2a80544e6c Mon Sep 17 00:00:00 2001 From: Reza Yazdani Aminabadi Date: Thu, 17 Nov 2022 13:43:40 -0800 Subject: [PATCH 01/10] add configurable quantization for enabling 4-bit inference --- csrc/transformer/inference/csrc/dequantize.cu | 62 ++++++++++++++++++- .../transformer/inference/csrc/pt_binding.cpp | 62 ++++++++++++------- deepspeed/inference/config.py | 2 + deepspeed/inference/engine.py | 1 + .../transformers/ds_transformer.py | 2 +- deepspeed/module_inject/replace_module.py | 16 +++-- deepspeed/ops/transformer/inference/config.py | 6 +- .../ops/transformer/inference/ds_attention.py | 14 +++-- deepspeed/ops/transformer/inference/ds_mlp.py | 7 ++- 9 files changed, 130 insertions(+), 42 deletions(-) diff --git a/csrc/transformer/inference/csrc/dequantize.cu b/csrc/transformer/inference/csrc/dequantize.cu index 959016bf10e3..bed87ec6e27b 100644 --- a/csrc/transformer/inference/csrc/dequantize.cu +++ b/csrc/transformer/inference/csrc/dequantize.cu @@ -158,6 +158,57 @@ __global__ void dequantize_kernel(__half* output, } } + + +__global__ void dequantize_kernel_4bits(float* output, + const int8_t* input, + const float* qscale, + int hidden_dim, + unsigned merge_hidden, + int cnt) +{ +} + +__global__ void dequantize_kernel_4bits(__half* output, + const int8_t* input, + const float* qscale, + unsigned hidden_dim, + unsigned merge_hidden, + int cnt) +{ + unsigned bid = blockIdx.x * gridDim.y + blockIdx.y; + unsigned tid = threadIdx.x; + + float local_scale = qscale[blockIdx.x]; + + const float* input_cast = reinterpret_cast(input); + float4* output_cast = reinterpret_cast(output); + + input_cast += bid * merge_hidden; + output_cast += bid * merge_hidden; + + for (int c = 0; c < cnt; c++) { + if (tid < merge_hidden) { + float q = input_cast[tid]; + uint8_t* q_int8 = (uint8_t*)&q; + + float4 q_f; + __half* q_h = (__half*)&q_f; + + q_h[0] = __float2half(local_scale * (float)((int8_t)(q_int8[0] & 0x7))); + q_h[1] = __float2half(local_scale * (float)((int8_t)(q_int8[0] >> 4))); + q_h[2] = __float2half(local_scale * (float)((int8_t)(q_int8[1] & 0x7))); + q_h[3] = __float2half(local_scale * (float)((int8_t)(q_int8[1] >> 4))); + q_h[4] = __float2half(local_scale * (float)((int8_t)(q_int8[2] & 0x7))); + q_h[5] = __float2half(local_scale * (float)((int8_t)(q_int8[2] >> 4))); + q_h[6] = __float2half(local_scale * (float)((int8_t)(q_int8[3] & 0x7))); + q_h[7] = __float2half(local_scale * (float)((int8_t)(q_int8[3] >> 4))); + output_cast[tid] = q_f; + tid += blockDim.x; + } + } +} + template void launch_dequantize(T* output, const int8_t* input, @@ -165,6 +216,7 @@ void launch_dequantize(T* output, unsigned output_size, unsigned hidden_dim, unsigned groups, + int q_bits, cudaStream_t stream) { unsigned threads = 1024; @@ -177,8 +229,12 @@ void launch_dequantize(T* output, dim3 block_dims(threads); dim3 grid_dims(groups, blocks); - dequantize_kernel<<>>( - output, input, qscale, hidden_dim, hid_cnt * hidden_dim, thd_cnt); + if (q_bits == 4) + dequantize_kernel_4bits<<>>( + output, input, qscale, hidden_dim, hid_cnt * hidden_dim, thd_cnt); + else + dequantize_kernel<<>>( + output, input, qscale, hidden_dim, hid_cnt * hidden_dim, thd_cnt); } template void launch_dequantize(float*, @@ -187,6 +243,7 @@ template void launch_dequantize(float*, unsigned, unsigned, unsigned, + int, cudaStream_t); template void launch_dequantize<__half>(__half*, const int8_t*, @@ -194,4 +251,5 @@ template void launch_dequantize<__half>(__half*, unsigned, unsigned, unsigned, + int, cudaStream_t); diff --git a/csrc/transformer/inference/csrc/pt_binding.cpp b/csrc/transformer/inference/csrc/pt_binding.cpp index 7c9ab0bfd3ea..2e67eac1f2d5 100644 --- a/csrc/transformer/inference/csrc/pt_binding.cpp +++ b/csrc/transformer/inference/csrc/pt_binding.cpp @@ -763,7 +763,8 @@ void quantized_gemm(void* output, at::Tensor& weight, at::Tensor& qscale, int groups, - int bsz) + int bsz, + int q_bits) { T* weight16 = (T*)Context::Instance().GetWorkSpace() + 12 * Context::Instance().GetMaxTokenLenght() * weight.size(1); @@ -774,14 +775,17 @@ void quantized_gemm(void* output, weight.size(0), weight.size(1), groups, + q_bits, Context::Instance().GetCurrentStream()); + int out_size = quantize ? weight.size(0) : weight.size(1); + if (q_bits == 4) out_size *= 2; float alpha = (T)1.0; float gemm_beta = (T)0.0; cublas_gemm_ex(Context::Instance().GetCublasHandle(), CUBLAS_OP_T, CUBLAS_OP_N, - weight.size(0), + out_size, bsz, weight.size(1), &alpha, @@ -806,15 +810,18 @@ at::Tensor qkv_unfused_cublas(at::Tensor& output, at::Tensor& beta, const float epsilon, bool add_bias, - bool q_int8) + bool quantize, + int q_bits) { int bsz = input.size(0) * input.size(1); T* workspace = (T*)Context::Instance().GetWorkSpace(); workspace += (3 * bsz * input.size(2)); ds_layer_norm_internal(workspace, input, gamma, beta, epsilon); - if (q_int8) { - quantized_gemm(output.data_ptr(), workspace, weight, q_scale, q_scale.size(0), bsz); + int out_size = quantize ? weight.size(0) : weight.size(1); + if (q_bits == 4) out_size *= 2; + if (quantize) { + quantized_gemm(output.data_ptr(), workspace, weight, q_scale, q_scale.size(0), bsz, q_bits); } else { float alpha = (T)1.0; float gemm_beta = (T)0.0; @@ -841,7 +848,7 @@ at::Tensor qkv_unfused_cublas(at::Tensor& output, if (add_bias) launch_bias_add((T*)output.data_ptr(), (T*)bias.data_ptr(), - q_int8 ? weight.size(0) : weight.size(1), + out_size, bsz, Context::Instance().GetCurrentStream()); return torch::from_blob(workspace, input.sizes(), input.options()); @@ -860,12 +867,13 @@ std::vector ds_qkv_gemm(at::Tensor& input, bool external_cache, unsigned mp_size, unsigned rank, - bool q_int8) + bool quantize, + int q_bits) { int bsz = input.size(0) * input.size(1); T* workspace = (T*)Context::Instance().GetWorkSpace(); - int out_size = q_int8 ? weight.size(0) : weight.size(1); - + int out_size = quantize ? weight.size(0) : weight.size(1); + if (q_bits == 4) out_size *= 2; auto options = at::TensorOptions() .dtype(input.options().dtype()) .layout(at::kStrided) @@ -874,7 +882,7 @@ std::vector ds_qkv_gemm(at::Tensor& input, auto output = at::from_blob(workspace, {input.size(0), input.size(1), out_size}, options); auto inp_norm = qkv_unfused_cublas( - output, input, weight, q_scale, bias, gamma, beta, epsilon, add_bias, q_int8); + output, input, weight, q_scale, bias, gamma, beta, epsilon, add_bias, quantize, q_bits); return {output, inp_norm}; } @@ -1195,7 +1203,8 @@ at::Tensor ds_vector_matmul(at::Tensor& input, at::Tensor& weight, bool async_op, at::Tensor& q_scale, - bool q_int8) + bool quantize, + int q_bits) { auto options = at::TensorOptions() .dtype(input.options().dtype()) @@ -1203,14 +1212,15 @@ at::Tensor ds_vector_matmul(at::Tensor& input, .device(at::kCUDA) .requires_grad(false); - int out_size = q_int8 ? weight.size(0) : weight.size(1); + int out_size = quantize ? weight.size(0) : weight.size(1); + if (q_bits == 4) out_size *= 2; int bsz = input.size(0) * input.size(1); T* workspace = (T*)Context::Instance().GetWorkSpace(); auto output = at::from_blob(workspace, {input.size(0), input.size(1), out_size}, options); - if (q_int8) { + if (quantize) { quantized_gemm( - output.data_ptr(), (T*)input.data_ptr(), weight, q_scale, q_scale.size(0), bsz); + output.data_ptr(), (T*)input.data_ptr(), weight, q_scale, q_scale.size(0), bsz, q_bits); } else { float alpha = (T)1.0; float gemm_beta = (T)0.0; @@ -1271,10 +1281,13 @@ at::Tensor mlp_unfused_cublas(at::Tensor& output, bool mlp_after_attn, at::Tensor& q_scale, at::Tensor& q_scale1, - bool q_int8, + bool quantize, + int q_bits, ActivationFuncType act_func_type) { int bsz = input.size(0) * input.size(1); + int out_size = quantize ? weight_out.size(0) : weight_out.size(1); + if (q_bits == 4) out_size *= 2; T* inp_norm = (T*)Context::Instance().GetWorkSpace() + torch::numel(input) + torch::numel(output); T* intermediate = inp_norm + torch::numel(input); @@ -1294,8 +1307,8 @@ at::Tensor mlp_unfused_cublas(at::Tensor& output, ds_layer_norm_internal(inp_norm, input, gamma, beta, epsilon); } - if (q_int8) { - quantized_gemm(intermediate, inp_norm, weight, q_scale, q_scale.size(0), bsz); + if (quantize) { + quantized_gemm(intermediate, inp_norm, weight, q_scale, q_scale.size(0), bsz, q_bits); } else { float alpha = (T)1.0; float gemm_beta = (T)0.0; @@ -1321,17 +1334,17 @@ at::Tensor mlp_unfused_cublas(at::Tensor& output, if (act_func_type == ActivationFuncType::GELU) { launch_bias_gelu(intermediate, (T*)bias.data_ptr(), - q_int8 ? weight.size(0) : weight.size(1), + out_size, bsz, Context::Instance().GetCurrentStream()); } else if (act_func_type == ActivationFuncType::ReLU) { launch_bias_relu(intermediate, (T*)bias.data_ptr(), - q_int8 ? weight.size(0) : weight.size(1), + out_size, bsz, Context::Instance().GetCurrentStream()); } - if (q_int8) { + if (quanitze) { quantized_gemm( output.data_ptr(), intermediate, weight1, q_scale1, q_scale1.size(0), bsz); } else { @@ -1374,7 +1387,8 @@ std::vector ds_mlp_gemm(at::Tensor& input, bool mlp_after_attn, at::Tensor& q_scale, at::Tensor& q_scale1, - bool q_int8, + bool quantize, + int q_bits, int activation_type) { auto options = at::TensorOptions() @@ -1383,7 +1397,8 @@ std::vector ds_mlp_gemm(at::Tensor& input, .device(at::kCUDA) .requires_grad(false); - int out_size = q_int8 ? weight_out.size(0) : weight_out.size(1); + int out_size = quantize ? weight_out.size(0) : weight_out.size(1); + if (q_bits == 4) out_size *= 2; auto output = at::from_blob((T*)Context::Instance().GetWorkSpace() + torch::numel(input), {input.size(0), input.size(1), out_size}, options); @@ -1404,7 +1419,8 @@ std::vector ds_mlp_gemm(at::Tensor& input, mlp_after_attn, q_scale, q_scale1, - q_int8, + quantize, + q_bits, act_func_type); return {output, res_add}; diff --git a/deepspeed/inference/config.py b/deepspeed/inference/config.py index 42facc9d3176..6e07ff9388e5 100644 --- a/deepspeed/inference/config.py +++ b/deepspeed/inference/config.py @@ -134,6 +134,8 @@ class DeepSpeedInferenceConfig(DeepSpeedConfigModel): Desired model data type, will convert model to this type. Supported target types: torch.half, torch.int8, torch.float """ + + quantization_bits: int = 8 tensor_parallel: DeepSpeedTPConfig = Field({}, alias="tp") """ diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py index e85f30852622..594ee41fb611 100755 --- a/deepspeed/inference/engine.py +++ b/deepspeed/inference/engine.py @@ -358,6 +358,7 @@ def _apply_injection_policy(self, config, client_module=None): training=False, return_tuple=config.return_tuple, quantize=(config.dtype == torch.int8), + quantization_bits=config.quantization_bits, quantize_settings=(self.quantization_scales, self.quantize_merge_count, self.mlp_extra_grouping, diff --git a/deepspeed/model_implementations/transformers/ds_transformer.py b/deepspeed/model_implementations/transformers/ds_transformer.py index 49b7c81698a1..f38995d91f8c 100644 --- a/deepspeed/model_implementations/transformers/ds_transformer.py +++ b/deepspeed/model_implementations/transformers/ds_transformer.py @@ -128,7 +128,7 @@ def forward( input = input[0] input_type = input.dtype - if (self.config.fp16 or self.config.q_int8) \ + if (self.config.fp16 or self.config.qunatize) \ and input.dtype == torch.float: input = input.half() diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py index 52f8af00f6f4..427c94a81729 100755 --- a/deepspeed/module_inject/replace_module.py +++ b/deepspeed/module_inject/replace_module.py @@ -137,9 +137,8 @@ def get_transformer_name(replaced_module): class GroupQuantizer: - def __init__(self, q_int8=True, num_groups=32, group_size=32, num_bits=8): + def __init__(self, q_int8=True, num_bits=8, num_groups=32): self.num_groups = num_groups - self.group_size = group_size self.num_bits = num_bits self.q_int8 = q_int8 @@ -156,7 +155,12 @@ def quantize(self, inputs, qkv=True, count=1, parallel_dim=0): scale = torch.max(input_min.abs(), input_max.abs()) * 2.0 / (q_range) input_flat = (input_flat / scale).round().clamp(-q_range // 2, q_range // 2 - 1) inputs_q = input_flat.reshape(inputs.shape).to(torch.int8).contiguous() - out = torch.nn.Parameter(inputs_q, requires_grad=False) + if self.num_bits == 4: + int4_data = torch.empty(inputs_q.size(0), inputs_q.size(1) // 2, dtype=torch.uint8, device=inputs_q.device) + int4_data = (inputs_q[:, 1::2].to(torch.uint8)) | inputs_q[:, ::2] + out = torch.nn.Parameter(int4_data, requires_grad=False) + else: + out = torch.nn.Parameter(inputs_q, requires_grad=False) #print(inputs.shape) inputs_split = inputs.split(inputs.shape[parallel_dim] // 2, dim=parallel_dim) input_flat = [ @@ -304,6 +308,7 @@ def replace_transformer_layer(orig_layer_impl, stochastic_mode=True, training=True, quantize=False, + quantization_bits=8, quantize_settings=None, triangular_masking=False, return_tuple=True, @@ -415,7 +420,7 @@ def replace_with_policy(child, #expert_mp_replace = ReplaceWithTensorSlicing(mp_group=expert_mp_group) - quantizer = GroupQuantizer(q_int8=quantize) + quantizer = GroupQuantizer(q_int8=quantize, num_bits=quantization_bits) if inference: scale_attn_by_inverse_layer_idx = config.scale_attn_by_inverse_layer_idx if hasattr( config, @@ -457,7 +462,8 @@ def replace_with_policy(child, fp16=fp16, pre_layer_norm=policy.pre_attn_norm, mp_size=mp_size, - q_int8=quantize, + quantize=quantize, + quantization_bits=quantization_bits, return_tuple=(return_tuple or (policy_cls is HFBertLayerPolicy)), triangular_masking=(policy_cls is not HFBertLayerPolicy), local_attention=((config.attention_layers[layer_id] == "local") diff --git a/deepspeed/ops/transformer/inference/config.py b/deepspeed/ops/transformer/inference/config.py index 64e931d9a9da..9203fd1dc63c 100644 --- a/deepspeed/ops/transformer/inference/config.py +++ b/deepspeed/ops/transformer/inference/config.py @@ -49,7 +49,8 @@ def __init__(self, local_rank=-1, mp_size=1, fp16=False, - q_int8=False, + quantize=False, + quantization_bits=8, pre_layer_norm=True, stochastic_mode=False, scale_attention=True, @@ -78,7 +79,8 @@ def __init__(self, self.stochastic_mode = stochastic_mode self.epsilon = layer_norm_eps self.mp_size = mp_size - self.q_int8 = q_int8 + self.quantize = quantize + self.quantization_bits = quantization_bits self.scale_attention = scale_attention self.triangular_masking = triangular_masking self.local_attention = local_attention diff --git a/deepspeed/ops/transformer/inference/ds_attention.py b/deepspeed/ops/transformer/inference/ds_attention.py index 610bd882ecf4..9177647054fb 100644 --- a/deepspeed/ops/transformer/inference/ds_attention.py +++ b/deepspeed/ops/transformer/inference/ds_attention.py @@ -182,7 +182,7 @@ def compute_attention(qkv_out, input_mask): head_size = (qkv_out.shape[-1] // 3 // num_attention_heads_per_partition) if no_masking: input_mask = torch.empty(1) - if merge_count > 0 and config.q_int8: + if merge_count > 0 and config.quantize: split_dim = (qkv_out.dim() - 1) qkv_split = torch.split(qkv_out, (qkv_out.shape[-1] // (2**merge_count)), @@ -324,13 +324,15 @@ def selfAttention_fp(): config.bigscience_bloom, config.mp_size, dist.get_rank() if dist.is_initialized() else 0, - config.q_int8) + config.quantize, + config.quantization_bits) context_layer, key_layer, value_layer = compute_attention(qkv_out[0] if isinstance(qkv_out, list) else qkv_out, input_mask) output = vector_matmul_func(context_layer, attn_ow, False, attn_ow.scale, - config.q_int8) + config.quantize, + config.quantization_bits) return output, key_layer, value_layer, context_layer, qkv_out[-1] @@ -390,14 +392,14 @@ def __init__(self, qkv_merging=False): super(DeepSpeedSelfAttention, self).__init__() self.config = config - data_type = torch.int8 if config.q_int8 else torch.half if config.fp16 else torch.float + data_type = torch.int8 if config.quantize else torch.half if config.fp16 else torch.float data_type_fp = torch.half if config.fp16 else torch.float self.config.layer_id = DeepSpeedSelfAttention.num_layers DeepSpeedSelfAttention.num_layers = DeepSpeedSelfAttention.num_layers + 1 device = torch.cuda.current_device() if config.bigscience_bloom else 'cpu' qkv_size_per_partition = (self.config.hidden_size // self.config.mp_size) * 3 self.attn_qkvw = nn.Parameter(torch.empty(self.config.hidden_size, - qkv_size_per_partition, + qkv_size_per_partition // 2 if self.config.quantization_bits==4 else qkv_size_per_partition, dtype=data_type, device=device), requires_grad=False) @@ -407,7 +409,7 @@ def __init__(self, requires_grad=False) out_size_per_partition = self.config.hidden_size // self.config.mp_size self.attn_ow = nn.Parameter(torch.empty(out_size_per_partition, - self.config.hidden_size, + self.config.hidden_size // 2 if self.config.quantization_bits==4 else self.config.hidden_size, dtype=data_type, device=device), requires_grad=False) diff --git a/deepspeed/ops/transformer/inference/ds_mlp.py b/deepspeed/ops/transformer/inference/ds_mlp.py index 4f1c705c55ea..ee7628745628 100644 --- a/deepspeed/ops/transformer/inference/ds_mlp.py +++ b/deepspeed/ops/transformer/inference/ds_mlp.py @@ -60,7 +60,8 @@ def forward(ctx, config.mlp_after_attn, inter_w.scale, output_w.scale, - config.q_int8, + config.quantize, + config.quantization_bits, config.mlp_act_func_type) residual = residual if config.pre_layer_norm else residual_add residual_add_func( @@ -107,7 +108,7 @@ def __init__(self, requires_grad=False) intm_size_per_partition = self.config.intermediate_size // self.config.mp_size self.inter_w = nn.Parameter(torch.empty(self.config.hidden_size, - intm_size_per_partition, + intm_size_per_partition // 2 if self.config.quantization_bits==4 else intm_size_per_partition, dtype=data_type, device=device), requires_grad=False) @@ -116,7 +117,7 @@ def __init__(self, device=device), requires_grad=False) self.output_w = nn.Parameter(torch.empty(intm_size_per_partition, - self.config.hidden_size, + self.config.hidden_size // 2 if self.config.quantization_bits==4 else self.config.hidden_size, dtype=data_type, device=device), requires_grad=False) From 2ce22d7dccd4427631e1cd7daabe6b8a5a2e88aa Mon Sep 17 00:00:00 2001 From: Reza Yazdani Date: Thu, 17 Nov 2022 22:38:38 +0000 Subject: [PATCH 02/10] fix a few bugs --- csrc/transformer/inference/csrc/dequantize.cu | 8 ++++---- csrc/transformer/inference/csrc/pt_binding.cpp | 10 +++++----- .../inference/includes/inference_cuda_layers.h | 3 ++- deepspeed/ops/transformer/inference/ds_mlp.py | 6 +++--- op_builder/transformer_inference.py | 10 +++++----- 5 files changed, 19 insertions(+), 18 deletions(-) diff --git a/csrc/transformer/inference/csrc/dequantize.cu b/csrc/transformer/inference/csrc/dequantize.cu index bed87ec6e27b..8b84a80f8b95 100644 --- a/csrc/transformer/inference/csrc/dequantize.cu +++ b/csrc/transformer/inference/csrc/dequantize.cu @@ -182,7 +182,7 @@ __global__ void dequantize_kernel_4bits(__half* output, float local_scale = qscale[blockIdx.x]; const float* input_cast = reinterpret_cast(input); - float4* output_cast = reinterpret_cast(output); + float4* output_cast = reinterpret_cast(output); input_cast += bid * merge_hidden; output_cast += bid * merge_hidden; @@ -210,7 +210,7 @@ __global__ void dequantize_kernel_4bits(__half* output, } template -void launch_dequantize(T* output, +void launch_dequantize_v2(T* output, const int8_t* input, const float* qscale, unsigned output_size, @@ -237,7 +237,7 @@ void launch_dequantize(T* output, output, input, qscale, hidden_dim, hid_cnt * hidden_dim, thd_cnt); } -template void launch_dequantize(float*, +template void launch_dequantize_v2(float*, const int8_t*, const float*, unsigned, @@ -245,7 +245,7 @@ template void launch_dequantize(float*, unsigned, int, cudaStream_t); -template void launch_dequantize<__half>(__half*, +template void launch_dequantize_v2<__half>(__half*, const int8_t*, const float*, unsigned, diff --git a/csrc/transformer/inference/csrc/pt_binding.cpp b/csrc/transformer/inference/csrc/pt_binding.cpp index 2e67eac1f2d5..4db7ba862972 100644 --- a/csrc/transformer/inference/csrc/pt_binding.cpp +++ b/csrc/transformer/inference/csrc/pt_binding.cpp @@ -769,7 +769,7 @@ void quantized_gemm(void* output, T* weight16 = (T*)Context::Instance().GetWorkSpace() + 12 * Context::Instance().GetMaxTokenLenght() * weight.size(1); - launch_dequantize(weight16, + launch_dequantize_v2(weight16, (int8_t*)weight.data_ptr(), (float*)qscale.data_ptr(), weight.size(0), @@ -778,7 +778,7 @@ void quantized_gemm(void* output, q_bits, Context::Instance().GetCurrentStream()); - int out_size = quantize ? weight.size(0) : weight.size(1); + int out_size = weight.size(0); if (q_bits == 4) out_size *= 2; float alpha = (T)1.0; float gemm_beta = (T)0.0; @@ -1286,7 +1286,7 @@ at::Tensor mlp_unfused_cublas(at::Tensor& output, ActivationFuncType act_func_type) { int bsz = input.size(0) * input.size(1); - int out_size = quantize ? weight_out.size(0) : weight_out.size(1); + int out_size = quantize ? weight.size(0) : weight.size(1); if (q_bits == 4) out_size *= 2; T* inp_norm = (T*)Context::Instance().GetWorkSpace() + torch::numel(input) + torch::numel(output); @@ -1344,9 +1344,9 @@ at::Tensor mlp_unfused_cublas(at::Tensor& output, bsz, Context::Instance().GetCurrentStream()); } - if (quanitze) { + if (quantize) { quantized_gemm( - output.data_ptr(), intermediate, weight1, q_scale1, q_scale1.size(0), bsz); + output.data_ptr(), intermediate, weight1, q_scale1, q_scale1.size(0), bsz, q_bits); } else { float alpha = (T)1.0; float gemm_beta = (T)0.0; diff --git a/csrc/transformer/inference/includes/inference_cuda_layers.h b/csrc/transformer/inference/includes/inference_cuda_layers.h index d6a97a463197..f00e0f112150 100644 --- a/csrc/transformer/inference/includes/inference_cuda_layers.h +++ b/csrc/transformer/inference/includes/inference_cuda_layers.h @@ -123,12 +123,13 @@ void launch_dequantize(T* output, cudaStream_t stream); template -void launch_dequantize(T* output, +void launch_dequantize_v2(T* output, const int8_t* input, const float* qscale, unsigned output_size, unsigned hidden_dim, unsigned groups, + int q_bits, cudaStream_t stream); template void launch_gptj_residual_add(T* input, diff --git a/deepspeed/ops/transformer/inference/ds_mlp.py b/deepspeed/ops/transformer/inference/ds_mlp.py index ee7628745628..ccfab8293bd5 100644 --- a/deepspeed/ops/transformer/inference/ds_mlp.py +++ b/deepspeed/ops/transformer/inference/ds_mlp.py @@ -95,7 +95,7 @@ def __init__(self, super(DeepSpeedMLP, self).__init__() self.config = config - data_type = torch.int8 if config.q_int8 else torch.half if config.fp16 else torch.float + data_type = torch.int8 if config.quantize else torch.half if config.fp16 else torch.float data_type_fp = torch.half if config.fp16 else torch.float device = torch.cuda.current_device() if config.bigscience_bloom else 'cpu' self.attn_nw = nn.Parameter(torch.empty(self.config.hidden_size, @@ -145,10 +145,10 @@ def __init__(self, self.fused_gemm_gelu = inference_cuda_module.fused_gemm_gelu_fp16 if config.fp16 else \ inference_cuda_module.fused_gemm_gelu_fp32 - self.bias_residual_func = inference_cuda_module.bias_residual_fp16 if config.fp16 or config.q_int8 else \ + self.bias_residual_func = inference_cuda_module.bias_residual_fp16 if config.fp16 or config.quantize else \ inference_cuda_module.bias_residual_fp32 - self.residual_add_func = inference_cuda_module.residual_add_bias_fp16 if config.fp16 or config.q_int8 else \ + self.residual_add_func = inference_cuda_module.residual_add_bias_fp16 if config.fp16 or config.quantize else \ inference_cuda_module.residual_add_bias_fp32 def forward(self, input, residual, residual_norm, bias): diff --git a/op_builder/transformer_inference.py b/op_builder/transformer_inference.py index 14ea975eb1cb..9bae6a68bea0 100755 --- a/op_builder/transformer_inference.py +++ b/op_builder/transformer_inference.py @@ -44,11 +44,11 @@ def sources(self): 'csrc/transformer/inference/csrc/transform.cu', ] - def extra_ldflags(self): - if not self.is_rocm_pytorch(): - return ['-lcurand'] - else: - return [] +# def extra_ldflags(self): +# if not self.is_rocm_pytorch(): +# return ['-lcurand'] +# else: +# return [] def include_paths(self): return ['csrc/transformer/inference/includes', 'csrc/includes'] From d2997bf49793ad068c60d15f77472512e74e2ad1 Mon Sep 17 00:00:00 2001 From: Reza Yazdani Aminabadi Date: Thu, 17 Nov 2022 15:05:19 -0800 Subject: [PATCH 03/10] add q-int4 on top of q-int8 --- deepspeed/module_inject/load_checkpoint.py | 4 ++-- deepspeed/module_inject/replace_module.py | 10 +++++++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/deepspeed/module_inject/load_checkpoint.py b/deepspeed/module_inject/load_checkpoint.py index f577a1a0e1bc..4cb2b52b3730 100644 --- a/deepspeed/module_inject/load_checkpoint.py +++ b/deepspeed/module_inject/load_checkpoint.py @@ -63,8 +63,8 @@ def load_parameters(module, prefix): transpose(tmp_data) if weight_quantizer. q_int8 else tmp_data) else: - p = torch.nn.parameter.Parameter(tmp_data, - requires_grad=False) + p = weight_quantizer.quantize(torch.nn.parameter.Parameter(tmp_data, + requires_grad=False), scale=scale) p.scale = scale setattr(module, n, p) else: diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py index 427c94a81729..2bb10957617d 100755 --- a/deepspeed/module_inject/replace_module.py +++ b/deepspeed/module_inject/replace_module.py @@ -142,11 +142,15 @@ def __init__(self, q_int8=True, num_bits=8, num_groups=32): self.num_bits = num_bits self.q_int8 = q_int8 - def quantize(self, inputs, qkv=True, count=1, parallel_dim=0): + def quantize(self, inputs, qkv=True, count=1, parallel_dim=0, scale=None): if not self.q_int8 or not qkv: inputs = torch.nn.Parameter(inputs, requires_grad=False) inputs.scale = torch.empty(1) return inputs + if scale is not None and inputs.dtype == torch.int8: + input_flat = inputs.reshape(self.num_groups, -1).contiguous() + input_flat = input_flat * scale.view(-1)[:32].unsqueeze(1) + inputs = input_flat.reshape(inputs.shape).to(torch.half).contiguous() q_range = 2**self.num_bits inputs = inputs.to(torch.cuda.current_device()) input_flat = inputs.reshape(self.num_groups, -1).contiguous() @@ -157,7 +161,7 @@ def quantize(self, inputs, qkv=True, count=1, parallel_dim=0): inputs_q = input_flat.reshape(inputs.shape).to(torch.int8).contiguous() if self.num_bits == 4: int4_data = torch.empty(inputs_q.size(0), inputs_q.size(1) // 2, dtype=torch.uint8, device=inputs_q.device) - int4_data = (inputs_q[:, 1::2].to(torch.uint8)) | inputs_q[:, ::2] + int4_data = (inputs_q[:, 1::2].to(torch.uint8) << 4) | inputs_q[:, ::2] out = torch.nn.Parameter(int4_data, requires_grad=False) else: out = torch.nn.Parameter(inputs_q, requires_grad=False) @@ -956,7 +960,7 @@ def replace_fn(child, _policy, layer_id=0): replace_fn=replace_fn, _replace_policy=policy) - quantizer = GroupQuantizer(q_int8=quantize) + quantizer = GroupQuantizer(q_int8=quantize, num_bits=quantization_bits) world_size = dist.get_world_size() if dist.is_initialized() else 1 rank = dist.get_rank() if dist.is_initialized() else 0 if checkpoint_dict is not None: From 53418114b84f160a74ed97c415f5a83a88d63d19 Mon Sep 17 00:00:00 2001 From: Reza Yazdani Date: Fri, 18 Nov 2022 17:57:43 +0000 Subject: [PATCH 04/10] fix a few issues to run inference on 2 A100-80G (4 A100-40G) --- csrc/transformer/inference/csrc/dequantize.cu | 110 ++++++++++++++++-- .../transformer/inference/csrc/pt_binding.cpp | 32 ++++- .../inference/includes/inference_context.h | 2 +- .../includes/inference_cuda_layers.h | 8 ++ deepspeed/inference/engine.py | 2 +- .../transformers/ds_transformer.py | 1 + deepspeed/module_inject/load_checkpoint.py | 36 +++--- deepspeed/module_inject/replace_module.py | 25 ++-- .../ops/transformer/inference/ds_attention.py | 5 + deepspeed/ops/transformer/inference/ds_mlp.py | 1 + 10 files changed, 183 insertions(+), 39 deletions(-) diff --git a/csrc/transformer/inference/csrc/dequantize.cu b/csrc/transformer/inference/csrc/dequantize.cu index 8b84a80f8b95..25b3caea396e 100644 --- a/csrc/transformer/inference/csrc/dequantize.cu +++ b/csrc/transformer/inference/csrc/dequantize.cu @@ -160,6 +160,7 @@ __global__ void dequantize_kernel(__half* output, + __global__ void dequantize_kernel_4bits(float* output, const int8_t* input, const float* qscale, @@ -169,6 +170,13 @@ __global__ void dequantize_kernel_4bits(float* output, { } + +struct PackedInt4 { + + int8_t low : 4; + int8_t high : 4; +}; + __global__ void dequantize_kernel_4bits(__half* output, const int8_t* input, const float* qscale, @@ -190,19 +198,19 @@ __global__ void dequantize_kernel_4bits(__half* output, for (int c = 0; c < cnt; c++) { if (tid < merge_hidden) { float q = input_cast[tid]; - uint8_t* q_int8 = (uint8_t*)&q; + PackedInt4* q_int8 = (PackedInt4*)&q; float4 q_f; __half* q_h = (__half*)&q_f; - q_h[0] = __float2half(local_scale * (float)((int8_t)(q_int8[0] & 0x7))); - q_h[1] = __float2half(local_scale * (float)((int8_t)(q_int8[0] >> 4))); - q_h[2] = __float2half(local_scale * (float)((int8_t)(q_int8[1] & 0x7))); - q_h[3] = __float2half(local_scale * (float)((int8_t)(q_int8[1] >> 4))); - q_h[4] = __float2half(local_scale * (float)((int8_t)(q_int8[2] & 0x7))); - q_h[5] = __float2half(local_scale * (float)((int8_t)(q_int8[2] >> 4))); - q_h[6] = __float2half(local_scale * (float)((int8_t)(q_int8[3] & 0x7))); - q_h[7] = __float2half(local_scale * (float)((int8_t)(q_int8[3] >> 4))); + q_h[0] = __float2half(local_scale * (float)((int8_t)(q_int8[0].low))); + q_h[1] = __float2half(local_scale * (float)((int8_t)(q_int8[0].high))); + q_h[2] = __float2half(local_scale * (float)((int8_t)(q_int8[1].low))); + q_h[3] = __float2half(local_scale * (float)((int8_t)(q_int8[1].high))); + q_h[4] = __float2half(local_scale * (float)((int8_t)(q_int8[2].low))); + q_h[5] = __float2half(local_scale * (float)((int8_t)(q_int8[2].high))); + q_h[6] = __float2half(local_scale * (float)((int8_t)(q_int8[3].low))); + q_h[7] = __float2half(local_scale * (float)((int8_t)(q_int8[3].high))); output_cast[tid] = q_f; tid += blockDim.x; } @@ -220,7 +228,7 @@ void launch_dequantize_v2(T* output, cudaStream_t stream) { unsigned threads = 1024; - hidden_dim /= 4; + hidden_dim = (q_bits == 4) ? hidden_dim / 8 : hidden_dim / 4; unsigned hid_cnt = threads / hidden_dim; unsigned thd_cnt = (hidden_dim - 1) / threads + 1; hid_cnt = hid_cnt > 0 ? hid_cnt : 1; @@ -253,3 +261,85 @@ template void launch_dequantize_v2<__half>(__half*, unsigned, int, cudaStream_t); + +__global__ void dequantize_kernel_4bits(float* output, + const int8_t* input, + int hidden_dim, + unsigned merge_hidden, + int cnt) +{ +} + +__global__ void dequantize_kernel_4bits(__half* output, + const int8_t* input, + unsigned hidden_dim, + unsigned merge_hidden, + int cnt) +{ + unsigned bid = blockIdx.x * gridDim.y + blockIdx.y; + unsigned tid = threadIdx.x; + + const float* input_cast = reinterpret_cast(input); + float4* output_cast = reinterpret_cast(output); + + input_cast += bid * merge_hidden; + output_cast += bid * merge_hidden; + + for (int c = 0; c < cnt; c++) { + if (tid < merge_hidden) { + float q = input_cast[tid]; + PackedInt4* q_int8 = (PackedInt4*)&q; + + + float4 q_f; + __half* q_h = (__half*)&q_f; + q_h[0] = __float2half((float)((int8_t)(q_int8[0].low))); + q_h[1] = __float2half((float)((int8_t)(q_int8[0].high))); + q_h[2] = __float2half((float)((int8_t)(q_int8[1].low))); + q_h[3] = __float2half((float)((int8_t)(q_int8[1].high))); + q_h[4] = __float2half((float)((int8_t)(q_int8[2].low))); + q_h[5] = __float2half((float)((int8_t)(q_int8[2].high))); + q_h[6] = __float2half((float)((int8_t)(q_int8[3].low))); + q_h[7] = __float2half((float)((int8_t)(q_int8[3].high))); + output_cast[tid] = q_f; + tid += blockDim.x; + } + } +} + +template +void launch_dequantize_v2(T* output, + const int8_t* input, + unsigned output_size, + unsigned hidden_dim, + int q_bits, + cudaStream_t stream) +{ + unsigned threads = 1024; + hidden_dim = (q_bits == 4) ? hidden_dim / 4 : hidden_dim / 4; + unsigned hid_cnt = threads / hidden_dim; + unsigned thd_cnt = (hidden_dim - 1) / threads + 1; + hid_cnt = hid_cnt > 0 ? hid_cnt : 1; + + unsigned blocks = (output_size + hid_cnt - 1) / (hid_cnt); + dim3 block_dims(threads); + dim3 grid_dims(1, blocks); + + if (q_bits == 4) + dequantize_kernel_4bits<<>>( + output, input, hidden_dim, hid_cnt * hidden_dim, thd_cnt); +} + +template void launch_dequantize_v2(float*, + const int8_t*, + unsigned, + unsigned, + int, + cudaStream_t); +template void launch_dequantize_v2<__half>(__half*, + const int8_t*, + unsigned, + unsigned, + int, + cudaStream_t); + diff --git a/csrc/transformer/inference/csrc/pt_binding.cpp b/csrc/transformer/inference/csrc/pt_binding.cpp index 4db7ba862972..b8bb1c22d515 100644 --- a/csrc/transformer/inference/csrc/pt_binding.cpp +++ b/csrc/transformer/inference/csrc/pt_binding.cpp @@ -757,6 +757,27 @@ std::vector ds_layer_norm_residual_store(at::Tensor& input, return {norm_output, res_output}; } +template +at::Tensor dequantize(at::Tensor& weight, + int q_bits) +{ + auto options = at::TensorOptions() + .dtype(at::kHalf) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + auto out = torch::empty( + {weight.size(0), weight.size(1) * 2}, options); + launch_dequantize_v2((T*)out.data_ptr(), + (int8_t*)weight.data_ptr(), + weight.size(0), + weight.size(1), + q_bits, + Context::Instance().GetCurrentStream()); + return out; +} + template void quantized_gemm(void* output, T* input, @@ -768,18 +789,18 @@ void quantized_gemm(void* output, { T* weight16 = (T*)Context::Instance().GetWorkSpace() + 12 * Context::Instance().GetMaxTokenLenght() * weight.size(1); - + int out_size = weight.size(0); + if (q_bits == 4) out_size *= 2; launch_dequantize_v2(weight16, (int8_t*)weight.data_ptr(), (float*)qscale.data_ptr(), - weight.size(0), + out_size, weight.size(1), groups, q_bits, Context::Instance().GetCurrentStream()); - int out_size = weight.size(0); - if (q_bits == 4) out_size *= 2; + float alpha = (T)1.0; float gemm_beta = (T)0.0; cublas_gemm_ex(Context::Instance().GetCublasHandle(), @@ -1741,4 +1762,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) m.def("allocate_workspace_fp16", &allocate_workspace<__half>, "DeepSpeed memory allocation for GPT inference with fp16 (CUDA)"); + + m.def("dequantize_fp16", &dequantize, "dequantize fp16"); + m.def("dequantize_fp32", &dequantize, "dequantize fp32"); } diff --git a/csrc/transformer/inference/includes/inference_context.h b/csrc/transformer/inference/includes/inference_context.h index b3851ca43b72..e20512f4acee 100644 --- a/csrc/transformer/inference/includes/inference_context.h +++ b/csrc/transformer/inference/includes/inference_context.h @@ -106,7 +106,7 @@ class Context { const int padded_head_size = head_size <= 32 ? 32 : (head_size <= 64 ? 64 : 128); const int effective_head_size = (head_size > 128) ? head_size : padded_head_size; - size_t activation_size = 16 * (num_heads * effective_head_size) * batch_size; + size_t activation_size = 64 * (num_heads * effective_head_size) * batch_size; // Other sequence length dimension is added when the final workSpaceSize is calculated size_t temp_size = batch_size * num_heads * max_out_tokens * 2; size_t cache_size = diff --git a/csrc/transformer/inference/includes/inference_cuda_layers.h b/csrc/transformer/inference/includes/inference_cuda_layers.h index f00e0f112150..1eecf08b225b 100644 --- a/csrc/transformer/inference/includes/inference_cuda_layers.h +++ b/csrc/transformer/inference/includes/inference_cuda_layers.h @@ -131,6 +131,14 @@ void launch_dequantize_v2(T* output, unsigned groups, int q_bits, cudaStream_t stream); + +template +void launch_dequantize_v2(T* output, + const int8_t* input, + unsigned output_size, + unsigned hidden_dim, + int q_bits, + cudaStream_t stream); template void launch_gptj_residual_add(T* input, T* output, diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py index 594ee41fb611..5e3a2e98b2be 100755 --- a/deepspeed/inference/engine.py +++ b/deepspeed/inference/engine.py @@ -122,7 +122,7 @@ def __init__(self, model, config): self._apply_injection_policy(config) device = torch.cuda.current_device() - self.module.to(device) + #self.module.to(device) if config.tensor_parallel.tp_size > 1: _rng_state = torch.cuda.get_rng_state().to(torch.cuda.current_device()) diff --git a/deepspeed/model_implementations/transformers/ds_transformer.py b/deepspeed/model_implementations/transformers/ds_transformer.py index f38995d91f8c..d842ffcdf395 100644 --- a/deepspeed/model_implementations/transformers/ds_transformer.py +++ b/deepspeed/model_implementations/transformers/ds_transformer.py @@ -157,6 +157,7 @@ def forward( self.config.epsilon) output = output.to(input_type) + #print(output.shape, presents[0].shape) if get_present: output = (output, presents) diff --git a/deepspeed/module_inject/load_checkpoint.py b/deepspeed/module_inject/load_checkpoint.py index 4cb2b52b3730..32d5f677514e 100644 --- a/deepspeed/module_inject/load_checkpoint.py +++ b/deepspeed/module_inject/load_checkpoint.py @@ -46,13 +46,15 @@ def load_parameters(module, prefix): if prefix + n in sd[0] and len(n.split('.')) == 1: if type(sd[0][prefix + n]) is list: tmp_data, scale = sd[0][prefix + n] - tmp_data = tmp_data + tmp_data = tmp_data.to(torch.cuda.current_device()) scale = scale.to(torch.cuda.current_device()) else: tmp_data = sd[0][prefix + n].to(torch.cuda.current_device()) scale = None src_shape = tmp_data.shape dst_shape = p.shape + if weight_quantizer.num_bits == 4 and len(dst_shape) > 1: + dst_shape = torch.Size((dst_shape[0], dst_shape[1] * 2)) inner_dim = 1 if tmp_data.dtype == torch.int8 else 0 outer_dim = 0 if tmp_data.dtype == torch.int8 else 1 if (len(src_shape) == 2 and len(dst_shape) == 2): @@ -63,9 +65,11 @@ def load_parameters(module, prefix): transpose(tmp_data) if weight_quantizer. q_int8 else tmp_data) else: + p = weight_quantizer.quantize(torch.nn.parameter.Parameter(tmp_data, requires_grad=False), scale=scale) - p.scale = scale + #p.scale = scale + setattr(module, n, p) else: dim = inner_dim if src_shape[inner_dim] != dst_shape[ @@ -85,9 +89,15 @@ def load_parameters(module, prefix): weight_quantizer.num_groups, -1).contiguous() else: - assert tmp_data.dtype != torch.int8, \ - '''Merging of the checkpoints are not supported when using INT8 checkpoint! \ - Please use a as many GPUs as TP-size for the checkpoint''' + #assert tmp_data.dtype != torch.int8, \ + # '''Merging of the checkpoints are not supported when using INT8 checkpoint! \ + # Please use a as many GPUs as TP-size for the checkpoint''' + if tmp_data.dtype == torch.int8: + for j in range(len(sd)): + inputs, scal = sd[j][prefix + n] + input_flat = inputs.to('cpu').reshape(scal.shape[0], -1).contiguous() + input_flat = input_flat * scal.view(-1)[:scal.shape[0]].unsqueeze(1) + sd[j][prefix + n] = input_flat.reshape(inputs.shape).to(torch.half).contiguous() all_data = [ sd[j][prefix + n] if type(sd[j][prefix + n]) is list else @@ -95,27 +105,23 @@ def load_parameters(module, prefix): for j in range(len(sd)) ] weight_partition = torch.cat([ - ad[0].to(torch.cuda.current_device()) + ad[0] if type(ad) is list else ad for ad in all_data ], dim=dim) - if tmp_data.dtype == torch.int8: - scale = torch.cat([ - ad[1].to(torch.cuda.current_device()) - for ad in all_data - ], - dim=dim) + scale = None if tmp_data.dtype != torch.int8: + weight_partition = weight_quantizer.quantize( transpose(weight_partition), \ parallel_dim=(0 if dim == 1 else 1)) if weight_quantizer.q_int8 else \ weight_quantizer.quantize(weight_partition) else: - weight_partition = torch.nn.parameter.Parameter( + weight_partition = weight_quantizer.quantize(torch.nn.parameter.Parameter( weight_partition, - requires_grad=False) - weight_partition.scale = scale + requires_grad=False), scale=scale) + #weight_partition.scale = scale setattr(module, n, weight_partition) else: if src_shape[0] == dst_shape[0]: diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py index 2bb10957617d..0a8a38fecba9 100755 --- a/deepspeed/module_inject/replace_module.py +++ b/deepspeed/module_inject/replace_module.py @@ -134,11 +134,11 @@ def get_transformer_name(replaced_module): break break return transformer_name - +iter1 = 1 class GroupQuantizer: def __init__(self, q_int8=True, num_bits=8, num_groups=32): - self.num_groups = num_groups + self.num_groups = 128 #num_groups self.num_bits = num_bits self.q_int8 = q_int8 @@ -147,10 +147,14 @@ def quantize(self, inputs, qkv=True, count=1, parallel_dim=0, scale=None): inputs = torch.nn.Parameter(inputs, requires_grad=False) inputs.scale = torch.empty(1) return inputs - if scale is not None and inputs.dtype == torch.int8: - input_flat = inputs.reshape(self.num_groups, -1).contiguous() + if scale is not None and inputs.dtype == torch.int8 and self.num_bits < 8: + input_flat = inputs.reshape(32, -1).contiguous() input_flat = input_flat * scale.view(-1)[:32].unsqueeze(1) inputs = input_flat.reshape(inputs.shape).to(torch.half).contiguous() + elif inputs.dtype != torch.half: + inputs = torch.nn.Parameter(inputs, requires_grad=False) + inputs.scale = scale + return inputs q_range = 2**self.num_bits inputs = inputs.to(torch.cuda.current_device()) input_flat = inputs.reshape(self.num_groups, -1).contiguous() @@ -158,11 +162,16 @@ def quantize(self, inputs, qkv=True, count=1, parallel_dim=0, scale=None): input_max = torch.max(input_flat, dim=1, keepdim=True)[0].float() scale = torch.max(input_min.abs(), input_max.abs()) * 2.0 / (q_range) input_flat = (input_flat / scale).round().clamp(-q_range // 2, q_range // 2 - 1) - inputs_q = input_flat.reshape(inputs.shape).to(torch.int8).contiguous() - if self.num_bits == 4: - int4_data = torch.empty(inputs_q.size(0), inputs_q.size(1) // 2, dtype=torch.uint8, device=inputs_q.device) - int4_data = (inputs_q[:, 1::2].to(torch.uint8) << 4) | inputs_q[:, ::2] + inputs_q = input_flat.reshape(inputs.shape).to(torch.int).contiguous() + if self.num_bits < 8: + #global iter1 + #iter1 += 1 + #int4_data = torch.empty(inputs_q.size(0), inputs_q.size(1) // 2, dtype=torch.uint8, device=inputs_q.device) + int4_data = ((inputs_q[:, 1::2].to(torch.uint8) << 4) | (inputs_q[:, ::2] & 0xf).to(torch.uint8)).to(torch.uint8).reshape(inputs.shape[0]//2, inputs.shape[1]) out = torch.nn.Parameter(int4_data, requires_grad=False) + #print(f'{inputs_q} \n ---------------------------- \n {out}') + #if iter1 == 4: + # exit() else: out = torch.nn.Parameter(inputs_q, requires_grad=False) #print(inputs.shape) diff --git a/deepspeed/ops/transformer/inference/ds_attention.py b/deepspeed/ops/transformer/inference/ds_attention.py index 9177647054fb..be7da4789e89 100644 --- a/deepspeed/ops/transformer/inference/ds_attention.py +++ b/deepspeed/ops/transformer/inference/ds_attention.py @@ -158,6 +158,7 @@ def backup_attention(mixed_x_layer, layer_past, alibi, input_mask, norm_factor): offset, config.mp_size) # change view [batch_size x num_heads, q_length, k_length] + attention_probs_reshaped = attention_probs.view(*matmul_result.shape) # matmul: [batch_size * num_heads, q_length, head_dim] @@ -172,6 +173,7 @@ def backup_attention(mixed_x_layer, layer_past, alibi, input_mask, norm_factor): context_layer = _transpose_for_context(context_layer) + return context_layer, presents ###################### End of HF modeling_bloom addition ######################## @@ -326,6 +328,7 @@ def selfAttention_fp(): dist.get_rank() if dist.is_initialized() else 0, config.quantize, config.quantization_bits) + context_layer, key_layer, value_layer = compute_attention(qkv_out[0] if isinstance(qkv_out, list) else qkv_out, input_mask) output = vector_matmul_func(context_layer, attn_ow, @@ -334,6 +337,7 @@ def selfAttention_fp(): config.quantize, config.quantization_bits) + return output, key_layer, value_layer, context_layer, qkv_out[-1] def selfAttention_int8(): @@ -356,6 +360,7 @@ def selfAttention_int8(): q_scales[0], (q_groups * (3 if qkv_merging else 1) * (2**merge_count)), (attn_qkvb is not None)) + context_layer, key_layer, value_layer = compute_attention(qkv_out) output = inference_cuda_module.vector_matmul_int8(context_layer, attn_ow, diff --git a/deepspeed/ops/transformer/inference/ds_mlp.py b/deepspeed/ops/transformer/inference/ds_mlp.py index ccfab8293bd5..722b0d49d6f8 100644 --- a/deepspeed/ops/transformer/inference/ds_mlp.py +++ b/deepspeed/ops/transformer/inference/ds_mlp.py @@ -63,6 +63,7 @@ def forward(ctx, config.quantize, config.quantization_bits, config.mlp_act_func_type) + residual = residual if config.pre_layer_norm else residual_add residual_add_func( output, # hidden state From e2f6fe91593e5d4b0c7480d960601f60f1b9783a Mon Sep 17 00:00:00 2001 From: Reza Yazdani Aminabadi Date: Fri, 18 Nov 2022 10:23:35 -0800 Subject: [PATCH 05/10] Add more config params for quantization --- csrc/transformer/inference/csrc/dequantize.cu | 1 - csrc/transformer/inference/csrc/pt_binding.cpp | 1 - deepspeed/inference/config.py | 8 ++++++++ deepspeed/inference/engine.py | 1 + deepspeed/module_inject/load_checkpoint.py | 3 --- deepspeed/module_inject/replace_module.py | 16 +++++----------- 6 files changed, 14 insertions(+), 16 deletions(-) diff --git a/csrc/transformer/inference/csrc/dequantize.cu b/csrc/transformer/inference/csrc/dequantize.cu index 25b3caea396e..20a3b95e1bea 100644 --- a/csrc/transformer/inference/csrc/dequantize.cu +++ b/csrc/transformer/inference/csrc/dequantize.cu @@ -172,7 +172,6 @@ __global__ void dequantize_kernel_4bits(float* output, struct PackedInt4 { - int8_t low : 4; int8_t high : 4; }; diff --git a/csrc/transformer/inference/csrc/pt_binding.cpp b/csrc/transformer/inference/csrc/pt_binding.cpp index b8bb1c22d515..46c59464488f 100644 --- a/csrc/transformer/inference/csrc/pt_binding.cpp +++ b/csrc/transformer/inference/csrc/pt_binding.cpp @@ -1762,7 +1762,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) m.def("allocate_workspace_fp16", &allocate_workspace<__half>, "DeepSpeed memory allocation for GPT inference with fp16 (CUDA)"); - m.def("dequantize_fp16", &dequantize, "dequantize fp16"); m.def("dequantize_fp32", &dequantize, "dequantize fp32"); } diff --git a/deepspeed/inference/config.py b/deepspeed/inference/config.py index 6e07ff9388e5..8dd54941ac85 100644 --- a/deepspeed/inference/config.py +++ b/deepspeed/inference/config.py @@ -136,6 +136,14 @@ class DeepSpeedInferenceConfig(DeepSpeedConfigModel): """ quantization_bits: int = 8 + """ + Number of bits used to quantize the model using ZeroQuant technology + """ + + num_groups: int = 256 + """ + Number of groups used to quantize the model using ZeroQuant technology + """ tensor_parallel: DeepSpeedTPConfig = Field({}, alias="tp") """ diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py index 5e3a2e98b2be..fd12bb3564d6 100755 --- a/deepspeed/inference/engine.py +++ b/deepspeed/inference/engine.py @@ -359,6 +359,7 @@ def _apply_injection_policy(self, config, client_module=None): return_tuple=config.return_tuple, quantize=(config.dtype == torch.int8), quantization_bits=config.quantization_bits, + num_groups=config.num_groups, quantize_settings=(self.quantization_scales, self.quantize_merge_count, self.mlp_extra_grouping, diff --git a/deepspeed/module_inject/load_checkpoint.py b/deepspeed/module_inject/load_checkpoint.py index 32d5f677514e..feeb572fc4cf 100644 --- a/deepspeed/module_inject/load_checkpoint.py +++ b/deepspeed/module_inject/load_checkpoint.py @@ -68,8 +68,6 @@ def load_parameters(module, prefix): p = weight_quantizer.quantize(torch.nn.parameter.Parameter(tmp_data, requires_grad=False), scale=scale) - #p.scale = scale - setattr(module, n, p) else: dim = inner_dim if src_shape[inner_dim] != dst_shape[ @@ -121,7 +119,6 @@ def load_parameters(module, prefix): weight_partition = weight_quantizer.quantize(torch.nn.parameter.Parameter( weight_partition, requires_grad=False), scale=scale) - #weight_partition.scale = scale setattr(module, n, weight_partition) else: if src_shape[0] == dst_shape[0]: diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py index 0a8a38fecba9..ebd563ab394c 100755 --- a/deepspeed/module_inject/replace_module.py +++ b/deepspeed/module_inject/replace_module.py @@ -134,11 +134,11 @@ def get_transformer_name(replaced_module): break break return transformer_name -iter1 = 1 + class GroupQuantizer: def __init__(self, q_int8=True, num_bits=8, num_groups=32): - self.num_groups = 128 #num_groups + self.num_groups = num_groups self.num_bits = num_bits self.q_int8 = q_int8 @@ -164,17 +164,10 @@ def quantize(self, inputs, qkv=True, count=1, parallel_dim=0, scale=None): input_flat = (input_flat / scale).round().clamp(-q_range // 2, q_range // 2 - 1) inputs_q = input_flat.reshape(inputs.shape).to(torch.int).contiguous() if self.num_bits < 8: - #global iter1 - #iter1 += 1 - #int4_data = torch.empty(inputs_q.size(0), inputs_q.size(1) // 2, dtype=torch.uint8, device=inputs_q.device) int4_data = ((inputs_q[:, 1::2].to(torch.uint8) << 4) | (inputs_q[:, ::2] & 0xf).to(torch.uint8)).to(torch.uint8).reshape(inputs.shape[0]//2, inputs.shape[1]) out = torch.nn.Parameter(int4_data, requires_grad=False) - #print(f'{inputs_q} \n ---------------------------- \n {out}') - #if iter1 == 4: - # exit() else: out = torch.nn.Parameter(inputs_q, requires_grad=False) - #print(inputs.shape) inputs_split = inputs.split(inputs.shape[parallel_dim] // 2, dim=parallel_dim) input_flat = [ inputs_split[i].reshape(self.num_groups, @@ -322,6 +315,7 @@ def replace_transformer_layer(orig_layer_impl, training=True, quantize=False, quantization_bits=8, + num_groups=256, quantize_settings=None, triangular_masking=False, return_tuple=True, @@ -433,7 +427,7 @@ def replace_with_policy(child, #expert_mp_replace = ReplaceWithTensorSlicing(mp_group=expert_mp_group) - quantizer = GroupQuantizer(q_int8=quantize, num_bits=quantization_bits) + quantizer = GroupQuantizer(q_int8=quantize, num_bits=quantization_bits, num_groups=num_groups) if inference: scale_attn_by_inverse_layer_idx = config.scale_attn_by_inverse_layer_idx if hasattr( config, @@ -969,7 +963,7 @@ def replace_fn(child, _policy, layer_id=0): replace_fn=replace_fn, _replace_policy=policy) - quantizer = GroupQuantizer(q_int8=quantize, num_bits=quantization_bits) + quantizer = GroupQuantizer(q_int8=quantize, num_bits=quantization_bits, num_groups=num_groups) world_size = dist.get_world_size() if dist.is_initialized() else 1 rank = dist.get_rank() if dist.is_initialized() else 0 if checkpoint_dict is not None: From dd03ae68b7e52df7f9e727e397f0647bd3021a66 Mon Sep 17 00:00:00 2001 From: Reza Yazdani Date: Fri, 18 Nov 2022 21:28:12 +0000 Subject: [PATCH 06/10] fix formating --- csrc/transformer/inference/csrc/dequantize.cu | 5 ++-- deepspeed/inference/config.py | 2 +- deepspeed/module_inject/load_checkpoint.py | 30 ++++++++++++------- deepspeed/module_inject/replace_module.py | 8 +++-- .../ops/transformer/inference/ds_attention.py | 22 +++++++------- deepspeed/ops/transformer/inference/ds_mlp.py | 20 ++++++++----- op_builder/transformer_inference.py | 1 + 7 files changed, 54 insertions(+), 34 deletions(-) diff --git a/csrc/transformer/inference/csrc/dequantize.cu b/csrc/transformer/inference/csrc/dequantize.cu index 20a3b95e1bea..6e282880c5c6 100644 --- a/csrc/transformer/inference/csrc/dequantize.cu +++ b/csrc/transformer/inference/csrc/dequantize.cu @@ -236,7 +236,7 @@ void launch_dequantize_v2(T* output, dim3 block_dims(threads); dim3 grid_dims(groups, blocks); - if (q_bits == 4) + if (q_bits == 4) dequantize_kernel_4bits<<>>( output, input, qscale, hidden_dim, hid_cnt * hidden_dim, thd_cnt); else @@ -324,7 +324,7 @@ void launch_dequantize_v2(T* output, dim3 block_dims(threads); dim3 grid_dims(1, blocks); - if (q_bits == 4) + if (q_bits == 4) dequantize_kernel_4bits<<>>( output, input, hidden_dim, hid_cnt * hidden_dim, thd_cnt); } @@ -341,4 +341,3 @@ template void launch_dequantize_v2<__half>(__half*, unsigned, int, cudaStream_t); - diff --git a/deepspeed/inference/config.py b/deepspeed/inference/config.py index 8dd54941ac85..368efdd119fa 100644 --- a/deepspeed/inference/config.py +++ b/deepspeed/inference/config.py @@ -134,7 +134,7 @@ class DeepSpeedInferenceConfig(DeepSpeedConfigModel): Desired model data type, will convert model to this type. Supported target types: torch.half, torch.int8, torch.float """ - + quantization_bits: int = 8 """ Number of bits used to quantize the model using ZeroQuant technology diff --git a/deepspeed/module_inject/load_checkpoint.py b/deepspeed/module_inject/load_checkpoint.py index feeb572fc4cf..2a41e7d002ae 100644 --- a/deepspeed/module_inject/load_checkpoint.py +++ b/deepspeed/module_inject/load_checkpoint.py @@ -66,8 +66,11 @@ def load_parameters(module, prefix): q_int8 else tmp_data) else: - p = weight_quantizer.quantize(torch.nn.parameter.Parameter(tmp_data, - requires_grad=False), scale=scale) + p = weight_quantizer.quantize( + torch.nn.parameter.Parameter( + tmp_data, + requires_grad=False), + scale=scale) setattr(module, n, p) else: dim = inner_dim if src_shape[inner_dim] != dst_shape[ @@ -93,9 +96,14 @@ def load_parameters(module, prefix): if tmp_data.dtype == torch.int8: for j in range(len(sd)): inputs, scal = sd[j][prefix + n] - input_flat = inputs.to('cpu').reshape(scal.shape[0], -1).contiguous() - input_flat = input_flat * scal.view(-1)[:scal.shape[0]].unsqueeze(1) - sd[j][prefix + n] = input_flat.reshape(inputs.shape).to(torch.half).contiguous() + input_flat = inputs.to('cpu').reshape( + scal.shape[0], + -1).contiguous() + input_flat = input_flat * scal.view( + -1)[:scal.shape[0]].unsqueeze(1) + sd[j][prefix + n] = input_flat.reshape( + inputs.shape).to( + torch.half).contiguous() all_data = [ sd[j][prefix + n] if type(sd[j][prefix + n]) is list else @@ -103,8 +111,8 @@ def load_parameters(module, prefix): for j in range(len(sd)) ] weight_partition = torch.cat([ - ad[0] - if type(ad) is list else ad for ad in all_data + ad[0] if type(ad) is list else ad + for ad in all_data ], dim=dim) scale = None @@ -116,9 +124,11 @@ def load_parameters(module, prefix): parallel_dim=(0 if dim == 1 else 1)) if weight_quantizer.q_int8 else \ weight_quantizer.quantize(weight_partition) else: - weight_partition = weight_quantizer.quantize(torch.nn.parameter.Parameter( - weight_partition, - requires_grad=False), scale=scale) + weight_partition = weight_quantizer.quantize( + torch.nn.parameter.Parameter( + weight_partition, + requires_grad=False), + scale=scale) setattr(module, n, weight_partition) else: if src_shape[0] == dst_shape[0]: diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py index ebd563ab394c..ccb98c648a3f 100755 --- a/deepspeed/module_inject/replace_module.py +++ b/deepspeed/module_inject/replace_module.py @@ -427,7 +427,9 @@ def replace_with_policy(child, #expert_mp_replace = ReplaceWithTensorSlicing(mp_group=expert_mp_group) - quantizer = GroupQuantizer(q_int8=quantize, num_bits=quantization_bits, num_groups=num_groups) + quantizer = GroupQuantizer(q_int8=quantize, + num_bits=quantization_bits, + num_groups=num_groups) if inference: scale_attn_by_inverse_layer_idx = config.scale_attn_by_inverse_layer_idx if hasattr( config, @@ -963,7 +965,9 @@ def replace_fn(child, _policy, layer_id=0): replace_fn=replace_fn, _replace_policy=policy) - quantizer = GroupQuantizer(q_int8=quantize, num_bits=quantization_bits, num_groups=num_groups) + quantizer = GroupQuantizer(q_int8=quantize, + num_bits=quantization_bits, + num_groups=num_groups) world_size = dist.get_world_size() if dist.is_initialized() else 1 rank = dist.get_rank() if dist.is_initialized() else 0 if checkpoint_dict is not None: diff --git a/deepspeed/ops/transformer/inference/ds_attention.py b/deepspeed/ops/transformer/inference/ds_attention.py index be7da4789e89..5b379d0f22bc 100644 --- a/deepspeed/ops/transformer/inference/ds_attention.py +++ b/deepspeed/ops/transformer/inference/ds_attention.py @@ -173,7 +173,6 @@ def backup_attention(mixed_x_layer, layer_past, alibi, input_mask, norm_factor): context_layer = _transpose_for_context(context_layer) - return context_layer, presents ###################### End of HF modeling_bloom addition ######################## @@ -337,7 +336,6 @@ def selfAttention_fp(): config.quantize, config.quantization_bits) - return output, key_layer, value_layer, context_layer, qkv_out[-1] def selfAttention_int8(): @@ -403,20 +401,24 @@ def __init__(self, DeepSpeedSelfAttention.num_layers = DeepSpeedSelfAttention.num_layers + 1 device = torch.cuda.current_device() if config.bigscience_bloom else 'cpu' qkv_size_per_partition = (self.config.hidden_size // self.config.mp_size) * 3 - self.attn_qkvw = nn.Parameter(torch.empty(self.config.hidden_size, - qkv_size_per_partition // 2 if self.config.quantization_bits==4 else qkv_size_per_partition, - dtype=data_type, - device=device), + self.attn_qkvw = nn.Parameter(torch.empty( + self.config.hidden_size, + qkv_size_per_partition // + 2 if self.config.quantization_bits == 4 else qkv_size_per_partition, + dtype=data_type, + device=device), requires_grad=False) self.attn_qkvb = nn.Parameter(torch.empty(qkv_size_per_partition, dtype=data_type_fp, device=device), requires_grad=False) out_size_per_partition = self.config.hidden_size // self.config.mp_size - self.attn_ow = nn.Parameter(torch.empty(out_size_per_partition, - self.config.hidden_size // 2 if self.config.quantization_bits==4 else self.config.hidden_size, - dtype=data_type, - device=device), + self.attn_ow = nn.Parameter(torch.empty( + out_size_per_partition, + self.config.hidden_size // + 2 if self.config.quantization_bits == 4 else self.config.hidden_size, + dtype=data_type, + device=device), requires_grad=False) self.attn_ob = nn.Parameter(torch.empty(self.config.hidden_size, diff --git a/deepspeed/ops/transformer/inference/ds_mlp.py b/deepspeed/ops/transformer/inference/ds_mlp.py index 722b0d49d6f8..4b2b8268a979 100644 --- a/deepspeed/ops/transformer/inference/ds_mlp.py +++ b/deepspeed/ops/transformer/inference/ds_mlp.py @@ -108,19 +108,23 @@ def __init__(self, device=device), requires_grad=False) intm_size_per_partition = self.config.intermediate_size // self.config.mp_size - self.inter_w = nn.Parameter(torch.empty(self.config.hidden_size, - intm_size_per_partition // 2 if self.config.quantization_bits==4 else intm_size_per_partition, - dtype=data_type, - device=device), + self.inter_w = nn.Parameter(torch.empty( + self.config.hidden_size, + intm_size_per_partition // + 2 if self.config.quantization_bits == 4 else intm_size_per_partition, + dtype=data_type, + device=device), requires_grad=False) self.inter_b = nn.Parameter(torch.empty(intm_size_per_partition, dtype=data_type_fp, device=device), requires_grad=False) - self.output_w = nn.Parameter(torch.empty(intm_size_per_partition, - self.config.hidden_size // 2 if self.config.quantization_bits==4 else self.config.hidden_size, - dtype=data_type, - device=device), + self.output_w = nn.Parameter(torch.empty( + intm_size_per_partition, + self.config.hidden_size // + 2 if self.config.quantization_bits == 4 else self.config.hidden_size, + dtype=data_type, + device=device), requires_grad=False) self.output_b = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, diff --git a/op_builder/transformer_inference.py b/op_builder/transformer_inference.py index 9bae6a68bea0..418d3068720e 100755 --- a/op_builder/transformer_inference.py +++ b/op_builder/transformer_inference.py @@ -44,6 +44,7 @@ def sources(self): 'csrc/transformer/inference/csrc/transform.cu', ] + # def extra_ldflags(self): # if not self.is_rocm_pytorch(): # return ['-lcurand'] From efab0aa00aada3c06fd38554b185574dca4442c8 Mon Sep 17 00:00:00 2001 From: Reza Yazdani Aminabadi Date: Fri, 18 Nov 2022 16:14:40 -0800 Subject: [PATCH 07/10] fix config --- deepspeed/inference/config.py | 12 +----------- deepspeed/inference/engine.py | 4 +--- deepspeed/module_inject/replace_module.py | 12 +++++------- 3 files changed, 7 insertions(+), 21 deletions(-) diff --git a/deepspeed/inference/config.py b/deepspeed/inference/config.py index e8e39143cc22..5ff87a1cce66 100644 --- a/deepspeed/inference/config.py +++ b/deepspeed/inference/config.py @@ -89,7 +89,7 @@ class BaseQuantConfig(DeepSpeedConfigModel): enabled = True num_bits = 8 q_type: QuantTypeEnum = QuantTypeEnum.sym - q_groups: int = 1 + q_groups: int = 256 class WeightQuantConfig(BaseQuantConfig): @@ -135,16 +135,6 @@ class DeepSpeedInferenceConfig(DeepSpeedConfigModel): Supported target types: `torch.half`, `torch.int8`, `torch.float` """ - quantization_bits: int = 8 - """ - Number of bits used to quantize the model using ZeroQuant technology - """ - - num_groups: int = 256 - """ - Number of groups used to quantize the model using ZeroQuant technology - """ - tensor_parallel: DeepSpeedTPConfig = Field({}, alias="tp") """ Configuration for tensor parallelism used to split the model across several diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py index fd12bb3564d6..dfe07125d953 100755 --- a/deepspeed/inference/engine.py +++ b/deepspeed/inference/engine.py @@ -357,9 +357,7 @@ def _apply_injection_policy(self, config, client_module=None): fp16=(config.dtype == torch.half) or (config.dtype == torch.int8), training=False, return_tuple=config.return_tuple, - quantize=(config.dtype == torch.int8), - quantization_bits=config.quantization_bits, - num_groups=config.num_groups, + quantize=(config.dtype == torch.int8) or (config.quant.enabled), quantize_settings=(self.quantization_scales, self.quantize_merge_count, self.mlp_extra_grouping, diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py index ccb98c648a3f..076acdab7465 100755 --- a/deepspeed/module_inject/replace_module.py +++ b/deepspeed/module_inject/replace_module.py @@ -314,8 +314,6 @@ def replace_transformer_layer(orig_layer_impl, stochastic_mode=True, training=True, quantize=False, - quantization_bits=8, - num_groups=256, quantize_settings=None, triangular_masking=False, return_tuple=True, @@ -428,8 +426,8 @@ def replace_with_policy(child, #expert_mp_replace = ReplaceWithTensorSlicing(mp_group=expert_mp_group) quantizer = GroupQuantizer(q_int8=quantize, - num_bits=quantization_bits, - num_groups=num_groups) + num_bits=config.quant.weight.num_bits, + num_groups=config.quant.weight.q_groups) if inference: scale_attn_by_inverse_layer_idx = config.scale_attn_by_inverse_layer_idx if hasattr( config, @@ -472,7 +470,7 @@ def replace_with_policy(child, pre_layer_norm=policy.pre_attn_norm, mp_size=mp_size, quantize=quantize, - quantization_bits=quantization_bits, + quantization_bits=config.quant.weight.num_bits, return_tuple=(return_tuple or (policy_cls is HFBertLayerPolicy)), triangular_masking=(policy_cls is not HFBertLayerPolicy), local_attention=((config.attention_layers[layer_id] == "local") @@ -966,8 +964,8 @@ def replace_fn(child, _policy, layer_id=0): _replace_policy=policy) quantizer = GroupQuantizer(q_int8=quantize, - num_bits=quantization_bits, - num_groups=num_groups) + num_bits=config.quant.weight.num_bits, + num_groups=config.quant.weight.q_groups) world_size = dist.get_world_size() if dist.is_initialized() else 1 rank = dist.get_rank() if dist.is_initialized() else 0 if checkpoint_dict is not None: From bc1d63ecf691c634604c7e56d9c2152f387b2f21 Mon Sep 17 00:00:00 2001 From: Reza Yazdani Aminabadi Date: Fri, 18 Nov 2022 16:31:01 -0800 Subject: [PATCH 08/10] fix quantization parameters --- deepspeed/inference/engine.py | 2 ++ deepspeed/module_inject/replace_module.py | 12 +++++++----- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py index dfe07125d953..04dfe98df6a1 100755 --- a/deepspeed/inference/engine.py +++ b/deepspeed/inference/engine.py @@ -358,6 +358,8 @@ def _apply_injection_policy(self, config, client_module=None): training=False, return_tuple=config.return_tuple, quantize=(config.dtype == torch.int8) or (config.quant.enabled), + quantization_bits=config.quant.weight.num_bits, + num_groups=config.quant.weight.q_groups, quantize_settings=(self.quantization_scales, self.quantize_merge_count, self.mlp_extra_grouping, diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py index 076acdab7465..ccb98c648a3f 100755 --- a/deepspeed/module_inject/replace_module.py +++ b/deepspeed/module_inject/replace_module.py @@ -314,6 +314,8 @@ def replace_transformer_layer(orig_layer_impl, stochastic_mode=True, training=True, quantize=False, + quantization_bits=8, + num_groups=256, quantize_settings=None, triangular_masking=False, return_tuple=True, @@ -426,8 +428,8 @@ def replace_with_policy(child, #expert_mp_replace = ReplaceWithTensorSlicing(mp_group=expert_mp_group) quantizer = GroupQuantizer(q_int8=quantize, - num_bits=config.quant.weight.num_bits, - num_groups=config.quant.weight.q_groups) + num_bits=quantization_bits, + num_groups=num_groups) if inference: scale_attn_by_inverse_layer_idx = config.scale_attn_by_inverse_layer_idx if hasattr( config, @@ -470,7 +472,7 @@ def replace_with_policy(child, pre_layer_norm=policy.pre_attn_norm, mp_size=mp_size, quantize=quantize, - quantization_bits=config.quant.weight.num_bits, + quantization_bits=quantization_bits, return_tuple=(return_tuple or (policy_cls is HFBertLayerPolicy)), triangular_masking=(policy_cls is not HFBertLayerPolicy), local_attention=((config.attention_layers[layer_id] == "local") @@ -964,8 +966,8 @@ def replace_fn(child, _policy, layer_id=0): _replace_policy=policy) quantizer = GroupQuantizer(q_int8=quantize, - num_bits=config.quant.weight.num_bits, - num_groups=config.quant.weight.q_groups) + num_bits=quantization_bits, + num_groups=num_groups) world_size = dist.get_world_size() if dist.is_initialized() else 1 rank = dist.get_rank() if dist.is_initialized() else 0 if checkpoint_dict is not None: From 92f7aab487d81b7f74a2a2ba27c048a79eb3b44c Mon Sep 17 00:00:00 2001 From: Ammar Ahmad Awan Date: Mon, 21 Nov 2022 10:01:45 -0800 Subject: [PATCH 09/10] rename the vars in quant config. --- deepspeed/inference/config.py | 4 ++-- deepspeed/inference/engine.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/deepspeed/inference/config.py b/deepspeed/inference/config.py index 5ff87a1cce66..f519469049dc 100644 --- a/deepspeed/inference/config.py +++ b/deepspeed/inference/config.py @@ -88,8 +88,8 @@ class QuantTypeEnum(str, Enum): class BaseQuantConfig(DeepSpeedConfigModel): enabled = True num_bits = 8 - q_type: QuantTypeEnum = QuantTypeEnum.sym - q_groups: int = 256 + type: QuantTypeEnum = QuantTypeEnum.sym + num_groups: int = 256 class WeightQuantConfig(BaseQuantConfig): diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py index 04dfe98df6a1..48b5c6bd9007 100755 --- a/deepspeed/inference/engine.py +++ b/deepspeed/inference/engine.py @@ -359,7 +359,7 @@ def _apply_injection_policy(self, config, client_module=None): return_tuple=config.return_tuple, quantize=(config.dtype == torch.int8) or (config.quant.enabled), quantization_bits=config.quant.weight.num_bits, - num_groups=config.quant.weight.q_groups, + num_groups=config.quant.weight.num_groups, quantize_settings=(self.quantization_scales, self.quantize_merge_count, self.mlp_extra_grouping, From 803bf2a4102a8bb69290855368c541e18e2fe0ac Mon Sep 17 00:00:00 2001 From: Ammar Ahmad Awan Date: Mon, 21 Nov 2022 10:02:16 -0800 Subject: [PATCH 10/10] format fixes. --- csrc/transformer/inference/csrc/dequantize.cu | 113 +++++++++--------- .../transformer/inference/csrc/pt_binding.cpp | 34 +++--- .../includes/inference_cuda_layers.h | 24 ++-- 3 files changed, 82 insertions(+), 89 deletions(-) diff --git a/csrc/transformer/inference/csrc/dequantize.cu b/csrc/transformer/inference/csrc/dequantize.cu index 6e282880c5c6..4c51b0ada7bc 100644 --- a/csrc/transformer/inference/csrc/dequantize.cu +++ b/csrc/transformer/inference/csrc/dequantize.cu @@ -158,30 +158,26 @@ __global__ void dequantize_kernel(__half* output, } } - - - __global__ void dequantize_kernel_4bits(float* output, - const int8_t* input, - const float* qscale, - int hidden_dim, - unsigned merge_hidden, - int cnt) + const int8_t* input, + const float* qscale, + int hidden_dim, + unsigned merge_hidden, + int cnt) { } - struct PackedInt4 { int8_t low : 4; int8_t high : 4; }; __global__ void dequantize_kernel_4bits(__half* output, - const int8_t* input, - const float* qscale, - unsigned hidden_dim, - unsigned merge_hidden, - int cnt) + const int8_t* input, + const float* qscale, + unsigned hidden_dim, + unsigned merge_hidden, + int cnt) { unsigned bid = blockIdx.x * gridDim.y + blockIdx.y; unsigned tid = threadIdx.x; @@ -218,13 +214,13 @@ __global__ void dequantize_kernel_4bits(__half* output, template void launch_dequantize_v2(T* output, - const int8_t* input, - const float* qscale, - unsigned output_size, - unsigned hidden_dim, - unsigned groups, - int q_bits, - cudaStream_t stream) + const int8_t* input, + const float* qscale, + unsigned output_size, + unsigned hidden_dim, + unsigned groups, + int q_bits, + cudaStream_t stream) { unsigned threads = 1024; hidden_dim = (q_bits == 4) ? hidden_dim / 8 : hidden_dim / 4; @@ -245,35 +241,35 @@ void launch_dequantize_v2(T* output, } template void launch_dequantize_v2(float*, - const int8_t*, - const float*, - unsigned, - unsigned, - unsigned, - int, - cudaStream_t); + const int8_t*, + const float*, + unsigned, + unsigned, + unsigned, + int, + cudaStream_t); template void launch_dequantize_v2<__half>(__half*, - const int8_t*, - const float*, - unsigned, - unsigned, - unsigned, - int, - cudaStream_t); + const int8_t*, + const float*, + unsigned, + unsigned, + unsigned, + int, + cudaStream_t); __global__ void dequantize_kernel_4bits(float* output, - const int8_t* input, - int hidden_dim, - unsigned merge_hidden, - int cnt) + const int8_t* input, + int hidden_dim, + unsigned merge_hidden, + int cnt) { } __global__ void dequantize_kernel_4bits(__half* output, - const int8_t* input, - unsigned hidden_dim, - unsigned merge_hidden, - int cnt) + const int8_t* input, + unsigned hidden_dim, + unsigned merge_hidden, + int cnt) { unsigned bid = blockIdx.x * gridDim.y + blockIdx.y; unsigned tid = threadIdx.x; @@ -289,7 +285,6 @@ __global__ void dequantize_kernel_4bits(__half* output, float q = input_cast[tid]; PackedInt4* q_int8 = (PackedInt4*)&q; - float4 q_f; __half* q_h = (__half*)&q_f; q_h[0] = __float2half((float)((int8_t)(q_int8[0].low))); @@ -308,11 +303,11 @@ __global__ void dequantize_kernel_4bits(__half* output, template void launch_dequantize_v2(T* output, - const int8_t* input, - unsigned output_size, - unsigned hidden_dim, - int q_bits, - cudaStream_t stream) + const int8_t* input, + unsigned output_size, + unsigned hidden_dim, + int q_bits, + cudaStream_t stream) { unsigned threads = 1024; hidden_dim = (q_bits == 4) ? hidden_dim / 4 : hidden_dim / 4; @@ -330,14 +325,14 @@ void launch_dequantize_v2(T* output, } template void launch_dequantize_v2(float*, - const int8_t*, - unsigned, - unsigned, - int, - cudaStream_t); + const int8_t*, + unsigned, + unsigned, + int, + cudaStream_t); template void launch_dequantize_v2<__half>(__half*, - const int8_t*, - unsigned, - unsigned, - int, - cudaStream_t); + const int8_t*, + unsigned, + unsigned, + int, + cudaStream_t); diff --git a/csrc/transformer/inference/csrc/pt_binding.cpp b/csrc/transformer/inference/csrc/pt_binding.cpp index 46c59464488f..5584437c5d73 100644 --- a/csrc/transformer/inference/csrc/pt_binding.cpp +++ b/csrc/transformer/inference/csrc/pt_binding.cpp @@ -758,8 +758,7 @@ std::vector ds_layer_norm_residual_store(at::Tensor& input, } template -at::Tensor dequantize(at::Tensor& weight, - int q_bits) +at::Tensor dequantize(at::Tensor& weight, int q_bits) { auto options = at::TensorOptions() .dtype(at::kHalf) @@ -767,14 +766,13 @@ at::Tensor dequantize(at::Tensor& weight, .device(at::kCUDA) .requires_grad(false); - auto out = torch::empty( - {weight.size(0), weight.size(1) * 2}, options); + auto out = torch::empty({weight.size(0), weight.size(1) * 2}, options); launch_dequantize_v2((T*)out.data_ptr(), - (int8_t*)weight.data_ptr(), - weight.size(0), - weight.size(1), - q_bits, - Context::Instance().GetCurrentStream()); + (int8_t*)weight.data_ptr(), + weight.size(0), + weight.size(1), + q_bits, + Context::Instance().GetCurrentStream()); return out; } @@ -792,14 +790,13 @@ void quantized_gemm(void* output, int out_size = weight.size(0); if (q_bits == 4) out_size *= 2; launch_dequantize_v2(weight16, - (int8_t*)weight.data_ptr(), - (float*)qscale.data_ptr(), - out_size, - weight.size(1), - groups, - q_bits, - Context::Instance().GetCurrentStream()); - + (int8_t*)weight.data_ptr(), + (float*)qscale.data_ptr(), + out_size, + weight.size(1), + groups, + q_bits, + Context::Instance().GetCurrentStream()); float alpha = (T)1.0; float gemm_beta = (T)0.0; @@ -842,7 +839,8 @@ at::Tensor qkv_unfused_cublas(at::Tensor& output, int out_size = quantize ? weight.size(0) : weight.size(1); if (q_bits == 4) out_size *= 2; if (quantize) { - quantized_gemm(output.data_ptr(), workspace, weight, q_scale, q_scale.size(0), bsz, q_bits); + quantized_gemm( + output.data_ptr(), workspace, weight, q_scale, q_scale.size(0), bsz, q_bits); } else { float alpha = (T)1.0; float gemm_beta = (T)0.0; diff --git a/csrc/transformer/inference/includes/inference_cuda_layers.h b/csrc/transformer/inference/includes/inference_cuda_layers.h index 1eecf08b225b..7b137fc9d539 100644 --- a/csrc/transformer/inference/includes/inference_cuda_layers.h +++ b/csrc/transformer/inference/includes/inference_cuda_layers.h @@ -124,21 +124,21 @@ void launch_dequantize(T* output, template void launch_dequantize_v2(T* output, - const int8_t* input, - const float* qscale, - unsigned output_size, - unsigned hidden_dim, - unsigned groups, - int q_bits, - cudaStream_t stream); + const int8_t* input, + const float* qscale, + unsigned output_size, + unsigned hidden_dim, + unsigned groups, + int q_bits, + cudaStream_t stream); template void launch_dequantize_v2(T* output, - const int8_t* input, - unsigned output_size, - unsigned hidden_dim, - int q_bits, - cudaStream_t stream); + const int8_t* input, + unsigned output_size, + unsigned hidden_dim, + int q_bits, + cudaStream_t stream); template void launch_gptj_residual_add(T* input, T* output,