diff --git a/docs/.nav.yml b/docs/.nav.yml index 6bf161b1b..e7157b573 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -56,6 +56,7 @@ nav: - Diffusion Features: - Overview: user_guide/diffusion_features.md - Feature Compatibility: user_guide/feature_compatibility.md + - Regional Compilation: user_guide/diffusion/regional_compilation.md - Cache Acceleration: - TeaCache: user_guide/diffusion/cache_acceleration/teacache.md - Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md diff --git a/docs/user_guide/diffusion/regional_compilation.md b/docs/user_guide/diffusion/regional_compilation.md new file mode 100644 index 000000000..469a62b2c --- /dev/null +++ b/docs/user_guide/diffusion/regional_compilation.md @@ -0,0 +1,38 @@ +# Regional Compilation + +Regional compilation applies `torch.compile` to the repeated transformer blocks +declared by a diffusion model. It is the default compilation scope when +diffusion inference runs without `--enforce-eager`. + +## Configuration + +Dynamic compilation is enabled by default so the compiled regions can handle +mixed resolutions. For a fixed-shape workload, disable it explicitly: + +```bash +vllm serve --omni --no-diffusion-compile-dynamic +``` + +The equivalent per-stage deploy configuration is: + +```yaml +stages: + - stage_id: 0 + diffusion_compile_granularity: regional + diffusion_compile_dynamic: false +``` + +For an experimental whole-transformer compile scope, set +`--diffusion-compile-granularity full` or use +`diffusion_compile_granularity: full` in the stage configuration. Full scope may +still contain graph breaks; it does not force one graph. It is rejected when +HSDP, sequence parallelism, CPU offload, or layerwise offload is enabled. Use +regional scope with those features. + +These settings control the generic model-runner compilation path. Pipelines +that provide their own `setup_compile()` implementation manage their compilation +policy independently. Compilation is lazy, so backend or graph errors can first +surface on the initial request. + +Use `--enforce-eager` to disable the model runner's generic compile setup. +Pipelines that compile internally define their own eager-mode behavior. diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 523542232..f8fb3ed7f 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -628,6 +628,47 @@ def test_from_pipeline_config_preserves_diffusion_parallel_mask_sp_padding(tmp_p assert stage.parallel_config.mask_sp_padding is True +def test_from_pipeline_config_routes_regional_compile_dynamic(tmp_path): + deploy_path = tmp_path / "dreamzero_compile.yaml" + deploy_path.write_text( + "\n".join( + [ + "pipeline: dreamzero", + "async_chunk: false", + "stages:", + " - stage_id: 0", + " diffusion_compile_granularity: regional", + " diffusion_compile_dynamic: false", + ] + ) + ) + + configured_stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0) + overridden_stage = _from_pipeline_key( + "dreamzero", + deploy_config_path=str(deploy_path), + cli_overrides={ + "diffusion_compile_granularity": "full", + "diffusion_compile_dynamic": True, + }, + ).stage_by_id(0) + + assert configured_stage.diffusion_config.diffusion_compile_granularity == "regional" + assert configured_stage.diffusion_config.diffusion_compile_dynamic is False + assert overridden_stage.diffusion_config.diffusion_compile_granularity == "full" + assert overridden_stage.diffusion_config.diffusion_compile_dynamic is True + + +def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic(): + with pytest.raises(ValidationError, match="diffusion_compile_dynamic"): + omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false") + + +def test_structured_diffusion_config_rejects_invalid_compile_granularity(): + with pytest.raises(ValidationError, match="diffusion_compile_granularity"): + omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block") + + def test_from_pipeline_config_matches_stage_config_to_omegaconf_behavior_for_representative_stage(): pipeline = _resolve_pipeline_or_skip("qwen3_tts") legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy(pipeline))[0] diff --git a/tests/diffusion/test_compile.py b/tests/diffusion/test_compile.py index ca898da95..447d8c75f 100644 --- a/tests/diffusion/test_compile.py +++ b/tests/diffusion/test_compile.py @@ -54,3 +54,23 @@ def test_regionally_compile_matches_wrapped_blocks_by_declared_container_attr(mo assert all(not block.compile_called for block in model.transformer_blocks) assert not model.other_blocks[0].compile_called assert model.transformer_blocks[0].forward("ok") == "compiled:ok" + + +def test_regionally_compile_does_not_partially_mutate_on_setup_failure(monkeypatch): + model = _ModelWithWrappedRepeatedBlocks() + original_forwards = [block.forward.__func__ for block in model.transformer_blocks] + compile_calls = 0 + + def _compile(fn, *args, **kwargs): + nonlocal compile_calls + compile_calls += 1 + if compile_calls == 2: + raise RuntimeError("compile setup failed") + return lambda *fn_args, **fn_kwargs: fn(*fn_args, **fn_kwargs) + + monkeypatch.setattr(compile_module.torch, "compile", _compile) + + with pytest.raises(RuntimeError, match="compile setup failed"): + regionally_compile(model, dynamic=True) + + assert [block.forward.__func__ for block in model.transformer_blocks] == original_forwards diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index 01c901d60..e44766feb 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -392,6 +392,57 @@ class TestRequestBatchCapability: warmup.assert_called_once_with() +class TestDiffusionCompileConfig: + pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + def test_config_defaults_to_dynamic_regional_compile(self) -> None: + config = OmniDiffusionConfig(model="test") + + assert config.diffusion_compile_granularity == "regional" + assert config.diffusion_compile_dynamic is True + + def test_from_kwargs_preserves_compile_controls(self) -> None: + config = OmniDiffusionConfig.from_kwargs( + model="test", + diffusion_compile_granularity="full", + diffusion_compile_dynamic=False, + ) + + assert config.diffusion_compile_granularity == "full" + assert config.diffusion_compile_dynamic is False + + def test_config_rejects_invalid_compile_granularity(self) -> None: + with pytest.raises(ValueError, match="diffusion_compile_granularity"): + OmniDiffusionConfig(model="test", diffusion_compile_granularity="block") + + def test_config_rejects_non_boolean_compile_dynamic(self) -> None: + with pytest.raises(TypeError, match="diffusion_compile_dynamic"): + OmniDiffusionConfig(model="test", diffusion_compile_dynamic="false") + + @pytest.mark.parametrize( + "kwargs, feature", + [ + ({"parallel_config": {"ulysses_degree": 2}, "num_gpus": 2}, "sequence parallelism"), + ( + { + "parallel_config": {"use_hsdp": True, "hsdp_shard_size": 2}, + "num_gpus": 2, + }, + "HSDP", + ), + ({"enable_cpu_offload": True}, "CPU offload"), + ({"enable_layerwise_offload": True}, "layerwise offload"), + ], + ) + def test_full_compile_rejects_incompatible_features(self, kwargs, feature) -> None: + with pytest.raises(ValueError, match=feature): + OmniDiffusionConfig( + model="test", + diffusion_compile_granularity="full", + **kwargs, + ) + + class TestRequestBatchAdmission: pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index cf9f0d3b4..a7f230d93 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -121,6 +121,15 @@ class _FinalOnlyStepPipeline: return DiffusionOutput(output=state.latents.clone()) +class _CompileTrackingModel: + def __init__(self): + self.compile_calls = [] + + def compile(self, *args, **kwargs): + self.compile_calls.append((args, kwargs)) + return self + + def _make_request(): sampling_params = SimpleNamespace( generator=None, @@ -177,10 +186,20 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar return runner -def _make_compile_runner(*, use_hsdp: bool): +def _make_compile_runner( + model=None, + *, + compile_granularity: str = "regional", + compile_dynamic: bool = True, + use_hsdp: bool = False, +): runner = object.__new__(DiffusionModelRunner) - runner.pipeline = SimpleNamespace(transformer=SimpleNamespace()) - runner.od_config = SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=use_hsdp)) + runner.pipeline = SimpleNamespace(transformer=model or SimpleNamespace()) + runner.od_config = SimpleNamespace( + diffusion_compile_granularity=compile_granularity, + diffusion_compile_dynamic=compile_dynamic, + parallel_config=SimpleNamespace(use_hsdp=use_hsdp), + ) return runner @@ -208,6 +227,65 @@ def test_compile_transformer_regionally_compiles_blocks(monkeypatch, use_hsdp): ] +@pytest.mark.core_model +@pytest.mark.cpu +def test_compile_transformer_uses_regional_dynamic_false_config(monkeypatch): + model = _CompileTrackingModel() + runner = _make_compile_runner(model, compile_dynamic=False) + compiled_model = object() + regional_calls = [] + + def _regionally_compile(target, **kwargs): + regional_calls.append((target, kwargs)) + return compiled_model + + monkeypatch.setattr(model_runner_module, "regionally_compile", _regionally_compile) + + DiffusionModelRunner._compile_transformer(runner, "transformer") + + assert regional_calls == [(model, {"dynamic": False})] + assert model.compile_calls == [] + assert runner.pipeline.transformer is compiled_model + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_compile_transformer_uses_full_granularity(monkeypatch): + model = _CompileTrackingModel() + runner = _make_compile_runner(model, compile_granularity="full", compile_dynamic=False) + regional_calls = [] + + monkeypatch.setattr( + model_runner_module, + "regionally_compile", + lambda *args, **kwargs: regional_calls.append((args, kwargs)), + ) + + DiffusionModelRunner._compile_transformer(runner, "transformer") + + assert model.compile_calls == [((), {"dynamic": False})] + assert regional_calls == [] + assert runner.pipeline.transformer is model + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_compile_transformer_falls_back_after_synchronous_setup_failure(monkeypatch, caplog): + model = _CompileTrackingModel() + runner = _make_compile_runner(model) + + def _regionally_compile(*args, **kwargs): + raise RuntimeError("compile setup failed") + + monkeypatch.setattr(model_runner_module, "regionally_compile", _regionally_compile) + + DiffusionModelRunner._compile_transformer(runner, "transformer") + + assert runner.pipeline.transformer is model + assert "failed before activation" in caplog.text + assert "lazy compilation errors" in caplog.text + + @pytest.mark.core_model @pytest.mark.cpu def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch): diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index d6b33dc89..8987ae1d8 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -85,6 +85,8 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs(): "tensor_parallel_size": None, "enable_expert_parallel": None, "enforce_eager": None, + "diffusion_compile_granularity": None, + "diffusion_compile_dynamic": None, } )[0] @@ -94,6 +96,8 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs(): assert parallel_config.tensor_parallel_size == 1 assert parallel_config.enable_expert_parallel is False assert stage_cfg["engine_args"]["enforce_eager"] is False + assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "regional" + assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is True def test_default_stage_config_propagates_ulysses_mode(): @@ -258,6 +262,32 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag(): assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True +def test_serve_cli_accepts_diffusion_compile_controls(): + """Ensure both compile controls reach the diffusion stage.""" + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="command") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args( + [ + "serve", + "Lightricks/LTX-Video-0.9.8-13B-distilled", + "--omni", + "--diffusion-compile-granularity", + "full", + "--no-diffusion-compile-dynamic", + ] + ) + + explicit_kwargs = args.get_explicit_kwargs_dict() + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + + assert args.diffusion_compile_granularity == "full" + assert args.diffusion_compile_dynamic is False + assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "full" + assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is False + + def test_serve_cli_accepts_diffusion_attention_backend(): """Ensure diffusion serve CLI exposes the shorthand backend flag.""" parser = TrackingArgumentParser() diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index cdcbafa79..14d27ca71 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -14,7 +14,7 @@ import copy from collections.abc import Mapping from dataclasses import dataclass, field, fields from pathlib import Path -from typing import Any, TypeAlias, TypedDict, cast +from typing import Any, Literal, TypeAlias, TypedDict, cast from pydantic import ConfigDict, Field from vllm.config.utils import config @@ -470,6 +470,8 @@ class _DiffusionConfigProjection: enable_cpu_offload: bool = False enable_layerwise_offload: bool = False pin_cpu_memory: bool = True + diffusion_compile_granularity: Literal["regional", "full"] = "regional" + diffusion_compile_dynamic: bool = Field(default=True, strict=True) vae_use_slicing: bool = False vae_use_tiling: bool = False mask_strategy_file_path: str | None = None diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 2f9749420..e275751fe 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -388,6 +388,8 @@ class StageDeployConfig: diffusion_attention_config: dict[str, Any] | None = None # Diffusion execution, cache, and VAE behavior. + diffusion_compile_granularity: str | None = None + diffusion_compile_dynamic: bool | None = None cache_backend: str | None = None cache_config: dict[str, Any] | None = None enable_cache_dit_summary: bool | None = None diff --git a/vllm_omni/diffusion/compile.py b/vllm_omni/diffusion/compile.py index 45216febc..227d4f330 100644 --- a/vllm_omni/diffusion/compile.py +++ b/vllm_omni/diffusion/compile.py @@ -53,23 +53,29 @@ def regionally_compile( repeated_block_attrs = getattr(model, "_layerwise_offload_blocks_attrs", []) - # Check if we have modules with the specified class names - has_compiled_region = False - compiled_region_count = 0 + # Build all compiled callables before mutating the model. This keeps setup + # failures atomic: callers can safely continue with the uncompiled model if + # torch.compile raises synchronously for any repeated block. + compiled_forwards: list[tuple[nn.Module, Any]] = [] for name, submod in model.named_modules(): if _matches_repeated_block(name, submod, repeated_blocks, repeated_block_attrs): # Compile the block compute while keeping nn.Module.__call__ hooks # outside the compiled graph. - submod.forward = torch.compile(submod.forward, *compile_args, **compile_kwargs) - has_compiled_region = True - compiled_region_count += 1 + compiled_forwards.append( + ( + submod, + torch.compile(submod.forward, *compile_args, **compile_kwargs), + ) + ) - if not has_compiled_region: + if not compiled_forwards: logger.warning(f"Regional compilation skipped because {repeated_blocks} classes are not found in the model.") else: + for submod, compiled_forward in compiled_forwards: + submod.forward = compiled_forward logger.info( "Regional compilation applied to %d module(s) for repeated blocks %s.", - compiled_region_count, + len(compiled_forwards), repeated_blocks, ) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 2d8ed2410..078354d96 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -684,6 +684,10 @@ class OmniDiffusionConfig: # Compilation enforce_eager: bool = False + # Controls the generic compilation path used when a pipeline does not + # provide its own setup_compile() implementation. + diffusion_compile_granularity: str = "regional" + diffusion_compile_dynamic: bool = True # Parallel weight loading (for faster diffusion model startup) enable_multithread_weight_load: bool = True @@ -880,6 +884,13 @@ class OmniDiffusionConfig: ) def __post_init__(self): + if self.diffusion_compile_granularity not in {"regional", "full"}: + raise ValueError( + "diffusion_compile_granularity must be 'regional' or 'full', " + f"got {self.diffusion_compile_granularity!r}" + ) + if not isinstance(self.diffusion_compile_dynamic, bool): + raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}") self.master_port = self._resolve_master_port() self.request_batch_max_wait_ms = float(self.request_batch_max_wait_ms or 0.0) if self.request_batch_max_wait_ms < 0: @@ -915,6 +926,23 @@ class OmniDiffusionConfig: f"num_gpus ({self.num_gpus}) < parallel_config.world_size ({self.parallel_config.world_size})" ) + if self.diffusion_compile_granularity == "full": + incompatible_features = [] + if self.parallel_config.use_hsdp: + incompatible_features.append("HSDP") + if self.parallel_config.sequence_parallel_size > 1: + incompatible_features.append("sequence parallelism") + if self.enable_cpu_offload: + incompatible_features.append("CPU offload") + if self.enable_layerwise_offload: + incompatible_features.append("layerwise offload") + if incompatible_features: + features = ", ".join(incompatible_features) + raise ValueError( + "diffusion_compile_granularity='full' is incompatible with " + f"{features}; use 'regional' compilation instead" + ) + # Convert string dtype to torch.dtype if needed if isinstance(self.dtype, str): dtype_map = { diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index ced51ba3d..ee1b61c21 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -151,15 +151,33 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): if model is None: return + compile_granularity = self.od_config.diffusion_compile_granularity + compile_dynamic = self.od_config.diffusion_compile_dynamic try: - setattr(self.pipeline, attr_name, regionally_compile(model, dynamic=True)) - logger.info("Model runner: %s compiled with torch.compile.", attr_name) + if compile_granularity == "full": + model.compile(dynamic=compile_dynamic) + compiled_model = model + else: + compiled_model = regionally_compile(model, dynamic=compile_dynamic) + setattr(self.pipeline, attr_name, compiled_model) except Exception as e: logger.warning( - "Model runner: torch.compile for %s failed: %s. Using eager mode.", + "Model runner: %s torch.compile setup for %s failed before activation: %s. " + "Continuing with the uncompiled model; lazy compilation errors can still " + "surface on the first request.", + compile_granularity, attr_name, e, ) + return + + logger.info( + "Model runner: %s configured for lazy %s torch.compile with dynamic=%s; " + "compilation errors may surface on the first request.", + attr_name, + compile_granularity, + compile_dynamic, + ) def load_model( self, diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 2d049d73f..4d7622b91 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -466,6 +466,8 @@ class OrchestratorArgs: hsdp_replicate_size: int = 1 diffusion_attention_backend: str | None = None diffusion_attention_config: str | None = None + diffusion_compile_granularity: str | None = None + diffusion_compile_dynamic: bool | None = None cache_backend: str = "none" cache_config: str | None = None enable_cache_dit_summary: bool = False diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 7740f11e7..9deabf6d3 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -999,6 +999,14 @@ class AsyncOmniEngine: "enable_cpu_offload": kwargs.get("enable_cpu_offload", False), "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False), "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"), + "diffusion_compile_granularity": ( + "regional" + if kwargs.get("diffusion_compile_granularity") is None + else kwargs["diffusion_compile_granularity"] + ), + "diffusion_compile_dynamic": ( + True if kwargs.get("diffusion_compile_dynamic") is None else kwargs["diffusion_compile_dynamic"] + ), "boundary_ratio": kwargs.get("boundary_ratio", None), "flow_shift": kwargs.get("flow_shift", None), "diffusion_load_format": kwargs.get("diffusion_load_format", "default"), diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index d807c22ed..98c086d86 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -423,6 +423,25 @@ class OmniServeCommand(CLISubcommand): "custom_pipeline, dummy, or diffusers for the HF diffusers adapter." ), ) + omni_config_group.add_argument( + "--diffusion-compile-granularity", + choices=["regional", "full"], + default=None, + help=( + "Compilation scope for the generic diffusion model runner. " + "'regional' compiles repeated blocks (default); 'full' compiles the whole transformer and is " + "incompatible with HSDP, sequence parallelism, CPU offload, and layerwise offload." + ), + ) + omni_config_group.add_argument( + "--diffusion-compile-dynamic", + action=argparse.BooleanOptionalAction, + default=None, + help=( + "Use dynamic shapes for the selected generic diffusion compile scope. " + "Disable for fixed-shape workloads with --no-diffusion-compile-dynamic." + ), + ) omni_config_group.add_argument( "--diffusers-load-kwargs", dest="diffusers_load_kwargs",