diff --git a/docs/mkdocs/hooks/search_index.py b/docs/mkdocs/hooks/search_index.py index 1b9e268d9..8cf23e869 100644 --- a/docs/mkdocs/hooks/search_index.py +++ b/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 diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index da8ffb452..fdb2cd4b9 100644 --- a/tests/config/test_omni_config.py +++ b/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 diff --git a/tests/diffusion/lora/test_loader.py b/tests/diffusion/lora/test_loader.py index deb5ec80d..4fd876937 100644 --- a/tests/diffusion/lora/test_loader.py +++ b/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 diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 565949ff6..c7a87ac34 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/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")) diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 250aff884..f20538284 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/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", diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 2b016be0d..f0752077b 100644 --- a/vllm_omni/config/omni_config.py +++ b/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", } diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 0e962b4cc..063a95b0e 100644 --- a/vllm_omni/config/stage_config.py +++ b/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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index d27cc36e1..7a7de3bb1 100644 --- a/vllm_omni/diffusion/data.py +++ b/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 diff --git a/vllm_omni/diffusion/lora/loader.py b/vllm_omni/diffusion/lora/loader.py index a13363c7a..eb3f35188 100644 --- a/vllm_omni/diffusion/lora/loader.py +++ b/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 diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 27774aa37..78247e7b9 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/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, diff --git a/vllm_omni/diffusion/model_loader/host_weight_plan.py b/vllm_omni/diffusion/model_loader/host_weight_plan.py index e6301b0ec..f0ea8024d 100644 --- a/vllm_omni/diffusion/model_loader/host_weight_plan.py +++ b/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): diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index ee5552751..20319de86 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/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 diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 598cbe893..f9297aa0c 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/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,