nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(compile): per-component compile + transformer_refine + prepare hook

Browse Source 
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>
This commit is contained in:
SolitaryThinker
2026-05-07 13:37:25 -07:00
co-authored by Junda Su Matthew Noto XOR-op Zhang Peiyuan
parent 82e532727f
commit d21676efa7
2 changed files with 66 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/models/encoders/gemma.py
+4
View File
@@ -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:
fastvideo/pipelines/composed_pipeline_base.py
+62 -13
View File
@@ -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...")