From d49831f9f68dbd3b5373c5b10e7e2d0ea053b5a9 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Mon, 31 Aug 2026 06:27:15 +0000 Subject: [PATCH] Synchronize the compute stream before offloading to disk `_offload_to_disk` releases the onloaded tensors at the end, returning their device memory to the allocator while the compute stream may still be reading them. The memory is then reused by the next onload, so the in-flight kernel reads foreign data and the output becomes NaN. `_offload_to_memory` already synchronizes for this reason, do the same on the disk path. Reproducer: tests/models/autoencoders/test_models_autoencoder_vidtok.py::TestAutoencoderVidTokMemory::test_group_offloading_with_disk[leaf_level-False] --- src/diffusers/hooks/group_offloading.py | 5 +++++ 1 file changed, 5 insertions(+) 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