Synchronize the compute stream before offloading to disk

[jiqing-feng]

What is fixed

Group offloading to disk produces all-NaN output with offload_type="leaf_level", use_stream=True and record_stream=False.

Reproducer on main:

pytest tests/models/autoencoders/test_models_autoencoder_vidtok.py::TestAutoencoderVidTokMemory::test_group_offloading_with_disk[leaf_level-False]

Why

_offload_to_disk ends with tensor_obj.data = torch.empty_like(tensor_obj.data, device=self.offload_device), which returns the on-device weight memory to the allocator. It runs on the host thread from post_forward, while kernels on the compute stream may still be reading those weights. The allocator hands the block to the next group's onload, so the in-flight kernel reads foreign data and the result becomes NaN. _offload_to_memory already guards this with current_stream().synchronize(); the disk path did not, so this PR adds the same guard.

Only that one combination fails because every other one removes the race: record_stream=True makes the allocator aware that the compute stream still owns the block, the memory path already synchronizes, and block_level groups are coarse enough that the offload never overlaps the next onload. Nothing here is device specific, the CUDA allocator just happens to pick a different block.

Validation

All eight offload_type x offload_to_disk x record_stream combinations of the VidTok autoencoder now match the non-offloaded output at atol=1e-5, and tests/hooks/test_group_offloading.py plus tests/models/autoencoders/test_models_autoencoder_vidtok.py pass (50 passed, 3 skipped for the memory tests, 38 passed for the hook tests). Tested on Intel Arc Pro B60 with torch 2.13.0+xpu.

[github-actions]

Hi @jiqing-feng, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.

[jiqing-feng]

Hi @sayakpaul . Would you please review this pr if you have bandwidth? Thanks!

[sayakpaul]

/diffusers-bot pytest tests/models -k "test_group_offloading_with_disk"

[github-actions]

⚠️pytest tests/models -k "test_group_offloading_with_disk" was cancelled on GPU — view logs.

[sayakpaul]

@jiqing-feng thanks for the PR. I just triggered the respective CI for this. Let's see how it pans out.

[jiqing-feng]

⚠️pytest tests/models -k "test_group_offloading_with_disk" was cancelled on GPU — view logs.

The job seems been cancelled.

[sayakpaul]

/diffusers-bot pytest tests/models -k "test_group_offloading_with_disk"

[github-actions]

✅ pytest tests/models -k "test_group_offloading_with_disk" passed on GPU — view logs.