| 32 | |
| 33 | |
| 34 | def get_quantize_weight_fn(quantizer: Quantizer, pre_quant_weight: nn.Parameter) -> Callable: |
| 35 | |
| 36 | def func() -> Tuple[nn.Parameter, Tensor, Tensor]: |
| 37 | quantized_weights, quant_scale, quant_min = quantizer.quantize(pre_quant_weight.data) |
| 38 | # A temporary hack as zero Zero3 assume all model weights has the same type. in all_gather_coalesced.get_only_unique_item |
| 39 | quantized_weights = quantized_weights.view(pre_quant_weight.dtype) |
| 40 | quant_scale = quant_scale.type(pre_quant_weight.dtype) |
| 41 | quant_min = quant_min.type(pre_quant_weight.dtype) |
| 42 | return quantized_weights, quant_scale, quant_min |
| 43 | |
| 44 | return func |
| 45 | |
| 46 | |
| 47 | class QuantizedLinear(nn.Linear): |