Quantized parameter class that implements weight quantization. Weights are stored in quantized form on GPUs, and can be dequantized on-the-fly when needed by the model. The weights are actually quantized during any `.to(device)`. Arguments: data (Tensor): parameter tensor.
| 16 | |
| 17 | |
| 18 | class QuantizedParameter(nn.Parameter): |
| 19 | """ |
| 20 | Quantized parameter class that implements weight quantization. Weights |
| 21 | are stored in quantized form on GPUs, and can be dequantized on-the-fly when |
| 22 | needed by the model. The weights are actually quantized during any `.to(device)`. |
| 23 | |
| 24 | Arguments: |
| 25 | data (Tensor): parameter tensor. |
| 26 | requires_grad (bool, optional): if the parameter requires gradient. Defaults |
| 27 | to False and is not supported to be True. Argument provided only for interface |
| 28 | compatibility with torch.nn.Parameter. |
| 29 | quantization_config (QuantizationConfig, optional): |
| 30 | quantizer (Quantizer, optional): Defaults to FP_Quantize but can be any quantizer |
| 31 | that implements deepspeed.ops.fp_quantizer.Quantizer. This argument is also |
| 32 | required since the quantizer is stashed in the Parameter itself, some models |
| 33 | may clone the Parameter by passing an attribute __dict__. For an example, see |
| 34 | tests/unit/linear/test_quant_param.py::TestQuantParam::test_hf_clone |
| 35 | """ |
| 36 | |
| 37 | def __new__( |
| 38 | cls, |
| 39 | data: Optional[torch.Tensor] = None, |
| 40 | requires_grad: bool = False, # quantized weights must be frozen |
| 41 | quantization_config: QuantizationConfig = None, |
| 42 | quantizer: Quantizer = None, |
| 43 | ): |
| 44 | if requires_grad: |
| 45 | raise ValueError("requires_grad=True is not supported with QuantizedParameter") |
| 46 | if data is None: |
| 47 | data = torch.empty(0) |
| 48 | self = torch.Tensor._make_subclass(cls, data, requires_grad) |
| 49 | self.quantization_config = QuantizationConfig() if quantization_config is None else quantization_config |
| 50 | if quantizer is not None: |
| 51 | self.quantizer = quantizer |
| 52 | else: |
| 53 | # if FPQuantizerBuilder is not compatible in this env this init will fail |
| 54 | self.quantizer = FP_Quantize(quantization_config=self.quantization_config) |
| 55 | self._ensure_quantized(self) |
| 56 | return self |
| 57 | |
| 58 | def _ensure_quantized(self, tensor: torch.Tensor): |
| 59 | # If the tensor is on the accelerator and is not quantized, then quantize it in-place. |
| 60 | if get_accelerator().on_accelerator(tensor) and tensor.dtype != self.quantization_config.q_dtype: |
| 61 | with get_accelerator().stream(get_accelerator().current_stream(tensor.device)): |
| 62 | tensor.data = self.quantizer.quantize(tensor.data, |
| 63 | q_bits=self.quantization_config.q_bits, |
| 64 | q_mantisa_bits=self.quantization_config.mantissa_bits) |
| 65 | assert tensor.dtype == self.quantization_config.q_dtype |
| 66 | |
| 67 | def dequantized(self) -> torch.Tensor: |
| 68 | """ |
| 69 | Return a tensor containing the dequantized weights of this parameter. |
| 70 | """ |
| 71 | if get_accelerator().on_accelerator(self.data) and self.data.dtype == self.quantization_config.q_dtype: |
| 72 | with get_accelerator().stream(get_accelerator().current_stream(self.data.device)): |
| 73 | return self.quantizer.dequantize(self.data, |
| 74 | q_bits=self.quantization_config.q_bits, |
| 75 | q_mantisa_bits=self.quantization_config.mantissa_bits) |
no outgoing calls