Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update Tests for Compression Config Change #2314

Merged
merged 2 commits into from
Jun 4, 2024
Merged

Update Tests for Compression Config Change #2314

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
15c995c
update tests
May 31, 2024
53ac12d
Merge branch 'main' into config_update_tests
Jun 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 7 additions & 9 deletions tests/sparseml/transformers/sparsification/test_compress_tensor_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,8 @@
from transformers import AutoConfig

import sparseml
from compressed_tensors import (
COMPRESSION_CONFIG_NAME,
QUANTIZATION_CONFIG_NAME,
SPARSITY_CONFIG_NAME,
)
from compressed_tensors import COMPRESSION_CONFIG_NAME
from compressed_tensors.compressors import ModelCompressor
from compressed_tensors.config import BitmaskConfig, DenseSparsityConfig
from compressed_tensors.quantization import (
QuantizationStatus,
Expand Down Expand Up @@ -96,7 +93,7 @@ def test_sparse_model_reload(compressed, config, dtype, tmp_path):

config = AutoConfig.from_pretrained(tmp_path / "compress_out")
compression_config = getattr(config, COMPRESSION_CONFIG_NAME, None)
sparsity_config = compression_config.get(SPARSITY_CONFIG_NAME, None)
sparsity_config = ModelCompressor.parse_sparsity_config(compression_config)
assert (
sparsity_config["format"] == "dense"
if (not compressed and config is None)
Expand Down Expand Up @@ -146,7 +143,8 @@ def test_dense_model_save(tmp_path, skip_compression_stats, save_compressed):

# for models with 0% sparsity no sparsity config is saved regardless
config = AutoConfig.from_pretrained(tmp_path / "dense_out")
sparsity_config = getattr(config, SPARSITY_CONFIG_NAME, None)
compression_config = getattr(config, COMPRESSION_CONFIG_NAME, None)
sparsity_config = ModelCompressor.parse_sparsity_config(compression_config)
assert sparsity_config is None

shutil.rmtree(tmp_path)
Expand Down Expand Up @@ -203,7 +201,7 @@ def test_quant_model_reload(format, dtype, tmp_path):

config = AutoConfig.from_pretrained(tmp_path / "compress_out")
compression_config = getattr(config, COMPRESSION_CONFIG_NAME, None)
quant_config = compression_config.get(QUANTIZATION_CONFIG_NAME, None)
quant_config = ModelCompressor.parse_quantization_config(compression_config)
assert quant_config["format"] == format

dense_model = SparseAutoModelForCausalLM.from_pretrained(
Expand Down Expand Up @@ -273,7 +271,7 @@ def test_quant_infer_format(status, expected_format, expected_dtype, tmp_path):

config = AutoConfig.from_pretrained(tmp_path / "compress_out")
compression_config = getattr(config, COMPRESSION_CONFIG_NAME, None)
quant_config = compression_config.get(QUANTIZATION_CONFIG_NAME, None)
quant_config = ModelCompressor.parse_quantization_config(compression_config)
assert quant_config["quantization_status"] == status.value
assert quant_config["format"] == expected_format

Expand Down
Loading