wehub-resource-sync
61 lines
2.2 KiB
Python
61 lines
2.2 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import pytest
|
|
import torch
|
|
import deepspeed
|
|
|
|
from deepspeed.accelerator import get_accelerator
|
|
from deepspeed.linear.quantization import QuantizedParameter
|
|
from deepspeed.linear.config import QuantizationConfig
|
|
|
|
from deepspeed.ops.op_builder import FPQuantizerBuilder
|
|
|
|
from unit.common import DistributedTest
|
|
|
|
if not deepspeed.ops.__compatible_ops__[FPQuantizerBuilder.NAME]:
|
|
pytest.skip("FPQuantizer op is not available on this system", allow_module_level=True)
|
|
|
|
|
|
class TestQuantParam(DistributedTest):
|
|
world_size = 1
|
|
|
|
@pytest.mark.parametrize('dtype', [torch.half, torch.float])
|
|
def test_unsupported_dtypes(self, dtype):
|
|
device = get_accelerator().current_device_name()
|
|
data = torch.rand(5, 5, device='cpu', dtype=dtype)
|
|
qp = QuantizedParameter(data)
|
|
with pytest.raises(AssertionError):
|
|
qp.to(device)
|
|
|
|
def test_requires_grad(self):
|
|
data = torch.rand(5, 5, dtype=torch.bfloat16)
|
|
with pytest.raises(ValueError):
|
|
QuantizedParameter(data, requires_grad=True)
|
|
|
|
def test_move_to_accelerator(self):
|
|
device = get_accelerator().current_device()
|
|
data = torch.rand(5, 5, device='cpu', dtype=torch.bfloat16)
|
|
quantization_config = QuantizationConfig()
|
|
quantization_config.q_dtype = FPQuantizerBuilder.get_default_quant_dtype()
|
|
qp = QuantizedParameter(data, quantization_config=quantization_config)
|
|
assert qp.device == torch.device('cpu')
|
|
qp = qp.to(get_accelerator().current_device_name())
|
|
assert qp.device == torch.device(device)
|
|
assert qp.dtype == quantization_config.q_dtype
|
|
|
|
def test_hf_clone(self):
|
|
device = get_accelerator().current_device_name()
|
|
data = torch.rand(5, 5, device=device, dtype=torch.bfloat16)
|
|
|
|
quantization_config = QuantizationConfig(q_bits=6)
|
|
qp = QuantizedParameter(data, quantization_config=quantization_config)
|
|
|
|
# should be able to clone parameter via dict, HF expects this to work
|
|
qp_copy = QuantizedParameter(qp.data, **qp.__dict__)
|
|
|
|
assert all(qp.data == qp_copy.data)
|
|
assert qp.quantization_config == qp_copy.quantization_config
|