Skip to content

[tests] Changes to the torch.compile() CI and tests #11508

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.

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

Merged
merged 13 commits into from
May 26, 2025
Merged

[tests] Changes to the torch.compile() CI and tests #11508

merged 13 commits into from
May 26, 2025

Conversation

sayakpaul
Copy link
Member

What does this PR do?

  • Removes the CUDA compile docker container
  • Consistency in how compilation cache is cleared and ensuring we start a compilation run with fresh cache
  • Removes unnecessary compilation tests from pipelines

Regarding the final point, I think we already test torch.compile() support for popular models and doing it at the pipeline-level makes little sense to me. But no strong opinions.

@sayakpaul sayakpaul requested a review from DN6 May 6, 2025 10:06
@sayakpaul
Copy link
Member Author

@DN6 a gentle ping.

Copy link
Collaborator

@DN6 DN6 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good 👍🏽 Thanks. Minor comments.

@@ -96,23 +92,8 @@ def test_gradient_checkpointing_is_applied(self):
expected_set = {"HunyuanVideoTransformer3DModel"}
super().test_gradient_checkpointing_is_applied(expected_set=expected_set)

@require_torch_gpu
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would make sense to add these decorators to the top of the Mixin no?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These decorators are present in the mixin:

@require_torch_gpu
@require_torch_2
@is_torch_compile
@slow

@@ -188,7 +188,7 @@ jobs:
group: aws-g4dn-2xlarge

container:
image: diffusers/diffusers-pytorch-compile-cuda
image: diffusers/diffusers-pytorch-cuda
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Small nit. We can remove the -k "compile" in the test runner step.

@sayakpaul sayakpaul merged commit 4af76d0 into main May 26, 2025
31 checks passed
@sayakpaul sayakpaul deleted the compile-ci branch May 26, 2025 15:31
@sayakpaul
Copy link
Member Author

Thanks for the reviews, @DN6! Hope to make our compile CI better and better.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants