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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.12.3a20251114
+Version: 0.12.3a20251212
 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.3a20251114 → compressed_tensors-0.12.3a20251212}/src/compressed_tensors/compressors/base.py

@@ -20,6 +20,11 @@ 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
 
@@ -185,10 +190,37 @@ class BaseCompressor(RegistryMixin, ABC):
         for name, parameter in module.named_parameters():
             compressed_data[name] = parameter
 
-        return self.decompress_weight(
+        # 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(
             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.3a20251114 → compressed_tensors-0.12.3a20251212}/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, QuantizationStrategy
+from compressed_tensors.quantization import QuantizationScheme
 from compressed_tensors.utils import (
     get_nested_mappings_from_state_dict,
     get_nested_weight_mappings,
@@ -85,6 +85,7 @@ class BaseQuantizationCompressor(BaseCompressor):
         """
         uncompressed_names = list(model_state.keys())
         compressed_dict = {}
+        compressed_param_names = set()
 
         # compress values
         desc = "Compressing with quantization"
@@ -119,54 +120,38 @@ class BaseQuantizationCompressor(BaseCompressor):
                     device=compression_device,
                 )
 
-                # update state dict
+                # update state dict and track which params were added
                 for key, value in compressed_values.items():
-                    compressed_dict[prefix + key] = value.to(compression_device)
+                    full_name = prefix + key
+                    compressed_dict[full_name] = value.to(compression_device)
+                    compressed_param_names.add(full_name)
 
             else:
-                # omit saving zero points for symmetric or packed quantization
-                if name.endswith("zero_point") and self._skip_zp(name, names_to_scheme):
+                # Skip qparams already added by compress_weight
+                if name in compressed_param_names:
                     continue
 
-                if name.endswith("weight_scale") and self._skip_scale():
-                    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
 
                 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.3a20251114 → compressed_tensors-0.12.3a20251212}/src/compressed_tensors/compressors/quantized_compressors/fp4_quantized.py

@@ -56,7 +56,6 @@ class NVFP4PackedCompressor(BaseQuantizationCompressor):
         return (
             "weight_packed",
             "weight_scale",
-            "weight_zero_point",
             "weight_global_scale",
         )
 
@@ -73,13 +72,20 @@ class NVFP4PackedCompressor(BaseQuantizationCompressor):
         :param quantization_args: quantization parameters for the weight
         :return: dictionary mapping compressed parameter names to shape and dtype
         """
-        output = {
+        return {
             "weight_packed": (
                 torch.Size((weight_shape[0], weight_shape[1] // 2)),
                 torch.uint8,
             ),
         }
-        return output
+
+    def compress_scale(
+        self,
+        scale: Tensor,
+        quantization_args: QuantizationArgs,
+    ) -> Dict[str, torch.Tensor]:
+        assert quantization_args.scale_dtype is not None
+        return scale.to(quantization_args.scale_dtype)
 
     def compress_weight(
         self,
@@ -103,7 +109,16 @@ class NVFP4PackedCompressor(BaseQuantizationCompressor):
         if device is not None:
             weight_packed = weight_packed.to(device)
         compressed_dict["weight_packed"] = weight_packed
-        compressed_dict["weight_scale"] = scale.to(quantization_args.scale_dtype)
+        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(
@@ -117,6 +132,12 @@ 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
         )
@@ -130,7 +151,21 @@ class MXFP4PackedCompressor(NVFP4PackedCompressor):
     Alias for mxfp4 quantized models
     """
 
-    pass
+    def compress_scale(
+        self,
+        scale: Tensor,
+        quantization_args: QuantizationArgs,
+    ) -> Dict[str, torch.Tensor]:
+        assert quantization_args.scale_dtype is not None
+        scale_exp = 127 + torch.floor(torch.log2(scale)).to(torch.int32) - 2
+        return scale_exp.to(quantization_args.scale_dtype)
+
+    def decompress_weight(
+        self,
+        compressed_data: Dict[str, Tensor],
+        quantization_args: Optional[QuantizationArgs] = None,
+    ) -> torch.Tensor:
+        raise NotImplementedError("MXFP4 Decompression is currently not supported")
 
 
 @torch.compile(fullgraph=True, dynamic=True)

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

@@ -64,25 +64,34 @@ 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),
         }
-        if not quantization_args.symmetric and quantization_args.strategy in [
+
+        # Add weight_scale - always needed for quantization
+        if quantization_args.strategy in [
             QuantizationStrategy.GROUP.value,
             QuantizationStrategy.CHANNEL.value,
         ]:
-            zp_factor = (
+            shape_factor = (
                 quantization_args.group_size
                 if quantization_args.strategy == QuantizationStrategy.GROUP.value
                 else weight_shape[-1]
             )
-
-            output["weight_zero_point"] = (
-                torch.Size((packed_size_zp, weight_shape[-1] // zp_factor)),
-                torch.int32,
+            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,
+                )
+
         return output
 
     def compress_weight(
@@ -175,6 +184,8 @@ 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
@@ -182,6 +193,20 @@ 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,
@@ -226,6 +251,9 @@ 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.3a20251114 → compressed_tensors-0.12.3a20251212}/src/compressed_tensors/config/format.py

@@ -50,6 +50,8 @@ def _get_quant_compression_format(
     is_weight_only = weight_args is not None and input_args is None
 
     if weight_args.num_bits == 4 and weight_args.type == QuantizationType.FLOAT.value:
+        if weight_args.group_size == 32:
+            return CompressionFormat.mxfp4_pack_quantized
         return CompressionFormat.nvfp4_pack_quantized
 
     if is_weight_only:  # w4a16 and w8a16

{compressed_tensors-0.12.3a20251114 → compressed_tensors-0.12.3a20251212}/src/compressed_tensors/quantization/quant_scheme.py

@@ -11,11 +11,11 @@
 # 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 warnings
 from copy import deepcopy
 from typing import List, Optional
 
+import torch
 from compressed_tensors.config import CompressionFormat
 from compressed_tensors.quantization.quant_args import (
     FP8_E4M3_DATA,
@@ -192,6 +192,43 @@ NVFP4 = dict(
     ),
 )
 
+MXFP4A16 = dict(
+    weights=QuantizationArgs(
+        num_bits=4,
+        type=QuantizationType.FLOAT,
+        strategy=QuantizationStrategy.GROUP,
+        symmetric=True,
+        dynamic=False,
+        group_size=32,
+        scale_dtype=torch.uint8,
+        zp_dtype=torch.uint8,
+    )
+)
+
+MXFP4 = dict(
+    weights=QuantizationArgs(
+        num_bits=4,
+        type=QuantizationType.FLOAT,
+        strategy=QuantizationStrategy.GROUP,
+        symmetric=True,
+        dynamic=False,
+        group_size=32,
+        scale_dtype=torch.uint8,
+        zp_dtype=torch.uint8,
+    ),
+    input_activations=QuantizationArgs(
+        num_bits=4,
+        type=QuantizationType.FLOAT,
+        strategy=QuantizationStrategy.GROUP,
+        dynamic=True,
+        symmetric=True,
+        group_size=32,
+        scale_dtype=torch.uint8,
+        zp_dtype=torch.uint8,
+    ),
+)
+
+
 # 8 bit integer weights and 8 bit activations quantization
 INT8_W8A8 = dict(
     weights=QuantizationArgs(
@@ -343,4 +380,6 @@ PRESET_SCHEMES = {
     "FP8_BLOCK": FP8_BLOCK,
     "NVFP4A16": NVFP4A16,
     "NVFP4": NVFP4,
+    "MXFP4A16": MXFP4A16,
+    "MXFP4": MXFP4,
 }

{compressed_tensors-0.12.3a20251114 → compressed_tensors-0.12.3a20251212}/src/compressed_tensors/quantization/utils/helpers.py

@@ -27,6 +27,11 @@ from compressed_tensors.quantization.quant_args import (
     round_to_quantized_type_dtype,
 )
 from compressed_tensors.quantization.quant_scheme import QuantizationScheme
+from compressed_tensors.quantization.utils.mxfp4_utils import (
+    generate_mxfp4_scales,
+    maybe_convert_from_mxfp4_exp,
+    should_generatre_mxfp4_scales,
+)
 from compressed_tensors.utils import deprecated
 from loguru import logger
 from torch import FloatTensor, IntTensor, Tensor
@@ -88,7 +93,10 @@ def calculate_qparams(
     # 1. Generate scale and zero-point
     if quantization_args.symmetric:
         max_val_pos = torch.max(torch.abs(min_vals), torch.abs(max_vals))
-        scales = max_val_pos / (float(bit_range) / 2)
+        if should_generatre_mxfp4_scales(args=quantization_args):
+            scales = generate_mxfp4_scales(x=max_val_pos)
+        else:
+            scales = max_val_pos / (float(bit_range) / 2)
         zero_points = torch.zeros(scales.shape, device=device, dtype=min_vals.dtype)
     else:
         if (
@@ -112,7 +120,10 @@ def calculate_qparams(
             scales, dtype=quantization_args.scale_dtype
         )
 
-    # 4. Update any 0s with small values to
+    # 4. Optionally remove exponent
+    scales = maybe_convert_from_mxfp4_exp(quantization_args, scales)
+
+    # 5. Update any 0s with small values to
     # prevent div by 0
     eps = _get_dtype_eps(
         dtype=quantization_args.scale_dtype
@@ -125,7 +136,7 @@ def calculate_qparams(
         scales,
     )
 
-    # 5. Round the zp to zp_dtype
+    # 6. Round the zp to zp_dtype
     zero_points = round_to_quantized_type_dtype(
         zero_points, dtype=quantization_args.zp_dtype, cast_to_original_dtype=False
     )

{compressed_tensors-0.12.3a20251114 → compressed_tensors-0.12.3a20251212}/src/compressed_tensors/quantization/utils/mxfp4_utils.py

@@ -13,16 +13,29 @@
 # limitations under the License.
 
 import torch
-from compressed_tensors.quantization.quant_args import BFLOAT16_DATA, FP4_E2M1_DATA
+from compressed_tensors.quantization.quant_args import (
+    BFLOAT16_DATA,
+    FP4_E2M1_DATA,
+    QuantizationArgs,
+)
 
 
-__all__ = ["convert_mxfp4_exp_scale", "generate_mxfp4_scales", "round_to_power_2"]
+__all__ = [
+    "maybe_convert_from_mxfp4_exp",
+    "generate_mxfp4_scales",
+    "round_to_power_2",
+    "should_generatre_mxfp4_scales",
+]
 
 # Reference: https://github.com/vllm-project/vllm/blob/main/tests/quantization/reference_mxfp4.py # noqa: E501
 
 
-def convert_mxfp4_exp_scale(
-    scale: torch.Tensor, dtype: torch.dtype = torch.bfloat16
+def should_generatre_mxfp4_scales(args: QuantizationArgs):
+    return args.num_bits == 4 and args.type == "float" and args.group_size == 32
+
+
+def maybe_convert_from_mxfp4_exp(
+    args: QuantizationArgs, scale: torch.Tensor
 ) -> torch.Tensor:
     """
     Converts mxfp4 scales. Scales are powers of 2, with the
@@ -32,10 +45,12 @@ def convert_mxfp4_exp_scale(
     :param scale: uint8 exponent scale
     :param dtype: dense dtype
     """
-    assert scale.dtype == torch.uint8
-    scale_exp = scale.to(torch.int32) - 127
-    scale = 2.00 ** (scale_exp.to(torch.float))
-    return scale.to(dtype)
+    original_dtype = scale.dtype
+    if should_generatre_mxfp4_scales(args):
+        scale_exp = scale.to(torch.int32) - 127
+        scale = 2.00 ** (scale_exp.to(torch.float))
+        return scale.to(original_dtype)
+    return scale
 
 
 def round_to_power_2(x: torch.Tensor) -> torch.Tensor:
@@ -77,21 +92,12 @@ def generate_mxfp4_scales(x: torch.Tensor) -> torch.Tensor:
     Generate mxfp4 scales. The scales require the following steps
     1. Round to the closest power of 2
     2. Convert to exponent
-    3. Store in uint8
 
     Called when calculating qparams using observers.
 
     :param x: tensor to round to closest power of 2
-    :returns uint8 scales as exponents
+    :returns scales as exponents
     """
     # Round to closest power of 2
     scale_power_2 = round_to_power_2(x)
-    # Convert to exponent
-    scale_exp = 127 + torch.floor(torch.log2(scale_power_2)).to(torch.int32) - 2
-    # Clamp and store in uint8, as expected by mxfp4
-    scale_exp = torch.clamp(
-        scale_exp,
-        max=torch.iinfo(torch.uint8).max,
-        min=torch.iinfo(torch.uint8).min,
-    )
-    return scale_exp.to(torch.uint8)
+    return 127 + torch.floor(torch.log2(scale_power_2)) - 2