-
Notifications
You must be signed in to change notification settings - Fork 6.3k
[tests] feat: add AoT compilation tests #12203
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
Conversation
@angelayi
angelayi
Aug 21, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not passing in a path should also automatically give you a path in the tmp dir!
@angelayi
angelayi
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for adding a test!
@angelayi
angelayi
Aug 21, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is there a particular reason why you're running it twice?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
To emulate the real scenario as the model is typically invoked more than once during the actual generation process.
Very nice! Ideally we should pass multiple inputs to the exported (and/or compiled) model in order to test proper dynamism (for instance in the case of QwenImageTransformer2DModel, export test might pass but we'll miss the fact that it will only work at a given image shape). Of course this probably means quite a (too) big refactor of prepare_init_args_and_inputs_for_common
Very nice! Ideally we should pass multiple inputs to the exported (and/or compiled) model in order to test proper dynamism (for instance in the case of QwenImageTransformer2DModel, export test might pass but we'll miss the fact that it will only work at a given image shape). Of course this probably means quite a (too) big refactor of prepare_init_args_and_inputs_for_common
We already have this kind of test:
diffusers/tests/models/test_modeling_common.py
Line 2119 in 91a151b
I think it's fine for now to just test for static shapes. If dynamic AoT becomes more popular, we can revisit it.
We already have this kind of test:
Very nice! I think that prepare_dummy_inputs could actually be useful outside of tests too (for torch.export.export input examples)
What does this PR do?
AoT compilation is exciting because it helps cut the framework overhead. It also helps realize similar benefits as JiT compilation for environments where JiT might not be feasible (ZeroGPU Spaces, for example). For example, with AoT-compilation, we were able to obtain 1.7x speedups in latency on a ZeroGPU Space. Wouldn't have been feasible, otherwise.
To run,
RUN_SLOW=1 RUN_COMPILE=yes pytest tests/models/ -k "aot".