Function func

deepspeed/inference/quantization/layers.py:36–42

()

def get_quantize_weight_fn(quantizer: Quantizer, pre_quant_weight: nn.Parameter) -> Callable:

    def func() -> Tuple[nn.Parameter, Tensor, Tensor]:
        quantized_weights, quant_scale, quant_min = quantizer.quantize(pre_quant_weight.data)
        # A temporary hack as zero Zero3 assume all model weights has the same type. in all_gather_coalesced.get_only_unique_item
        quantized_weights = quantized_weights.view(pre_quant_weight.dtype)
        quant_scale = quant_scale.type(pre_quant_weight.dtype)
        quant_min = quant_min.type(pre_quant_weight.dtype)
        return quantized_weights, quant_scale, quant_min

    return func