[rfc][compile] compile method for DiffusionPipeline #11705
Conversation
anijain2305
force-pushed
the
compile_utils
branch
from
59f1f8c to
e825022
Compare
| @@ -286,6 +286,8 @@ def __init__( | |||
|
|
|||
| self.gradient_checkpointing = False | |||
|
|
|||
| self.compile_region_classes = (FluxTransformerBlock, FluxSingleTransformerBlock) | |||
There was a problem hiding this comment.
Should this be a class-level attribute like this?
There was a problem hiding this comment.
Yes, thats better.
| @@ -2027,6 +2027,40 @@ def _maybe_raise_error_if_group_offload_active( | |||
| return True | |||
| return False | |||
|
|
|||
| def compile( | |||
There was a problem hiding this comment.
We could make it a part of ModelMixin rather than DiffusionPipeline I believe:
For most cases, users want to just do pipe.transformer.compile(). So, perhaps easier with this being added to ModelMixin?
There was a problem hiding this comment.
Do you expect any improvements from compile on the pre/post transformer blocks? If not, then yeah, moving it inside makes sense to me.
There was a problem hiding this comment.
With this change, do we have to do pipe.transformer.compile()? Hmm, this could be confusing because torch.nn.Module also has a method named compile. If we move in the ModelMixin, we might want a different method name.
There was a problem hiding this comment.
Well, to tackle those cases, we could do something like this:
diffusers/src/diffusers/models/modeling_utils.py
Line 1306 in 62cbde8
So, we include all the acceptable args and kwargs in the ModelMixin compile()` method but additionally include the regional/hierarchical compilation related kwargs. Would that work?
There was a problem hiding this comment.
Nice! Let me try this out.
There was a problem hiding this comment.
I feel like there are some advantages of Pipeline.compile:
- It captures the user intent better, e.g., "I just want to accelerate this pipeline" (the word "compile" is also great because it signals "expect some initial delay").
- It's more future proof (diffuser devs get to tweak things to accommodate changes from diffusers/pytorch versions, or even future hardware and pipeline architecture/perf-characteristics).
- It's more user friendly (small things like
.transformerdoes make a noticeable effect on user experience).
(2) and (3) are effectively consequences of (1). What do you think? @sayakpaul
There was a problem hiding this comment.
Sayak suggested to phase this out.
Have pipeline.transformer.compile as the first PR
And then pipeline.compile in future which can internally call pipeline.transformer.compile. Its actually better this way because then the compile region classes is also hidden from the pipe.
There was a problem hiding this comment.
Yes, we will eventually ship pipe.compile() :)
| compile_regions_for_transformer: bool = True, | ||
| transformer_module_name: str = "transformer", |
There was a problem hiding this comment.
These two will have to have changed if we proceed to have it inside ModelMixin.
|
thanks for the PR ! |
This is something that we will have to do manually. I deliberately made that choice. It should be be fairly easy to find those classes, basically the transformer classes. In absence of compile_region_classes, we can fallback to the full model compilation with a warning. I think this is a fair compromise because it gives control to the model author, and with little guidance it should be very easy. |
This PR adds a compile method to the DiffusionPipeline to make it really easy for diffusers users to apply torch.compile.
The default setting enables regional compilation, which for diffusion models give as good speedup as full model compilation, with 8x-10x smaller compilation time. More data to support regional compilation as default is below
To onboard a new model, we just have to define the
compile_region_classesattribute in the models, and that should allow regional compilation.cc @sayakpaul