[New] [Offload] [1/2] Disambiguate synchronous __setitem__/offload from update_offload

src/compressed_tensors/offload/__init__.py

@@ -110,6 +110,7 @@ def disable_onloading():
         yield
 
 
+@torch.no_grad()
 def update_offload_parameter(module: torch.nn.Module, name: str, data: torch.Tensor):
     """
     Update the offload and onload data of an existing parameter/buffer. Supports both
@@ -127,26 +128,27 @@ def update_offload_parameter(module: torch.nn.Module, name: str, data: torch.Ten
     :param data: tensor to update parameter with
     """
     if isinstance(module._parameters, OffloadCache):
-        # | Component | Update Implementation       |
-        # | --------- | --------------------------- |
-        # | CPU       | Copy into shared cpu memory |
-        # | Disk      | Write file to disk          |
-        # | Device    | Copy into local device      |
-        # | --------- | --------------------------- |
-        # all implementations update onloaded data if applicable
         if name in module._parameters:
             cache = module._parameters
         elif name in module._buffers:
             cache = module._buffers
         else:
             raise AttributeError(f"{type(module)} has no attribute {name}")
 
-        # triggers update if shapes match
-        cache[name] = data
+        # when onloading is disabled, parameters can be access and assigned directly
+        if cache.onloading_disabled:
+            cache.offloaded_values[name] = data
+            return
+
+        # get offloaded value for updating
+        offloaded = cache.offloaded_values[name]
+        if offloaded is None:
+            raise ValueError(f"Cannot update offload value `None` for param {name}")
+
+        cache.update_offload(offloaded, data)
 
     else:
-        with torch.no_grad():
-            getattr(module, name).copy_(data)
+        getattr(module, name).copy_(data)
 
 
 def get_execution_device(

src/compressed_tensors/offload/cache/base.py

@@ -215,18 +215,11 @@ def __setitem__(self, key: Hashable, value: torch.Tensor | None):
             self.offloaded_values[key] = value
             return
 
-        # if the key already exists, update with the new value
-        offloaded = self.offloaded_values.get(key, None)
-        if offloaded is not None and torch.is_same_size(offloaded, value):
-            self.update_offload(offloaded, value)
+        if key in self:
+            del self[key]
 
-            onloaded = self.keep_onloaded_values.get(offloaded, None)
-            if onloaded is not None and onloaded is not offloaded:
-                onloaded.copy_(value)
-
-        # if the key does not exist (or the value is None), offload the new value
-        else:
-            self.offloaded_values[key] = self.offload(value)
+        # synchronously offload value
+        self.offloaded_values[key] = self.offload(value)
 
     def __delitem__(self, key: Hashable):
         """

src/compressed_tensors/offload/cache/disk.py

@@ -113,7 +113,7 @@ def offload(
         }
 
         assert self._is_ct_file_path(file_path), f"Attempted to write to {file_path}"
-        save_file({"weight": tensor}, file_path)
+        save_file({"weight": tensor.contiguous()}, file_path)
         return offloaded
 
     def __delitem__(self, key: str):
@@ -154,7 +154,8 @@ def update_offload(self, offloaded: torch.Tensor, data: torch.Tensor | None):
 
         # save with data using original weight_name
         assert self._is_ct_file_path(file_path), f"Attempted to write to {file_path}"
-        save_file({weight_name: data.reshape_as(offloaded).to(dtype=dtype)}, file_path)
+        data = data.reshape_as(offloaded).to(dtype=dtype).contiguous()
+        save_file({weight_name: data}, file_path)
 
     @classmethod
     def create_checkpoint_symlink(

tests/test_compressors/distributed/conftest.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+
+import pytest
+import torch.distributed as dist
+from compressed_tensors.offload.dist_utils import is_distributed
+
+
+@pytest.fixture
+def offload_folder(tmp_path) -> str:
+    offload_path = tmp_path / "offload_dir"
+
+    if not is_distributed() or dist.get_rank() == 0:
+        os.makedirs(offload_path, exist_ok=True)
+
+    if is_distributed():
+        broadcast_object = [str(offload_path)]
+        dist.broadcast_object_list(broadcast_object, src=0)
+        offload_path = broadcast_object[0]
+
+    return str(offload_path)

tests/test_compressors/distributed/test_distributed_compression.py

@@ -172,6 +172,39 @@ def test_distributed_compression_with_offload():
     assert hasattr(model.layer2, "weight_packed")
 
 
+@pytest.mark.unit
+@requires_gpu(2)
+@torchrun(world_size=2, init_dist=True)
+def test_distributed_compression_with_disk_offload(offload_folder):
+    """Test distributed compression with offloaded modules."""
+    model = TwoLayerModel()
+    setup_quantized_model(model)
+
+    # Offload model to CPU
+    offload_module(
+        model.layer1,
+        onload_device="cuda",
+        offload_device="disk",
+        offload_dir=offload_folder,
+    )
+    offload_module(
+        model.layer2,
+        onload_device="cuda",
+        offload_device="disk",
+        offload_dir=offload_folder,
+    )
+
+    q_config = create_quantization_config(bits=4, format="pack-quantized")
+    compressor = ModelCompressor(quantization_config=q_config)
+
+    # Compress the model
+    compressor.compress_model(model)
+
+    # Verify compression happened even with offloading
+    assert hasattr(model.layer1, "weight_packed")
+    assert hasattr(model.layer2, "weight_packed")
+
+
 @pytest.mark.unit
 @requires_gpu(2)
 @torchrun(world_size=2, init_dist=True)

tests/test_offload/cache/test_disk.py

@@ -112,7 +112,7 @@ def test_files(tmp_path):
         read_tensor = file.get_tensor("weight")
         assert_tensor_equal(read_tensor, tensor)
 
-    # modify
+    # rewrite
     tensor = torch.ones(10)
     cache["weight"] = tensor
 

tests/test_offload/cache/test_dist_cpu.py

@@ -158,10 +158,14 @@ def test_distributed_async_update(onload_device):
     rank = dist.get_rank()
     if rank == 0:
         # Rank 0 updates tensor_0
-        cache[f"tensor_{rank}"] = torch.ones(10, device=onload_device) * 1.0
+        with disable_onloading():
+            offloaded = cache[f"tensor_{rank}"]
+        cache.update_offload(offloaded, torch.ones(10, device=onload_device) * 1.0)
     elif rank == 1:
         # Rank 1 updates tensor_1
-        cache[f"tensor_{rank}"] = torch.ones(10, device=onload_device) * 2.0
+        with disable_onloading():
+            offloaded = cache[f"tensor_{rank}"]
+        cache.update_offload(offloaded, torch.ones(10, device=onload_device) * 2.0)
 
     # Synchronize to ensure all updates are complete
     dist.barrier()

tests/test_offload/cache/test_dist_disk.py

@@ -172,7 +172,9 @@ def test_distributed_files(tmp_path):
     # modify on one rank
     tensor = torch.ones(10)
     if dist.get_rank() == 0:
-        cache["weight"] = tensor
+        with disable_onloading():
+            offloaded = cache["weight"]
+        cache.update_offload(offloaded, tensor)
 
     assert len(DiskCache.index) == 1
     if dist.get_rank() == 0:  # only rank0 bc `tmp_path` is not shared between ranks
@@ -218,10 +220,14 @@ def test_distributed_async_update(tmp_path):
     rank = dist.get_rank()
     if rank == 0:
         # Rank 0 updates tensor_0
-        cache[f"tensor_{rank}"] = torch.ones(10, device=onload_device) * 1.0
+        with disable_onloading():
+            offloaded = cache[f"tensor_{rank}"]
+        cache.update_offload(offloaded, torch.ones(10, device=onload_device) * 1.0)
     elif rank == 1:
         # Rank 1 updates tensor_1
-        cache[f"tensor_{rank}"] = torch.ones(10, device=onload_device) * 2.0
+        with disable_onloading():
+            offloaded = cache[f"tensor_{rank}"]
+        cache.update_offload(offloaded, torch.ones(10, device=onload_device) * 2.0)
 
     # Synchronize to ensure all updates are complete
     dist.barrier()

tests/test_offload/convert/test_from_accelerate.py

@@ -13,6 +13,7 @@
     disable_onloading,
     from_accelerate,
     load_offloaded_model,
+    update_offload_parameter,
 )
 from compressed_tensors.offload.cache import CPUCache, DeviceCache, DiskCache
 from compressed_tensors.offload.convert.from_accelerate import (
@@ -169,10 +170,11 @@ def test_dist_disk_safetensors_update(tmp_path):
         rank_1_module = model.model.layers[-1].self_attn.k_proj
         rank = dist.get_rank()
         if rank == 0:
-            rank_0_module.weight *= 0
+            new_weight = rank_0_module.weight * 0
+            update_offload_parameter(rank_0_module, "weight", new_weight)
         elif rank == 1:
-            rank_1_module.weight *= 0
-            rank_1_module.weight += 1
+            new_weight = rank_1_module.weight * 0 + 1
+            update_offload_parameter(rank_1_module, "weight", new_weight)
         dist.barrier()
 
         # Check that onloaded values are updated across ranks

tests/test_offload/test_interface.py

@@ -79,7 +79,6 @@ def test_update_offload_parameter(linear: torch.nn.Linear, cache, offload):
     linear.weight = torch.nn.Parameter(init_data, requires_grad=False)
     if offload:
         offload_module(linear, ONLOAD_DEVICE, OFFLOAD_DEVICE)
-
     assert linear.weight == 0
 
     update_offload_parameter(linear, "weight", torch.tensor(1))
@@ -96,6 +95,32 @@ def test_update_offload_parameter(linear: torch.nn.Linear, cache, offload):
     assert linear.weight == 3
 
 
+@pytest.mark.unit
+@requires_gpu
+def test_update_offload_parameter_only(offloaded_linear: torch.nn.Linear):
+    # Use tensors matching the weight shape and explicitly on the offload device
+    offloaded_linear.weight = torch.nn.Parameter(
+        torch.zeros(5, 5, device=OFFLOAD_DEVICE), requires_grad=False
+    )
+
+    with disable_offloading():
+        # Access weight to onload it and cache it
+        _ = offloaded_linear.weight
+        # Update only the offloaded value, not the onloaded cache
+        update_offload_parameter(
+            offloaded_linear, "weight", torch.ones(5, 5, device=OFFLOAD_DEVICE)
+        )
+
+        # Verify offloaded value was updated to ones
+        with disable_onloading():
+            offload = offloaded_linear.weight
+        assert torch.all(offload == 1).item()
+
+        # Verify onloaded cached value is still zeros (not updated)
+        onload = offloaded_linear.weight
+        assert torch.all(onload == 0).item()
+
+
 @pytest.mark.unit
 def test_update_offload_parameter_with_grad(linear: torch.nn.Linear):
     zeros = torch.nn.Parameter(torch.zeros(5, 5), requires_grad=True)

tests/test_offload/test_module.py

@@ -204,16 +204,18 @@ def test_forward_signature(linear: torch.nn.Linear, cache):
 
 
 @pytest.mark.unit
+@requires_gpu
 def test_set_item(offloaded_linear: torch.nn.Linear):
-    # update
+    # update with same size - should alias when already on offload device
     update = torch.nn.Parameter(
         torch.rand(5, 5, device=OFFLOAD_DEVICE), requires_grad=False
     )
     offloaded_linear.weight = update
     with disable_onloading():
-        assert offloaded_linear.weight is not update
+        # When value is already on offload device, __setitem__ can alias it
+        assert offloaded_linear.weight is update
 
-    # overwrite with different size
+    # overwrite with different size - should also alias
     overwrite = torch.nn.Parameter(
         torch.rand(6, 6, device=OFFLOAD_DEVICE), requires_grad=False
     )
@@ -223,14 +225,18 @@ def test_set_item(offloaded_linear: torch.nn.Linear):
 
 
 @pytest.mark.unit
+@requires_gpu
 def test_set_item_buffers(offloaded_linear: torch.nn.Linear):
     # common case: registering buffers of difference sizes twice
     new = torch.rand(5)
     offloaded_linear.register_buffer("buffer", new, persistent=False)
     with disable_onloading():
         assert offloaded_linear.buffer is new
 
-    overwrite = torch.rand(6)
-    offloaded_linear.register_buffer("buffer", overwrite, persistent=False)
-    with disable_onloading():
-        assert offloaded_linear.buffer is overwrite
+    for size in (5, 6):
+        overwrite = torch.rand(size, device=ONLOAD_DEVICE)
+        offloaded_linear.register_buffer("buffer", overwrite, persistent=False)
+        with disable_onloading():
+            offloaded = offloaded_linear.buffer
+            assert offloaded.device == OFFLOAD_DEVICE
+            assert torch.equal(offloaded.to(overwrite), overwrite)