Skip to content

Synchronize the compute stream before offloading to disk#14657

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jiqing-feng wants to merge 2 commits into
base:main
Choose a base branch
from
+5 −0
Open

Synchronize the compute stream before offloading to disk #14657

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d49831f
Synchronize the compute stream before offloading to disk
jiqing-feng Aug 31, 2026
bb47795
Merge branch 'main' into fix-group-offload-disk-stream-sync
jiqing-feng Aug 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions src/diffusers/hooks/group_offloading.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand Down
Loading