wehub-resource-sync
135 lines
4.8 KiB
Python
135 lines
4.8 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import pytest
|
|
import torch
|
|
from deepspeed.linear import QuantizationConfig
|
|
|
|
import deepspeed
|
|
|
|
from deepspeed.ops.fp_quantizer import FP_Quantize
|
|
from deepspeed.ops.op_builder import FPQuantizerBuilder
|
|
from deepspeed.accelerator import get_accelerator
|
|
|
|
if not deepspeed.ops.__compatible_ops__[FPQuantizerBuilder.NAME]:
|
|
pytest.skip("FPQuantizer op is not available on this system", allow_module_level=True)
|
|
|
|
# warning: this import silently JIT builds a set of kernels and may take a minute
|
|
from qtorch.quant import float_quantize
|
|
|
|
|
|
def qtorch_quantize(input, exp_bits=4, man_bits=3, rounding="nearest", group_size=1024):
|
|
ori_dt = input.dtype
|
|
ori_shape = input.shape
|
|
last_dim = group_size
|
|
input = input.view(-1, last_dim)
|
|
|
|
q_bits = exp_bits + man_bits + 1
|
|
q_range = FPQuantizerBuilder.get_quant_range(q_bits)
|
|
input_to_float = input.float()
|
|
input_max = input_to_float.abs().amax(dim=-1, keepdim=True)
|
|
|
|
return ((float_quantize(input_to_float / input_max * q_range, exp_bits, man_bits, rounding=rounding) * \
|
|
input_max / q_range).to(ori_dt)).reshape(ori_shape)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bf16"])
|
|
def test_fp_quant_meta(dtype):
|
|
device_name = get_accelerator().device_name()
|
|
group_size = 128
|
|
q_bits = 8
|
|
exp_bits = 4
|
|
man_bits = 3
|
|
|
|
quant_config = QuantizationConfig()
|
|
quant_config.q_dtype = FPQuantizerBuilder.get_default_quant_dtype()
|
|
quant_config.group_size = group_size
|
|
fpq = FP_Quantize(quantization_config=quant_config)
|
|
|
|
for i in range(10):
|
|
x = torch.rand(4, 1024, dtype=dtype)
|
|
|
|
ds_x = x.clone().to(device_name)
|
|
x_quantized, meta_tensor = fpq.quantize(ds_x, q_bits=q_bits, return_meta_tensor=True)
|
|
x_dequantized = fpq.dequantize(x_quantized, q_bits=q_bits, scale=meta_tensor)
|
|
|
|
qtorch_out = qtorch_quantize(x, exp_bits=exp_bits, man_bits=man_bits, group_size=group_size)
|
|
qtorch_error = (qtorch_out - x).abs().sum() / x.numel()
|
|
ds_error = (x_dequantized - ds_x).abs().sum() / x.numel()
|
|
|
|
assert 0.0004 > abs(qtorch_error.item() - ds_error.item()), f"failed on iteration {i}"
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bf16"])
|
|
def test_fp_quant_selective(dtype):
|
|
group_size = 128
|
|
q_bits = 8
|
|
exp_bits = 4
|
|
man_bits = 3
|
|
|
|
device_name = get_accelerator().device_name()
|
|
|
|
quant_config = QuantizationConfig()
|
|
quant_config.q_dtype = FPQuantizerBuilder.get_default_quant_dtype()
|
|
quant_config.group_size = group_size
|
|
fpq = FP_Quantize(quantization_config=quant_config)
|
|
|
|
indexes = torch.zeros(2, dtype=torch.int32, device=device_name)
|
|
indexes[0] = 1
|
|
indexes[1] = 3
|
|
for i in range(10):
|
|
x = torch.rand(4, 1024, dtype=dtype, device=device_name)
|
|
|
|
x = x.reshape(4, 1, x.shape[-1])
|
|
ds_x = x.clone()
|
|
x_quantized = fpq.quantize(ds_x, q_bits=q_bits)
|
|
x_dequantized = fpq.selective_dequantize(x_quantized, indexes, q_bits=q_bits)
|
|
|
|
qtorch_out = qtorch_quantize(x.index_select(0, indexes),
|
|
exp_bits=exp_bits,
|
|
man_bits=man_bits,
|
|
group_size=group_size)
|
|
qtorch_error = (qtorch_out - x.index_select(0, indexes)).abs().sum() / x.numel()
|
|
ds_error = (x_dequantized - x.index_select(0, indexes)).abs().sum() / x.numel()
|
|
|
|
assert 0.0004 > abs(qtorch_error.item() - ds_error.item()), f"failed on iteration {i}"
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bf16"])
|
|
@pytest.mark.parametrize("q_bits", [8, 6, 12], ids=["qbits8", "qbits6", "qbits12"])
|
|
def test_fp_quant(dtype, q_bits):
|
|
device_name = get_accelerator().device_name()
|
|
|
|
quant_config = QuantizationConfig()
|
|
quant_config.q_dtype = FPQuantizerBuilder.get_default_quant_dtype()
|
|
quant_config.group_size = 128
|
|
fpq = FP_Quantize(quantization_config=quant_config)
|
|
|
|
for i in range(10):
|
|
x = torch.rand(4, 1024, dtype=dtype)
|
|
|
|
ds_x = x.clone().to(device_name)
|
|
x_quantized = fpq.quantize(ds_x, q_bits=q_bits)
|
|
x_dequantized = fpq.dequantize(x_quantized, q_bits=q_bits)
|
|
|
|
if q_bits == 8:
|
|
exp_bits = 4
|
|
man_bits = 3
|
|
elif q_bits == 6:
|
|
exp_bits = 3
|
|
man_bits = 2
|
|
elif q_bits == 12:
|
|
exp_bits = 4
|
|
man_bits = 7
|
|
else:
|
|
raise ValueError(f"unknown {q_bits=}")
|
|
|
|
qtorch_out = qtorch_quantize(x, exp_bits=exp_bits, man_bits=man_bits, group_size=quant_config.group_size)
|
|
|
|
qtorch_error = (qtorch_out - x).abs().sum() / x.numel()
|
|
ds_error = (x_dequantized - ds_x).abs().sum() / x.numel()
|
|
|
|
assert 0.0004 > abs(qtorch_error.item() - ds_error.item()), f"failed on iteration {i}"
|