From d21676efa728308a748e38aa7fc4833a07bb5a6a Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 26 Apr 2026 17:37:40 -0700 Subject: [PATCH] feat(compile): per-component compile + transformer_refine + prepare hook MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bring the public ``ComposedPipelineBase.post_init`` compile dispatch to feature parity with FastVideo-internal: * Compile ``transformer_refine`` alongside ``transformer`` and ``transformer_2`` whenever the DiT compile flag is on. Without this the LTX-2 stage-2 refine pass silently runs eager while the main transformer is compiled — different inductor fusion choices vs. eager execution can produce last-bit divergences in bf16. * Drive the new per-component compile flags from ``FastVideoArgs``: text encoder (``enable_torch_compile_text_encoder``), VAE (``enable_torch_compile_vae``), audio VAE (``enable_torch_compile_audio_vae``). Each picks its own kwargs dict (``torch_compile_kwargs_*``) and falls back to the master ``torch_compile_kwargs`` when empty. * Call ``module.prepare_for_compile()`` on each compiled submodule before invoking ``torch.compile`` so model-specific external state (e.g. lazy HF loads) lands outside Dynamo's tracer. ``Gemma3`` implements this hook to materialize the Gemma weights up front. Co-authored-by: Junda (David) Su <90978028+Davids048@users.noreply.github.com> Co-authored-by: Matthew Noto <99706358+RandNMR73@users.noreply.github.com> Co-authored-by: XOR-op <17672363+XOR-op@users.noreply.github.com> Co-authored-by: Zhang Peiyuan <42993249+jzhang38@users.noreply.github.com> --- fastvideo/models/encoders/gemma.py | 4 + fastvideo/pipelines/composed_pipeline_base.py | 75 +++++++++++++++---- 2 files changed, 66 insertions(+), 13 deletions(-) diff --git a/fastvideo/models/encoders/gemma.py b/fastvideo/models/encoders/gemma.py index 6741a3473..59d13a6d7 100644 --- a/fastvideo/models/encoders/gemma.py +++ b/fastvideo/models/encoders/gemma.py @@ -361,6 +361,10 @@ class LTX2GemmaTextEncoderModel(TextEncoder): continue yield name, param + def prepare_for_compile(self) -> None: + # Load Gemma outside Dynamo so torch.compile does not trace HF file-system checks. + _ = self.gemma_model + @property def gemma_model(self) -> Gemma3ForConditionalGeneration: if self._gemma_model is None: diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index 079ea169a..91aa0f3c4 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -131,6 +131,11 @@ class ComposedPipelineBase(ABC): ) return + prepare_for_compile = getattr(module, "prepare_for_compile", None) + if callable(prepare_for_compile): + logger.info("Running prepare_for_compile for %s", module_name) + prepare_for_compile() + compiled_count = self._compile_with_conditions(module, compile_kwargs) if compiled_count > 0: logger.info( @@ -159,7 +164,11 @@ class ComposedPipelineBase(ABC): self.initialize_validation_pipeline(self.training_args) self.initialize_pipeline(self.fastvideo_args) - if self.fastvideo_args.enable_torch_compile: + compile_transformer = self.fastvideo_args.enable_torch_compile + compile_text_encoder = (self.fastvideo_args.enable_torch_compile_text_encoder) + compile_vae = self.fastvideo_args.enable_torch_compile_vae + compile_audio_vae = self.fastvideo_args.enable_torch_compile_audio_vae + if (compile_transformer or compile_text_encoder or compile_vae or compile_audio_vae): if self.fastvideo_args.training_mode: logger.info("Torch Compile enabled via FSDP loader for training; skipping additional pipeline compile") else: @@ -170,18 +179,58 @@ class ComposedPipelineBase(ABC): except Exception: # pragma: no cover - FSDP not always available fsdp_module_cls = None - compile_kwargs = self.fastvideo_args.torch_compile_kwargs or {} - self._maybe_compile_pipeline_module( - module_name="transformer", - fsdp_module_cls=fsdp_module_cls, - compile_kwargs=compile_kwargs, - ) - self._maybe_compile_pipeline_module( - module_name="transformer_2", - fsdp_module_cls=fsdp_module_cls, - compile_kwargs=compile_kwargs, - ) - logger.info("Torch Compile enabled for DiT") + global_compile_kwargs = (self.fastvideo_args.torch_compile_kwargs or {}) + dit_compile_kwargs = (self.fastvideo_args.torch_compile_kwargs_dit or global_compile_kwargs) + text_compile_kwargs = (self.fastvideo_args.torch_compile_kwargs_text_encoder or global_compile_kwargs) + vae_compile_kwargs = (self.fastvideo_args.torch_compile_kwargs_vae or global_compile_kwargs) + audio_vae_compile_kwargs = (self.fastvideo_args.torch_compile_kwargs_audio_vae or global_compile_kwargs) + + if compile_transformer: + self._maybe_compile_pipeline_module( + module_name="transformer", + fsdp_module_cls=fsdp_module_cls, + compile_kwargs=dit_compile_kwargs, + ) + self._maybe_compile_pipeline_module( + module_name="transformer_refine", + fsdp_module_cls=fsdp_module_cls, + compile_kwargs=dit_compile_kwargs, + ) + self._maybe_compile_pipeline_module( + module_name="transformer_2", + fsdp_module_cls=fsdp_module_cls, + compile_kwargs=dit_compile_kwargs, + ) + logger.info("Torch Compile enabled for DiT") + + if compile_text_encoder: + self._maybe_compile_pipeline_module( + module_name="text_encoder", + fsdp_module_cls=fsdp_module_cls, + compile_kwargs=text_compile_kwargs, + ) + self._maybe_compile_pipeline_module( + module_name="text_encoder_2", + fsdp_module_cls=fsdp_module_cls, + compile_kwargs=text_compile_kwargs, + ) + logger.info("Torch Compile enabled for text encoder") + + if compile_vae: + self._maybe_compile_pipeline_module( + module_name="vae", + fsdp_module_cls=fsdp_module_cls, + compile_kwargs=vae_compile_kwargs, + ) + logger.info("Torch Compile enabled for VAE") + + if compile_audio_vae: + self._maybe_compile_pipeline_module( + module_name="audio_vae", + fsdp_module_cls=fsdp_module_cls, + compile_kwargs=audio_vae_compile_kwargs, + ) + logger.info("Torch Compile enabled for audio VAE") if not self.fastvideo_args.training_mode: logger.info("Creating pipeline stages...")