huggingface / huggingface/diffusers
[tests] help us test `torch.compile()` for impactful models
- Dominant language
- Python
- Stars
- 34.5k
- Forks
- 7.3k
- Avg merge
- 3d 3h
- Merged PRs (30d)
- 91
Description
https://github.com/huggingface/diffusers/pull/11085 added a test for checking if there's any graph break or recompilation issue for `torch.compile`d model.
We should add this test to the most impactful models to ensure our code is `torch.compile` friendly and has the potential to benefit from it. So far, we test it for `FluxTransformer2DModel`. Below are some models I have in mind where this test should be added:
- [ ] `HiDreamImageTransformer2DModel` (~currently doesn't have a test class like other models from `src/diffusers/models`~)
- [x] [`HunyuanVideoTransformer3DTests`](https://github.com/huggingface/diffusers/blob/a7e9f85e21dde12f2f2489702ea82db40ebb31d2/tests/models/transformers/test_models_transformer_hunyuan_video.py#L28C7-L28C37)
- [x] [`WanTransformer3DTests`](https://github.com/huggingface/diffusers/blob/a7e9f85e21dde12f2f2489702ea82db40ebb31d2/tests/models/transformers/test_models_transformer_wan.py#L28C7-L28C28)
- [x] [`UNet2DConditionModelTests`](https://github.com/huggingface/diffusers/blob/a7e9f85e21dde12f2f2489702ea82db40ebb31d2/tests/models/unets/test_models_unet_2d_condition.py#L353C7-L353C32)
- [x] LTXVideoTransformer3DModel
## Steps to contribute
1. Refer to https://github.com/huggingface/diffusers/pull/11085 to under the changes needed in the test classes of the respective models.
2. To be specific, always attempt to just add `TorchCompileTesterMixin` to the respective model testing file. Like: https://github.com/huggingface/diffusers/blob/fb29132b98abdd218bacb6dbaab372f5bb177a2e/tests/models/transformers/test_models_transformer_flux.py#L81-L81
3. Make the changes.
4. Run the tests locally with `RUN_SLOW=1 RUN_COMPILE=1 pytest tests/ -k "test_torch_compile_recompilation_and_graph_break"`. Make sure the machine has a GPU. Please confirm that this test is passing when opening the PR. In case it doesn't pass let us know.
5. Open the PR and mention this issue (#11430). Tag @sayakpaul and @DN6 for a review. Don't hesitate to ask for guidance/help if needed.
@DN6 any model classes I am missing?
Contributor guide
Assessment
This issue has not been assessed yet.