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 for unknown reasons. This causes the offload_to_memory function to update param.data without waiting for the default stream (current_stream) to complete (the forward pass).

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]

As for why record_stream fails, I suspect it’s because calling record_stream on Tensor.data is unsafe — Tensor.data is somewhat deprecated, though I still need confirmation on this.

[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"?

[sayakpaul]

Why do we have to condition on param.device.type == 'cuda'?

[KimbingNg]

Because Tensors on CPU can not call record_stream

[sayakpaul]

We can do something like if self.record_stream and param.device.type != "cpu":, right?

[KimbingNg]

Yes, I updated the code as you recommended in the latest commit

[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.