[Perf][Diffusion] Optimize HSDP startup via Rank-0 shared weight loading and accelerated LoRA delta computation (#7005)

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

docs/mkdocs/hooks/search_index.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Reduce browser search-index memory for generated API reference pages."""
 
 from material.plugins.search.plugin import SearchIndex

tests/config/test_omni_config.py

@@ -714,6 +714,7 @@ def test_sub_config_fields_match_structured_scopes():
         "max_cudagraph_capture_size",
         "enable_flashinfer_autotune",
         "enable_multithread_weight_load",
+        "enable_broadcast_weight_load",
         "num_weight_load_threads",
         "disable_autocast",
         # Per-stage checkpoint resolution for repos whose stages live in

tests/diffusion/lora/test_loader.py

@@ -509,3 +509,15 @@ def test_lora_loader_mixin_lora_is_fused_property(mocker: MockerFixture):
     assert pipeline2.lora_is_fused is False
     pipeline2.lora_is_fused = True
     assert pipeline2.lora_is_fused is True
+
+
+def test_prepare_lora_delta_with_compute_device():
+    base_key = "blocks.0.attn.to_q"
+    state_dict = {
+        f"{base_key}.lora_A.weight": torch.ones(RANK, HEAD_DIM),
+        f"{base_key}.lora_B.weight": torch.ones(HEAD_DIM, RANK),
+    }
+    delta, used_keys = _prepare_lora_delta(state_dict, base_key, compute_device=torch.device("cpu"))
+    assert delta is not None
+    assert delta.shape == (HEAD_DIM, HEAD_DIM)
+    assert len(used_keys) == 2

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -6,9 +6,11 @@ Tests for the DiffusersPipelineLoader.
 """
 
 import json
+import os
+import tempfile
 from pathlib import Path
 from types import SimpleNamespace
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
@@ -260,17 +262,16 @@ 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():
+def test_maybe_fuse_distilled_lora_skips_when_lora_path_unset():
     cfg = SimpleNamespace(
         lora_backend="distill",
-        lora_path="/path/to/lora.safetensors",
+        lora_path=None,
         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()
@@ -278,7 +279,7 @@ def test_maybe_fuse_distilled_lora_skips_when_hwr_warm_snapshot_present():
     loader._maybe_fuse_distilled_lora(model)
 
     model.load_lora_weights.assert_not_called()
-    assert getattr(model, "lora_is_fused", False) is True
+    assert getattr(model, "lora_is_fused", False) is False
 
 
 def test_maybe_fuse_distilled_lora_fuses_when_not_warm_snapshot():
@@ -1053,6 +1054,34 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
     assert loader.take_host_weight_plan() is None
 
 
+def test_dlo_mmap_plan_with_distilled_lora_falls_back_to_ordinary_loader(monkeypatch):
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        lora_backend=LoRABackend.DISTILL,
+        lora_path="/fake/lora.safetensors",
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.load_lora_weights = MagicMock()
+    calls: list[str] = []
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: calls.append("load")  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: calls.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+
+    assert loader.load_model(load_device="cpu") is model
+    assert calls == ["load", "process"]
+    assert model.load_lora_weights.call_count == 1
+    assert loader.take_host_weight_plan() is None
+
+
 def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch):
     from vllm.model_executor.layers.quantization.online.fp8 import (
         Fp8PerTensorOnlineLinearMethod,
@@ -1487,3 +1516,413 @@ def test_load_model(prefetch_helios_model, mock_tp_group):
     )
     model = loader.load_model(load_device="cpu")
     assert isinstance(model, HeliosPipeline)
+
+
+def test_hsdp_broadcast_weight_load_rank0(mocker):
+    """Ensure that on rank 0 with enable_broadcast_weight_load=True, weights are loaded, fused, and broadcast."""
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    mocker.patch("torch.distributed.is_initialized", return_value=True)
+    mocker.patch("torch.distributed.get_world_size", return_value=4)
+    mocker.patch("torch.distributed.get_rank", return_value=0)
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=4,
+        ),
+        enable_broadcast_weight_load=True,
+        lora_backend=LoRABackend.DISTILL,
+        lora_path="/fake/lora.safetensors",
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    events: list[str] = []
+
+    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
+    loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load"))  # type: ignore[method-assign]
+    loader._maybe_fuse_distilled_lora = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model: events.append("fuse")
+    )
+    loader._broadcast_model_weights = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model, **_kwargs: events.append("broadcast")
+    )
+    loader._process_weights_after_loading = mocker.Mock(  # type: ignore[method-assign]
+        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", "broadcast", "process", "shard"]
+
+
+def test_hsdp_broadcast_weight_load_rank_nonzero(mocker):
+    """Ensure that on rank > 0 with enable_broadcast_weight_load=True, loading is skipped and broadcast is received."""
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    mocker.patch("torch.distributed.is_initialized", return_value=True)
+    mocker.patch("torch.distributed.get_world_size", return_value=4)
+    mocker.patch("torch.distributed.get_rank", return_value=1)
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=4,
+        ),
+        enable_broadcast_weight_load=True,
+        lora_backend=LoRABackend.DISTILL,
+        lora_path="/fake/lora.safetensors",
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.load_lora_weights = mocker.Mock()  # type: ignore[assignment]
+    events: list[str] = []
+
+    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
+    loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load"))  # type: ignore[method-assign]
+    loader._maybe_fuse_distilled_lora = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model: events.append("fuse")
+    )
+    loader._broadcast_model_weights = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model, **_kwargs: events.append("broadcast")
+    )
+    loader._process_weights_after_loading = mocker.Mock(  # type: ignore[method-assign]
+        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"))
+
+    # Rank 1 should NOT call load_weights or _maybe_fuse_distilled_lora from disk
+    assert "load" not in events
+    assert "fuse" not in events
+    assert events == ["broadcast", "process", "shard"]
+    assert getattr(model, "lora_is_fused", False) is True
+
+
+def test_broadcast_model_weights_invokes_dist_broadcast(mocker):
+    broadcast_calls = []
+    mocker.patch("torch.distributed.is_initialized", return_value=True)
+    mocker.patch("torch.distributed.get_world_size", return_value=2)
+    mocker.patch("torch.distributed.get_rank", return_value=0)
+    mocker.patch("torch.distributed.get_backend", return_value="gloo")
+    mocker.patch("torch.distributed.broadcast", side_effect=lambda tensor, src: broadcast_calls.append((tensor, src)))
+    mocker.patch("torch.distributed.barrier")
+
+    model = nn.Module()
+    model.linear = nn.Linear(2, 2)
+    model.register_buffer("buf", torch.tensor([1.0, 2.0]))
+    model.register_buffer("int_buf", torch.tensor([1, 2, 3], dtype=torch.long))
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=True),
+        enable_broadcast_weight_load=True,
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    loader._broadcast_model_weights(model, target_device=torch.device("cpu"), src_rank=0)
+
+    # float32 tensors (linear.weight, linear.bias, buf) are coalesced into 1 bucket,
+    # and int64 tensor (int_buf) is in a separate bucket -> 2 coalesced broadcast calls
+    assert len(broadcast_calls) == 2
+    assert all(src == 0 for _, src in broadcast_calls)
+
+
+def test_broadcast_model_weights_receiver_copies_data(mocker):
+    """Ensure receiver rank copies broadcasted flat tensor data into its model parameters/buffers."""
+    mocker.patch("torch.distributed.is_initialized", return_value=True)
+    mocker.patch("torch.distributed.get_world_size", return_value=2)
+    mocker.patch("torch.distributed.get_rank", return_value=1)
+    mocker.patch("torch.distributed.get_backend", return_value="gloo")
+    mocker.patch("torch.distributed.barrier")
+
+    # Simulate broadcast by filling the flat tensor with source values
+    def mock_broadcast(tensor, src):
+        if tensor.dtype == torch.float32:
+            tensor.fill_(3.14)
+        elif tensor.dtype == torch.long:
+            tensor.fill_(42)
+
+    mocker.patch("torch.distributed.broadcast", side_effect=mock_broadcast)
+
+    model = nn.Module()
+    model.linear = nn.Linear(2, 2)
+    model.register_buffer("buf", torch.zeros(2))
+    model.register_buffer("int_buf", torch.zeros(3, dtype=torch.long))
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=True),
+        enable_broadcast_weight_load=True,
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    loader._broadcast_model_weights(model, target_device=torch.device("cpu"), src_rank=0)
+
+    assert torch.allclose(model.linear.weight, torch.tensor(3.14))
+    assert torch.allclose(model.linear.bias, torch.tensor(3.14))
+    assert torch.allclose(model.buf, torch.tensor(3.14))
+    assert torch.equal(model.int_buf, torch.tensor([42, 42, 42], dtype=torch.long))
+
+
+def test_broadcast_model_weights_mixed_cpu_and_accelerator_devices(mocker):
+    """Ensure that when some pipeline modules are on CPU (transformer) and others on accelerator (VAE/encoder),
+    broadcasting normalizes tensors to CPU before concatenation and succeeds without device mismatch.
+    """
+    dev = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+    mocker.patch("torch.distributed.is_initialized", return_value=True)
+    mocker.patch("torch.distributed.get_world_size", return_value=2)
+    mocker.patch("torch.distributed.get_rank", return_value=0)
+    mocker.patch("torch.distributed.get_backend", return_value="nccl" if dev.type == "cuda" else "gloo")
+    mocker.patch("torch.distributed.broadcast")
+    mocker.patch("torch.distributed.barrier")
+
+    model = nn.Module()
+    # Transformer parameters on CPU
+    model.transformer = nn.Linear(4, 4)
+    # VAE / Encoder parameters on accelerator
+    model.vae = nn.Linear(4, 4).to(dev)
+    model.text_encoder = nn.Linear(4, 4).to(dev)
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=True),
+        enable_broadcast_weight_load=True,
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    # Should not raise RuntimeError: Tensors must have same device
+    loader._broadcast_model_weights(model, target_device=dev, src_rank=0)
+
+
+def test_broadcast_model_weights_mixed_devices_receiver(mocker):
+    """Ensure receiver rank with mixed CPU/accelerator tensors correctly receives broadcast data."""
+    dev = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+    mocker.patch("torch.distributed.is_initialized", return_value=True)
+    mocker.patch("torch.distributed.get_world_size", return_value=2)
+    mocker.patch("torch.distributed.get_rank", return_value=1)
+    mocker.patch("torch.distributed.get_backend", return_value="nccl" if dev.type == "cuda" else "gloo")
+    mocker.patch("torch.distributed.barrier")
+
+    def mock_broadcast(tensor, src):
+        tensor.fill_(7.0)
+
+    mocker.patch("torch.distributed.broadcast", side_effect=mock_broadcast)
+
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2)  # CPU
+    model.vae = nn.Linear(2, 2).to(dev)  # Accelerator
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=True),
+        enable_broadcast_weight_load=True,
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    loader._broadcast_model_weights(model, target_device=dev, src_rank=0)
+
+    assert torch.allclose(model.transformer.weight, torch.tensor(7.0))
+    assert torch.allclose(model.vae.weight, torch.tensor(7.0, device=dev))
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_hsdp_broadcast_weight_load_falls_back_when_online_quantization_enabled(mocker, rank):
+    """Ensure that when online quantization is enabled, HSDP weight loading falls back to
+    ordinary per-rank loading on all ranks rather than rank-0 broadcast.
+    """
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    mocker.patch("torch.distributed.is_initialized", return_value=True)
+    mocker.patch("torch.distributed.get_world_size", return_value=4)
+    mocker.patch("torch.distributed.get_rank", return_value=rank)
+
+    mock_quant_config = SimpleNamespace(
+        is_checkpoint_quantized=False,
+    )
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=4,
+        ),
+        enable_broadcast_weight_load=True,
+        lora_backend=None,
+        lora_path=None,
+        quantization_config=mock_quant_config,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    events: list[str] = []
+
+    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
+    loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load"))  # type: ignore[method-assign]
+    loader._broadcast_model_weights = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model, **_kwargs: events.append("broadcast")
+    )
+    loader._process_weights_after_loading = mocker.Mock(  # type: ignore[method-assign]
+        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"))
+
+    # Both rank 0 and non-zero ranks MUST execute ordinary load_weights, not broadcast
+    assert "load" in events
+    assert "broadcast" not in events
+    assert events == ["load", "process", "shard"]
+
+
+def _mp_worker_online_quant(rank: int, world_size: int, rendezvous: str, temp_dir: str) -> None:
+    torch.set_num_threads(1)
+    import torch.distributed as dist
+
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=world_size)
+    try:
+
+        class _MockOnlineQuantModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.transformer = nn.Linear(4, 4)
+                self.load_called = False
+
+            def load_weights(self, weights):
+                self.load_called = True
+                return {"transformer.weight", "transformer.bias"}
+
+        mock_quant_config = SimpleNamespace(is_checkpoint_quantized=False)
+        od_config = SimpleNamespace(
+            dtype=torch.float32,
+            parallel_config=SimpleNamespace(
+                use_hsdp=True,
+                hsdp_replicate_size=1,
+                hsdp_shard_size=world_size,
+            ),
+            enable_broadcast_weight_load=True,
+            lora_backend=None,
+            lora_path=None,
+            quantization_config=mock_quant_config,
+        )
+        loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+        model = _MockOnlineQuantModel()
+        loader._init_from_load_format = lambda *args, **kwargs: model  # type: ignore[method-assign]
+        loader._process_weights_after_loading = lambda *args, **kwargs: None  # type: ignore[method-assign]
+
+        with (
+            patch.object(
+                loader_module.ModuleDiscovery,
+                "discover",
+                return_value=PipelineModules(
+                    dits=[model.transformer],
+                    dit_names=["transformer"],
+                    vaes=[],
+                    encoders=[],
+                    encoder_names=[],
+                    resident_modules=[],
+                    resident_names=[],
+                ),
+            ),
+            patch.object(loader_module, "apply_hsdp_to_model", return_value=None),
+        ):
+            loader._load_model_with_hsdp(torch.device("cpu"))
+
+        if model.load_called:
+            with open(os.path.join(temp_dir, f"rank_{rank}_success.flag"), "w") as f:
+                f.write("ok")
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.skipif(
+    not torch.distributed.is_available() or not torch.distributed.is_gloo_available(),
+    reason="requires torch.distributed gloo",
+)
+def test_hsdp_broadcast_weight_load_online_quant_multiprocess():
+    """Multi-process regression test verifying that with online quantization and broadcast enabled,
+    all ranks fall back to independent load_weights in parallel processes without hanging.
+    """
+    import torch.multiprocessing as mp
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        rendezvous = f"file://{os.path.join(temp_dir, 'gloo-rendezvous')}"
+        world_size = 2
+        with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}):
+            mp.spawn(
+                _mp_worker_online_quant,
+                args=(world_size, rendezvous, temp_dir),
+                nprocs=world_size,
+                join=True,
+            )
+        assert os.path.exists(os.path.join(temp_dir, "rank_0_success.flag"))
+        assert os.path.exists(os.path.join(temp_dir, "rank_1_success.flag"))

tests/engine/test_stage_engine_args.py

@@ -82,6 +82,7 @@ _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
         "custom_voice_dir",
         "devices",
         "disable_autocast",
+        "enable_broadcast_weight_load",
         "enable_multithread_weight_load",
         "env",
         "has_sampling_extra_args",

vllm_omni/config/omni_config.py

@@ -165,6 +165,7 @@ class _ModelEngineOverrides(TypedDict, total=False):
     max_cudagraph_capture_size: int
     enable_flashinfer_autotune: bool
     enable_multithread_weight_load: bool
+    enable_broadcast_weight_load: bool
     num_weight_load_threads: int
     disable_autocast: bool
 
@@ -435,6 +436,7 @@ class OmniStageModelConfig(_TrackExplicitConfigFields):
     max_cudagraph_capture_size: int | None = Field(default=None, ge=0)
     enable_flashinfer_autotune: bool | None = None
     enable_multithread_weight_load: bool = True
+    enable_broadcast_weight_load: bool = False
     num_weight_load_threads: int = Field(default=4, ge=1)
     disable_autocast: bool = False
     # Per-stage checkpoint/tokenizer subdirectories under the model root
@@ -1028,6 +1030,7 @@ _DIFFUSION_MOVED_SHARED_FIELDS = frozenset(
         "enable_sleep_mode",
         "enforce_eager",
         "enable_multithread_weight_load",
+        "enable_broadcast_weight_load",
         "num_weight_load_threads",
         "disable_autocast",
     }

vllm_omni/config/stage_config.py

@@ -472,6 +472,7 @@ class StageDeployConfig:
 
     # Runtime optimizations used by diffusion loading/execution.
     enable_multithread_weight_load: bool | None = None
+    enable_broadcast_weight_load: bool | None = None
     num_weight_load_threads: int | None = None
     diffusion_offload_config: dict[str, Any] | None = None
     # Compatibility aliases for existing callers and model-specific stage

vllm_omni/diffusion/data.py

@@ -920,6 +920,7 @@ class OmniDiffusionConfig:
 
     # Parallel weight loading (for faster diffusion model startup)
     enable_multithread_weight_load: bool = True
+    enable_broadcast_weight_load: bool = False
     num_weight_load_threads: int = 4
 
     # Enable sleep mode

vllm_omni/diffusion/lora/loader.py

@@ -43,6 +43,7 @@ def _prepare_lora_delta(
     lora_a_suffix: str = "lora_A.weight",
     lora_b_suffix: str = "lora_B.weight",
     lora_bias_suffix: str = "bias",
+    compute_device: torch.device | None = None,
 ):
     used_keys: set[str] = set()
     # stacked_params_mapping                       param_to_weight_names
@@ -70,6 +71,8 @@ def _prepare_lora_delta(
                 if lora_bias_key not in lora_state_dict:
                     return None, used_keys
                 delta = lora_state_dict[lora_bias_key]
+                if compute_device is not None and compute_device.type != "cpu":
+                    delta = delta.to(compute_device)
                 stacked_deltas.append(delta)
                 used_keys.add(lora_bias_key)
         else:
@@ -80,6 +83,9 @@ def _prepare_lora_delta(
                     return None, used_keys
                 a = lora_state_dict[lora_a_key]
                 b = lora_state_dict[lora_b_key]
+                if compute_device is not None and compute_device.type != "cpu":
+                    a = a.to(compute_device)
+                    b = b.to(compute_device)
                 delta = torch.matmul(b, a)
                 stacked_deltas.append(delta)
                 used_keys.add(lora_a_key)
@@ -94,7 +100,10 @@ def _prepare_lora_delta(
         if lora_bias_key not in lora_state_dict:
             return None, used_keys
         used_keys.add(lora_bias_key)
-        return lora_state_dict[lora_bias_key], used_keys
+        delta = lora_state_dict[lora_bias_key]
+        if compute_device is not None and compute_device.type != "cpu":
+            delta = delta.to(compute_device)
+        return delta, used_keys
 
     lora_a_key = f"{base_key}.{lora_a_suffix}"
     lora_b_key = f"{base_key}.{lora_b_suffix}"
@@ -102,9 +111,13 @@ def _prepare_lora_delta(
         return None, used_keys
     a = lora_state_dict[lora_a_key]
     b = lora_state_dict[lora_b_key]
+    if compute_device is not None and compute_device.type != "cpu":
+        a = a.to(compute_device)
+        b = b.to(compute_device)
+    delta = torch.matmul(b, a)
     used_keys.add(lora_a_key)
     used_keys.add(lora_b_key)
-    return torch.matmul(b, a), used_keys
+    return delta, used_keys
 
 
 def _load_lora_state_dict(
@@ -267,6 +280,7 @@ class LoraLoaderMixin:
             else:
                 continue
 
+            compute_dev = params.device if params.device.type != "cpu" else None
             delta, used_keys = _prepare_lora_delta(
                 state_dict,
                 base_key,
@@ -275,6 +289,7 @@ class LoraLoaderMixin:
                 lora_a_suffix,
                 lora_b_suffix,
                 lora_bias_suffix,
+                compute_device=compute_dev,
             )
             if delta is None:
                 continue
@@ -319,6 +334,7 @@ class LoraLoaderMixin:
             else:
                 continue
 
+            compute_dev = param.device if param.device.type != "cpu" else None
             delta, used_keys = _prepare_lora_delta(
                 state_dict,
                 base_key,
@@ -327,6 +343,7 @@ class LoraLoaderMixin:
                 lora_a_suffix,
                 lora_b_suffix,
                 lora_bias_suffix,
+                compute_device=compute_dev,
             )
             if delta is None:
                 continue

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -509,11 +509,6 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
         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
@@ -531,6 +526,111 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
         else:
             logger.warning("Pipeline %s does not support loading distilled LoRA weights.", model.__class__.__name__)
 
+    def _broadcast_model_weights(
+        self,
+        model: nn.Module,
+        target_device: torch.device,
+        src_rank: int = 0,
+        bucket_cap_bytes: int = 256 * 1024 * 1024,
+    ) -> None:
+        """Broadcast model parameters and buffers from src_rank to all other ranks using coalesced bucketing."""
+        if not torch.distributed.is_initialized() or torch.distributed.get_world_size() <= 1:
+            return
+
+        rank = torch.distributed.get_rank()
+        world_size = torch.distributed.get_world_size()
+        backend = torch.distributed.get_backend()
+        is_cpu_backend = backend == "gloo"
+
+        logger.info(
+            "Worker %d: %s full model weights via %s across %d ranks",
+            rank,
+            "Broadcasting" if rank == src_rank else "Receiving",
+            backend,
+            world_size,
+        )
+        t0 = time.perf_counter()
+
+        all_tensors = [p.data for _, p in model.named_parameters() if p.numel() > 0] + [
+            b.data for _, b in model.named_buffers() if b.numel() > 0
+        ]
+        if not all_tensors:
+            return
+
+        count = len(all_tensors)
+        total_bytes = sum(t.numel() * t.element_size() for t in all_tensors)
+
+        # Group tensors into buckets by dtype and maximum byte size
+        buckets: list[list[torch.Tensor]] = []
+        curr_bucket: list[torch.Tensor] = []
+        curr_bytes = 0
+
+        for tensor in all_tensors:
+            t_bytes = tensor.numel() * tensor.element_size()
+            if curr_bucket and (curr_bucket[0].dtype != tensor.dtype or curr_bytes + t_bytes > bucket_cap_bytes):
+                buckets.append(curr_bucket)
+                curr_bucket = []
+                curr_bytes = 0
+            curr_bucket.append(tensor)
+            curr_bytes += t_bytes
+
+        if curr_bucket:
+            buckets.append(curr_bucket)
+
+        for bucket in buckets:
+            tot_numel = sum(t.numel() for t in bucket)
+            dtype = bucket[0].dtype
+
+            if is_cpu_backend:
+                if rank == src_rank:
+                    flat_cpu = torch.cat([t.detach().to("cpu").reshape(-1) for t in bucket])
+                else:
+                    flat_cpu = torch.empty(tot_numel, dtype=dtype)
+                torch.distributed.broadcast(flat_cpu, src=src_rank)
+                if rank != src_rank:
+                    offset = 0
+                    for t in bucket:
+                        n = t.numel()
+                        t.copy_(flat_cpu[offset : offset + n].view_as(t))
+                        offset += n
+                del flat_cpu
+            else:
+                if rank == src_rank:
+                    flat_cpu = torch.cat([t.detach().to("cpu").reshape(-1) for t in bucket])
+                    dev_flat = flat_cpu.to(target_device, non_blocking=False)
+                    del flat_cpu
+                    torch.distributed.broadcast(dev_flat, src=src_rank)
+                    del dev_flat
+                else:
+                    dev_flat = torch.empty(tot_numel, dtype=dtype, device=target_device)
+                    torch.distributed.broadcast(dev_flat, src=src_rank)
+                    flat_cpu = dev_flat.to("cpu", non_blocking=False)
+                    del dev_flat
+                    offset = 0
+                    for t in bucket:
+                        n = t.numel()
+                        t.copy_(flat_cpu[offset : offset + n].view_as(t), non_blocking=False)
+                        offset += n
+                    del flat_cpu
+
+        if not is_cpu_backend and target_device.type != "cpu":
+            from vllm_omni.platforms import current_omni_platform
+
+            current_omni_platform.synchronize()
+        torch.distributed.barrier()
+
+        elapsed = time.perf_counter() - t0
+        gb = total_bytes / 1e9
+        logger.info(
+            "Worker %d: Shared weight broadcast complete (%d tensors, %d buckets, %.2f GB) in %.2fs (%.2f GB/s)",
+            rank,
+            count,
+            len(buckets),
+            gb,
+            elapsed,
+            gb / max(elapsed, 1e-6),
+        )
+
     def load_model(
         self,
         load_device: str,
@@ -654,6 +754,10 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
                 if hwr_active and hwr_state is not None:
                     self.host_weight_plan = cast(HostWeightPlan | None, hwr_state.get("plan"))
                 if dit_distributed_offload and not hwr_active and not self._force_canonical_load:
+                    lora_backend = getattr(self.od_config, "lora_backend", None)
+                    has_distilled_lora = lora_backend in (LoRABackend.DISTILL, "distill") and bool(
+                        getattr(self.od_config, "lora_path", None)
+                    )
                     plan_result = build_checkpoint_mmap_plan(
                         model,
                         dit_modules=tuple(zip(modules.dit_names, modules.dits)),
@@ -662,6 +766,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
                         tensor_parallel_size=tensor_parallel_size,
                         use_hsdp=use_hsdp,
                         online_quantization=any(self._has_online_quant(dit) for dit in modules.dits),
+                        has_distilled_lora=has_distilled_lora,
                     )
                     self.host_weight_plan = plan_result.plan
 
@@ -689,7 +794,6 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
                             sources=ordinary_sources,
                             planned_weights=host_weight_plan.bindings,
                         )
-                    self._maybe_fuse_distilled_lora(model)
                 else:
                     if dit_distributed_offload and plan_result is not None:
                         logger.info(
@@ -1097,8 +1201,39 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
         if load_format == "diffusers":
             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)
+        world_size = 1
+        rank = 0
+        if torch.distributed.is_initialized():
+            world_size = torch.distributed.get_world_size()
+            rank = torch.distributed.get_rank()
+
+        has_online_quant = self._has_online_quant(model) or (
+            self.quant_config is not None and not getattr(self.quant_config, "is_checkpoint_quantized", False)
+        )
+        enable_broadcast = bool(getattr(self.od_config, "enable_broadcast_weight_load", False)) and world_size > 1
+
+        if enable_broadcast and has_online_quant:
+            logger.info(
+                "Worker %d: Online quantization detected; falling back to ordinary per-rank weight loading for HSDP",
+                rank,
+            )
+            enable_broadcast = False
+
+        if enable_broadcast:
+            if rank == 0:
+                self.load_weights(model)
+                self._maybe_fuse_distilled_lora(model)
+            self._broadcast_model_weights(model, target_device=target_device, src_rank=0)
+            if (
+                rank != 0
+                and getattr(self.od_config, "lora_backend", None) in (LoRABackend.DISTILL, "distill")
+                and getattr(self.od_config, "lora_path", None)
+                and hasattr(model, "load_lora_weights")
+            ):
+                setattr(model, "lora_is_fused", True)
+        else:
+            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_loader/host_weight_plan.py

@@ -322,6 +322,7 @@ def build_checkpoint_mmap_plan(
     tensor_parallel_size: int,
     use_hsdp: bool,
     online_quantization: bool,
+    has_distilled_lora: bool = False,
 ) -> HostWeightPlanResult:
     """Build a complete direct-checkpoint plan or return a fallback reason."""
     if tensor_parallel_size != 1:
@@ -330,6 +331,8 @@ def build_checkpoint_mmap_plan(
         return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
     if online_quantization:
         return HostWeightPlanResult(None, "online quantization requires the ordinary loader")
+    if has_distilled_lora:
+        return HostWeightPlanResult(None, "distilled LoRA requires the ordinary loader")
 
     remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None)
     if not callable(remap_fn):

vllm_omni/engine/arg_utils.py

@@ -572,6 +572,7 @@ class OrchestratorArgs:
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False
     enable_multithread_weight_load: bool = True
+    enable_broadcast_weight_load: bool = False
     num_weight_load_threads: int = 4
     diffusion_offload_config: dict[str, Any] | None = None
     # Compatibility aliases for existing callers and model-specific stage

vllm_omni/entrypoints/cli/serve.py

@@ -726,6 +726,13 @@ class OmniServeCommand(CLISubcommand):
             default=True,
             help="Disable multi-threaded safetensors loading (default: enabled with 4 threads).",
         )
+        omni_config_group.add_argument(
+            "--enable-broadcast-weight-load",
+            action="store_true",
+            dest="enable_broadcast_weight_load",
+            default=False,
+            help="Enable Rank-0 shared weight broadcast across workers for HSDP (default: disabled).",
+        )
         omni_config_group.add_argument(
             "--num-weight-load-threads",
             type=int,