diff --git a/src/diffusers/hooks/group_offloading.py b/src/diffusers/hooks/group_offloading.py index 10d3f0c245a1..50d9b7fe8ec4 100644 --- a/src/diffusers/hooks/group_offloading.py +++ b/src/diffusers/hooks/group_offloading.py @@ -175,8 +175,9 @@ def __init__( @staticmethod def _to_cpu(tensor, low_cpu_mem_usage): # For TorchAO tensors, `.data` returns an incomplete wrapper without internal attributes - # (e.g. `.qdata`, `.scale`), so we must call `.cpu()` on the tensor directly. - t = tensor.cpu() if _is_torchao_tensor(tensor) else tensor.data.cpu() + # (e.g. `.qdata`, `.scale`), so we must call `.to(..., copy=True)` on the tensor directly. `tensor.cpu()` can + # return `tensor` itself when it is already on CPU, which would alias the cached CPU copy with the live parameter. + t = tensor.to("cpu", copy=True) if _is_torchao_tensor(tensor) else tensor.data.cpu() return t if low_cpu_mem_usage else t.pin_memory() def _init_cpu_param_dict(self): diff --git a/tests/hooks/test_group_offloading.py b/tests/hooks/test_group_offloading.py index a903186aa6b4..41c0335e3fa7 100644 --- a/tests/hooks/test_group_offloading.py +++ b/tests/hooks/test_group_offloading.py @@ -21,6 +21,7 @@ from diffusers import AutoencoderKL from diffusers.hooks import HookRegistry, ModelHook +from diffusers.hooks.group_offloading import ModuleGroup, _restore_torchao_tensor, _swap_torchao_tensor from diffusers.models import ModelMixin from diffusers.pipelines.pipeline_utils import DiffusionPipeline from diffusers.utils import logging as diffusers_logging @@ -372,6 +373,30 @@ def test_error_raised_if_group_offloading_applied_on_sequential_offloaded_module with pytest.raises(ValueError, match="Cannot apply group offloading"): pipe.model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=3) + def test_torchao_cpu_cache_does_not_alias_live_parameter(self): + try: + from torchao.quantization import Int8WeightOnlyConfig, quantize_ + except ImportError: + pytest.skip("test requires torchao") + + linear = torch.nn.Linear(16, 16, bias=False, dtype=torch.bfloat16) + quantize_(linear, Int8WeightOnlyConfig(version=2)) + weight = linear.weight + + cpu_copy = ModuleGroup._to_cpu(weight, low_cpu_mem_usage=True) + assert cpu_copy is not weight + + moved = weight.to("meta") + _swap_torchao_tensor(weight, moved) + + tensor_data_names = getattr(cpu_copy.__class__, "tensor_data_names") + for attr_name in tensor_data_names: + assert getattr(cpu_copy, attr_name).device.type == "cpu" + + _restore_torchao_tensor(weight, cpu_copy) + for attr_name in tensor_data_names: + assert getattr(weight, attr_name).device.type == "cpu" + def test_block_level_stream_with_invocation_order_different_from_initialization_order(self): if torch.device(torch_device).type not in ["cuda", "xpu"]: return