[Bugfix][Diffusion] Fuse distilled LoRA weights before HSDP sharding (#6948)

Signed-off-by: samithuang <285365963@qq.com>
Signed-off-by: Samit <285365963@qq.com>

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).

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

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(

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

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):

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,

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,

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",
     ),
 }
 

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: