compressed-tensors 0.9.5a20250424__tar.gz → 0.9.5a20250428__tar.gz

{compressed_tensors-0.9.5a20250424 → compressed_tensors-0.9.5a20250428}/.github/workflows/build.yml

@@ -76,7 +76,7 @@ jobs:
 
             - name: build
               id: build
-              uses: neuralmagic/nm-actions/actions/build-ml-whl@fix-whl-checks
+              uses: neuralmagic/nm-actions/actions/build-ml-whl@c7e5a66c382104e1beadcb7dadf429f8ab15b344  # v1.20.0
               with:
                   dev: false
                   release: ${{ inputs.wf_category == 'RELEASE' }}

{compressed_tensors-0.9.5a20250424 → compressed_tensors-0.9.5a20250428}/.github/workflows/test-check.yaml

@@ -12,16 +12,16 @@ jobs:
   python-tests:
     runs-on: ubuntu-24.04
     steps:
-        - uses: actions/setup-python@v4
+        - uses: actions/setup-python@v5
           with:
             python-version: '3.10'
-        - uses: actions/checkout@v3
+        - uses: actions/checkout@v4
+          with:
+            fetch-depth: 0
+            fetch-tags: true
         - name: Set Env
           run: |
             pip3 install --upgrade pip && pip3 install --upgrade setuptools
-            pip3 install virtualenv
-            virtualenv venv
-            source venv/bin/activate
         - name: "⚙️ Install dependencies"
           run: pip3 install .[dev,accelerate]
         - name: "🔬 Running tests"

{compressed_tensors-0.9.5a20250424/src/compressed_tensors.egg-info → compressed_tensors-0.9.5a20250428}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.9.5a20250424
+Version: 0.9.5a20250428
 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.9.5a20250428/pyproject.toml

@@ -0,0 +1,7 @@
+[build-system]
+requires = ["setuptools", "wheel", "setuptools_scm==8.2.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.black]
+line-length = 88
+target-version = ['py36']

{compressed_tensors-0.9.5a20250424 → compressed_tensors-0.9.5a20250428}/src/compressed_tensors/quantization/lifecycle/initialize.py

@@ -31,7 +31,7 @@ from compressed_tensors.quantization.quant_scheme import QuantizationScheme
 from compressed_tensors.quantization.utils import is_kv_cache_quant_scheme
 from compressed_tensors.utils import (
     disable_hf_hook,
-    has_offloaded_params,
+    get_execution_device,
     register_offload_parameter,
 )
 from torch.nn import Module, Parameter
@@ -148,11 +148,8 @@ def _initialize_scale_zero_point(
     if quantization_args.dynamic:
         return
 
-    # begin on the same device as other parameters or cpu if offloaded.
-    # in the offloaded case, there's no point moving tensors to the execution device
-    # if they're going to be immediately offloaded by `register_offload_parameter`
-    params_device = next(module.parameters()).device
-    device = "cpu" if has_offloaded_params(module) else params_device
+    # initialize on execution device to avoid performing quantized ops on cpu
+    device = get_execution_device(module)
 
     # infer expected scale/zero point shape
     if quantization_args.strategy == QuantizationStrategy.TOKEN:

{compressed_tensors-0.9.5a20250424 → compressed_tensors-0.9.5a20250428}/src/compressed_tensors/utils/offload.py

@@ -28,7 +28,7 @@ Utilities associated with offloading functionality provided by `accelerate`.
 import contextlib
 import warnings
 from functools import wraps
-from typing import Any, Callable, Dict, Literal, Optional, Union
+from typing import Any, Callable, Dict, Iterable, Literal, Optional, Union
 
 import torch
 
@@ -67,6 +67,8 @@ __all__ = [
     "delete_offload_parameter",
     "has_offloaded_params",
     "disable_hf_hook",
+    "disable_offload",
+    "align_modules",
     "align_module_device",
 ]
 
@@ -344,6 +346,43 @@ def delete_from_weights_map(
         )
 
 
+@contextlib.contextmanager
+def disable_offload(module: torch.nn.Module):
+    """
+    Context manager to disable module onloading and offloading. Parameters will stay on
+    their current device
+
+    :param module: module to disable offloading for
+    """
+    if has_offloaded_params(module):
+        module._hf_hook.offload = False
+        yield
+        module._hf_hook.offload = True
+    else:
+        yield
+
+
+@contextlib.contextmanager
+def align_modules(
+    modules: Union[torch.nn.Module, Iterable[torch.nn.Module]],
+    execution_device: Optional[torch.device] = None,
+):
+    """
+    Context manager for onloading modules to a device, and disabling onload and offload
+    attempts triggered by forward calls. Used for sequential onloading of layers
+
+    :param modules: `torch.nn.Module` or iterable of `torch.nn.Module`s to onload
+    :param execution_device: device to onload to
+    """
+    modules = (modules,) if isinstance(modules, torch.nn.Module) else modules
+
+    with contextlib.ExitStack() as stack:
+        for module in modules:
+            stack.enter_context(align_module_device(module, execution_device))
+            stack.enter_context(disable_offload(module))  # disable redundant onloading
+        yield
+
+
 """ Upstreamed Functions """
 
 

{compressed_tensors-0.9.5a20250424 → compressed_tensors-0.9.5a20250428}/src/compressed_tensors/version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.9.5.a20250424'
+__version__ = version = '0.9.5.a20250428'
 __version_tuple__ = version_tuple = (0, 9, 5)

{compressed_tensors-0.9.5a20250424 → compressed_tensors-0.9.5a20250428/src/compressed_tensors.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.9.5a20250424
+Version: 0.9.5a20250428
 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.9.5a20250424 → compressed_tensors-0.9.5a20250428}/tests/test_utils/test_offload.py

@@ -15,6 +15,7 @@ import pytest
 import torch
 from compressed_tensors.utils import (
     align_module_device,
+    align_modules,
     delete_offload_parameter,
     disable_hf_hook,
     get_execution_device,
@@ -248,6 +249,35 @@ def test_disable_hf_hook_model_recurse():
     assert hasattr(module2, "_hf_hook")
 
 
+@requires_accelerate()
+def test_align_modules():
+    from accelerate.hooks import attach_align_device_hook
+
+    module0 = ExampleModule()
+    module1 = ExampleModule()
+    module2 = ExampleModule()
+    model = torch.nn.Sequential(module0, torch.nn.Sequential(module1, module2))
+    attach_align_device_hook(
+        model,
+        execution_device=torch.device("cpu"),
+        offload=True,
+        weights_map=model.state_dict(),
+    )
+
+    assert module0.a.device == torch.device("meta")
+    assert module1.a.device == torch.device("meta")
+    assert module2.a.device == torch.device("meta")
+
+    with align_modules((module0, module1)):
+        assert module0.a.device != torch.device("meta")
+        assert module1.a.device != torch.device("meta")
+        assert module2.a.device == torch.device("meta")
+
+    assert module0.a.device == torch.device("meta")
+    assert module1.a.device == torch.device("meta")
+    assert module2.a.device == torch.device("meta")
+
+
 @requires_accelerate()
 def test_offload_to_weights_map():
     from accelerate.utils import OffloadedWeightsLoader, PrefixedDataset

compressed_tensors-0.9.5a20250424/pyproject.toml

@@ -1,3 +0,0 @@
-[tool.black]
-line-length = 88
-target-version = ['py36']