From 8fd2c050674d51d202e2fcfed8978c22ca70b826 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Wed, 3 Sep 2025 19:23:14 -0400 Subject: [PATCH 1/6] apply transform config Signed-off-by: Kyle Sayers --- .../quantizer_compressed_tensors.py | 12 ++---- src/transformers/utils/import_utils.py | 8 ++-- src/transformers/utils/quantization_config.py | 16 +++++++- .../test_compressed_tensors.py | 37 +++++++++++-------- 4 files changed, 44 insertions(+), 29 deletions(-) diff --git a/src/transformers/quantizers/quantizer_compressed_tensors.py b/src/transformers/quantizers/quantizer_compressed_tensors.py index 803c55775214..7332eff488f1 100644 --- a/src/transformers/quantizers/quantizer_compressed_tensors.py +++ b/src/transformers/quantizers/quantizer_compressed_tensors.py @@ -62,21 +62,17 @@ def validate_environment(self, *args, **kwargs): # torch already should be installed as part of compressed tensors raise ImportError("torch is required for using compressed-tensors quantization") - def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype": - if dtype is None: - logger.info("Loading model using torch.float16 for compressed-tensors quantization") - dtype = torch.float16 - elif dtype != torch.float16: - logger.info("We suggest you to set `dtype=torch.float16` for better efficiency with compressed_tensors.") - return dtype - def _process_model_before_weight_loading(self, model, **kwargs): from compressed_tensors.quantization import apply_quantization_config + from compressed_tensors.transform import apply_transform_config ct_quantization_config = self.compressor.quantization_config + ct_transform_config = self.quantization_config.transform_config # Always initialize compressed wrappers to match the checkpoint apply_quantization_config(model, ct_quantization_config, self.run_compressed) + if ct_transform_config is not None: + apply_transform_config(model, ct_transform_config) if ( self.quantization_config.is_quantization_compressed or self.quantization_config.is_sparsification_compressed diff --git a/src/transformers/utils/import_utils.py b/src/transformers/utils/import_utils.py index a0957ea30e23..e33abccd9484 100644 --- a/src/transformers/utils/import_utils.py +++ b/src/transformers/utils/import_utils.py @@ -122,6 +122,7 @@ def _is_package_available(pkg_name: str, return_version: bool = False) -> Union[ TORCHAO_MIN_VERSION = "0.4.0" AUTOROUND_MIN_VERSION = "0.5.0" TRITON_MIN_VERSION = "1.0.0" +COMPRESSED_TENSORS_MIN_VERSION = "0.11.0" _accelerate_available, _accelerate_version = _is_package_available("accelerate", return_version=True) _apex_available = _is_package_available("apex") @@ -191,8 +192,7 @@ def _is_package_available(pkg_name: str, return_version: bool = False) -> Union[ _is_optimum_quanto_available = True except importlib.metadata.PackageNotFoundError: _is_optimum_quanto_available = False -# For compressed_tensors, only check spec to allow compressed_tensors-nightly package -_compressed_tensors_available = importlib.util.find_spec("compressed_tensors") is not None +_compressed_tensors_available, _compressed_tensors_version = _is_package_available("compressed_tensors", return_version=True) _pandas_available = _is_package_available("pandas") _peft_available = _is_package_available("peft") _phonemizer_available = _is_package_available("phonemizer") @@ -1364,8 +1364,8 @@ def is_qutlass_available() -> Union[tuple[bool, str], bool]: return _qutlass_available -def is_compressed_tensors_available() -> bool: - return _compressed_tensors_available +def is_compressed_tensors_available(min_version: str = COMPRESSED_TENSORS_MIN_VERSION) -> bool: + return _compressed_tensors_available and version.parse(_compressed_tensors_version) >= version.parse(min_version) def is_auto_gptq_available() -> Union[tuple[bool, str], bool]: diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index 037bf3ed73d4..5caa64696a73 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -1319,21 +1319,22 @@ def __init__( sparsity_config: Optional[dict[str, Any]] = None, quant_method: str = "compressed-tensors", run_compressed: bool = True, + transform_config: Optional[dict[str, Any]] = None, **kwargs, ): if is_compressed_tensors_available(): from compressed_tensors.config import SparsityCompressionConfig from compressed_tensors.quantization import QuantizationConfig + from compressed_tensors.transform import TransformConfig else: raise ImportError( "compressed_tensors is not installed and is required for compressed-tensors quantization. Please install it with `pip install compressed-tensors`." ) - self.quantization_config = None - self.sparsity_config = None self.run_compressed = run_compressed # parse from dict to load nested QuantizationScheme objects + self.quantization_config = None if config_groups or kv_cache_scheme: self.quantization_config = QuantizationConfig.model_validate( { @@ -1348,11 +1349,16 @@ def __init__( } ) + self.sparsity_config = None if sparsity_config: self.sparsity_config = SparsityCompressionConfig.load_from_registry( sparsity_config.get("format"), **sparsity_config ) + self.transform_config = None + if transform_config: + self.transform_config = TransformConfig.model_validate(transform_config) + self.quant_method = QuantizationMethod.COMPRESSED_TENSORS def post_init(self): @@ -1392,6 +1398,7 @@ def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs): if "quantization_config" in config_dict: config_dict = dict( sparsity_config=config_dict.get("sparsity_config"), + transform_config=config_dict.get("transform_config"), **config_dict["quantization_config"], ) @@ -1415,6 +1422,11 @@ def to_dict(self) -> dict[str, Any]: else: quantization_config["sparsity_config"] = {} + if self.transform_config is not None: + quantization_config["transform_config"] = self.transform_config.model_dump() + else: + quantization_config["transform_config"] = {} + return quantization_config def to_diff_dict(self) -> dict[str, Any]: diff --git a/tests/quantization/compressed_tensors_integration/test_compressed_tensors.py b/tests/quantization/compressed_tensors_integration/test_compressed_tensors.py index d44e560fff0d..825d9290d891 100644 --- a/tests/quantization/compressed_tensors_integration/test_compressed_tensors.py +++ b/tests/quantization/compressed_tensors_integration/test_compressed_tensors.py @@ -17,8 +17,9 @@ class CompressedTensorsTest(unittest.TestCase): tinyllama_w4a16 = "nm-testing/tinyllama-w4a16-compressed-hf-quantizer" tinyllama_w8a8 = "nm-testing/tinyllama-w8a8-compressed-hf-quantizer" llama3_8b_fp8 = "nm-testing/Meta-Llama-3-8B-Instruct-fp8-hf_compat" + quip_w4a16 = "nm-testing/Llama-3.2-1B-Instruct-quipv-w4a16" - prompt = "Paris is the capital of which country?" + prompt = "The capital of France is Paris, the capital of Germany is Berlin" def tearDown(self): gc.collect() @@ -47,30 +48,32 @@ def test_config_to_from_dict(self): self.assertIsInstance(config_from_dict.sparsity_config, SparsityCompressionConfig) def test_tinyllama_w8a8(self): - expected_out = " Paris is the capital of which country?\n\n 1. Paris is the capital of which country?\n\n 1. Paris is the capital of which country?\n\n 1. Paris is the capital of which country?\n\n" - self._test_quantized_model(self.tinyllama_w8a8, expected_out) + self._test_quantized_model(self.tinyllama_w8a8, 30.0) def test_tinyllama_w4a16(self): - expected_out = " Paris is the capital of which country?\nAnswer: Paris is the capital of France.\nQuestion: Which country is the capital of which city?\nAnswer: The capital of the city of New York is New York.\nQuestion: Which" - self._test_quantized_model(self.tinyllama_w4a16, expected_out) + self._test_quantized_model(self.tinyllama_w4a16, 20.0) def test_tinyllama_w8a16(self): - expected_out = " Paris is the capital of which country?\nA. France\nB. Germany\nC. Spain\nD. Italy\nE. Switzerland\nQ10. Which of the following is not a country in the European Union?\nA." - self._test_quantized_model(self.tinyllama_w8a16, expected_out) + self._test_quantized_model(self.tinyllama_w8a16, 20.0) def test_llama_8b_fp8(self): - expected_out = "<|begin_of_text|>Paris is the capital of which country? France\nWhat is the name of the famous museum in Paris that is home to the Mona Lisa? The Louvre\nWhat is the name of the famous bridge in Paris that is often associated with the city" - self._test_quantized_model(self.llama3_8b_fp8, expected_out) + self._test_quantized_model(self.llama3_8b_fp8, 10.0) - def _test_quantized_model(self, model_name: str, expected_output: str): - """Carry out generation""" + def test_quip_w4a16(self): + self._test_quantized_model(self.quip_w4a16, 10.0) + + def _test_quantized_model(self, model_name: str, expected_perplexity: float): + # load model quantized_model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto") tokenizer = AutoTokenizer.from_pretrained(model_name) device = quantized_model.device + + # check config self.assertIsNotNone( quantized_model.config.quantization_config, "quantization_config should not be None", ) + # check scales self.assertTrue( any( key @@ -79,9 +82,13 @@ def _test_quantized_model(self, model_name: str, expected_output: str): ), "quantized model should load a non-trivial scale into the state dict", ) + + # compute outputs with loss inputs = tokenizer(self.prompt, return_tensors="pt").to(device) - generated_ids = quantized_model.generate(**inputs, max_length=50, do_sample=False) - outputs = tokenizer.batch_decode(generated_ids) + labels = inputs["input_ids"] + with torch.no_grad(): + outputs = quantized_model(**inputs, labels=labels) - self.assertIsNotNone(outputs) - self.assertEqual(outputs[0], expected_output) + # check perplexity + perplexity = torch.exp(outputs.loss) + self.assertLessEqual(perplexity, expected_perplexity) From da8a76aa57d688201f8c9c448a523c37705c15f0 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Wed, 3 Sep 2025 19:35:22 -0400 Subject: [PATCH 2/6] remove unused code Signed-off-by: Kyle Sayers --- src/transformers/quantizers/quantizer_compressed_tensors.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/transformers/quantizers/quantizer_compressed_tensors.py b/src/transformers/quantizers/quantizer_compressed_tensors.py index 7332eff488f1..cd9483807090 100644 --- a/src/transformers/quantizers/quantizer_compressed_tensors.py +++ b/src/transformers/quantizers/quantizer_compressed_tensors.py @@ -17,10 +17,6 @@ from ..utils.quantization_config import CompressedTensorsConfig from .base import HfQuantizer - -if is_torch_available(): - import torch - logger = logging.get_logger(__name__) From 49b14f026daf60fd8e3bd25ceb07359e16592d34 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Wed, 3 Sep 2025 20:30:00 -0400 Subject: [PATCH 3/6] reorder application Signed-off-by: Kyle Sayers --- .../quantizers/quantizer_compressed_tensors.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/transformers/quantizers/quantizer_compressed_tensors.py b/src/transformers/quantizers/quantizer_compressed_tensors.py index cd9483807090..7e7164aabe07 100644 --- a/src/transformers/quantizers/quantizer_compressed_tensors.py +++ b/src/transformers/quantizers/quantizer_compressed_tensors.py @@ -65,10 +65,14 @@ def _process_model_before_weight_loading(self, model, **kwargs): ct_quantization_config = self.compressor.quantization_config ct_transform_config = self.quantization_config.transform_config - # Always initialize compressed wrappers to match the checkpoint - apply_quantization_config(model, ct_quantization_config, self.run_compressed) + # apply transform config first (before modules are converted to run compressed) if ct_transform_config is not None: apply_transform_config(model, ct_transform_config) + + # apply quantization config (potentially convert to run compressed) + apply_quantization_config(model, ct_quantization_config, self.run_compressed) + + # compress meta model to match checkpoint format if ( self.quantization_config.is_quantization_compressed or self.quantization_config.is_sparsification_compressed From 655149af55ed0b8b95d9a832d6007399d4ee1f8c Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Wed, 3 Sep 2025 22:14:59 -0400 Subject: [PATCH 4/6] fix style Signed-off-by: Kyle Sayers --- src/transformers/quantizers/quantizer_compressed_tensors.py | 1 + src/transformers/utils/import_utils.py | 4 +++- src/transformers/utils/quantization_config.py | 2 ++ 3 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/transformers/quantizers/quantizer_compressed_tensors.py b/src/transformers/quantizers/quantizer_compressed_tensors.py index 7e7164aabe07..8e5503e55a05 100644 --- a/src/transformers/quantizers/quantizer_compressed_tensors.py +++ b/src/transformers/quantizers/quantizer_compressed_tensors.py @@ -17,6 +17,7 @@ from ..utils.quantization_config import CompressedTensorsConfig from .base import HfQuantizer + logger = logging.get_logger(__name__) diff --git a/src/transformers/utils/import_utils.py b/src/transformers/utils/import_utils.py index e33abccd9484..912664fce9c0 100644 --- a/src/transformers/utils/import_utils.py +++ b/src/transformers/utils/import_utils.py @@ -192,7 +192,9 @@ def _is_package_available(pkg_name: str, return_version: bool = False) -> Union[ _is_optimum_quanto_available = True except importlib.metadata.PackageNotFoundError: _is_optimum_quanto_available = False -_compressed_tensors_available, _compressed_tensors_version = _is_package_available("compressed_tensors", return_version=True) +_compressed_tensors_available, _compressed_tensors_version = _is_package_available( + "compressed_tensors", return_version=True +) _pandas_available = _is_package_available("pandas") _peft_available = _is_package_available("peft") _phonemizer_available = _is_package_available("phonemizer") diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index 5caa64696a73..d82dde8f67fa 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -1306,6 +1306,8 @@ class CompressedTensorsConfig(QuantizationConfigMixin): do not override, should be compressed-tensors run_compressed (`bool`, *optional*, defaults to `True`): alter submodules (usually linear) in order to emulate compressed model execution if True, otherwise use default submodule + transform_config (`typing.dict[str, typing.Any]`, *optional*): + configuration for online and offline transforms to improve accuracy recovery """ def __init__( From ac720f41ce875875589301376cfbc606a0a475f8 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Sun, 7 Sep 2025 20:44:22 -0400 Subject: [PATCH 5/6] patch_tie_weights_fn Signed-off-by: Kyle Sayers --- src/transformers/modeling_utils.py | 6 +- .../quantizer_compressed_tensors.py | 61 +++++++++++++++++-- 2 files changed, 60 insertions(+), 7 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 973ee405cb3a..0cd40b2c507d 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -3038,11 +3038,13 @@ def _init_weights(self, module): std = getattr(self.config.get_text_config(), "initializer_range", 0.02) if isinstance(module, (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d, nn.ConvTranspose1d, nn.ConvTranspose2d)): - module.weight.data.normal_(mean=0.0, std=std) + if hasattr(module, "weight"): + module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): - module.weight.data.normal_(mean=0.0, std=std) + if hasattr(module, "weight"): + module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() elif isinstance(module, nn.MultiheadAttention): diff --git a/src/transformers/quantizers/quantizer_compressed_tensors.py b/src/transformers/quantizers/quantizer_compressed_tensors.py index 8e5503e55a05..a326fa95d871 100644 --- a/src/transformers/quantizers/quantizer_compressed_tensors.py +++ b/src/transformers/quantizers/quantizer_compressed_tensors.py @@ -13,10 +13,17 @@ # limitations under the License. +from typing import List, Tuple, TYPE_CHECKING +from collections import defaultdict +import torch + from ..utils import is_compressed_tensors_available, is_torch_available, logging from ..utils.quantization_config import CompressedTensorsConfig from .base import HfQuantizer +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + logger = logging.get_logger(__name__) @@ -58,8 +65,16 @@ def validate_environment(self, *args, **kwargs): if not is_torch_available(): # torch already should be installed as part of compressed tensors raise ImportError("torch is required for using compressed-tensors quantization") - - def _process_model_before_weight_loading(self, model, **kwargs): + + def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype": + if dtype is None: + logger.info("Loading model using torch.float16 for compressed-tensors quantization") + dtype = torch.float16 + elif dtype != torch.float16: + logger.info("We suggest you to set `dtype=torch.float16` for better efficiency with compressed_tensors.") + return dtype + + def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwargs): from compressed_tensors.quantization import apply_quantization_config from compressed_tensors.transform import apply_transform_config @@ -80,13 +95,45 @@ def _process_model_before_weight_loading(self, model, **kwargs): ): self.compressor.compress_model(model=model) - def _process_model_after_weight_loading(self, model, **kwargs): - """Decompress loaded model if necessary - need for qat""" + # update tied weights to include added transforms (_dynamic_tied_weights_keys) + self.patch_tie_weights_fn(model) + def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): + """Decompress loaded model if necessary - need for qat""" if ( self.quantization_config.is_quantization_compressed and not self.run_compressed ) or self.quantization_config.is_sparsification_compressed: - self.compressor.decompress_model(model=model) + self.dequantize(model) + + def patch_tie_weights_fn(self, model: "PreTrainedModel"): + # record shared tensors before weight loading + shared_weights: defaultdict[int, List[Tuple[torch.nn.Module, str]]] = defaultdict(list) + for module in model.modules(): + shared_keys = getattr(module, "_dynamic_tied_weights_keys", set()) + for key in shared_keys: + weight = getattr(module, key) + shared_weights[id(weight)].append((module, key)) + + original_fn = model.tie_weights.__func__ + + # this function is called after weight loading but before dispatch + def tie_weights(self: "PreTrainedModel"): + # broadcast loaded weight to other shared weights + for modules_keys in shared_weights.values(): + weights = [getattr(module, key) for module, key in modules_keys] + loaded_weights = [weight for weight in weights if weight.device.type != "meta"] + if len(loaded_weights) <= 0: + raise ValueError("Failed to load shared weight") + if len(loaded_weights) >= 2: + raise ValueError("Loaded too many shared weights") + + loaded_weight = loaded_weights[0] + for module, key in modules_keys: + module.load_state_dict({key: loaded_weight}, strict=False, assign=True) + + original_fn(self) + + model.tie_weights = tie_weights.__get__(model) def update_tp_plan(self, config): additional_plan = { @@ -113,3 +160,7 @@ def is_qat_trainable(self) -> bool: def is_serializable(self, safe_serialization=None) -> bool: """Models quantized using compressed tensors can be saved to disk""" return True + + def dequantize(self, model: "PreTrainedModel"): + """Decompress model""" + self.compressor.decompress_model(model=model) \ No newline at end of file From 02d5f25c8a10a90173d2dc155f03ae5a1d2e0e48 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Wed, 1 Oct 2025 13:15:48 -0400 Subject: [PATCH 6/6] minor updates Signed-off-by: Kyle Sayers --- .../quantizers/quantizer_compressed_tensors.py | 7 ++++--- .../test_compressed_tensors.py | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/transformers/quantizers/quantizer_compressed_tensors.py b/src/transformers/quantizers/quantizer_compressed_tensors.py index a326fa95d871..32c5aa047bdc 100644 --- a/src/transformers/quantizers/quantizer_compressed_tensors.py +++ b/src/transformers/quantizers/quantizer_compressed_tensors.py @@ -109,10 +109,11 @@ def patch_tie_weights_fn(self, model: "PreTrainedModel"): # record shared tensors before weight loading shared_weights: defaultdict[int, List[Tuple[torch.nn.Module, str]]] = defaultdict(list) for module in model.modules(): - shared_keys = getattr(module, "_dynamic_tied_weights_keys", set()) + shared_keys = getattr(module, "_dynamic_tied_weights_keys", list()) for key in shared_keys: - weight = getattr(module, key) - shared_weights[id(weight)].append((module, key)) + weight = getattr(module, key, None) + if weight is not None: + shared_weights[id(weight)].append((module, key)) original_fn = model.tie_weights.__func__ diff --git a/tests/quantization/compressed_tensors_integration/test_compressed_tensors.py b/tests/quantization/compressed_tensors_integration/test_compressed_tensors.py index 825d9290d891..d99bbab40981 100644 --- a/tests/quantization/compressed_tensors_integration/test_compressed_tensors.py +++ b/tests/quantization/compressed_tensors_integration/test_compressed_tensors.py @@ -17,7 +17,7 @@ class CompressedTensorsTest(unittest.TestCase): tinyllama_w4a16 = "nm-testing/tinyllama-w4a16-compressed-hf-quantizer" tinyllama_w8a8 = "nm-testing/tinyllama-w8a8-compressed-hf-quantizer" llama3_8b_fp8 = "nm-testing/Meta-Llama-3-8B-Instruct-fp8-hf_compat" - quip_w4a16 = "nm-testing/Llama-3.2-1B-Instruct-quipv-w4a16" + quip_w4a16 = "nm-testing/Llama-3.2-1B-Instruct-quip-w4a16" prompt = "The capital of France is Paris, the capital of Germany is Berlin"