compressed-tensors 0.10.2a20250616__tar.gz → 0.10.2a20250620__tar.gz

{compressed_tensors-0.10.2a20250616 → compressed_tensors-0.10.2a20250620}/.github/workflows/report.yml

@@ -120,7 +120,7 @@ jobs:
               shell: bash
 
             - name: report to reportportal
-              uses: neuralmagic/nm-actions/actions/reportportal_submit_execution_results@v1.15.0
+              uses: neuralmagic/nm-actions/actions/reportportal_submit_execution_results@v1.22.0
               with:
                 droute_username: ${{ secrets.DROUTE_USERNAME }}
                 droute_password: ${{ secrets.DROUTE_PASSWORD }}

{compressed_tensors-0.10.2a20250616/src/compressed_tensors.egg-info → compressed_tensors-0.10.2a20250620}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.10.2a20250616
+Version: 0.10.2a20250620
 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.10.2a20250616 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/utils/offload.py

@@ -85,6 +85,7 @@ __all__ = [
     "delete_offload_module",
     "offloaded_dispatch",
     "disable_offloading",
+    "remove_dispatch",
 ]
 
 
@@ -170,22 +171,22 @@ def update_parameter_data(
 
 def get_execution_device(module: torch.nn.Module) -> torch.device:
     """
-    Get the device which inputs should be moved to before module execution
+    Get the device which inputs should be moved to before module execution.
+    Assume that modules execute in the same order as returned by `model.modules()`
 
     :param module: module to check, may be offloaded
     :return: onload device of module
     """
-    if has_offloaded_params(module):
-        return module._hf_hook.execution_device
+    for submodule in module.modules():
+        if has_offloaded_params(submodule):
+            return submodule._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")
+        param = next(submodule.parameters(recurse=False), None)
+        if param is not None:
+            return param.device
 
-    return first_param.device
+    warnings.warn(f"Unable to get execution device of {module}, falling back to CPU")
+    return torch.device("cpu")
 
 
 def register_offload_parameter(
@@ -514,6 +515,9 @@ def offloaded_dispatch(
     if offload_device == "disk":
         raise NotImplementedError("Disk offloading is not currently supported")
 
+    # remove any existing hooks
+    remove_dispatch(module)
+
     # create weights map
     state_dict = module.state_dict()
     state_dict = {key: val.to(offload_device) for key, val in state_dict.items()}
@@ -535,6 +539,33 @@ def offloaded_dispatch(
         weights_map=weights_map,
         tied_params_map=tied_params_map,
     )
+
+    # when saving a model, `PretrainedModel.save_pretrained` will only
+    # onload weights if the following requirements are met
+    # if (
+    #     hasattr(self, "hf_device_map")
+    #     and len(set(self.hf_device_map.values())) > 1
+    #     and ("cpu" in self.hf_device_map.values()
+    #          or "disk" in self.hf_device_map.values())
+    # ):
+    # because this function always offloads, disregard actual devices and
+    # always use `cpu` and `cuda:0` to guarantee this condition passes
+    setattr(module, "hf_device_map", {"fake_offload": "cpu", "fake_exec": "cuda:0"})
+
+    return module
+
+
+def remove_dispatch(module: torch.nn.Module) -> torch.nn.Module:
+    """
+    Remove any existing dispatches from module
+
+    :param module: module which may be dispatched with hf hooks
+    :return: module without dispatch
+    """
+    remove_hook_from_module(module, recurse=True)
+    if hasattr(module, "hf_device_map"):
+        delattr(module, "hf_device_map")
+
     return module
 
 
@@ -563,7 +594,7 @@ def disable_offloading():
     # update any parameters which may have changed
     for module, (hook, offload) in onloaded_modules.items():
         hook.offload = offload
-        for name, param in module.named_parameters():
+        for name, param in module.named_parameters(recurse=False):
             update_offload_parameter(module, name, param.data)
         hook.post_forward(module, None)
 

{compressed_tensors-0.10.2a20250616 → compressed_tensors-0.10.2a20250620}/src/compressed_tensors/version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.10.2.a20250616'
+__version__ = version = '0.10.2.a20250620'
 __version_tuple__ = version_tuple = (0, 10, 2)

{compressed_tensors-0.10.2a20250616 → compressed_tensors-0.10.2a20250620/src/compressed_tensors.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.10.2a20250616
+Version: 0.10.2a20250620
 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.10.2a20250616 → compressed_tensors-0.10.2a20250620}/tests/test_utils/test_offload.py

@@ -102,6 +102,25 @@ def test_get_execution_device():
         assert get_execution_device(module) == torch.device("cuda:0")
 
 
+@requires_gpu
+@requires_accelerate()
+def test_get_execution_device_model():
+    class Model(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.a = torch.nn.Linear(1, 2)
+            self.b = torch.nn.Linear(2, 2, device="cuda:0")
+
+        def forward(self, x):
+            return self.b(self.a(x).to("cuda:0"))
+
+    model = Model()
+    assert get_execution_device(model) == torch.device("cpu")
+
+    offloaded_dispatch(model.a, torch.device("cuda:0"))
+    assert get_execution_device(model) == torch.device("cuda:0")
+
+
 @requires_accelerate()
 def test_register_offload_parameter():
     from accelerate import init_empty_weights