compressed-tensors 0.11.1a20250929__tar.gz → 0.12.1__tar.gz

{compressed_tensors-0.11.1a20250929 → compressed_tensors-0.12.1}/.github/actions/test/action.yml

@@ -9,8 +9,7 @@ inputs:
     required: true
   code_coverage:
     description: whether to collect code coverage metrics during test run
-    type: boolean
-    default: false
+    default: 'false'
 outputs:
   status:
     description: "final status from test"
@@ -51,7 +50,7 @@ runs:
 
           if [[ "${ENABLE_COVERAGE}" == "true" ]]; then
             echo "::group::Installing code coverage requirements via pip"
-            pip install bashlex https://github.com/neuralmagic/pytest-nm-releng/archive/v0.4.0.tar.gz
+            pip install https://github.com/neuralmagic/pytest-nm-releng/archive/v0.4.0.tar.gz
             pip install coverage pytest-cov
 
             # Adding Code coverage to the tests
@@ -76,7 +75,7 @@ runs:
             fi
             echo "::endgroup::"
           fi
-          
+
           deactivate
           exit ${SUCCESS}
       shell: bash

compressed_tensors-0.12.1/.github/workflows/post-release-nightly-build.yml

@@ -0,0 +1,15 @@
+name: Post-release nightly build & release
+
+on:
+  release:
+    types: [released]
+
+jobs:
+  BUILD-TEST-NIGHTLY:
+    name: Post-release nightly build & release
+    uses: ./.github/workflows/trigger-all.yml
+    with:
+      wf_category: NIGHTLY
+      push_to_pypi: true
+      gitref: main
+    secrets: inherit

{compressed_tensors-0.11.1a20250929 → compressed_tensors-0.12.1}/.github/workflows/test-check.yaml

@@ -12,7 +12,10 @@ on:
 
 jobs:
   python-tests:
-    runs-on: ubuntu-24.04
+    runs-on: k8s-util
+    env:
+        HF_HOME: /model-cache
+        HF_TOKEN: ${{ secrets.NM_HF_TOKEN_READ_ONLY }}
     steps:
         - uses: actions/setup-python@v5
           with:
@@ -21,11 +24,13 @@ jobs:
           with:
             fetch-depth: 0
             fetch-tags: true
+        - name: Install system dependencies
+          run: |-
+            sudo apt-get update
+            sudo apt-get install -y --no-install-recommends g++ gcc make
         - name: Set Env
-          run: |
-            pip3 install --upgrade pip && pip3 install --upgrade setuptools
+          run: pip3 install --upgrade pip setuptools
         - name: "⚙️ Install dependencies"
           run: pip3 install .[dev,accelerate]
         - name: "🔬 Running tests"
           run: make test
-

{compressed_tensors-0.11.1a20250929 → compressed_tensors-0.12.1}/.github/workflows/test.yml

@@ -71,6 +71,9 @@ jobs:
             contents: 'read'
             id-token: 'write'
             pages: 'write'
+        env:
+            HF_HOME: /model-cache
+            HF_TOKEN: ${{ secrets.NM_HF_TOKEN_READ_ONLY }}
         environment:
             name: github-pages
             url: ${{ steps.coverage.outputs.page_url }}
@@ -90,13 +93,9 @@ jobs:
                   python-version: ${{ inputs.python }}
 
             - name: install system dependencies
-              run: |
-                if command -v g++ >/dev/null 2>&1; then
-                  echo "found g++ compiler"
-                else
-                   echo "installing g++ etc compilers..."
-                   sudo apt update && sudo apt install -y g++ gcc
-                fi
+              run: |-
+                  sudo apt-get update
+                  sudo apt-get install -y --no-install-recommends g++ gcc make
               shell: bash
 
             - name: checkout code
@@ -172,7 +171,7 @@ jobs:
               uses: actions/upload-artifact@v4
               if: success() || failure()
               with:
-                  name: report-${{ inputs.test_label }}.xml
+                  name: report-${{ inputs.python }}-${{ inputs.test_label }}.xml
                   path: test-results/report.xml
                   retention-days: 5
 

{compressed_tensors-0.11.1a20250929 → compressed_tensors-0.12.1}/.github/workflows/trigger-all.yml

@@ -5,6 +5,21 @@ on:
       # * is a special character in YAML so you have to quote this string
       - cron: '30 0 * * *'  # nightly run
 
+    workflow_call:
+        inputs:
+            wf_category:
+                description: "workflow category, must be 'NIGHTLY' or 'RELEASE' (default: NIGHTLY)"
+                type: string
+                default: NIGHTLY
+            push_to_pypi:
+                description: "when set and tests pass, then '.whl' & '.tar.gz' will be pushed to public pypi"
+                type: boolean
+                default: false
+            gitref:
+                description: "git commit hash or tag name"
+                type: string
+                default: 'main'
+
     workflow_dispatch:
         inputs:
             wf_category:
@@ -19,7 +34,7 @@ on:
                 type: boolean
                 default: false
             gitref:
-                description: "git commit hash or tag name" 
+                description: "git commit hash or tag name"
                 type: string
                 default: 'main'
 
@@ -32,9 +47,8 @@ jobs:
             wf_category: ${{ inputs.wf_category || 'NIGHTLY' }}
             gitref: ${{ inputs.gitref || 'main' }}
             push_to_pypi: ${{ (github.event.schedule == '30 0 * * *') || inputs.push_to_pypi || false }}
-            test_configs: '[{"python":"3.11.4","label":"ubuntu-24.04","timeout":"40","code_coverage":true},
-                            {"python":"3.10.12","label":"ubuntu-22.04","timeout":"40"},
+            test_configs: '[{"python":"3.11.4","label":"k8s-util","timeout":"40","code_coverage":true},
+                            {"python":"3.10.12","label":"k8s-util","timeout":"40"},
                             {"python":"3.9.17","label":"k8s-h100-solo","timeout":"40"},
                             {"python":"3.12.6","label":"k8s-a100-duo","timeout":"40"}]'
-
         secrets: inherit

{compressed_tensors-0.11.1a20250929 → compressed_tensors-0.12.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.11.1a20250929
+Version: 0.12.1
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/neuralmagic/compressed-tensors
 Author: Neuralmagic, Inc.

{compressed_tensors-0.11.1a20250929 → compressed_tensors-0.12.1}/src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -344,6 +344,55 @@ class ModelCompressor:
                     format, config=quantization_config
                 )
 
+    def get_missing_module_keys(self, model: Module) -> List[str]:
+        """
+        Identifies the expected missing weight keys in the compressed state_dict.
+
+        When a model undergoes sparsity or quantization compression, certain
+        weight tensors may be absent from the checkpoint by virtue of compression.
+        This function determines which weight keys are missing based on the
+        applied compression techniques.
+
+        :param model: The PyTorch model to check for missing keys.
+        :return: A list of missing keys expected in the compressed state_dict.
+        """
+        missing_keys = set()
+
+        # Determine missing keys due to sparsity compression
+        if (
+            self.sparsity_compressor
+            and self.sparsity_config.format != CompressionFormat.dense.value
+        ):
+            sparse_targets = match_named_modules(
+                model=model,
+                targets=self.sparsity_config.targets,
+                ignore=self.sparsity_config.ignore,
+            )
+
+            missing_keys.update(
+                merge_names(target_name, "weight")
+                for target_name, _module in sparse_targets
+            )
+
+        # Determine missing keys due to pack quantization
+        if (
+            self.quantization_compressor
+            and self.quantization_config.format
+            == CompressionFormat.pack_quantized.value
+        ):
+            for scheme in self.quantization_config.config_groups.values():
+                quant_targets = match_named_modules(
+                    model=model,
+                    targets=scheme.targets,
+                    ignore=self.quantization_config.ignore,
+                )
+                missing_keys.update(
+                    merge_names(target_name, "weight")
+                    for target_name, _module in quant_targets
+                )
+
+        return list(missing_keys)
+
     def get_unexpected_file_keys(self, model: Module) -> List[str]:
         """
         Identifies extra keys introduced by the compression process in the

{compressed_tensors-0.11.1a20250929 → compressed_tensors-0.12.1}/src/compressed_tensors/transform/factory/hadamard.py

@@ -52,7 +52,7 @@ class HadamardFactory(TransformFactory):
         :param args: defines how the transform will be applied to the module
         """
         assert hasattr(module, "weight")
-        size = get_transform_size(module, args.location, self.scheme.block_size)
+        size = get_transform_size(module, args.location, self.scheme.head_dim)
         exec_device = get_execution_device(module)
         device = get_offloaded_device(module)
         precision = self.scheme.precision if args.is_online() else torch.float64

{compressed_tensors-0.11.1a20250929 → compressed_tensors-0.12.1}/src/compressed_tensors/transform/factory/matrix_multiply.py

@@ -51,7 +51,7 @@ class RandomMatrixFactory(TransformFactory):
         :param args: defines how the transform will be applied to the module
         """
         assert hasattr(module, "weight")
-        size = get_transform_size(module, args.location, self.scheme.block_size)
+        size = get_transform_size(module, args.location, self.scheme.head_dim)
         device = get_offloaded_device(module)
         precision = self.scheme.precision if args.is_online() else torch.float64
 

{compressed_tensors-0.11.1a20250929 → compressed_tensors-0.12.1}/src/compressed_tensors/transform/transform_scheme.py

@@ -17,7 +17,7 @@ from typing import List, Optional
 import torch
 from compressed_tensors.transform import TransformArgs
 from compressed_tensors.utils import TorchDtype
-from pydantic import BaseModel, ConfigDict, Field, model_validator
+from pydantic import BaseModel, ConfigDict, Field
 
 
 __all__ = ["TransformScheme"]
@@ -36,8 +36,11 @@ class TransformScheme(BaseModel):
     :param randomize: True if uniquely randomized transform weights should be used,
         otherwise use identical transform weights where applicable
     :param requires_grad: True if weights include gradients for training
-    :param block_size: If set, the transform matrix will be block diagonal, with each
-        block being a square matrix of this size.
+    :param head_dim: If set, the transform matrix will be block diagonal with each
+        block being a square matrix of this size. The name head_dim was chosen because
+        some rotations need to be block-diagonal with block size equal to the head_dim,
+        but research has shown value in applying some rotations with smaller block size,
+        irrespective of head_dim.
     :param precision: Precision at which this transform should be applied during online
         rotations. Fused (offline) rotations are always performed in float64
     """
@@ -46,21 +49,7 @@ class TransformScheme(BaseModel):
     apply: List[TransformArgs] = Field(default_factory=list)
     randomize: bool = Field(default=False)
     requires_grad: bool = Field(default=False)
-    block_size: Optional[int] = Field(default=None)
-    head_dim: Optional[int] = Field(
-        default=None, deprecated="head_dim is deprecated, use block_size instead"
-    )
+    head_dim: Optional[int] = Field(default=None)
     precision: TorchDtype = Field(default=torch.float32)
 
-    @model_validator(mode="after")
-    def validate_model_after(model: "TransformScheme") -> "TransformScheme":
-        """
-        If head_dim is used instead of block_size, set block_size to head_dim
-        and remove head_dim
-        """
-        if model.block_size is None and model.head_dim is not None:
-            model.block_size = model.head_dim
-            model.head_dim = None
-        return model
-
     model_config = ConfigDict(extra="forbid")

{compressed_tensors-0.11.1a20250929 → compressed_tensors-0.12.1}/src/compressed_tensors/version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.11.1.a20250929'
-__version_tuple__ = version_tuple = (0, 11, 1)
+__version__ = version = '0.12.1'
+__version_tuple__ = version_tuple = (0, 12, 1)

{compressed_tensors-0.11.1a20250929 → compressed_tensors-0.12.1}/src/compressed_tensors.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.11.1a20250929
+Version: 0.12.1
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/neuralmagic/compressed-tensors
 Author: Neuralmagic, Inc.

{compressed_tensors-0.11.1a20250929 → compressed_tensors-0.12.1}/src/compressed_tensors.egg-info/SOURCES.txt

@@ -10,6 +10,7 @@ setup.py
 .github/scripts/step-status
 .github/workflows/build-test.yml
 .github/workflows/build.yml
+.github/workflows/post-release-nightly-build.yml
 .github/workflows/quality-check.yaml
 .github/workflows/report.yml
 .github/workflows/test-check.yaml

{compressed_tensors-0.11.1a20250929 → compressed_tensors-0.12.1}/tests/test_compressors/model_compressors/test_model_compressor.py

@@ -450,6 +450,26 @@ def test_multiple_quant_compressors():
     assert all(format in compressor.quantization_compressor for format in formats)
 
 
+@pytest.mark.parametrize(
+    "model, sparsity_config, quantization_config, expected",
+    [
+        (
+            TwoLayerModel(),
+            get_bitmask_sparsity_config(targets=["re:.*layer1$"]),
+            create_quantization_config(bits=8, type="int", strategy="channel"),
+            {"layer1.weight"},
+        )
+    ],
+)
+def test_get_missing_keys(model, sparsity_config, quantization_config, expected):
+    model_compressor = ModelCompressor(
+        sparsity_config=sparsity_config, quantization_config=quantization_config
+    )
+
+    actual = model_compressor.get_missing_module_keys(model)
+    assert len(actual) == len(expected) and all(key in actual for key in expected)
+
+
 @pytest.mark.parametrize(
     "model, sparsity_config, quantization_config, expected",
     [

{compressed_tensors-0.11.1a20250929 → compressed_tensors-0.12.1}/tests/test_transform/factory/test_correctness.py

@@ -33,7 +33,7 @@ from tests.testing_utils import requires_accelerate, requires_gpu
 def test_correctness_linear(type, randomize, head_dim, input_batch_size):
     size = (4, 8)
     module = torch.nn.Linear(*size, bias=False)
-    scheme = TransformScheme(type=type, randomize=randomize, block_size=head_dim)
+    scheme = TransformScheme(type=type, randomize=randomize, head_dim=head_dim)
     factory = TransformFactory.from_scheme(scheme, name="")
 
     input_tfm = factory.create_transform(
@@ -150,7 +150,7 @@ def test_correctness_attention_heads(type, randomize, head_dim, input_batch_size
             "": TransformScheme(
                 type=type,
                 randomize=randomize,
-                block_size=head_dim,
+                head_dim=head_dim,
                 apply=[
                     TransformArgs(targets="v_proj", location="weight_output"),
                     TransformArgs(

{compressed_tensors-0.11.1a20250929 → compressed_tensors-0.12.1}/tests/test_transform/test_transform_scheme.py

@@ -72,28 +72,3 @@ def test_multiple_groups():
     assert not scheme.randomize
     assert scheme.type == "hadamard"
     assert len(scheme.apply) == 20
-
-
-def test_transform_scheme_block_size():
-    """
-    Ensure json with (deprecated) `head_dim` or `block_size`
-    both load up correctly and save with `block_size` field
-    """
-
-    old_scheme = TransformScheme.model_validate_json(
-        '{"type": "hadamard", "head_dim": 128}'
-    )
-    assert old_scheme.block_size == 128
-    assert old_scheme.model_dump()["block_size"] == 128
-    old_scheme = TransformScheme(type="hadamard", head_dim=64)
-    assert old_scheme.block_size == 64
-    assert old_scheme.model_dump()["block_size"] == 64
-
-    new_scheme = TransformScheme.model_validate_json(
-        '{"type": "hadamard", "block_size": 128}'
-    )
-    assert new_scheme.block_size == 128
-    assert new_scheme.model_dump()["block_size"] == 128
-    new_scheme = TransformScheme(type="hadamard", block_size=64)
-    assert new_scheme.block_size == 64
-    assert new_scheme.model_dump()["block_size"] == 64

{compressed_tensors-0.11.1a20250929 → compressed_tensors-0.12.1}/tests/test_utils/test_helpers.py

@@ -36,12 +36,11 @@ def tensors():
 
 
 @pytest.fixture
-def llama_model(tmp_path):
-    model_name = "neuralmagic/llama2.c-stories110M-pruned50"
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name, torch_dtype="auto", cache_dir=tmp_path
+def llama_model():
+    return AutoModelForCausalLM.from_pretrained(
+        "RedHatAI/llama2.c-stories110M-pruned50",
+        torch_dtype="auto",
     )
-    yield model
 
 
 def test_save_compressed_sparse_bitmask(tmp_path, tensors):
@@ -120,9 +119,9 @@ def test_load_compressed_dense(tmp_path, tensors):
 
 
 def test_load_compressed_sharded(tmp_path, llama_model):
-    sharded_model_path = tmp_path / "shared_model"
+    sharded_model_path = tmp_path / "sharded_model"
     llama_model.save_pretrained(sharded_model_path, max_shard_size="2MB")
-    # make sure that model is shared on disk
+    # make sure that model is sharded on disk
     assert len(os.listdir(sharded_model_path)) > 1
     loaded_state_dict = dict(load_compressed(sharded_model_path))
     for key, value in llama_model.state_dict().items():