[tests] add tests to check for graph breaks and recompilation in pipelines during torch.compile()

[sayakpaul]

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.

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

@DN6@hlky a gentle ping.

[hlky]

Looks good, thanks for working on this.

[sayakpaul]

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

[sayakpaul]

~@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 functionforwardin /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 functionforwardin /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 functionforwardin /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 functionforwardin /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 functionforwardin /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 functionforwardin /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 functionforwardin /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 functionforwardin /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 functionforwardin /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 supportedFAILED 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 supportedFAILED 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 supportedFAILED 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 functionforwardin /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 typefor 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 functionforwardin /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 functionforwardin /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 functionforwardin /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 functionforwardin /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 functionforwardin /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 functionforwardin /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 functionforwardin /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 functionforwardin /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?~

[sayakpaul]

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

[sayakpaul]

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]

@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]

Failing test is unrelated.