From edd517e3fa76d4876860f9040e3cab41b9ca8571 Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Sat, 7 Nov 2026 15:55:36 -0800 Subject: [PATCH] [bugfix] Make torch.compile reduce-overhead work in the shared denoising loop (Wan T2V) (#2047) Co-authored-by: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> --- docs/inference/optimizations.md | 17 +- fastvideo/fastvideo_args.py | 30 +++ fastvideo/models/loader/fsdp_load.py | 4 + fastvideo/pipelines/composed_pipeline_base.py | 16 +- fastvideo/pipelines/stages/denoising.py | 9 +- .../stages/test_denoising_cudagraph_step.py | 250 ++++++++++++++++++ 6 files changed, 321 insertions(+), 5 deletions(-) create mode 100644 fastvideo/tests/stages/test_denoising_cudagraph_step.py diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index 2ad6edfe0..e926a249c 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -501,17 +501,28 @@ MS-SSIM gate on *your* config, especially when combining path in the regional support envelope; other VSA paths and the FA3 grad-enabled path remain outside it. Use the default inductor mode shown above unless your exact configuration has its own gate. +- **Generic compile with `mode="reduce-overhead"`**: when the DiT compile + records CUDA graphs, the shared denoising loop (`DenoisingStage`, used by + Wan T2V) calls `torch.compiler.cudagraph_mark_step_begin()` once per step, + as the MiniMax-H3 loop does, so graph replays do not overwrite tensors the + current step still reads. Turn off layerwise offload + (`dit_layerwise_offload=False`) as well: offload copies the DiT weights to + new GPU memory on every forward, so the graphs are recorded again each + step instead of replayed, which is slower than the default mode. The win + is largest when the DiT is launch-bound (small shapes, few tokens); at + 480x832 the step is GPU-bound and matches the default mode. Extra `torch.compile` options are passed through `torch_compile_kwargs` (a dict), accepted by `VideoGenerator.from_pretrained(...)` and by the -CLI as a JSON string via `--torch-compile-kwargs`. Example (currently -**not** recommended — see the CUDA-graphs caveat above): +CLI as a JSON string via `--torch-compile-kwargs`. Example (see the +CUDA-graphs notes above): ```python VideoGenerator.from_pretrained( "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", enable_torch_compile=True, - torch_compile_kwargs={"mode": "reduce-overhead"}, # may error today + torch_compile_kwargs={"mode": "reduce-overhead"}, + dit_layerwise_offload=False, ) ``` diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 7024fe716..133470b06 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -1252,6 +1252,36 @@ def get_current_fastvideo_args() -> FastVideoArgs: return _current_fastvideo_args +# Stamped on a DiT by whichever code path actually attaches the compile: +# ``ComposedPipelineBase._compile_pipeline_module_instance`` for the +# pipeline-level compile, ``_compile_model_regions`` for the regional one. +DIT_COMPILE_CUDAGRAPHS_ATTR = "_fastvideo_compile_records_cudagraphs" + + +def compile_kwargs_record_cudagraphs(compile_kwargs: dict[str, Any] | None) -> bool: + """Whether these ``torch.compile`` kwargs make inductor record CUDA graphs.""" + kwargs = compile_kwargs or {} + if kwargs.get("disable"): + return False + options = kwargs.get("options") or {} + return kwargs.get("mode") in ("reduce-overhead", "max-autotune") or bool(options.get("triton.cudagraphs")) + + +def dit_compile_uses_cudagraphs(module: Any) -> bool: + """Whether the compile attached to ``module`` records CUDA graphs (cudagraph trees). + + Reads the stamp left by the code path that compiled this module, so the + answer never depends on which flag asked for the compile. A request flag + cannot answer it: ``_compile_pipeline_module_instance`` skips the compile + only when the module *is* an ``FSDPModule``, so ``use_fsdp_inference=True`` + still compiles a transformer handed in through + ``from_pretrained(loaded_modules=...)``, or one whose model declares no + shard conditions; and ``inference_torch_compile`` reaches cudagraphs + through ``torch_compile_kwargs["options"]["triton.cudagraphs"]``. + """ + return bool(getattr(module, DIT_COMPILE_CUDAGRAPHS_ATTR, False)) + + @dataclasses.dataclass class TrainingArgs(FastVideoArgs): """ diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 4165d5c8e..ca8435cb6 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -865,6 +865,10 @@ def _compile_model_regions(model: nn.Module, compile_kwargs: dict[str, Any], *, if compiled_count == 0: raise ValueError(f"No submodules in {type(model).__name__} matched _compile_conditions") + # Regional compile rejects ``mode`` above but merges the caller's + # ``options``, so {"triton.cudagraphs": True} records graphs here too. + from fastvideo.fastvideo_args import DIT_COMPILE_CUDAGRAPHS_ATTR, compile_kwargs_record_cudagraphs + setattr(model, DIT_COMPILE_CUDAGRAPHS_ATTR, compile_kwargs_record_cudagraphs(kwargs)) logger.info( "Enabled regional torch.compile for %d submodules in %s after FSDP setup with kwargs=%s", compiled_count, diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index faea0dc96..24d579adf 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -21,7 +21,8 @@ from fastvideo.distributed import ( maybe_init_distributed_environment_and_model_parallel, ) from fastvideo.distributed.communication_op import (warmup_sequence_parallel_communication) -from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs +from fastvideo.fastvideo_args import (DIT_COMPILE_CUDAGRAPHS_ATTR, FastVideoArgs, TrainingArgs, + compile_kwargs_record_cudagraphs) from fastvideo.hooks.activation_trace import attach_activation_trace, detach_activation_trace from fastvideo.logger import init_logger from fastvideo.profiler import get_or_create_profiler @@ -222,6 +223,11 @@ class ComposedPipelineBase(ABC): logger.info("Running prepare_for_compile for %s", module_name) prepare_for_compile() + # The FSDP check above is an instance test, so use_fsdp_inference=True + # still reaches here for a transformer that was never wrapped. Stamp + # before any wrapper goes on, so the loop reads the compile off the DiT. + setattr(module, DIT_COMPILE_CUDAGRAPHS_ATTR, compile_kwargs_record_cudagraphs(compile_kwargs)) + compiled_count = ComposedPipelineBase._compile_with_conditions(module, compile_kwargs) if compiled_count > 0: logger.info( @@ -327,6 +333,14 @@ class ComposedPipelineBase(ABC): ) if any(name in self.modules for name in ("transformer", "transformer_refine", "transformer_2")): logger.info("Torch Compile enabled for DiT") + if (self.fastvideo_args.dit_layerwise_offload and compile_kwargs_record_cudagraphs(dit_compile_kwargs) + and not getattr(self, "_warned_cudagraph_offload", False)): + self._warned_cudagraph_offload = True + logger.warning( + "The DiT torch.compile records CUDA graphs (%s), but dit_layerwise_offload copies the DiT " + "weights to new GPU memory on every forward, so the graphs are recorded again each step " + "instead of replayed (slower than the default mode). Set dit_layerwise_offload=False.", + dit_compile_kwargs) if compile_text_encoder and any(_want(name) for name in ("text_encoder", "text_encoder_2")): for name in ("text_encoder", "text_encoder_2"): diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index acdc478ae..4e8faa0ee 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -16,7 +16,7 @@ import fastvideo.envs as envs from fastvideo.attention import get_attn_backend from fastvideo.attention.selector import component_attention_backend from fastvideo.distributed import (get_local_torch_device, get_world_group) -from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.fastvideo_args import FastVideoArgs, dit_compile_uses_cudagraphs from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger from fastvideo.models.loader.component_loader import TransformerLoader @@ -309,6 +309,13 @@ class DenoisingStage(PipelineStage): continue current_model, current_guidance_scale = self.select_model(t, batch, fastvideo_args, state) + # Without a mark, cudagraph trees start a new generation on + # every compiled call, so a replay can overwrite outputs the + # step still reads. Asked of the selected model, not the args: + # Wan 2.2 swaps DiTs mid-loop. select_model moves devices but + # never calls the model, so the mark still precedes the call. + if dit_compile_uses_cudagraphs(current_model): + torch.compiler.cudagraph_mark_step_begin() latent_model_input = self.prepare_model_input(latents, batch, target_dtype, state) assert not torch.isnan(latent_model_input).any(), "latent_model_input contains nan" t_expand = self.prepare_timestep(t, latent_model_input, batch, state) diff --git a/fastvideo/tests/stages/test_denoising_cudagraph_step.py b/fastvideo/tests/stages/test_denoising_cudagraph_step.py new file mode 100644 index 000000000..2cc199d8f --- /dev/null +++ b/fastvideo/tests/stages/test_denoising_cudagraph_step.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 +"""DenoisingStage under torch.compile(mode="reduce-overhead"); no weights required. + +The CPU tests check that each denoising step opens a new cudagraph-trees step +only when the DiT compile records CUDA graphs, and that such a compile warns +under layerwise offload. The CUDA test compiles each block of a tiny DiT the +way ``_compile_conditions`` compiles Wan blocks and runs the shared CFG loop. +""" + +import pytest +import torch + +from fastvideo.tests.stages._denoising_fixtures import ( + NullProgressBar, TinyDenoiser, TinyScheduler, _patch_denoising_module, _tiny_args, _tiny_batch, +) + + +def _compile_args(args, compile_kwargs, **overrides): + args.enable_torch_compile = compile_kwargs is not None + args.torch_compile_kwargs = compile_kwargs or {} + args.training_mode = False + args.inference_torch_compile = False + args.use_fsdp_inference = False + for key, value in overrides.items(): + setattr(args, key, value) + return args + + +class _FakeFSDPModule: + """Stands in for torch.distributed.fsdp.FSDPModule, which needs no GPU here.""" + + +def _compiled_dit(monkeypatch, compile_kwargs, *, fsdp_wrapped=False): + """A DiT put through the real pipeline compile path, so it carries its stamp. + + Going through ``_compile_pipeline_module_instance`` rather than setting + flags is the point of these cases: the mark has to follow the compile that + was actually attached, and that function decides from an ``isinstance`` + check, not from ``use_fsdp_inference``. + """ + from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase + + monkeypatch.setattr(torch, "compile", lambda target, **kwargs: target) + model = TinyDenoiser() + if fsdp_wrapped: + model.__class__ = type("FSDPTinyDenoiser", (TinyDenoiser, _FakeFSDPModule), {}) + return ComposedPipelineBase._compile_pipeline_module_instance("transformer", model, _FakeFSDPModule, + compile_kwargs or {}) + + +@pytest.mark.parametrize("compile_kwargs, fsdp_wrapped, expected_marks", [ + ({"mode": "reduce-overhead"}, False, [0, 2, 4, 6]), + ({"mode": "max-autotune"}, False, [0, 2, 4, 6]), + ({"options": {"triton.cudagraphs": True}}, False, [0, 2, 4, 6]), + ({"mode": "default"}, False, []), + ({"mode": "max-autotune-no-cudagraphs"}, False, []), + ({"mode": "reduce-overhead", "disable": True}, False, []), + # Really FSDP-wrapped: the pipeline skips the compile, so nothing to mark. + ({"mode": "reduce-overhead"}, True, []), +]) +def test_denoising_marks_one_cudagraph_step_per_timestep(monkeypatch, env_overrides, compile_kwargs, fsdp_wrapped, + expected_marks): + denoising, _ = _patch_denoising_module(monkeypatch, env_overrides, None) + marks = [] + model = _compiled_dit(monkeypatch, compile_kwargs, fsdp_wrapped=fsdp_wrapped) + monkeypatch.setattr(torch.compiler, "cudagraph_mark_step_begin", lambda: marks.append(len(model.calls))) + stage = denoising.DenoisingStage(model, TinyScheduler()) + stage.progress_bar = lambda iterable=None, total=None: NullProgressBar() + + stage.forward(_tiny_batch(), _compile_args(_tiny_args(), compile_kwargs)) + + # With CUDA graphs: one mark before each step's cond+uncond pair, never + # between them. Without: no mark, so other cudagraph users keep + # automatic step boundaries. + assert model.calls == ["cond", "uncond"] * 4 + assert marks == expected_marks + + +def test_marks_follow_the_attached_compile_not_the_fsdp_flag(monkeypatch, env_overrides): + """Regression: use_fsdp_inference=True on a DiT the pipeline still compiled. + + ``from_pretrained(loaded_modules=...)`` hands the pipeline a plain module, + and a model that declares no shard conditions is never wrapped either. In + both cases the FSDP instance check does not fire, the DiT is compiled with + CUDA graphs, and the flag says otherwise. Marking has to happen anyway, or + the crash this guards against comes back. + """ + denoising, _ = _patch_denoising_module(monkeypatch, env_overrides, None) + marks = [] + model = _compiled_dit(monkeypatch, {"mode": "reduce-overhead"}, fsdp_wrapped=False) + monkeypatch.setattr(torch.compiler, "cudagraph_mark_step_begin", lambda: marks.append(len(model.calls))) + stage = denoising.DenoisingStage(model, TinyScheduler()) + stage.progress_bar = lambda iterable=None, total=None: NullProgressBar() + + args = _compile_args(_tiny_args(), {"mode": "reduce-overhead"}) + args.use_fsdp_inference = True # the flag says no compile; the module says otherwise + + stage.forward(_tiny_batch(), args) + + assert marks == [0, 2, 4, 6] + + +class _Block(torch.nn.Module): + + def __init__(self, dim): + super().__init__() + self.proj = torch.nn.Linear(dim, dim) + + def forward(self, hidden_states, context, temb): + return hidden_states + torch.tanh(self.proj(hidden_states) + context.mean(dim=1, keepdim=True) + temb) + + +def _is_tiny_block(name: str, submodule: torch.nn.Module) -> bool: + return isinstance(submodule, _Block) + + +class _TinyBlockDiT(torch.nn.Module): + """Eager embed/head around per-block compiled blocks, like Wan.""" + hidden_size = 1 + num_attention_heads = 1 + # Declared so _compile_with_conditions compiles the blocks rather than the + # whole module, which is the shape this test is about. + _compile_conditions = [_is_tiny_block] + + def __init__(self, dim=32, num_blocks=3): + super().__init__() + self.config = type("Cfg", (), {"use_meanflow": False})() + self.embed = torch.nn.Linear(2, dim) + self.time = torch.nn.Linear(1, dim) + self.blocks = torch.nn.ModuleList(_Block(dim) for _ in range(num_blocks)) + self.head = torch.nn.Linear(dim, 2) + + def forward(self, latent_model_input, prompt_embeds, timestep, guidance=None): + b, c, f, h, w = latent_model_input.shape + x = self.embed(latent_model_input.flatten(2).transpose(1, 2).float()) + temb = self.time(timestep.float().reshape(b, 1, 1) / 1000) + context = prompt_embeds[0].float() + for block in self.blocks: + x = block(x, context, temb) + return self.head(x).transpose(1, 2).reshape(b, c, f, h, w) + + +def _run_compiled_stage(monkeypatch, env_overrides, mode): + from fastvideo.tests.stages._denoising_fixtures import _args, _batch + + denoising, _ = _patch_denoising_module(monkeypatch, env_overrides, None) + device = torch.device("cuda") + monkeypatch.setattr(denoising, "get_local_torch_device", lambda: device) + torch.manual_seed(0) + model = _TinyBlockDiT().to(device) + # Compile through the pipeline's own entry point rather than calling + # torch.compile on the blocks here. That is what production does, and it is + # also what stamps the DiT with what the compile records; hand-rolling the + # per-block compile skips the stamp, so the loop would not mark and this + # test would reproduce the very crash it guards against. + from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase + + model = ComposedPipelineBase._compile_pipeline_module_instance( + "transformer", model, None, {"mode": mode, "dynamic": False} + ) + stage = denoising.DenoisingStage(model, TinyScheduler()) + stage.progress_bar = lambda iterable=None, total=None: NullProgressBar() + outputs = [] + for _ in range(3): + batch = _batch(steps=6, cfg=True) + batch.latents = batch.latents.to(device) + batch.prompt_embeds = [torch.full((1, 4, 32), 0.5, device=device)] + batch.negative_prompt_embeds = [torch.full((1, 4, 32), -0.5, device=device)] + with torch.inference_mode(): + outputs.append(stage.forward(batch, _compile_args(_args(), {"mode": mode})).latents.clone()) + return outputs + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graphs need a CUDA device") +def test_reduce_overhead_cfg_loop_matches_default_compile(monkeypatch, env_overrides): + torch._dynamo.reset() + try: + reference = _run_compiled_stage(monkeypatch, env_overrides, "default") + torch._dynamo.reset() + graphed = _run_compiled_stage(monkeypatch, env_overrides, "reduce-overhead") + finally: + torch._dynamo.reset() + + for ref, out in zip(reference, graphed, strict=True): + assert torch.equal(ref, out) + + +@pytest.mark.parametrize("mode, offload, warned", [ + ("reduce-overhead", True, True), + ("max-autotune", True, True), + ("reduce-overhead", False, False), + ("default", True, False), +]) +def test_cudagraph_compile_warns_when_layerwise_offload_moves_weights(monkeypatch, mode, offload, warned): + from types import SimpleNamespace + + import fastvideo.pipelines.composed_pipeline_base as base + from fastvideo.tests.stages._denoising_fixtures import RecordingLogger + + class _Pipeline(base.ComposedPipelineBase): + + def __init__(self, fastvideo_args): + self.fastvideo_args = fastvideo_args + self.modules = {"transformer": torch.nn.Linear(2, 2)} + + def create_pipeline_stages(self, fastvideo_args): + raise NotImplementedError + + logger = RecordingLogger() + monkeypatch.setattr(base, "logger", logger) + monkeypatch.setattr(torch, "compile", lambda target, **kwargs: target) + args = SimpleNamespace(enable_torch_compile=True, enable_torch_compile_text_encoder=False, + enable_torch_compile_vae=False, enable_torch_compile_audio_vae=False, training_mode=False, + inference_torch_compile=False, torch_compile_kwargs={}, torch_compile_kwargs_dit={"mode": mode}, + torch_compile_kwargs_text_encoder={}, torch_compile_kwargs_vae={}, + torch_compile_kwargs_audio_vae={}, dit_layerwise_offload=offload) + + pipeline = _Pipeline(args) + pipeline._apply_inference_compile(("vae", )) + assert logger.warnings == [] + # MiniMax-H3 sequential loading calls this again for late modules; warn once. + pipeline._apply_inference_compile() + pipeline._apply_inference_compile(("transformer", )) + + assert len([msg for msg in logger.warnings if "dit_layerwise_offload" in msg]) == (1 if warned else 0) + + +@pytest.mark.parametrize("kwargs, records", [ + ({"options": {"triton.cudagraphs": True}}, True), + ({"fullgraph": True}, False), +]) +def test_regional_compile_stamps_what_it_recorded(monkeypatch, kwargs, records): + """The regional path reaches cudagraphs through options, so it must stamp too. + + _compile_model_regions rejects `mode` but merges the caller's `options`, + which is how {"triton.cudagraphs": True} gets to inductor on the + inference_torch_compile path. Before the stamp, the helper read False there + and the loop never marked. + """ + from fastvideo.fastvideo_args import DIT_COMPILE_CUDAGRAPHS_ATTR, dit_compile_uses_cudagraphs + from fastvideo.models.loader.fsdp_load import _compile_model_regions + + monkeypatch.setattr(torch, "compile", lambda target, **kw: target) + model = _TinyBlockDiT() + + compiled = _compile_model_regions(model, dict(kwargs)) + + assert compiled == 3, "the tiny DiT's blocks should match _compile_conditions" + assert getattr(model, DIT_COMPILE_CUDAGRAPHS_ATTR) is records + assert dit_compile_uses_cudagraphs(model) is records