| 94 | |
| 95 | |
| 96 | class DeQuantizer: |
| 97 | |
| 98 | def __init__(self, config: Dict, dtype: torch.dtype) -> None: |
| 99 | self.config = config |
| 100 | self.dtype = dtype |
| 101 | assert self.config['num_bits'] == 4 or self.config[ |
| 102 | 'num_bits'] == 8, 'Only INT4 and INT8 quantization is supported.' |
| 103 | assert self.config['symmetric'] == False, 'Only asymmetric quantization is supported at this moment.' |
| 104 | |
| 105 | def dequantize(self, tensor: Tensor, quant_scale: Tensor, quant_min: Tensor) -> Tensor: |
| 106 | # Use customized CUDA quantization kernel if possible. |
| 107 | if self.config['group_size'] % 8 == 0 and \ |
| 108 | (self.config['num_bits'] == 4 or self.config['num_bits'] == 8) and \ |
| 109 | self.config['group_dim'] == len(tensor.shape) - 1 and \ |
| 110 | self.dtype == torch.float16 and device == get_accelerator().device_name(): |
| 111 | |
| 112 | last_dimension_size = self.config['group_size'] |
| 113 | if self.config['num_bits'] == 4: |
| 114 | last_dimension_size = last_dimension_size // 2 |
| 115 | quantized_tensor = get_quantizer_module().dequantize_int4_to_half_experimental( |
| 116 | tensor.reshape(-1, last_dimension_size), quant_scale, quant_min, |
| 117 | tensor.numel() // last_dimension_size, self.config['group_size']) |
| 118 | shape = list(tensor.shape) |
| 119 | shape[-1] = shape[-1] * 2 |
| 120 | elif self.config['num_bits'] == 8: |
| 121 | # last_dimension_size = last_dimension_size // 2 |
| 122 | quantized_tensor = get_quantizer_module().dequantize_int8_to_half_experimental( |
| 123 | tensor.reshape(-1, last_dimension_size), quant_scale, quant_min, |
| 124 | tensor.numel() // last_dimension_size, self.config['group_size']) |
| 125 | shape = list(tensor.shape) |
| 126 | |
| 127 | return quantized_tensor.reshape(shape) |
| 128 | |
| 129 | if self.config['num_bits'] == 4: |
| 130 | tensor = self._decompress_uint4_to_uint8(tensor) |
| 131 | elif self.config['num_bits'] != 8: |
| 132 | assert False, 'Unsupported quantization bits {}'.format(self.config['num_bits']) |
| 133 | |
| 134 | shape = tensor.shape |
| 135 | num_groups = shape[self.config['group_dim']] // self.config['group_size'] |
| 136 | new_shape = (shape[:self.config['group_dim']] + (num_groups, self.config['group_size']) + |
| 137 | shape[self.config['group_dim'] + 1:]) |
| 138 | tensor = tensor.view(new_shape) |
| 139 | |
| 140 | dequantized_tensor = self._dequantize_int8(tensor, quant_scale, quant_min).view(shape) |
| 141 | return dequantized_tensor |
| 142 | |
| 143 | def _dequantize_int8(self, tensor: Tensor, quant_scale: Tensor, quant_min: Tensor) -> Tensor: |
| 144 | assert tensor.dtype == torch.uint8 |
| 145 | data = torch.zeros_like(tensor, dtype=self.dtype, device=tensor.device) |
| 146 | data = data.copy_(tensor) |
| 147 | data = data.div_(quant_scale).add_(quant_min) |
| 148 | |
| 149 | return data |
| 150 | |
| 151 | def _decompress_uint4_to_uint8(self, tensor: Tensor) -> Tensor: |
| 152 | new_data_shape = list(tensor.shape) |
| 153 | new_data_shape[-1] = new_data_shape[-1] * 2 |
no outgoing calls