[tests] add tests to check for graph breaks, recompilation, cuda syncs in pipelines during torch.compile() #11085

sayakpaul · 2025-03-17T09:00:14Z

What does this PR do?

Diffusers prides itself in being performant and also in providing good support for torch.compile() related optimizations. For the latter to be sufficiently effective, we need to ensure our pipelines don't cause:

graph breaks when fullgraph=True is requested
recompilation
CUDA syncs causing overheads (yet to add a test for this)

This PR adds a test suite for this. It is done through a test called test_torch_compile_recompilation_and_graph_break(). We decided to separate this out in a class called TorchCompileTesterMixin to follow #11085 (comment).

I think we should enable this for the most impactful diffusion models while skipping the ones that don't have much usage. I have hence started this by adding the test to test_models_transformer_flux.py as discussed with @DN6.

@DN6 @hlky if you could provide your thoughts on this direction of testing, that would be helpful.

tests/pipelines/test_pipelines_common.py

anijain2305

This is good. Just to reiterate - graph breaks and recompilation are different/orthogonal concepts.

if you want to ensure that your model has no graph breaks, fullgraph=True is enough.

if you want to ensure that your model does not recompile but has graph breaks, you can use

torch._dynamo.config.recompile_limit = 1
torch._dynamo.config.fail_on_recompile_limit_hit = True

If you want to ensure that your model has no graph breaks and no recompilations, you can use

model = torch.compile(model, fullgraph=True)
torch._dynamo.config.recompile_limit = 1

Here, fullgraph=True internally ensures that it raises an error if the total number of compilations exceed recompile_limit.

What you have in this PR is also fine. You want no graph break and no recompilations. So you are using fullgraph=True and

with torch._dynamo.config.patch(error_on_recompile=True):

This works too.

sayakpaul · 2025-04-09T06:56:22Z

@DN6 @hlky a gentle ping.

hlky

Looks good, thanks for working on this.

sayakpaul · 2025-04-09T09:59:53Z

I will run the test on all pipelines so that we can decide which ones should be skipped.

sayakpaul · 2025-04-14T07:17:32Z

~@DN6 I ran the full test suite (with ) and got the following

Unfold

========================================================= short test summary info =========================================================
FAILED tests/pipelines/cogview4/test_cogview4.py::CogView4PipelineFastTests::test_torch_compile_recompilation_and_graph_break - ImportError: This modeling file requires the following packages that were not found in your environment: tiktoken. Run `pip install ti...
FAILED tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py::FluxControlNetImg2ImgPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/transformers/transformer_...
FAILED tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py::FluxControlNetInpaintPipelineTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/transformers/transformer_...
FAILED tests/pipelines/controlnet_sd3/test_controlnet_sd3.py::StableDiffusion3ControlNetPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/transformers/transformer_...
FAILED tests/pipelines/flux/test_pipeline_flux_control.py::FluxControlPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/transformers/transformer_...
FAILED tests/pipelines/flux/test_pipeline_flux_control_img2img.py::FluxControlImg2ImgPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/transformers/transformer_...
FAILED tests/pipelines/flux/test_pipeline_flux_control_inpaint.py::FluxControlInpaintPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/transformers/transformer_...
FAILED tests/pipelines/flux/test_pipeline_flux_fill.py::FluxFillPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/transformers/transformer_...
FAILED tests/pipelines/flux/test_pipeline_flux_img2img.py::FluxImg2ImgPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/transformers/transformer_...
FAILED tests/pipelines/flux/test_pipeline_flux_inpaint.py::FluxInpaintPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/transformers/transformer_...
FAILED tests/pipelines/hidream/test_pipeline_hidream.py::HiDreamImagePipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.Unsupported: dynamic shape operator: aten.bincount.default; to enable, set torch._dynamo.config.capture_dynamic_outp...
FAILED tests/pipelines/hunyuan_video/test_hunyuan_image2video.py::HunyuanVideoImageToVideoPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.Unsupported: Dynamic slicing on data-dependent value is not supported
FAILED tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py::HunyuanSkyreelsImageToVideoPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.Unsupported: Dynamic slicing on data-dependent value is not supported
FAILED tests/pipelines/hunyuan_video/test_hunyuan_video.py::HunyuanVideoPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.Unsupported: Dynamic slicing on data-dependent value is not supported
FAILED tests/pipelines/kandinsky/test_kandinsky_prior.py::KandinskyPriorPipelineFastTests::test_torch_compile_recompilation_and_graph_break - AttributeError: 'KandinskyPriorPipeline' object has no attribute 'transformer'
FAILED tests/pipelines/kandinsky2_2/test_kandinsky_prior.py::KandinskyV22PriorPipelineFastTests::test_torch_compile_recompilation_and_graph_break - AttributeError: 'KandinskyV22PriorPipeline' object has no attribute 'transformer'
FAILED tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py::KandinskyV22PriorEmb2EmbPipelineFastTests::test_torch_compile_recompilation_and_graph_break - AttributeError: 'KandinskyV22PriorEmb2EmbPipeline' object has no attribute 'transformer'
FAILED tests/pipelines/lumina2/test_pipeline_lumina2.py::Lumina2PipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.UserError: Could not guard on data-dependent expression Eq(u0 + 4, 0) (unhinted: Eq(u0 + 4, 0)).  (Size-like symbols...
FAILED tests/pipelines/mochi/test_mochi.py::MochiPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.Unsupported: dynamic shape operator: aten.nonzero.default; to enable, set torch._dynamo.config.capture_dynamic_outpu...
FAILED tests/pipelines/omnigen/test_pipeline_omnigen.py::OmniGenPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.UserError: Dynamic control flow is not supported at the moment. Please use functorch.experimental.control_flow.cond ...
FAILED tests/pipelines/pag/test_pag_sd3_img2img.py::StableDiffusion3PAGImg2ImgPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/transformers/transformer_...
FAILED tests/pipelines/sana/test_sana_controlnet.py::SanaControlNetPipelineFastTests::test_torch_compile_recompilation_and_graph_break - RuntimeError: Expected a 'cpu' device type for generator but found 'cuda'
FAILED tests/pipelines/shap_e/test_shap_e.py::ShapEPipelineFastTests::test_torch_compile_recompilation_and_graph_break - AttributeError: 'ShapEPipeline' object has no attribute 'transformer'
FAILED tests/pipelines/shap_e/test_shap_e_img2img.py::ShapEImg2ImgPipelineFastTests::test_torch_compile_recompilation_and_graph_break - AttributeError: 'ShapEImg2ImgPipeline' object has no attribute 'transformer'
FAILED tests/pipelines/stable_cascade/test_stable_cascade_combined.py::StableCascadeCombinedPipelineFastTests::test_torch_compile_recompilation_and_graph_break - AttributeError: 'StableCascadeCombinedPipeline' object has no attribute 'transformer'
FAILED tests/pipelines/stable_cascade/test_stable_cascade_decoder.py::StableCascadeDecoderPipelineFastTests::test_torch_compile_recompilation_and_graph_break - AttributeError: 'StableCascadeDecoderPipeline' object has no attribute 'transformer'
FAILED tests/pipelines/stable_cascade/test_stable_cascade_prior.py::StableCascadePriorPipelineFastTests::test_torch_compile_recompilation_and_graph_break - AttributeError: 'StableCascadePriorPipeline' object has no attribute 'transformer'
FAILED tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py::StableDiffusionAttendAndExcitePipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.UserError: Dynamic control flow is not supported at the moment. Please use functorch.experimental.control_flow.cond ...
FAILED tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py::StableDiffusion3Img2ImgPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/transformers/transformer_...
FAILED tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py::StableDiffusion3InpaintPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/transformers/transformer_...
FAILED tests/pipelines/stable_diffusion_sag/test_stable_diffusion_sag.py::StableDiffusionSAGPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/unets/unet_2d_condition.p...
FAILED tests/pipelines/text_to_video_synthesis/test_text_to_video_zero_sdxl.py::TextToVideoZeroSDXLPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/unets/unet_2d_condition.p...
FAILED tests/pipelines/unclip/test_unclip.py::UnCLIPPipelineFastTests::test_torch_compile_recompilation_and_graph_break - AttributeError: 'UnCLIPPipeline' object has no attribute 'transformer'
FAILED tests/pipelines/unclip/test_unclip_image_variation.py::UnCLIPImageVariationPipelineFastTests::test_torch_compile_recompilation_and_graph_break - AttributeError: 'UnCLIPImageVariationPipeline' object has no attribute 'transformer'
FAILED tests/pipelines/unidiffuser/test_unidiffuser.py::UniDiffuserPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/pipelines/unidiffuser/modeling_u...
FAILED tests/pipelines/wan/test_wan.py::WanPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/transformers/transformer_...
FAILED tests/pipelines/wan/test_wan_image_to_video.py::WanImageToVideoPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/transformers/transformer_...
FAILED tests/pipelines/wan/test_wan_video_to_video.py::WanVideoToVideoPipelineFastTests::test_torch_compile_recompilation_and_graph_break - torch._dynamo.exc.RecompileError: Recompiling function forward in /home/sayak/diffusers/src/diffusers/models/transformers/transformer_...
FAILED tests/pipelines/wuerstchen/test_wuerstchen_combined.py::WuerstchenCombinedPipelineFastTests::test_torch_compile_recompilation_and_graph_break - AttributeError: 'WuerstchenCombinedPipeline' object has no attribute 'transformer'
FAILED tests/pipelines/wuerstchen/test_wuerstchen_decoder.py::WuerstchenDecoderPipelineFastTests::test_torch_compile_recompilation_and_graph_break - AttributeError: 'WuerstchenDecoderPipeline' object has no attribute 'transformer'
FAILED tests/pipelines/wuerstchen/test_wuerstchen_prior.py::WuerstchenPriorPipelineFastTests::test_torch_compile_recompilation_and_graph_break - AttributeError: 'WuerstchenPriorPipeline' object has no attribute 'transformer'

41 failed, 137 passed, 7036 deselected, 69 warnings in 3360.64s (0:56:00)

I can fix some of the obvious ones. But for the rest, should we skip and revisit in a follow-up?~

tests/pipelines/test_pipelines_common.py

sayakpaul · 2025-04-14T07:59:07Z

@anijain2305 sorry to be bothering you again, but does this PR clear the compiler cache appropriately?

sayakpaul · 2025-04-14T11:39:29Z

Had a talk with @DN6 and we agreed to have a common TorchCompileMixinTests class that will be added to a select few model test classes that have usage. We will then reorganize our existing torch.compile() tests a bit and make them a part of the TorchCompileMixinTests class in later PR.

sayakpaul · 2025-04-18T04:16:03Z

@DN6 I think we could also create an issue thread after this PR is merged to get help from the community in testing the most important model classes.

sayakpaul · 2025-04-28T00:36:29Z

Failing test is unrelated.

sayakpaul added 2 commits March 17, 2025 14:03

test for better torch.compile stuff.

de30cba

fixes

f389a4d

sayakpaul requested review from DN6 and hlky March 17, 2025 09:00

sayakpaul added 2 commits March 18, 2025 21:04

Merge branch 'main' into test-better-torch-compile

6b05db6

Merge branch 'main' into test-better-torch-compile

e5543dc

anijain2305 reviewed Mar 20, 2025

View reviewed changes

tests/pipelines/test_pipelines_common.py Outdated Show resolved Hide resolved

anijain2305 reviewed Mar 20, 2025

View reviewed changes

tests/pipelines/test_pipelines_common.py Outdated Show resolved Hide resolved

anijain2305 approved these changes Mar 20, 2025

View reviewed changes

sayakpaul added 4 commits March 21, 2025 08:53

recompilation and graph break.

6791037

Merge branch 'main' into test-better-torch-compile

abd1f6c

Merge branch 'main' into test-better-torch-compile

1f797b4

Merge branch 'main' into test-better-torch-compile

d669340

hlky approved these changes Apr 9, 2025

View reviewed changes

Merge branch 'main' into test-better-torch-compile

c49a855

sayakpaul mentioned this pull request Apr 12, 2025

Rewrite AuraFlowPatchEmbed.pe_selection_index_based_on_dim to be torch.compile compatible #11297

Merged

6 tasks

sayakpaul added 2 commits April 14, 2025 11:47

Merge branch 'main' into test-better-torch-compile

c060ba0

Merge branch 'main' into test-better-torch-compile

e75a9de

sayakpaul commented Apr 14, 2025

View reviewed changes

tests/pipelines/test_pipelines_common.py Outdated Show resolved Hide resolved

clear compilation cache.

c7f153a

Merge branch 'main' into test-better-torch-compile

c74c9a8

sayakpaul added 4 commits April 14, 2025 17:14

Merge branch 'main' into test-better-torch-compile

1a934b2

change to modeling level test.

e0566e6

Merge branch 'main' into test-better-torch-compile

38c1d0d

allow running compilation tests during nightlies.

87d957d

sayakpaul added the torch.compile label Apr 15, 2025

sayakpaul added 2 commits April 15, 2025 19:32

Merge branch 'main' into test-better-torch-compile

a8184ef

Merge branch 'main' into test-better-torch-compile

fae8b6c

sayakpaul changed the title ~~[WIP][tests] add tests to check for graph breaks, recompilation, cuda syncs in pipelines during torch.compile()~~ [tests] add tests to check for graph breaks, recompilation, cuda syncs in pipelines during torch.compile() Apr 18, 2025

sayakpaul added 2 commits April 21, 2025 14:37

Merge branch 'main' into test-better-torch-compile

1749955

Merge branch 'main' into test-better-torch-compile

a07c63b

DN6 approved these changes Apr 27, 2025

View reviewed changes

Merge branch 'main' into test-better-torch-compile

f71c8f6

sayakpaul marked this pull request as ready for review April 27, 2025 15:11

sayakpaul merged commit aa5f5d4 into main Apr 28, 2025
28 of 29 checks passed

sayakpaul deleted the test-better-torch-compile branch April 28, 2025 00:36

sayakpaul mentioned this pull request Apr 28, 2025

[tests] help us test torch.compile() for impactful models #11430

Open

5 tasks

yao-matrix mentioned this pull request Apr 28, 2025

[tests] fix import. #11434

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[tests] add tests to check for graph breaks, recompilation, cuda syncs in pipelines during torch.compile() #11085

[tests] add tests to check for graph breaks, recompilation, cuda syncs in pipelines during torch.compile() #11085

sayakpaul commented Mar 17, 2025 •

edited

Loading

anijain2305 left a comment •

edited

Loading

sayakpaul commented Apr 9, 2025

hlky left a comment

sayakpaul commented Apr 9, 2025

sayakpaul commented Apr 14, 2025 •

edited

Loading

sayakpaul commented Apr 14, 2025

sayakpaul commented Apr 14, 2025

sayakpaul commented Apr 18, 2025

sayakpaul commented Apr 28, 2025

[tests] add tests to check for graph breaks, recompilation, cuda syncs in pipelines during torch.compile() #11085

[tests] add tests to check for graph breaks, recompilation, cuda syncs in pipelines during torch.compile() #11085

Conversation

sayakpaul commented Mar 17, 2025 • edited Loading

What does this PR do?

anijain2305 left a comment • edited Loading

Choose a reason for hiding this comment

sayakpaul commented Apr 9, 2025

hlky left a comment

Choose a reason for hiding this comment

sayakpaul commented Apr 9, 2025

sayakpaul commented Apr 14, 2025 • edited Loading

sayakpaul commented Apr 14, 2025

sayakpaul commented Apr 14, 2025

sayakpaul commented Apr 18, 2025

sayakpaul commented Apr 28, 2025

sayakpaul commented Mar 17, 2025 •

edited

Loading

anijain2305 left a comment •

edited

Loading

sayakpaul commented Apr 14, 2025 •

edited

Loading