unit tests

tests/quantization/test_quant_args.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+from pydantic import ValidationError
+from sparsetensors.quantization import (
+    QuantizationArgs,
+    QuantizationStrategy,
+    QuantizationType,
+)
+
+
+def test_defaults():
+    default = QuantizationArgs()
+
+    assert default.num_bits == 8
+    assert default.type == QuantizationType.INT
+    assert default.symmetric
+    assert default.strategy == QuantizationStrategy.TENSOR
+    assert default.group_size is None
+    assert default.block_structure is None
+
+
+def test_group():
+    kwargs = {"strategy": "group", "group_size": 128}
+
+    group = QuantizationArgs(**kwargs)
+    assert group.strategy == QuantizationStrategy.GROUP
+    assert group.group_size == kwargs["group_size"]
+
+
+def test_block():
+    kwargs = {"strategy": "block", "block_structure": "2x4"}
+
+    block = QuantizationArgs(**kwargs)
+    assert block.strategy == QuantizationStrategy.BLOCK
+    assert block.block_structure == kwargs["block_structure"]
+
+
+def test_invalid():
+    with pytest.raises(ValidationError):
+        _ = QuantizationArgs(type="invalid")
+    with pytest.raises(ValidationError):
+        _ = QuantizationArgs(strategy="invalid")

tests/quantization/test_quant_config.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytest
+from pydantic import ValidationError
+from sparsetensors.quantization import (
+    QuantizationConfig,
+    QuantizationScheme,
+    QuantizationStatus,
+)
+
+
+def test_basic_config():
+    config_groups = {"group_1": QuantizationScheme(targets=[])}
+    config = QuantizationConfig(config_groups=config_groups)
+
+    assert config.config_groups == config_groups
+    assert config.quant_method == "sparseml"
+    assert config.format == "fakequant"
+    assert config.quantization_status == QuantizationStatus.INITIALIZED
+    assert config.global_compression_ratio is None
+    assert config.ignore is None
+
+
+def test_full_config():
+    config_groups = {
+        "group_1": QuantizationScheme(targets=[]),
+        "group_2": QuantizationScheme(targets=[]),
+    }
+    global_compression_ratio = 3.5
+    ignore = ["model.layers.0"]
+    quantization_status = "compressed"
+
+    config = QuantizationConfig(
+        config_groups=config_groups,
+        global_compression_ratio=global_compression_ratio,
+        ignore=ignore,
+        quantization_status=quantization_status,
+    )
+    assert config.config_groups == config_groups
+    assert config.global_compression_ratio == global_compression_ratio
+    assert config.ignore == ignore
+    assert config.quantization_status == QuantizationStatus.COMPRESSED
+
+
+def test_need_config_groups():
+    with pytest.raises(ValidationError):
+        _ = QuantizationScheme()

tests/quantization/test_quant_scheme.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+from pydantic import ValidationError
+from sparsetensors.quantization import QuantizationArgs, QuantizationScheme
+
+
+def test_basic_scheme():
+    targets = ["model.layer.0", "model.layer.3"]
+    weights = QuantizationArgs()
+
+    scheme = QuantizationScheme(targets=targets, weights=weights)
+    assert scheme.targets == targets
+    assert scheme.weights == weights
+    assert scheme.input_activations is None
+    assert scheme.output_activations is None
+
+
+def test_full_scheme():
+    targets = ["Linear"]
+    weights = QuantizationArgs()
+    input_activations = QuantizationArgs(num_bits=4)
+    output_activations = QuantizationArgs(num_bits=8, type="float", symmetric=False)
+
+    scheme = QuantizationScheme(
+        targets=targets,
+        weights=weights,
+        input_activations=input_activations,
+        output_activations=output_activations,
+    )
+    assert scheme.targets == targets
+    assert scheme.weights == weights
+    assert scheme.input_activations == input_activations
+    assert scheme.output_activations == output_activations
+
+
+def test_needs_targets():
+    with pytest.raises(ValidationError):
+        _ = QuantizationScheme()