diff --git a/src/compressed_tensors/offload/__init__.py b/src/compressed_tensors/offload/__init__.py index 51bba4fb2..c2fac355c 100644 --- a/src/compressed_tensors/offload/__init__.py +++ b/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,13 +128,6 @@ 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: @@ -141,12 +135,20 @@ def update_offload_parameter(module: torch.nn.Module, name: str, data: torch.Ten 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( diff --git a/src/compressed_tensors/offload/cache/base.py b/src/compressed_tensors/offload/cache/base.py index 51821dc15..a1abd6f08 100644 --- a/src/compressed_tensors/offload/cache/base.py +++ b/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): """ diff --git a/src/compressed_tensors/offload/cache/disk.py b/src/compressed_tensors/offload/cache/disk.py index 4857d00a1..a9bfab8c7 100644 --- a/src/compressed_tensors/offload/cache/disk.py +++ b/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( diff --git a/tests/test_compressors/distributed/conftest.py b/tests/test_compressors/distributed/conftest.py new file mode 100644 index 000000000..54486056d --- /dev/null +++ b/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) diff --git a/tests/test_compressors/distributed/test_distributed_compression.py b/tests/test_compressors/distributed/test_distributed_compression.py index 1e7f59ef4..d7e370ab5 100644 --- a/tests/test_compressors/distributed/test_distributed_compression.py +++ b/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) diff --git a/tests/test_offload/cache/test_disk.py b/tests/test_offload/cache/test_disk.py index 6ebdd2986..e53048d5f 100644 --- a/tests/test_offload/cache/test_disk.py +++ b/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 diff --git a/tests/test_offload/cache/test_dist_cpu.py b/tests/test_offload/cache/test_dist_cpu.py index c0c8c78c5..cfb39528a 100644 --- a/tests/test_offload/cache/test_dist_cpu.py +++ b/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() diff --git a/tests/test_offload/cache/test_dist_disk.py b/tests/test_offload/cache/test_dist_disk.py index 62b95ac83..0d16eb37d 100644 --- a/tests/test_offload/cache/test_dist_disk.py +++ b/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() diff --git a/tests/test_offload/convert/test_from_accelerate.py b/tests/test_offload/convert/test_from_accelerate.py index f9b755e04..8faeb6663 100644 --- a/tests/test_offload/convert/test_from_accelerate.py +++ b/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 diff --git a/tests/test_offload/test_interface.py b/tests/test_offload/test_interface.py index 784bd4bf8..300ded2a9 100644 --- a/tests/test_offload/test_interface.py +++ b/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) diff --git a/tests/test_offload/test_module.py b/tests/test_offload/test_module.py index 855b311e4..3623ae71a 100644 --- a/tests/test_offload/test_module.py +++ b/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,6 +225,7 @@ 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) @@ -230,7 +233,10 @@ def test_set_item_buffers(offloaded_linear: torch.nn.Linear): 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)