compressed-tensors 0.9.4a20250421__tar.gz → 0.9.5a20250424__tar.gz

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

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

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

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

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

{compressed_tensors-0.9.4a20250421 → compressed_tensors-0.9.5a20250424}/setup.py

@@ -101,6 +101,7 @@ setup(
     use_scm_version={
         "version_scheme": version_func,
         "local_scheme": localversion_func,
+        "version_file": "src/compressed_tensors/version.py",
     },
     author="Neuralmagic, Inc.",
     author_email="support@neuralmagic.com",

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

@@ -94,22 +94,6 @@ def is_module_offloaded(module: torch.nn.Module) -> bool:
     return has_offloaded_params(module)
 
 
-def get_execution_device(module: torch.nn.Module) -> torch.device:
-    """
-    :param module: module to check
-    :return: device module is loaded onto during forward pass
-    """
-    if has_offloaded_params(module):
-        return module._hf_hook.execution_device
-    device = next(module.parameters()).device
-
-    # offload only gets set for leaf modules, fallback to checking for device type
-    if device.type == "meta":
-        return module._hf_hook.execution_device
-
-    return device
-
-
 def get_offloaded_device(module: torch.nn.Module) -> torch.device:
     """
     :param module: module to check
@@ -158,6 +142,26 @@ def update_parameter_data(
 """ Candidates for Upstreaming """
 
 
+def get_execution_device(module: torch.nn.Module) -> torch.device:
+    """
+    Get the device which inputs should be moved to before module execution
+
+    :param module: module to check, may be offloaded
+    :return: onload device of module
+    """
+    if has_offloaded_params(module):
+        return module._hf_hook.execution_device
+
+    first_param = next(module.parameters(), None)
+    if first_param is None:
+        warnings.warn(
+            f"Unable able to infer execution device of {module}, falling back to CPU"
+        )
+        return torch.device("cpu")
+
+    return first_param.device
+
+
 def register_offload_parameter(
     module: torch.nn.Module,
     name: str,

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

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

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

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

@@ -17,12 +17,13 @@ from compressed_tensors.utils import (
     align_module_device,
     delete_offload_parameter,
     disable_hf_hook,
+    get_execution_device,
     has_offloaded_params,
     register_offload_parameter,
     update_offload_parameter,
 )
 from compressed_tensors.utils.offload import offload_to_weights_map
-from tests.testing_utils import requires_accelerate
+from tests.testing_utils import requires_accelerate, requires_gpu
 
 
 class ExampleModule(torch.nn.Module):
@@ -55,8 +56,46 @@ def test_has_offloaded_params():
     assert has_offloaded_params(module)
 
 
+@requires_gpu
+@requires_accelerate()
+def test_get_execution_device():
+    from accelerate import init_empty_weights
+    from accelerate.big_modeling import attach_align_device_hook
+
+    # no offloading
+    module = ExampleModule()
+    assert get_execution_device(module) == torch.device("cpu")
+
+    # with offloading
+    attach_align_device_hook(module, torch.device("cuda:0"))
+    assert get_execution_device(module) == torch.device("cuda:0")
+
+    # in meta context
+    with torch.device("meta"):
+        module = ExampleModule()
+        assert get_execution_device(module) == torch.device("meta")
+
+    # offloaded in meta context
+    module = ExampleModule()
+    attach_align_device_hook(module, torch.device("cuda:0"))
+    with torch.device("meta"):
+        assert get_execution_device(module) == torch.device("cuda:0")
+
+    # in empty weights context
+    with init_empty_weights():
+        module = ExampleModule()
+        assert get_execution_device(module) == torch.device("meta")
+
+    # offloaded in empty weights context
+    module = ExampleModule()
+    attach_align_device_hook(module, torch.device("cuda:0"))
+    with init_empty_weights():
+        assert get_execution_device(module) == torch.device("cuda:0")
+
+
 @requires_accelerate()
 def test_register_offload_parameter():
+    from accelerate import init_empty_weights
     from accelerate.hooks import attach_align_device_hook
 
     module = ExampleModule()
@@ -94,6 +133,12 @@ def test_register_offload_parameter():
         assert module.f.device == torch.device("cpu")
     assert module._hf_hook.weights_map["f"].device == torch.device("cpu")
 
+    # parameters registered in the empty init context are still empty
+    with init_empty_weights():
+        module = ExampleModule()
+        register_offload_parameter(module, "c", parameter)
+    assert module.a.device == module.b.device == module.c.device == torch.device("meta")
+
 
 @requires_accelerate()
 def test_update_offload_parameter():

compressed_tensors-0.9.4a20250421/pyproject.toml

@@ -1,10 +0,0 @@
-[build-system]
-requires = ["setuptools", "wheel", "setuptools_scm>8"]
-build-backend = "setuptools.build_meta"
-
-[tool.setuptools_scm]
-version_file = "src/compressed_tensors/version.py"
-
-[tool.black]
-line-length = 88
-target-version = ['py36']