compressed-tensors 0.12.3a20251212__tar.gz → 0.12.3a20251215__tar.gz

{compressed_tensors-0.12.3a20251212/src/compressed_tensors.egg-info → compressed_tensors-0.12.3a20251215}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.12.3a20251212
+Version: 0.12.3a20251215
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/vllm-project/compressed-tensors
 Author: Neuralmagic, Inc.

{compressed_tensors-0.12.3a20251212 → compressed_tensors-0.12.3a20251215}/examples/quantize_and_pack_int4.ipynb

@@ -15,7 +15,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -25,8 +25,7 @@
     "from compressed_tensors.quantization import (\n",
     "    QuantizationConfig,\n",
     "    QuantizationStatus,\n",
-    "    apply_quantization_config,\n",
-    "    compress_quantized_weights\n",
+    "    apply_quantization_config\n",
     ")\n",
     "from compressed_tensors.compressors import ModelCompressor\n",
     "from transformers import AutoModelForCausalLM, AutoTokenizer, DefaultDataCollator\n",
@@ -37,51 +36,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c883cdc8ecd04866bd01d61796b81c26",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "config.json:   0%|          | 0.00/560 [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "32b18b14b6774ce7b61d2854a1ed5f49",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "model.safetensors:   0%|          | 0.00/4.40G [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "370c6d18521a4b65833a411728be1ed7",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "generation_config.json:   0%|          | 0.00/129 [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
     {
      "data": {
       "text/plain": [
@@ -113,7 +70,7 @@
        ")"
       ]
      },
-     "execution_count": 3,
+     "execution_count": 13,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -122,7 +79,7 @@
     "# load a dense, unquantized tiny llama model\n",
     "device = \"cuda:0\"\n",
     "model_name = \"TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T\"\n",
-    "model = AutoModelForCausalLM.from_pretrained(model_name, device_map=device, torch_dtype=\"auto\")\n",
+    "model = AutoModelForCausalLM.from_pretrained(model_name, device_map=device, torch_dtype=torch.bfloat16)\n",
     "model"
    ]
   },
@@ -139,7 +96,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -164,7 +121,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -177,7 +134,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 16,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -198,14 +155,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": 17,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Running calibration: 512it [00:33, 15.42it/s]\n"
+      "Running calibration: 512it [00:58,  8.82it/s]\n"
      ]
     }
    ],
@@ -233,20 +190,24 @@
     "\n",
     "Notice that at this point, the weight itself is still a floating point and has not been quantized. \n",
     "\n",
-    "To convert the weights to an integer type, we need to apply the `compress_quantized_weights` function. After compressing the weights, a forward pass of the model can no longer be run in PyTorch"
+    "To convert the weights to an integer type, we need to apply the `compress_model` function. After compressing the weights, a forward pass of the model can no longer be run in PyTorch.\n",
+    "\n",
+    "After compressing the quantized model with the `pack-quantized` format, weights are represented as logical int4 values packed into int32 containers ( `weight_packed` ), with the original shape recorded in `weight_shape`.\n",
+    "\n",
+    "This packed representation is what gets saved to disk when using ModelCompressor.compress_model(model)."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": 18,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Scale: tensor([17296.], device='cuda:4', dtype=torch.float16), Zero Point: tensor([0], device='cuda:4', dtype=torch.int8)\n",
-      "Weight min: -1.587890625 max: 1.0283203125 dtype: torch.float16\n"
+      "Scale: tensor([-3.0465e+26], device='cuda:0', dtype=torch.bfloat16), Zero Point: tensor([0], device='cuda:0', dtype=torch.int8)\n",
+      "Weight min: -1.5859375 max: 1.03125 dtype: torch.bfloat16\n"
      ]
     }
    ],
@@ -262,64 +223,62 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": 19,
    "metadata": {},
    "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Compressing model: 154it [00:02, 59.75it/s]"
+     ]
+    },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Scale: tensor([17296.], device='cuda:4', dtype=torch.float16), Zero Point: tensor([0], device='cuda:4', dtype=torch.int8)\n",
-      "Weight min: 0 max: 0 dtype: torch.int8\n"
+      "Compressed weight scale: tensor([-3.0465e+26], device='cuda:0', dtype=torch.bfloat16), zero point: tensor([0], device='cuda:0', dtype=torch.int8)\n",
+      "Compressed weight  dtype: torch.int32\n",
+      "Compressed weight shape: torch.Size([2048, 256])\n",
+      "Uncompressed weight shape: tensor([2048, 2048], device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
      ]
     }
    ],
    "source": [
     "# convert quantized weights to integers\n",
-    "model.apply(compress_quantized_weights)\n",
+    "compressor = ModelCompressor(quantization_config=config)\n",
+    "compressor.compress_model(model)\n",
     "\n",
     "state_dict = model.state_dict()\n",
     "example_layer = \"model.layers.0.self_attn.q_proj.weight\"\n",
     "scale = state_dict[example_layer + \"_scale\"]\n",
     "zero_point = state_dict[example_layer + \"_zero_point\"]\n",
-    "weight = state_dict[example_layer]\n",
-    "print(f\"Scale: {scale}, Zero Point: {zero_point}\")\n",
-    "print(f\"Weight min: {torch.min(weight)} max: {torch.max(weight)} dtype: {weight.dtype}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "After compressing the quantized model, the weight matrix has a range of int4 but is stored in an int8. \n",
-    "\n",
-    "We can further compress the model on disk using the `pack-quantized` format we specified in the config. This compression format will pack the int4 weights into int32"
+    "weight = state_dict[example_layer + \"_packed\"]\n",
+    "shape = state_dict[example_layer + \"_shape\"]\n",
+    "print(f\"Compressed weight scale: {scale}, zero point: {zero_point}\")\n",
+    "print(f\"Compressed weight  dtype: {weight.dtype}\")\n",
+    "print(f\"Compressed weight shape: {weight.shape}\")\n",
+    "print(f\"Uncompressed weight shape: {shape}\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 31,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Compression format: pack-quantized\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Quantized Compression: 100%|██████████| 509/509 [00:03<00:00, 153.70it/s]\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Size of the model's weights on disk using safetensors: 712.23 MB\n"
+      "Compression format: pack-quantized\n",
+      "Size of the model's weights on disk using safetensors: 712.25 MB\n"
      ]
     }
    ],
@@ -330,9 +289,8 @@
     "compression_format = config.format\n",
     "print(f\"Compression format: {compression_format}\")\n",
     "\n",
-    "compressor = ModelCompressor(quantization_config=config)\n",
-    "compressed_state_dict = compressor.compress(model)\n",
-    "model.save_pretrained(output_dir, state_dict=compressed_state_dict)\n",
+    "\n",
+    "model.save_pretrained(output_dir, state_dict=model.state_dict())\n",
     "compressor.update_config(output_dir)\n",
     "\n",
     "compressed_size_on_disk_mb = os.path.getsize(os.path.join(output_dir, \"model.safetensors\")) / 1024 / 1024\n",
@@ -356,7 +314,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.12"
+   "version": "3.12.12"
   }
  },
  "nbformat": 4,

{compressed_tensors-0.12.3a20251212 → compressed_tensors-0.12.3a20251215}/src/compressed_tensors/compressors/base.py

@@ -20,11 +20,6 @@ from compressed_tensors.config import SparsityCompressionConfig
 from compressed_tensors.quantization import QuantizationArgs, QuantizationConfig
 from compressed_tensors.registry import RegistryMixin
 from compressed_tensors.utils import has_offloaded_params
-from compressed_tensors.utils.offload import (
-    delete_offload_parameter,
-    get_offloaded_device,
-    register_offload_parameter,
-)
 from torch import Tensor
 from torch.nn import Module
 
@@ -190,37 +185,10 @@ class BaseCompressor(RegistryMixin, ABC):
         for name, parameter in module.named_parameters():
             compressed_data[name] = parameter
 
-        # Save references to original parameters before decompression
-        original_scale = compressed_data.get("weight_scale")
-        original_zp = compressed_data.get("weight_zero_point")
-
-        # NOTE: decompress_weight may modify compressed_data dict in-place
-        # This is subtle but allows us to update the module's qparams with
-        # the unpacked values.
-        # TODO: Consider refactoring to return modified qparams explicitly
-        result = self.decompress_weight(
+        return self.decompress_weight(
             compressed_data=compressed_data, quantization_args=quantization_args
         ).to(device)
 
-        # Update module's parameters only if they were modified
-        for param_name, original_param in [
-            ("weight_scale", original_scale),
-            ("weight_zero_point", original_zp),
-        ]:
-            if (
-                param_name in compressed_data
-                and compressed_data[param_name] is not original_param
-            ):
-                # Delete the old parameter and register the updated one
-                delete_offload_parameter(module, param_name)
-                offload_device = get_offloaded_device(module)
-                param = torch.nn.Parameter(
-                    compressed_data[param_name], requires_grad=False
-                )
-                register_offload_parameter(module, param_name, param, offload_device)
-
-        return result
-
     def decompress_weight(
         self, compressed_data: Dict[str, Tensor], **kwargs
     ) -> torch.Tensor:

{compressed_tensors-0.12.3a20251212 → compressed_tensors-0.12.3a20251215}/src/compressed_tensors/compressors/quantized_compressors/base.py

@@ -18,7 +18,7 @@ from typing import Any, Dict, Generator, Tuple, Union
 
 import torch
 from compressed_tensors.compressors.base import BaseCompressor
-from compressed_tensors.quantization import QuantizationScheme
+from compressed_tensors.quantization import QuantizationScheme, QuantizationStrategy
 from compressed_tensors.utils import (
     get_nested_mappings_from_state_dict,
     get_nested_weight_mappings,
@@ -85,7 +85,6 @@ class BaseQuantizationCompressor(BaseCompressor):
         """
         uncompressed_names = list(model_state.keys())
         compressed_dict = {}
-        compressed_param_names = set()
 
         # compress values
         desc = "Compressing with quantization"
@@ -120,38 +119,54 @@ class BaseQuantizationCompressor(BaseCompressor):
                     device=compression_device,
                 )
 
-                # update state dict and track which params were added
+                # update state dict
                 for key, value in compressed_values.items():
-                    full_name = prefix + key
-                    compressed_dict[full_name] = value.to(compression_device)
-                    compressed_param_names.add(full_name)
+                    compressed_dict[prefix + key] = value.to(compression_device)
 
             else:
-                # Skip qparams already added by compress_weight
-                if name in compressed_param_names:
+                # omit saving zero points for symmetric or packed quantization
+                if name.endswith("zero_point") and self._skip_zp(name, names_to_scheme):
                     continue
 
-                # for symmetric quantization, omit zero_point
-                # manually because it wasn't handled in compress_weight
-                if name.endswith("weight_zero_point"):
-                    module_path = name.rsplit(".", 1)[0]
-                    if (
-                        module_path in names_to_scheme
-                        and names_to_scheme[module_path].weights.symmetric
-                    ):
-                        continue
-                    # Call compress_zp if available (for PackedQuantizationCompressor)
-                    if module_path in names_to_scheme and hasattr(self, "compress_zp"):
-                        value = self.compress_zp(
-                            value, names_to_scheme[module_path].weights
-                        )
-                        if value is None:
-                            continue
+                if name.endswith("weight_scale") and self._skip_scale():
+                    continue
 
                 compressed_dict[name] = value.to(compression_device)
 
         return compressed_dict
 
+    def _skip_scale(self):
+        from compressed_tensors.compressors import NVFP4PackedCompressor
+
+        return isinstance(self, NVFP4PackedCompressor)
+
+    def _skip_zp(
+        self, name: str, names_to_scheme: Dict[str, QuantizationScheme]
+    ) -> bool:
+        from compressed_tensors.compressors import PackedQuantizationCompressor
+
+        module_name, zp_name = name.rsplit(".", 1) if "." in name else ("", name)
+        scheme = names_to_scheme[module_name]
+
+        if zp_name == "weight_zero_point":
+            args = scheme.weights
+        if zp_name == "input_zero_point":
+            args = scheme.input_activations
+        if zp_name == "output_zero_point":
+            args = scheme.output_activations
+
+        symmetric = args.symmetric
+        packable_strategies = [
+            QuantizationStrategy.GROUP.value,
+            QuantizationStrategy.CHANNEL.value,
+        ]
+        packed = (
+            isinstance(self, PackedQuantizationCompressor)
+            and args.strategy in packable_strategies
+        )
+
+        return symmetric or packed
+
     def decompress(
         self,
         path_to_model_or_tensors: Union[str, Path, Dict[str, Any]],

{compressed_tensors-0.12.3a20251212 → compressed_tensors-0.12.3a20251215}/src/compressed_tensors/compressors/quantized_compressors/fp4_quantized.py

@@ -56,6 +56,7 @@ class NVFP4PackedCompressor(BaseQuantizationCompressor):
         return (
             "weight_packed",
             "weight_scale",
+            "weight_zero_point",
             "weight_global_scale",
         )
 
@@ -72,12 +73,13 @@ class NVFP4PackedCompressor(BaseQuantizationCompressor):
         :param quantization_args: quantization parameters for the weight
         :return: dictionary mapping compressed parameter names to shape and dtype
         """
-        return {
+        output = {
             "weight_packed": (
                 torch.Size((weight_shape[0], weight_shape[1] // 2)),
                 torch.uint8,
             ),
         }
+        return output
 
     def compress_scale(
         self,
@@ -112,13 +114,6 @@ class NVFP4PackedCompressor(BaseQuantizationCompressor):
         compressed_dict["weight_scale"] = self.compress_scale(
             scale=scale, quantization_args=quantization_args
         )
-
-        if global_scale is None:
-            raise ValueError(
-                "NVFP4 quantization requires global_scale (TENSOR_GROUP strategy). "
-                "Use TENSOR_GROUP strategy instead of GROUP for FP4 quantization."
-            )
-
         return compressed_dict
 
     def decompress_weight(
@@ -132,12 +127,6 @@ class NVFP4PackedCompressor(BaseQuantizationCompressor):
         m, n = weight.shape
         # TODO: use a user provided dequant dtype
         unpacked = unpack_fp4_from_uint8(weight, m, n * 2)
-
-        # cast scale dtype to match unpacked dtype for dequantization
-        if scale.dtype != unpacked.dtype:
-            scale = scale.to(unpacked.dtype)
-            compressed_data["weight_scale"] = scale
-
         decompressed_weight = dequantize(
             x_q=unpacked, scale=scale, global_scale=global_scale, dtype=unpacked.dtype
         )

{compressed_tensors-0.12.3a20251212 → compressed_tensors-0.12.3a20251215}/src/compressed_tensors/compressors/quantized_compressors/pack_quantized.py

@@ -64,34 +64,25 @@ class PackedQuantizationCompressor(BaseQuantizationCompressor):
         """
         pack_factor = 32 // quantization_args.num_bits
         packed_size = math.ceil(weight_shape[1] / pack_factor)
+        packed_size_zp = math.ceil(weight_shape[0] / pack_factor)
         output = {
             "weight_packed": (torch.Size((weight_shape[0], packed_size)), torch.int32),
             "weight_shape": (torch.Size((2,)), torch.int32),
         }
-
-        # Add weight_scale - always needed for quantization
-        if quantization_args.strategy in [
+        if not quantization_args.symmetric and quantization_args.strategy in [
             QuantizationStrategy.GROUP.value,
             QuantizationStrategy.CHANNEL.value,
         ]:
-            shape_factor = (
+            zp_factor = (
                 quantization_args.group_size
                 if quantization_args.strategy == QuantizationStrategy.GROUP.value
                 else weight_shape[-1]
             )
-            scale_cols = math.ceil(weight_shape[-1] / shape_factor)
-            output["weight_scale"] = (
-                torch.Size((weight_shape[0], scale_cols)),
-                quantization_args.scale_dtype,
-            )
-
-            # Add weight_zero_point for asymmetric quantization
-            if not quantization_args.symmetric:
-                output["weight_zero_point"] = (
-                    torch.Size((math.ceil(weight_shape[0] / pack_factor), scale_cols)),
-                    torch.int32,
-                )
 
+            output["weight_zero_point"] = (
+                torch.Size((packed_size_zp, weight_shape[-1] // zp_factor)),
+                torch.int32,
+            )
         return output
 
     def compress_weight(
@@ -184,8 +175,6 @@ class PackedQuantizationCompressor(BaseQuantizationCompressor):
             zero_point = unpack_from_int32(
                 zero_point, num_bits, original_zp_shape, packed_dim=0
             )
-            # Update the compressed_data dict with the unpacked zero_point
-            compressed_data["weight_zero_point"] = zero_point
 
         decompressed_weight = dequantize(
             x_q=unpacked, scale=scale, zero_point=zero_point, g_idx=g_idx
@@ -193,20 +182,6 @@ class PackedQuantizationCompressor(BaseQuantizationCompressor):
 
         return decompressed_weight
 
-    def compress_zp(
-        self, zero_point: Tensor, quantization_args: Optional[QuantizationArgs] = None
-    ) -> Optional[Tensor]:
-        if zero_point is None or quantization_args.symmetric:
-            return None
-        if zero_point.dtype == torch.int32:
-            return zero_point
-        if quantization_args.strategy in [
-            QuantizationStrategy.GROUP.value,
-            QuantizationStrategy.CHANNEL.value,
-        ]:
-            return pack_to_int32(zero_point, quantization_args.num_bits, packed_dim=0)
-        return zero_point
-
 
 def pack_to_int32(
     value: torch.Tensor,
@@ -251,9 +226,6 @@ def pack_to_int32(
     if packed_dim == 0:
         value = value.transpose(0, 1)
 
-    # Ensure contiguous memory for .view() operation
-    value = value.contiguous()
-
     rows, cols = value.shape
     padded_cols = math.ceil(cols / pack_factor) * pack_factor
     pad_len = padded_cols - cols

{compressed_tensors-0.12.3a20251212 → compressed_tensors-0.12.3a20251215}/src/compressed_tensors/version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.12.3.a20251212'
+__version__ = version = '0.12.3.a20251215'
 __version_tuple__ = version_tuple = (0, 12, 3)

{compressed_tensors-0.12.3a20251212 → compressed_tensors-0.12.3a20251215/src/compressed_tensors.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.12.3a20251212
+Version: 0.12.3a20251215
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/vllm-project/compressed-tensors
 Author: Neuralmagic, Inc.