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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.10.2a20250612
+Version: 0.10.2a20250616
 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.2a20250612 → compressed_tensors-0.10.2a20250616}/setup.py

@@ -113,5 +113,6 @@ setup(
     extras_require=_setup_extras(),
     install_requires=_setup_install_requires(),
     package_dir={"": "src"},
+    package_data={"": ["transform/utils/hadamards.safetensors"]},
     packages=_setup_packages(),
 )

{compressed_tensors-0.10.2a20250612 → compressed_tensors-0.10.2a20250616}/src/compressed_tensors/transform/factory/hadamard.py

@@ -59,7 +59,7 @@ class HadamardFactory(TransformFactory):
         return HadamardTransform(weight, args)
 
     def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
-        data = deterministic_hadamard_matrix(size)
+        data = deterministic_hadamard_matrix(size, dtype, device)
         data = data.to(dtype=dtype, device=device)
         return Parameter(data, requires_grad=self.scheme.requires_grad)
 

{compressed_tensors-0.10.2a20250612 → compressed_tensors-0.10.2a20250616}/src/compressed_tensors/transform/factory/random_hadamard.py

@@ -29,6 +29,6 @@ class RandomHadamardFactory(HadamardFactory):
     """
 
     def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
-        data = random_hadamard_matrix(size, self.generator)
+        data = random_hadamard_matrix(size, dtype, device, self.generator)
         data = data.to(dtype=dtype, device=device)
         return Parameter(data, requires_grad=self.scheme.requires_grad)

compressed_tensors-0.10.2a20250616/src/compressed_tensors/transform/utils/hadamard.py

@@ -0,0 +1,160 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from pathlib import Path
+from typing import Optional
+
+import torch
+from safetensors import safe_open
+
+
+REPO_PATH = Path(__file__).parent / "hadamards.safetensors"
+
+
+__all__ = ["random_hadamard_matrix", "deterministic_hadamard_matrix", "is_pow2"]
+
+
+# note that hadamard matrix multiplication can be accelerated using a library such as
+# https://github.com/Dao-AILab/fast-hadamard-transform/tree/master
+
+
+def deterministic_hadamard_matrix(
+    size: int,
+    dtype: torch.dtype = torch.bfloat16,
+    device: torch.device = torch.device("cpu"),
+) -> torch.Tensor:
+    """
+    Construct an n-by-n Hadamard matrix, using Sylvester's construction.
+    `n` must be a power of 2.
+
+    Adapated from https://github.com/scipy/scipy/blob/v1.15.2/scipy/linalg/_special_matrices.py  # noqa: E501
+
+    :param size: order of the matrix, must be a power of 2
+    :param dtype: data type of matrix
+    :param device: device to construct matrix on
+    :return: hadamard matrix of size `size`
+    """
+    if size <= 0:
+        raise ValueError("Cannot construct deterministic hadamard of size <= 0")
+
+    log2 = int(math.log2(size))
+    if size != 2**log2:
+        raise ValueError("Cannot construct deterministic hadamard of size != 2^n")
+
+    H = torch.tensor([[1]], dtype=dtype, device=device)
+
+    # Sylvester's construction
+    for _ in range(log2):
+        H = torch.vstack((torch.hstack((H, H)), torch.hstack((H, -H))))
+
+    return H / math.sqrt(size)
+
+
+def random_hadamard_matrix(
+    size: int,
+    dtype: torch.dtype = torch.bfloat16,
+    device: torch.device = torch.device("cpu"),
+    gen: Optional[torch.Generator] = None,
+) -> torch.Tensor:
+    """
+    Produces a randomly generated Hadamard matrix. Differs from
+    `deterministic_hadamard_matrix` in that this function supports non powers of 2
+    and randomization using a seeded generator
+
+    Adapated from https://github.com/facebookresearch/SpinQuant/blob/main/utils/hadamard_utils.py  # noqa: E501
+    Known matrices were retrieved from N. J. A. Sloane's Library of Hadamard Matrices http://www.neilsloane.com/hadamard/  # noqa: E501
+
+    :param size: The dimension of the hamadard matrix
+    :param dtype: data type of matrix
+    :param device: device to construct matrix on
+    :param gen: Optional generator random values
+    :return: randomly generated hadamard matrix
+    """
+    Q = torch.randint(low=0, high=2, size=(size,), generator=gen, dtype=dtype)  # cpu
+    Q = Q.to(device=device)
+    Q = Q * 2 - 1
+    Q = torch.diag(Q)
+    return _matmul_hadU(Q) / math.sqrt(size)
+
+
+def is_pow2(n: int) -> bool:
+    """
+    Check if a number is a power of 2
+
+    :param n: number to check
+    :return: True iff `n` is a power of 2
+    """
+    return n > 0 and (n & (n - 1) == 0)
+
+
+def _fetch_hadamard_divisor(
+    n: int,
+    dtype: torch.dtype,
+    device: torch.device = torch.device("cpu"),
+    file_path: str = REPO_PATH,
+) -> Optional[torch.Tensor]:
+    """
+    Fetch a known hadamard matrix from the given file path. The returned matrix will
+    be of of size `k` such that `n / k` is a power of two. Return None if no such
+    matrix exists.
+
+    Note: This function reopens the safetensors file every time it is called.
+    This is technically inefficient, but a very small runtime cost and simpler
+    than forcing callers to manage the file open context
+
+    :param n: size of known hadamard matrix
+    :return: a known hadamard matrix of size `n` if one exists, else None
+    """
+    with safe_open(file_path, framework="pt", device=str(device)) as file:
+        divisors = sorted((int(key) for key in file.keys()), reverse=True)
+        for divisor in divisors:
+            if n % divisor == 0 and is_pow2(n // divisor):
+                return file.get_tensor(str(divisor)).to(dtype=dtype)
+
+    return None
+
+
+def _matmul_hadU(X: torch.Tensor) -> torch.Tensor:
+    size = X.size(0)
+    dtype = X.dtype
+    device = X.device
+
+    # Check if we have the determined hadamard matrix
+    hadK = _fetch_hadamard_divisor(size, dtype, device=device)
+    if hadK is None:
+        raise ValueError(f"Cannot construct random hadamard matrix of size {size}")
+    K = hadK.size(0)
+
+    # Reshape diag matrix with randomized -1/+1
+    input = X.clone().view(-1, size, 1)
+    output = input.clone()
+    while input.shape[1] > K:
+        input = input.view(input.shape[0], input.shape[1] // 2, 2, input.shape[2])
+        output = output.view(input.shape)
+        output[:, :, 0, :] = input[:, :, 0, :] + input[:, :, 1, :]
+        output[:, :, 1, :] = input[:, :, 0, :] - input[:, :, 1, :]
+        output = output.view(input.shape[0], input.shape[1], -1)
+        (input, output) = (output, input)
+    assert input.shape[1] == K
+    del output
+
+    # Do not explicitly repeat - OOM
+    # input = torch.bmm(
+    #     hadK.repeat(len(input), 1, 1).to(input.device).to(input.dtype), input)
+    # Use bcast instead
+    input = hadK.view(1, K, K).to(input) @ input
+
+    # normalize
+    return input.view(X.shape)

{compressed_tensors-0.10.2a20250612 → compressed_tensors-0.10.2a20250616}/src/compressed_tensors/utils/offload.py

@@ -31,9 +31,10 @@ import contextlib
 import warnings
 from functools import wraps
 from operator import attrgetter
-from typing import Any, Callable, Dict, Iterable, Literal, Optional, Union
+from typing import Any, Callable, Dict, Iterable, Literal, Optional, Tuple, Union
 
 import torch
+from compressed_tensors.utils import patch_attr
 
 
 try:
@@ -83,6 +84,7 @@ __all__ = [
     "register_offload_module",
     "delete_offload_module",
     "offloaded_dispatch",
+    "disable_offloading",
 ]
 
 
@@ -204,9 +206,24 @@ def register_offload_parameter(
     has_onload = any(p.device != torch.device("meta") for p in module.parameters())
     module.register_parameter(name, parameter)
 
+    # do everything AlignDevicesHook.init_hook does
+    # https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L281
     if has_offloaded_params(module):
-        weights_map = module._hf_hook.weights_map
-        offload_to_weights_map(weights_map, name, parameter.data, offload_device)
+        hook: AlignDevicesHook = module._hf_hook
+        assert hook.weights_map is not None
+
+        # append to original_devices
+        hook.original_devices[name] = parameter.device
+
+        # append to weights map
+        offload_to_weights_map(hook.weights_map, name, parameter.data, offload_device)
+
+        # append to tied_params_map
+        offloaded = hook.weights_map[name]
+        if hook.tied_params_map is not None:
+            hook.tied_params_map[offloaded.data_ptr()] = {}  # (1)
+
+        # perform offloading
         if not has_onload:
             set_module_tensor_to_device(module, name, "meta")
 
@@ -214,7 +231,7 @@ def register_offload_parameter(
 def update_offload_parameter(
     module: torch.nn.Module,
     name: str,
-    data: Optional[torch.Tensor],
+    data: torch.Tensor,
     offload_device: Optional[Union[torch.device, Literal["disk"]]] = None,
 ):
     """
@@ -227,7 +244,7 @@ def update_offload_parameter(
     :param offload_device: device on which weight will be offloaded to. If None is
         provided, then infer device from parameters on module
     """
-    param = getattr(module, name)
+    param: torch.nn.Parameter = getattr(module, name)
     if param.data.shape != data.shape:
         warnings.warn(
             f"Shape of parameter being updated {param.data.shape} does not match shape "
@@ -235,7 +252,7 @@ def update_offload_parameter(
         )
 
     # copy data into onloaded parameter if applicable
-    if param.device != torch.device("meta"):
+    if param.device != torch.device("meta") and data is not param.data:
         param.data.copy_(data)
 
     # update offload dict
@@ -420,7 +437,6 @@ def register_offload_module(base: torch.nn.Module, name: str, module: torch.nn.M
         hook: AlignDevicesHook = base._hf_hook
         assert hook.offload
         assert hook.weights_map is not None
-        assert hook.tied_params_map is not None
 
         # offloading kwargs for submodule
         place_submodules = False
@@ -435,7 +451,8 @@ def register_offload_module(base: torch.nn.Module, name: str, module: torch.nn.M
             module, include_buffers=offload_buffers, recurse=place_submodules
         ):
             offloaded = param.to(offload_device)
-            hook.tied_params_map[offloaded.data_ptr()] = {}  # (1)
+            if hook.tied_params_map is not None:
+                hook.tied_params_map[offloaded.data_ptr()] = {}  # (1)
             offload_to_weights_map(hook.weights_map, f"{name}.{param_name}", offloaded)
 
             # if the parent places submodules, offload here
@@ -463,9 +480,6 @@ def register_offload_module(base: torch.nn.Module, name: str, module: torch.nn.M
 
     base.register_module(name, module)
 
-    # (1): Since we cannot know which pointers are shared when we add parameters in an
-    # online way, assume that all pointers are shared. This comes at no runtime cost
-
 
 def delete_offload_module(base: torch.nn.Module, name: str):
     """
@@ -501,7 +515,9 @@ def offloaded_dispatch(
         raise NotImplementedError("Disk offloading is not currently supported")
 
     # create weights map
-    weights_map = OffloadedWeightsLoader(state_dict=module.state_dict(), device="cpu")
+    state_dict = module.state_dict()
+    state_dict = {key: val.to(offload_device) for key, val in state_dict.items()}
+    weights_map = OffloadedWeightsLoader(state_dict=state_dict, device=offload_device)
 
     # create tied params map
     tied_params = find_tied_parameters(module)
@@ -522,6 +538,36 @@ def offloaded_dispatch(
     return module
 
 
+@contextlib.contextmanager
+def disable_offloading():
+    """
+    Keep modules onloaded and disable offloading until this context exits.
+    Affects modules which have been hooked with accelerate's `AlignDevicesHook`
+    """
+    original_pre_forward = AlignDevicesHook.pre_forward
+    onloaded_modules: Dict[torch.nn.Module, Tuple[AlignDevicesHook, bool]] = dict()
+
+    # onload once and disable any future onloading/offloading steps
+    def keep_onload_pre_forward(self: AlignDevicesHook, module, *args, **kwargs):
+        ret = original_pre_forward(self, module, *args, **kwargs)
+        if module not in onloaded_modules:
+            onloaded_modules[module] = (self, self.offload)
+            self.offload = False
+        return ret
+
+    # use the patched pre_forward function within the context
+    with patch_attr(AlignDevicesHook, "pre_forward", keep_onload_pre_forward):
+        yield
+
+    # manually offload all modules that were onloaded
+    # 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():
+            update_offload_parameter(module, name, param.data)
+        hook.post_forward(module, None)
+
+
 """ Upstreamed Functions """
 
 
@@ -589,3 +635,7 @@ def align_module_device(
 
     else:
         yield
+
+
+# (1): Since we cannot know which pointers are shared when we add parameters in an
+# online way, assume that all pointers are shared. This has virtually no runtime cost

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

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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.10.2a20250612
+Version: 0.10.2a20250616
 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.2a20250612 → compressed_tensors-0.10.2a20250616}/src/compressed_tensors.egg-info/SOURCES.txt

@@ -82,6 +82,7 @@ src/compressed_tensors/transform/factory/matrix_multiply.py
 src/compressed_tensors/transform/factory/random_hadamard.py
 src/compressed_tensors/transform/utils/__init__.py
 src/compressed_tensors/transform/utils/hadamard.py
+src/compressed_tensors/transform/utils/hadamards.safetensors
 src/compressed_tensors/transform/utils/utils.py
 src/compressed_tensors/utils/__init__.py
 src/compressed_tensors/utils/helpers.py

{compressed_tensors-0.10.2a20250612 → compressed_tensors-0.10.2a20250616}/tests/test_transform/utils/test_hadamard.py

@@ -13,46 +13,48 @@
 # limitations under the License.
 
 
-import numpy
 import pytest
 import torch
 from compressed_tensors.transform.utils.hadamard import (
-    _get_had12,
-    _get_had20,
     deterministic_hadamard_matrix,
+    is_pow2,
     random_hadamard_matrix,
 )
+from tests.testing_utils import requires_gpu
 
 
-@pytest.mark.parametrize(
-    "had_func",
-    [
-        _get_had12,
-        _get_had20,
-    ],
-)
-def test_packed_hadamard_compliant(had_func):
-    had_matrix = had_func()
-    size = had_matrix.size(0)
-    # HH.T == nI
-    product = had_matrix @ had_matrix.T
-    assert torch.equal(product, size * torch.eye(size))
+_sizes_to_test = [
+    768,  # gpt2 small
+    1024,  # gpt2 medium
+    1280,  # qwen_2_5_vl vision
+    1600,  # gpt2 xl
+    2048,  # gpt3 small
+    3584,  # qwen_2_5_vl
+    3840,  # qwen_2_5_vl vision qkv
+    4096,  # llama3
+    7168,  # deepseek_v3
+    14336,  # llama3 intermediate
+    18432,  # deepseek_v3 intermediate
+    18944,  # qwen_2_5_vl intermediate
+]
+_atol = 1e-1  # bfloat16 is low precision for large matrices
 
 
-@pytest.mark.parametrize(
-    "size",
-    [4096, 2048],
-)
+@requires_gpu
+@pytest.mark.parametrize("size", _sizes_to_test)
 def test_random_hadamard_matrix_compliant(size):
-    had_matrix = random_hadamard_matrix(size)
-    product = torch.round(had_matrix @ had_matrix.T)
-    assert torch.equal(product, torch.eye(size))
+    # (H / sqrt(n))(H.T / sqrt(n)) == I
+    matrix = random_hadamard_matrix(size, device="cuda")
+    product = matrix @ matrix.T
+    eye = torch.eye(size, dtype=product.dtype, device="cuda")
+    assert torch.allclose(product, eye, atol=_atol)
 
 
 def test_random_hadamard_generator():
+    # check that generation is deterministic with a seed
     generator = torch.Generator().manual_seed(42)
-    one = random_hadamard_matrix(2048, generator)
-    two = random_hadamard_matrix(2048, generator)
+    one = random_hadamard_matrix(2048, gen=generator)
+    two = random_hadamard_matrix(2048, gen=generator)
 
     one_true = torch.tensor(
         [
@@ -73,12 +75,16 @@ def test_random_hadamard_generator():
     assert torch.all(two[:3, :3].sign() == two_true.sign())
 
 
-@pytest.mark.parametrize(
-    "size",
-    [1024],
-)
+@requires_gpu
+@pytest.mark.parametrize("size", _sizes_to_test)
 def test_deterministic_hadamard_compliant(size):
-    had_matrix = deterministic_hadamard_matrix(size)
+    if not is_pow2(size):
+        with pytest.raises(ValueError):
+            matrix = deterministic_hadamard_matrix(size, device="cuda")
+        return
+
     # (H / sqrt(n))(H.T / sqrt(n)) == I
-    product = had_matrix @ had_matrix.T
-    assert numpy.array_equal(product, numpy.eye(size))
+    matrix = deterministic_hadamard_matrix(size, device="cuda")
+    product = matrix @ matrix.T
+    eye = torch.eye(size, dtype=product.dtype, device="cuda")
+    assert torch.allclose(product, eye, atol=_atol)

{compressed_tensors-0.10.2a20250612 → compressed_tensors-0.10.2a20250616}/tests/test_utils/test_offload.py

@@ -19,6 +19,7 @@ from compressed_tensors.utils import (
     delete_offload_module,
     delete_offload_parameter,
     disable_hf_hook,
+    disable_offloading,
     get_execution_device,
     has_offloaded_params,
     offloaded_dispatch,
@@ -148,6 +149,47 @@ def test_register_offload_parameter():
     assert module.a.device == module.b.device == module.c.device == torch.device("meta")
 
 
+@requires_accelerate()
+@requires_gpu
+def test_register_offload_parameter_hook_replacement():
+    module = ExampleModule()
+    parameter_c = torch.nn.Parameter(torch.tensor(1.0, device="cuda"))
+    parameter_d = torch.nn.Parameter(torch.tensor(1.0, device="cpu"))
+
+    offloaded_dispatch(module, "cuda")
+    register_offload_parameter(module, "c", parameter_c)
+    register_offload_parameter(module, "d", parameter_d)
+
+    with disable_hf_hook(module):
+        assert module.a.device == torch.device("cpu")
+        assert module.b.device == torch.device("cpu")
+        assert module.c.device == torch.device("cuda:0")
+        assert module.d.device == torch.device("cpu")
+
+    assert module.a.device == torch.device("meta")
+    assert module.b.device == torch.device("meta")
+    assert module.c.device == torch.device("meta")
+    assert module.d.device == torch.device("meta")
+    assert module._hf_hook.weights_map["a"].device == torch.device("cpu")
+    assert module._hf_hook.weights_map["b"].device == torch.device("cpu")
+    assert module._hf_hook.weights_map["c"].device == torch.device("cpu")
+    assert module._hf_hook.weights_map["d"].device == torch.device("cpu")
+
+
+@requires_accelerate()
+@requires_gpu
+def test_register_offload_parameter_shared():
+    module = ExampleModule()
+    parameter = torch.nn.Parameter(torch.tensor(1.0))
+
+    offloaded_dispatch(module, "cuda")
+    register_offload_parameter(module, "c", parameter)
+    register_offload_parameter(module, "d", parameter)
+
+    with align_module_device(module):
+        assert module.c is module.d
+
+
 @requires_accelerate()
 def test_update_offload_parameter():
     from accelerate.hooks import attach_align_device_hook
@@ -397,15 +439,23 @@ def test_delete_offload_module(exec_device):
 
 @requires_gpu
 @requires_accelerate()
-@pytest.mark.parametrize("exec_device", [torch.device("cpu"), torch.device("cuda")])
-def test_offloaded_dispatch(exec_device):
+@pytest.mark.parametrize(
+    "exec_device,offload_device",
+    [
+        (torch.device("cpu"), torch.device("cpu")),
+        (torch.device("cpu"), torch.device("cuda:0")),
+        (torch.device("cuda:0"), torch.device("cpu")),
+        (torch.device("cuda:0"), torch.device("cuda:0")),
+    ],
+)
+def test_offloaded_dispatch(exec_device, offload_device):
     # single module
-    module = torch.nn.Linear(1, 2)
-    module = offloaded_dispatch(module, exec_device)
+    module = torch.nn.Linear(1, 2, device=offload_device)
+    module = offloaded_dispatch(module, exec_device, offload_device)
     assert has_offloaded_params(module)
     assert module._hf_hook.offload
     assert module.weight.device == torch.device("meta")
-    assert "weight" in module._hf_hook.weights_map
+    assert module._hf_hook.weights_map["weight"].device == offload_device
     assert module._hf_hook.tied_params_map is not None
 
     # can run
@@ -413,13 +463,13 @@ def test_offloaded_dispatch(exec_device):
 
     # model
     model = ExampleModel()
-    model = offloaded_dispatch(model, exec_device)
+    model = offloaded_dispatch(model, exec_device, offload_device)
     assert not has_offloaded_params(model)
 
     assert has_offloaded_params(model.linear)
     assert model.linear._hf_hook.offload
     assert model.linear.weight.device == torch.device("meta")
-    assert "weight" in model.linear._hf_hook.weights_map
+    assert model.linear._hf_hook.weights_map["weight"].device == offload_device
     assert model.linear._hf_hook.tied_params_map is not None
 
     # can run
@@ -429,4 +479,43 @@ def test_offloaded_dispatch(exec_device):
     parameter = torch.nn.Parameter(torch.tensor(1.0))
     register_offload_parameter(module, "new_param", parameter)
     assert module.new_param.device == torch.device("meta")
-    assert module._hf_hook.weights_map["new_param"].device == torch.device("cpu")
+    assert module._hf_hook.weights_map["new_param"].device == offload_device
+
+
+@requires_gpu
+@requires_accelerate()
+@pytest.mark.parametrize(
+    "exec_device,offload_device",
+    [
+        (torch.device("cpu"), torch.device("cpu")),
+        (torch.device("cpu"), torch.device("cuda:0")),
+        (torch.device("cuda:0"), torch.device("cpu")),
+        (torch.device("cuda:0"), torch.device("cuda:0")),
+    ],
+)
+def test_disable_offloading(exec_device, offload_device):
+    module = torch.nn.Linear(1, 2, device=exec_device)
+
+    # non-offloaded modules are unaffected
+    with disable_offloading():
+        output = module(torch.empty(1, device=exec_device))
+        assert module.weight.device == exec_device
+        assert output.device == exec_device
+
+    # offloaded modules stay on device until context exit
+    offloaded_dispatch(module, exec_device, offload_device)
+    assert module.weight.device == torch.device("meta")
+    assert module._hf_hook.weights_map["weight"].device == offload_device
+
+    with disable_offloading():
+        assert module.weight.device == torch.device("meta")
+        output = module(torch.empty(1, device=exec_device))
+        assert module.weight.device == exec_device
+        assert output.device == exec_device
+
+        output = module(torch.empty(1, device=exec_device))
+        assert module.weight.device == exec_device
+        assert output.device == exec_device
+
+    assert module.weight.device == torch.device("meta")
+    assert module._hf_hook.weights_map["weight"].device == offload_device