Function quant_weight_int8

gpu/convert_checkpoint.py:23–27

(weight)

    print(f"Model config {config.__dict__}")

    def quant_weight_int8(weight):
        s = 1.0 / weight.abs().mean().clamp_(min=1e-5)
        new_weight = (weight * s).round().clamp(-1, 1).to(torch.int8)
        new_scale = (1.0 / s).to(torch.bfloat16)
        return new_weight, new_scale.reshape(1)

    def quant_weight_fp16(weight):
        s = 1.0 / weight.abs().mean().clamp_(min=1e-5)