| |
| |
|
|
| |
|
|
| import torch |
| import abc |
| from abc import ABC |
|
|
| import gc |
| from deepspeed.ops.op_builder import FPQuantizerBuilder |
| from deepspeed.accelerator import get_accelerator |
|
|
| fp_quant_module = None |
|
|
|
|
| class Quantizer(ABC): |
| """ |
| Abstract Quantizer class that implements quantize/dequantize methods. |
| |
| Arguments: |
| group_size (int, optional): number of values or elements that are grouped |
| together for the quantization process. |
| """ |
|
|
| def __init__(self, group_size=512) -> None: |
| self.group_size = group_size |
|
|
| @abc.abstractmethod |
| def quantize(self, |
| input, |
| q_bits=8, |
| q_mantisa_bits=3, |
| stochastic_mode=False, |
| return_meta_tensor=False) -> torch.Tensor: |
| ... |
|
|
| @abc.abstractmethod |
| def dequantize(self, input_q, fp_out=None, q_bits=8, q_mantisa_bits=3, scale=None) -> torch.Tensor: |
| ... |
|
|
|
|
| class FP_Quantize(Quantizer): |
|
|
| def __init__(self, quantization_config) -> None: |
| global fp_quant_module |
| super().__init__(group_size=quantization_config.group_size) |
| if fp_quant_module is None: |
| fp_quant_module = FPQuantizerBuilder().load() |
| self.cuda_impl = getattr(fp_quant_module, "CUDA_IMPL", True) |
| self.q_config = quantization_config |
|
|
| self.orig_dtype = None |
| self.num_groups = None |
| self.input_q = None |
| self.scale = None |
|
|
| def quantize(self, |
| input, |
| q_bits=8, |
| q_mantisa_bits=3, |
| stochastic_mode=False, |
| return_meta_tensor=False) -> torch.Tensor: |
| assert input.dtype == torch.bfloat16, "only support bf16 for now" |
| if return_meta_tensor: |
| assert q_bits == 8, "meta tensor is only supported with q_bit=8" |
|
|
| self.orig_dtype = input.dtype |
| self.orig_shape = input.shape |
|
|
| if q_bits == 8: |
| pass |
| elif q_bits == 12: |
| q_mantisa_bits = 4 |
| elif q_bits == 6: |
| q_mantisa_bits = 2 |
| elif q_bits == 4: |
| q_mantisa_bits = 1 |
| else: |
| assert (0), \ |
| f"Missing {q_bits}-quantization, please add the template arguments for the kernel to support this precision!" |
| self.num_groups = input.numel() // self.group_size |
| self.input_q = torch.ones(self.num_groups, |
| int(self.group_size * q_bits) // 8 + 4, |
| dtype=torch.uint8, |
| device=input.device) |
| out = fp_quant_module.quantize(self.input_q, input, self.group_size, stochastic_mode, q_bits, q_mantisa_bits) |
| if return_meta_tensor: |
| data, self.scale = out.split(self.group_size, dim=-1) |
| data = data.contiguous().reshape(input.shape) |
| self.scale = self.scale.contiguous() |
| del self.input_q |
| del out |
| gc.collect() |
| get_accelerator().empty_cache() |
| return data, self.scale |
|
|
| return out |
|
|
| def to(self, *args, **kwargs): |
| |
| if hasattr(self, 'input_q'): |
| self.input_q = self.input_q.to(*args, **kwargs) |
| if hasattr(self, 'scale'): |
| self.scale = self.scale.to(*args, **kwargs) |
|
|
| def get_scales(self): |
| return fp_quant_module.get_scales(self.scale, self.num_groups) |
|
|
| def dequantize(self, input_q, fp_out=None, q_bits=8, q_mantisa_bits=3, scale=None) -> torch.Tensor: |
| assert (self.orig_dtype is not None), \ |
| "[De-quantization Error]: you need to call quantize before dequantizing!" |
| fp_out = torch.empty(self.orig_shape, dtype=self.orig_dtype, |
| device=input_q.device) if fp_out is None else fp_out |
| if q_bits == 8: |
| pass |
| elif q_bits == 12: |
| q_mantisa_bits = 4 |
| elif q_bits == 6: |
| q_mantisa_bits = 2 |
| elif q_bits == 4: |
| q_mantisa_bits = 1 |
| else: |
| assert (0), \ |
| f"Missing {q_bits}-dequantization, please add the template arguments for the kernel to support this precision!" |
|
|
| if scale is not None: |
| assert input_q.numel() == fp_out.numel(), \ |
| f'[De-quantization Error]: quantized data should have the same size as original tensor when scale is not None!' |
| input_q = torch.cat([input_q.reshape(-1, self.group_size), scale], dim=-1).contiguous() |
| fp_quant_module.dequantize(fp_out, input_q, self.group_size, q_mantisa_bits, q_bits - q_mantisa_bits - 1) |
| return fp_out |
|
|
| def selective_dequantize(self, |
| input_q, |
| indexes, |
| fp_out=None, |
| q_bits=8, |
| q_mantisa_bits=3, |
| scale=None) -> torch.Tensor: |
| assert (not hasattr(self, 'orig_shape') or len(self.orig_shape) == 3), \ |
| "Selective-Dequantization works on 3d tensor only! Please reshape the tensor before calling dequantize function." |
| assert (self.orig_dtype is not None), \ |
| "[De-quantization Error]: you need to call quantize before dequantizing!" |
| fp_out = torch.empty( |
| (indexes.shape[0], |
| *self.orig_shape[1:]), dtype=self.orig_dtype, device=input_q.device) if fp_out is None else fp_out |
| if q_bits == 8: |
| pass |
| elif q_bits == 12: |
| q_mantisa_bits = 4 |
| elif q_bits == 6: |
| q_mantisa_bits = 2 |
| elif q_bits == 4: |
| q_mantisa_bits = 1 |
| else: |
| assert (0), \ |
| f"Missing {q_bits}-dequantization, please add the template arguments for the kernel to support this precision!" |
|
|
| if scale is not None: |
| assert input_q.numel() == fp_out.numel(), \ |
| f'[De-quantization Error]: quantized data should have the same size as original tensor when scale is not None!' |
| input_q = torch.cat([input_q.reshape(-1, self.group_size), scale], dim=-1).contiguous() |
|
|
| fp_quant_module.selective_dequantize(fp_out, input_q, indexes, self.group_size, q_mantisa_bits, |
| q_bits - q_mantisa_bits - 1) |
| return fp_out |
|
|
