Fixes #12673. record_stream in group offloading is not working properly

[KimbingNg]

What does this PR do?

Fixes # (12673) and Tencent-Hunyuan/HunyuanVideo-1.5#10 I identified the root cause: record_stream isn’t functioning as expected. This causes the offload_to_memory function to update param.data without waiting for the default stream (current_stream) to complete (the forward pass).

Updates [Nov 27]: The root cause!!
The current_stream() call inside _transfer_tensor_to_device is under with torch.cuda.stream(self.stream). So it is not returning the correct default stream to call record_stream on !

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• [s] Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

cc @asomoza @sayakpaul @a-r-r-o-w @DN6

[KimbingNg]

Also relevant to Tencent-Hunyuan/HunyuanVideo-1.5#10

[sayakpaul]

Thanks for your PR!

Could we also run the test_group_offloading tests?

pytest tests/models -k "test_group_offloading"?

[KimbingNg]

pytest tests/models -k "test_group_offloading"

Running pytest:

=============================================== 365 passed, 37 skipped, 3604 deselected, 757 warnings in 100.34s (0:01:40) ================================================

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[DN6]

Thanks @KimbingNg I think we might also need to make a similar change to _offload_to_disk.

[KimbingNg]

@DN6 , I found the root cause!!!!

The current_stream() call inside _transfer_tensor_to_device is under with torch.cuda.stream(self.stream). So it is not returning the correct default stream to call record_stream on !

As for the onload_from_disk case, the current_stream is obtained outside the with context, so it works fine.

I force push a new commit, please help reviewing that again, thanks!

[DN6]

@KimbingNg Ah great catch! 👍🏽

[sayakpaul]

We can also remove the following lines:

if self.__class__.__name__ == "AutoencoderKLCosmosTests" and offload_type == "leaf_level":
            pytest.skip("With `leaf_type` as the offloading type, it fails. Needs investigation.")

in the test_group_offloading_with_disk() test implementation.

[KimbingNg]

We can also remove the following lines:

if self.__class__.__name__ == "AutoencoderKLCosmosTests" and offload_type == "leaf_level":
            pytest.skip("With `leaf_type` as the offloading type, it fails. Needs investigation.")

in the test_group_offloading_with_disk() test implementation.

I ran pytest tests/models -k "test_group_offloading" after removing these two lines:

=============================================== 366 passed, 36 skipped, 3604 deselected, 1177 warnings in 123.01s (0:02:03) ===============================================

Should I push a commit that removes these two lines?

[sayakpaul]

Yes

[KimbingNg]

The two lines are removed