Quantizes a model by converting all weights to int8. Args: model: The model to quantize. qmode: The quantization mode, e.g. int8, 8da4w. computation_dtype: The dtype that ops are performed in (the resulting dtype of dequantization). Also the dtype of the
( # noqa C901
model: torch.nn.Module,
qmode: str,
computation_dtype: Optional[DType] = None,
checkpoint_dtype: Optional[DType] = None,
checkpoint_path: Optional[Path] = None,
# following arguments only available when setting int4 or gptq quantization.
group_size: Optional[int] = None,
# following arguments are only used for GPTQ
calibration_tasks: Optional[list] = None,
calibration_limit: Optional[int] = None,
calibration_seq_length: Optional[int] = None,
pad_calibration_inputs: bool = False,
percdamp: float = 0.01,
blocksize: int = 128,
tokenizer_path: Optional[Path] = None,
verbose: bool = False,
quantize_with_hqq: bool = True,
)
| 18 | |
| 19 | |
| 20 | def quantize( # noqa C901 |
| 21 | model: torch.nn.Module, |
| 22 | qmode: str, |
| 23 | computation_dtype: Optional[DType] = None, |
| 24 | checkpoint_dtype: Optional[DType] = None, |
| 25 | checkpoint_path: Optional[Path] = None, |
| 26 | # following arguments only available when setting int4 or gptq quantization. |
| 27 | group_size: Optional[int] = None, |
| 28 | # following arguments are only used for GPTQ |
| 29 | calibration_tasks: Optional[list] = None, |
| 30 | calibration_limit: Optional[int] = None, |
| 31 | calibration_seq_length: Optional[int] = None, |
| 32 | pad_calibration_inputs: bool = False, |
| 33 | percdamp: float = 0.01, |
| 34 | blocksize: int = 128, |
| 35 | tokenizer_path: Optional[Path] = None, |
| 36 | verbose: bool = False, |
| 37 | quantize_with_hqq: bool = True, |
| 38 | ) -> torch.nn.Module: |
| 39 | """ |
| 40 | Quantizes a model by converting all weights to int8. |
| 41 | |
| 42 | Args: |
| 43 | model: The model to quantize. |
| 44 | qmode: The quantization mode, e.g. int8, 8da4w. |
| 45 | computation_dtype: The dtype that ops are performed in (the resulting dtype of dequantization). |
| 46 | Also the dtype of the rest of the non-quantized compoents of the model. |
| 47 | checkpoint_dtype: The dtype of the checkpoint, this arg exists since it is more accurate to |
| 48 | quantize the weight in its original dtype. |
| 49 | |
| 50 | Returns: |
| 51 | A quantized model. |
| 52 | """ |
| 53 | if computation_dtype: |
| 54 | computation_torch_dtype = computation_dtype.to_torch_dtype() |
| 55 | else: |
| 56 | computation_torch_dtype = torch.float32 |
| 57 | |
| 58 | if not checkpoint_dtype: |
| 59 | checkpoint_torch_dtype = computation_torch_dtype |
| 60 | else: |
| 61 | checkpoint_torch_dtype = checkpoint_dtype.to_torch_dtype() |
| 62 | |
| 63 | if qmode == "int8": |
| 64 | # Add quantization mode options here: group size, bit width, etc. |
| 65 | return WeightOnlyInt8QuantHandler( |
| 66 | model, precision=checkpoint_torch_dtype |
| 67 | ).quantized_model() |
| 68 | elif qmode.startswith("torchao:fpa"): |
| 69 | pattern = r"torchao:fpa(\d+)w" |
| 70 | matches = re.findall(pattern, qmode) |
| 71 | assert len(matches) == 1, f"Expected 1 match for pattern but got {len(matches)}" |
| 72 | bitwidth = int(matches[0][0]) |
| 73 | _load_torchao_aten_lib(libname="libtorchao_ops_mps_aten") |
| 74 | from torchao.experimental.quant_api import UIntxWeightOnlyLinearQuantizer |
| 75 | |
| 76 | with torch.no_grad(): |
| 77 | # This quantize() is currently doing a model.to(self.precision) so cannot |
no test coverage detected