Fix TorchAO group offload CPU cache aliasing

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):

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