diff --git a/src/diffusers/hooks/group_offloading.py b/src/diffusers/hooks/group_offloading.py index 10d3f0c245a1..e6c8d8732113 100644 --- a/src/diffusers/hooks/group_offloading.py +++ b/src/diffusers/hooks/group_offloading.py @@ -294,6 +294,11 @@ def _onload_from_memory(self): def _offload_to_disk(self): self._check_disk_offload_torchao() + # Releasing the onloaded tensors below frees their device memory, which the compute stream may still be + # reading. `record_stream` already prevents the allocator from reusing it too early. + if self.stream is not None and not self.record_stream: + self._torch_accelerator_module.current_stream().synchronize() + # TODO: we can potentially optimize this code path by checking if the _all_ the desired # safetensor files exist on the disk and if so, skip this step entirely, reducing IO # overhead. Currently, we just check if the given `safetensors_file_path` exists and if not