diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 559bbbb50..be18141e6 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -467,7 +467,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--hsdp-shard-size", type=int, - default=1, + default=-1, help="Number of GPUs to shard weights across for HSDP.", ) parser.add_argument( @@ -575,6 +575,10 @@ def main(): lora_path = lora_path[0] omni_kwargs["lora_path"] = lora_path omni_kwargs["lora_backend"] = args.lora_backend + if args.use_hsdp: + omni_kwargs["use_hsdp"] = args.use_hsdp + omni_kwargs["hsdp_shard_size"] = args.hsdp_shard_size + omni_kwargs["hsdp_replicate_size"] = args.hsdp_replicate_size # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails # gate is an engine-level config (offline analog of the server's --no-guardrails). diff --git a/tests/diffusion/lora/test_loader.py b/tests/diffusion/lora/test_loader.py index f13c0f6c3..deb5ec80d 100644 --- a/tests/diffusion/lora/test_loader.py +++ b/tests/diffusion/lora/test_loader.py @@ -490,3 +490,22 @@ def test_prepare_lora_delta_unchanged_for_existing_fused_layouts(fused, shards): delta, used_keys = _prepare_lora_delta(state_dict, base_key, param_to_weight_names) assert delta.shape == (len(shards) * HEAD_DIM, HEAD_DIM) assert len(used_keys) == 2 * len(shards) + + +def test_lora_loader_mixin_lora_is_fused_property(mocker: MockerFixture): + pipeline = DummyQwenImagePipeline(1, HEAD_DIM) + assert pipeline.lora_is_fused is False + + lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer) + mocker.patch( + "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline", + return_value=_convert_non_diffusers_qwen_lora_to_diffusers, + ) + pipeline.load_lora_weights(lora_state_dict, "adapter0") + assert pipeline.lora_is_fused is True + + # Test explicit setter + pipeline2 = DummyQwenImagePipeline(1, HEAD_DIM) + assert pipeline2.lora_is_fused is False + pipeline2.lora_is_fused = True + assert pipeline2.lora_is_fused is True diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index bd59ee94b..4304f5d2e 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -8,6 +8,7 @@ Tests for the DiffusersPipelineLoader. import json from pathlib import Path from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch @@ -18,6 +19,7 @@ from vllm.config.load import LoadConfig from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig +from vllm_omni.diffusion.lora.manager import LoRABackend from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlan, @@ -219,6 +221,48 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading( assert len(tuple((store_root / "source-digests-v1" / "entries").glob("*.json"))) == 1 +def test_maybe_fuse_distilled_lora_skips_when_hwr_warm_snapshot_present(): + cfg = SimpleNamespace( + lora_backend="distill", + lora_path="/path/to/lora.safetensors", + lora_scale=1.0, + dtype=torch.bfloat16, + quantization_config=None, + parallel_config=SimpleNamespace(use_hsdp=False), + ) + loader = DiffusersPipelineLoader(LoadConfig(), cfg) + loader._hwr_state = {"warm_snapshot": {"transformer.weight": (1234, "hash")}} + + model = nn.Module() + model.load_lora_weights = MagicMock() + + loader._maybe_fuse_distilled_lora(model) + + model.load_lora_weights.assert_not_called() + assert getattr(model, "lora_is_fused", False) is True + + +def test_maybe_fuse_distilled_lora_fuses_when_not_warm_snapshot(): + cfg = SimpleNamespace( + lora_backend=LoRABackend.DISTILL, + lora_path="/path/to/lora.safetensors", + lora_scale=1.0, + dtype=torch.bfloat16, + quantization_config=None, + parallel_config=SimpleNamespace(use_hsdp=False), + ) + loader = DiffusersPipelineLoader(LoadConfig(), cfg) + loader._hwr_state = None + + model = nn.Module() + model.load_lora_weights = MagicMock() + + loader._maybe_fuse_distilled_lora(model) + + model.load_lora_weights.assert_called_once_with("/path/to/lora.safetensors") + assert getattr(model, "lora_is_fused", False) is True + + def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_path: Path, monkeypatch): from vllm_omni.diffusion.model_loader import diffusers_loader as loader_module @@ -1038,6 +1082,156 @@ def test_hsdp_processes_quantized_weights_before_sharding(mocker): assert events == ["load", "process", "prepare", "shard"] +def test_hsdp_fuses_distilled_lora_before_sharding(mocker): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + from vllm_omni.diffusion.offloader.module_collector import PipelineModules + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace( + use_hsdp=True, + hsdp_replicate_size=1, + hsdp_shard_size=2, + ), + quantization_config=None, + lora_backend="distill", + lora_path="/path/to/lora.safetensors", + lora_scale=1.0, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + loader.quant_config = None + + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + events: list[str] = [] + + model.load_lora_weights = mocker.Mock(side_effect=lambda path: events.append(f"fuse_lora:{path}")) + + loader._init_from_load_format = mocker.Mock(return_value=model) + loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load")) + loader._process_weights_after_loading = mocker.Mock(side_effect=lambda _model, _device: events.append("process")) + mocker.patch.object( + loader_mod.ModuleDiscovery, + "discover", + return_value=PipelineModules( + dits=[model.transformer], + dit_names=["transformer"], + vaes=[], + encoders=[], + encoder_names=[], + resident_modules=[], + resident_names=[], + ), + ) + mocker.patch.object( + loader_mod, + "apply_hsdp_to_model", + side_effect=lambda *_args, **_kwargs: events.append("shard"), + ) + + loader._load_model_with_hsdp(torch.device("cpu")) + + assert events == ["load", "fuse_lora:/path/to/lora.safetensors", "process", "shard"] + assert getattr(model, "lora_is_fused", False) is True + + +def test_hsdp_fuses_multi_file_distilled_lora_wan22(mocker): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + from vllm_omni.diffusion.offloader.module_collector import PipelineModules + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace( + use_hsdp=True, + hsdp_replicate_size=1, + hsdp_shard_size=2, + ), + quantization_config=None, + lora_backend="distill", + lora_path=["/path/to/high.safetensors", "/path/to/low.safetensors"], + lora_scale=1.0, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + loader.quant_config = None + + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + model.transformer_2 = nn.Linear(2, 2, bias=False) + events: list[str] = [] + + model.load_lora_weights = mocker.Mock(side_effect=lambda paths: events.append(f"fuse_lora:{paths}")) + + loader._init_from_load_format = mocker.Mock(return_value=model) + loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load")) + loader._process_weights_after_loading = mocker.Mock(side_effect=lambda _model, _device: events.append("process")) + mocker.patch.object( + loader_mod.ModuleDiscovery, + "discover", + return_value=PipelineModules( + dits=[model.transformer, model.transformer_2], + dit_names=["transformer", "transformer_2"], + vaes=[], + encoders=[], + encoder_names=[], + resident_modules=[], + resident_names=[], + ), + ) + mocker.patch.object( + loader_mod, + "apply_hsdp_to_model", + side_effect=lambda *_args, **_kwargs: events.append("shard"), + ) + + loader._load_model_with_hsdp(torch.device("cpu")) + + assert events == [ + "load", + "fuse_lora:['/path/to/high.safetensors', '/path/to/low.safetensors']", + "process", + "shard", + "shard", + ] + assert getattr(model, "lora_is_fused", False) is True + + +def test_load_model_fuses_distilled_lora_non_hsdp(mocker): + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace( + use_hsdp=False, + tensor_parallel_size=1, + data_parallel_size=1, + sequence_parallel_size=1, + ), + enable_cpu_offload=False, + quantization_config=None, + lora_backend="distill", + lora_path="/path/to/lora.safetensors", + lora_scale=1.0, + host_weight_runtime_mode="disabled", + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + loader._force_canonical_load = True + + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + events: list[str] = [] + + model.load_lora_weights = mocker.Mock(side_effect=lambda path: events.append(f"fuse_lora:{path}")) + + loader._init_from_load_format = mocker.Mock(return_value=model) + loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load")) + loader._process_weights_after_loading = mocker.Mock(side_effect=lambda _model, _device: events.append("process")) + loader._apply_skip_softmax_calibration = mocker.Mock() + + res = loader.load_model(load_device="cpu", device=torch.device("cpu")) + + assert res is model + assert events == ["load", "fuse_lora:/path/to/lora.safetensors", "process"] + assert getattr(model, "lora_is_fused", False) is True + + def test_get_all_weights(prefetch_helios_model, mock_tp_group): """Ensure that get all weights on a tiny model resolves to nonempty weights.""" od_config = OmniDiffusionConfig( diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index 52717fb6f..a27cd2cc8 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -312,3 +312,41 @@ class TestDiffusionWorkerWakeUp: mock_buffer2.data.copy_.assert_not_called() assert result is True + + +class TestDiffusionWorkerInitLoraManager: + """Test DiffusionWorker.init_lora_manager method.""" + + def test_skips_when_lora_already_fused(self, mocker: MockerFixture, mock_gpu_worker): + """When pipeline.lora_is_fused is True, init_lora_manager returns immediately.""" + from vllm_omni.diffusion.lora.manager import LoRABackend + + mock_gpu_worker.od_config.lora_backend = LoRABackend.DISTILL + mock_gpu_worker.od_config.lora_path = "/path/to/lora.safetensors" + mock_gpu_worker.od_config.lora_scale = 1.0 + + pipeline = mock_gpu_worker.model_runner.pipeline + pipeline.lora_is_fused = True + pipeline.load_lora_weights = mocker.Mock() + + mock_gpu_worker.init_lora_manager() + + pipeline.load_lora_weights.assert_not_called() + assert mock_gpu_worker.lora_manager is None + + def test_distill_backend_fuses_and_marks_flag_when_not_pre_fused(self, mocker: MockerFixture, mock_gpu_worker): + """When lora_is_fused is False, distill backend calls load_lora_weights and sets flag.""" + from vllm_omni.diffusion.lora.manager import LoRABackend + + mock_gpu_worker.od_config.lora_backend = LoRABackend.DISTILL + mock_gpu_worker.od_config.lora_path = "/path/to/lora.safetensors" + mock_gpu_worker.od_config.lora_scale = 1.0 + + pipeline = mock_gpu_worker.model_runner.pipeline + pipeline.lora_is_fused = False + pipeline.load_lora_weights = mocker.Mock() + + mock_gpu_worker.init_lora_manager() + + pipeline.load_lora_weights.assert_called_once_with("/path/to/lora.safetensors") + assert pipeline.lora_is_fused is True diff --git a/tests/entrypoints/openai_api/test_video_pipeline_capability.py b/tests/entrypoints/openai_api/test_video_pipeline_capability.py index 530c6e216..ed4d9db86 100644 --- a/tests/entrypoints/openai_api/test_video_pipeline_capability.py +++ b/tests/entrypoints/openai_api/test_video_pipeline_capability.py @@ -61,6 +61,9 @@ def test_video_pipeline_rejects_non_video_final_outputs(stage_configs): "WanDMDPipeline", "LingBotWorldCausalDMDPipeline", "LongCatVideoAvatarPipeline", + "SanaVideoPipeline", + "SanaImageToVideoPipeline", + "SanaWmPipeline", ], ) def test_registered_video_aliases_declare_video_output(model_class_name): diff --git a/vllm_omni/diffusion/lora/loader.py b/vllm_omni/diffusion/lora/loader.py index 676e8d514..3f74b4853 100644 --- a/vllm_omni/diffusion/lora/loader.py +++ b/vllm_omni/diffusion/lora/loader.py @@ -209,6 +209,7 @@ class LoraLoaderMixin: # Set by the pipeline this mixin is combined with; annotated, not assigned. transformer: torch.nn.Module _lora_loaded: dict[str | None, dict[str, torch.Tensor]] + _lora_is_fused: bool # Lazy initialization to avoid MRO issues: __init__ may not be called # when mixin with nn.Module @@ -222,6 +223,18 @@ class LoraLoaderMixin: def lora_loaded(self, value): self._lora_loaded = value + @property + def lora_is_fused(self) -> bool: + """True when LoRA weights are fused into base weights. + + Mixin loads are in-place fusions, so non-empty `_lora_loaded` implies fused. + """ + return getattr(self, "_lora_is_fused", False) or bool(getattr(self, "_lora_loaded", None)) + + @lora_is_fused.setter + def lora_is_fused(self, value: bool): + self._lora_is_fused = value + @classmethod def load_lora_into_module( cls, diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 7069e5012..30d0f0610 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -31,6 +31,7 @@ from vllm.utils.torch_utils import set_default_torch_dtype from vllm_omni.diffusion.config import set_current_diffusion_config from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp_to_model +from vllm_omni.diffusion.lora.manager import LoRABackend from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( get_checkpoint_adapter, ) @@ -465,6 +466,39 @@ class DiffusersPipelineLoader(HWRLoaderMixin): return all_parameter_names return {name for name in all_parameter_names if name.startswith(source_prefixes)} + def _maybe_fuse_distilled_lora(self, model: nn.Module) -> None: + """Fuse distilled LoRA weights into the model before sharding or quantization.""" + if self.od_config is None: + return + lora_backend = getattr(self.od_config, "lora_backend", None) + if lora_backend != LoRABackend.DISTILL and lora_backend != "distill": + return + + if getattr(model, "lora_is_fused", False): + return + + # A warm HWR restore already brings back the fused final-layout weights. + if self._hwr_state is not None and self._hwr_state.get("warm_snapshot") is not None: + setattr(model, "lora_is_fused", True) + return + + lora_path = getattr(self.od_config, "lora_path", None) + if not lora_path: + return + + if isinstance(lora_path, list) and len(lora_path) == 1: + lora_path = lora_path[0] + + if hasattr(model, "load_lora_weights"): + lora_scale = getattr(self.od_config, "lora_scale", 1.0) + if lora_scale > 1.0: + logger.warning("lora_scale > 1.0 may not take any effect when using distilled LoRA backend.") + logger.info("Fusing distilled LoRA weights from %s into %s", lora_path, model.__class__.__name__) + model.load_lora_weights(lora_path) + setattr(model, "lora_is_fused", True) + else: + logger.warning("Pipeline %s does not support loading distilled LoRA weights.", model.__class__.__name__) + def load_model( self, load_device: str, @@ -587,6 +621,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin): sources=ordinary_sources, planned_weights=host_weight_plan.bindings, ) + self._maybe_fuse_distilled_lora(model) else: if _dist_offload and _use_ag and _dlo_group_size > 1 and _has_online_quant: # An effective DLO group size of one performs no weight @@ -625,6 +660,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin): self.load_weights(model, stream_online_quant_to_cpu=True) else: self.load_weights(model) + self._maybe_fuse_distilled_lora(model) self._process_weights_after_loading(model, target_device) # A warm final-layout hit has already completed all @@ -1019,6 +1055,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin): raise ValueError("HSDP is not supported with the diffusers adapter load format") model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True) self.load_weights(model) + self._maybe_fuse_distilled_lora(model) # Quantization methods must finish while parameters are ordinary local # tensors. Some post-load transforms use operations (for example, diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py index 6b51e5753..67c58fbec 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/vllm_omni/diffusion/model_metadata.py @@ -97,6 +97,7 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = { "SanaWmPipeline": DiffusionModelMetadata( supports_multimodal_inputs=True, max_multimodal_image_inputs=1, + final_output_type="video", ), } diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index b65498d75..cc6adf362 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -565,6 +565,7 @@ class DiffusionWorker: if self.od_config.lora_scale > 1.0: logger.warning("lora_scale > 1.0 may not take any effect when using distilled LoRA backend.") pipeline.load_lora_weights(lora_path) + pipeline.lora_is_fused = True else: logger.warning("Pipeline does not support loading distilled LoRA weights for now.") else: