From a99cae3800708e33669e1e4896f96ff0733d5550 Mon Sep 17 00:00:00 2001 From: binchengxiong Date: Tue, 21 Jul 2026 11:12:56 +0800 Subject: [PATCH] feat: add compiled_region option to torch.compile the VAE encoder/decoder Add a 'compiled_region' knob to CompileConfig ('language' | 'all'). 'language' keeps the existing behavior (compile only the MoT TransformerBlocks); 'all' additionally applies torch.compile to the WanVAE_ Encoder3d/Decoder3d (CausalConv3d layers), which the existing apply_compile path never touches since it only wraps the _encode_vision/_decode_vision post-processing functions. The VAE is compiled with the torch.compile defaults (fullgraph=False, dynamic=None): the causal-conv feat_cache list mutation creates graph breaks, so fullgraph must stay off; and with automatic dynamic shapes the chunk temporal size only takes two values (1 frame for prime, temporal_window frames for steady state), so compilation stabilizes after at most a couple of recompiles without forcing dynamic=True. --- .../configs/toml_config/sft_config.py | 10 +++- .../model/generator/omni_mot_model.py | 58 +++++++++++++++++++ examples/toml/sft_config/vision_sft_nano.toml | 1 + 3 files changed, 68 insertions(+), 1 deletion(-) diff --git a/cosmos_framework/configs/toml_config/sft_config.py b/cosmos_framework/configs/toml_config/sft_config.py index 04d0efbc..45bc5a56 100644 --- a/cosmos_framework/configs/toml_config/sft_config.py +++ b/cosmos_framework/configs/toml_config/sft_config.py @@ -11,7 +11,7 @@ from __future__ import annotations from pathlib import Path -from typing import Any, Optional +from typing import Any, Literal, Optional import tomllib from pydantic import BaseModel, ConfigDict, Field @@ -201,6 +201,14 @@ class CompileConfig(BaseModel): "callback's progressive warmup." ), ) + compiled_region: Literal["language", "all"] = Field( + default="language", + description=( + "Which part of the network to torch.compile. " + "'language' compiles only the MoT TransformerBlocks; " + "'all' also compiles the VAE encoder/decoder heads (VFM only)." + ), + ) class ActivationCheckpointingConfig(BaseModel): diff --git a/cosmos_framework/model/generator/omni_mot_model.py b/cosmos_framework/model/generator/omni_mot_model.py index 1a7b38d3..48b0d833 100644 --- a/cosmos_framework/model/generator/omni_mot_model.py +++ b/cosmos_framework/model/generator/omni_mot_model.py @@ -380,10 +380,68 @@ def set_up_model(self): self.net_ema_worker.copy_to(src_model=self.net, tgt_model=self.net_ema) + # Compile VAE 3D-Conv encoder/decoder when compiled_region="all". + # Unlike apply_compile (which wraps _encode_vision post-processing), + # this directly compiles the CausalConv3d layers inside WanVAE_. + if config.compile.enabled and config.compile.compiled_region == "all": + self._compile_vae() + self.set_up_memory() torch.cuda.empty_cache() + def _compile_vae(self) -> None: + """Apply torch.compile to the VAE 3D-Conv encoder and decoder. + + The existing ``apply_compile`` in ``parallelize_vfm_network`` wraps + ``_encode_vision`` / ``_decode_vision`` — which are post-processing + functions (patchify + Linear + pos_embed) that contain **no** 3D Conv. + The real VAE encoder/decoder (``WanVAE_.encoder`` / ``WanVAE_.decoder``) + with ``CausalConv3d`` layers is never touched by that path. + + This method directly wraps ``Encoder3d`` and ``Decoder3d`` with + ``torch.compile``. + + Key differences from ``apply_compile``: + - ``fullgraph=False`` (default): the causal-conv caching logic uses + Python list mutation (``feat_cache[idx] = cache_x``) which creates + graph breaks; ``fullgraph=False`` allows dynamo to compile the + sub-graphs between breaks. + - ``dynamic=None`` (default, automatic dynamic shapes): the first + call compiles a static-shape kernel; when a different temporal + size shows up, dynamo recompiles once with dynamic shapes. Chunk + temporal sizes here only take two values (1 frame for prime, + ``temporal_window`` frames for steady state), so compilation + stabilizes after at most a couple of recompiles — no need to + force ``dynamic=True``. + - VAE is ``@torch.no_grad()`` and ``requires_grad_(False)``, so no + backward graph is involved. + """ + vae = self.tokenizer_vision_gen # Wan2pt2VAEInterface + wan_vae = getattr(vae, "model", None) # WanVAE (plain class) + if wan_vae is None: + log.warning("VAE tokenizer has no .model attribute, skipping VAE compile") + return + inner = getattr(wan_vae, "model", None) # WanVAE_ (nn.Module) + if inner is None: + log.warning("WanVAE has no .model (WanVAE_) attribute, skipping VAE compile") + return + + compiled_parts = [] + if hasattr(inner, "encoder"): + log.info("Applying torch.compile to VAE encoder (Encoder3d)") + inner.encoder = torch.compile(inner.encoder) + compiled_parts.append("encoder") + if hasattr(inner, "decoder"): + log.info("Applying torch.compile to VAE decoder (Decoder3d)") + inner.decoder = torch.compile(inner.decoder) + compiled_parts.append("decoder") + + if compiled_parts: + log.info(f"VAE torch.compile applied to: {', '.join(compiled_parts)}") + else: + log.warning("VAE model has no .encoder/.decoder, skipping VAE compile") + def install_attention_dispatch(self, net: torch.nn.Module) -> None: """Install a custom attention dispatch function on the network. diff --git a/examples/toml/sft_config/vision_sft_nano.toml b/examples/toml/sft_config/vision_sft_nano.toml index d6fd8b33..2b695a6a 100644 --- a/examples/toml/sft_config/vision_sft_nano.toml +++ b/examples/toml/sft_config/vision_sft_nano.toml @@ -30,6 +30,7 @@ data_parallel_replicate_degree = 1 [model.compile] enabled = true # was [model.parallelism].use_torch_compile compile_dynamic = true +compiled_region = "all" [model.activation_checkpointing] mode = "full"