Move files in compressed_tensors.compressor.utils compressed_tensors.utils to avoid name mangling (#102)

Author: rahul-tuli

src/compressed_tensors/compressors/marlin_24.py

@@ -18,15 +18,16 @@
 import numpy as np
 import torch
 from compressed_tensors.compressors import Compressor
-from compressed_tensors.compressors.utils import (
+from compressed_tensors.config import CompressionFormat
+from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy
+from compressed_tensors.quantization.lifecycle.forward import quantize
+from compressed_tensors.utils import (
     get_permutations_24,
+    is_quantization_param,
+    merge_names,
     sparse_semi_structured_from_dense_cutlass,
     tensor_follows_mask_structure,
 )
-from compressed_tensors.config import CompressionFormat
-from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy
-from compressed_tensors.quantization.lifecycle.forward import quantize
-from compressed_tensors.utils import is_quantization_param, merge_names
 from torch import Tensor
 from tqdm import tqdm
 

src/compressed_tensors/compressors/utils/__init__.py

@@ -13,4 +13,7 @@
 # limitations under the License.
 # flake8: noqa
 
+from .helpers import *
+from .permutations_24 import *
 from .safetensors_load import *
+from .semi_structured_conversions import *

src/compressed_tensors/compressors/utils/helpers.py

@@ -14,10 +14,15 @@
 
 from typing import Optional
 
+import torch
 from transformers import AutoConfig
 
 
-__all__ = ["infer_compressor_from_model_config", "fix_fsdp_module_name"]
+__all__ = [
+    "infer_compressor_from_model_config",
+    "fix_fsdp_module_name",
+    "tensor_follows_mask_structure",
+]
 
 FSDP_WRAPPER_NAME = "_fsdp_wrapped_module"
 
@@ -60,3 +65,28 @@ def fix_fsdp_module_name(name: str) -> str:
     return name.replace(FSDP_WRAPPER_NAME + ".", "").replace(
         "." + FSDP_WRAPPER_NAME, ""
     )
+
+
+def tensor_follows_mask_structure(tensor, mask: str = "2:4") -> bool:
+    """
+    :param tensor: tensor to check
+    :param mask: mask structure to check for, in the format "n:m"
+    :return: True if the tensor follows the mask structure, False otherwise.
+        Note, some weights can incidentally be zero, so we check for
+        atleast n zeros in each chunk of size m
+    """
+
+    n, m = tuple(map(int, mask.split(":")))
+    # Reshape the tensor into chunks of size m
+    tensor = tensor.view(-1, m)
+
+    # Count the number of zeros in each chunk
+    zero_counts = (tensor == 0).sum(dim=1)
+
+    # Check if the number of zeros in each chunk atleast n
+    # Greater than sign is needed as some weights can incidentally
+    # be zero
+    if not torch.all(zero_counts >= n).item():
+        raise ValueError()
+
+    return True

src/compressed_tensors/utils/__init__.py

@@ -21,7 +21,6 @@
     Marlin24Compressor,
     map_modules_to_quant_args,
 )
-from compressed_tensors.compressors.utils import mask_creator
 from compressed_tensors.config import CompressionFormat
 from compressed_tensors.quantization import (
     QuantizationArgs,
@@ -32,7 +31,7 @@
     apply_quantization_config,
     apply_quantization_status,
 )
-from compressed_tensors.utils import merge_names
+from compressed_tensors.utils import mask_creator, merge_names
 from torch.nn.modules import Linear, Sequential