vllm-omni/tests/diffusion/model_loader/test_diffusers_loader.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""
Tests for the DiffusersPipelineLoader.
"""

import json
import os
import tempfile
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest
import torch
import torch.nn as nn
from safetensors.torch import save_file
from vllm.config.load import LoadConfig

import vllm_omni.diffusion.model_loader.diffusers_loader as loader_module
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,
    HostWeightPlanResult,
    TensorBinding,
    build_checkpoint_binding_plan,
)
from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
from vllm_omni.diffusion.models.helios import HeliosPipeline
from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
from vllm_omni.diffusion.registry import initialize_model
from vllm_omni.quantization.component_config import ComponentQuantizationConfig
from vllm_omni.transformers_utils.repo_utils import hf_api

pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]

model_path = "hf-internal-testing/tiny-helios-modular-pipe"


@pytest.fixture(scope="module")
def prefetch_helios_model():
    """Downloads the tiny helios model prior to running a test."""
    hf_api().snapshot_download(model_path)


@pytest.fixture(scope="function")
def mock_tp_group(mocker):
    """Mocks the tensor parallel group; this is needed to initialize the Helios model."""
    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
    mock_group = mocker.MagicMock()
    mock_group.world_size = 1
    mock_group.rank_in_group = 0
    mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)


class _DummyPipelineModel(nn.Module):
    def __init__(self, *, source_prefix: str):
        super().__init__()
        self.transformer = nn.Linear(2, 2, bias=False)
        self.vae = nn.Linear(2, 2, bias=False)
        self.weights_sources = [
            DiffusersPipelineLoader.ComponentSource(
                model_or_path="dummy",
                subfolder="transformer",
                revision=None,
                prefix=source_prefix,
                fall_back_to_pt=True,
            )
        ]

    def load_weights(self, weights):
        params = dict(self.named_parameters())
        loaded: set[str] = set()
        for name, tensor in weights:
            if name not in params:
                continue
            params[name].data.copy_(tensor.to(dtype=params[name].dtype))
            loaded.add(name)
        return loaded


class _HWRTransformer(nn.Module):
    host_weight_restore_contract = FinalLayoutModelContract(
        implementation_id="test-hwr-transformer",
        version="1",
    )

    def __init__(self):
        super().__init__()
        self.weight = nn.Parameter(torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2))

    def validate_restored_host_weights(self):
        assert self.weight.dtype is torch.bfloat16


class _HWRPipeline(nn.Module):
    def __init__(self, source_root: Path):
        super().__init__()
        self.transformer = _HWRTransformer()
        self.load_count = 0
        self.weights_sources = [
            DiffusersPipelineLoader.ComponentSource(
                model_or_path=str(source_root),
                subfolder=None,
                revision=None,
                prefix="transformer.",
                fall_back_to_pt=False,
            )
        ]

    def load_weights(self, weights):
        loaded: set[str] = set()
        params = dict(self.named_parameters())
        for name, tensor in weights:
            if name in params:
                params[name].data.copy_(tensor.to(dtype=params[name].dtype))
                loaded.add(name)
        self.load_count += 1
        return loaded


def _hwr_config(model: str | Path, root: Path, *, mode: str = "preferred") -> SimpleNamespace:
    return SimpleNamespace(
        model=str(model),
        dtype=torch.bfloat16,
        host_weight_runtime_mode=mode,
        host_weight_runtime_root=str(root),
        enable_distributed_layerwise_offload=True,
        dlo_use_allgather=False,
        lora_path=None,
        quantization_config=None,
        diffusion_attention_config=None,
        parallel_config=SimpleNamespace(
            use_hsdp=False,
            data_parallel_size=1,
            sequence_parallel_size=1,
            tensor_parallel_size=1,
            pipeline_parallel_size=1,
            cfg_parallel_size=1,
            enable_expert_parallel=False,
            ulysses_degree=1,
            ring_degree=1,
            allgather_degree=1,
            ulysses_mode="strict",
        ),
    )


def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoader:
    od_config = SimpleNamespace(
        dtype=torch.float32,
        parallel_config=SimpleNamespace(use_hsdp=False),
        quantization_config=None,
    )
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)

    loader.counter_before_loading_weights = 0.0
    loader.counter_after_loading_weights = 0.0

    def _iter_weights(_model):
        for name in weight_names:
            yield name, torch.zeros((2, 2))

    loader.get_all_weights = _iter_weights  # type: ignore[assignment]
    return loader


def test_serialized_torchao_component_uses_pytorch_iterator(mocker):
    torchao_config = SimpleNamespace(
        get_name=lambda: "torchao",
        is_checkpoint_torchao_serialized=True,
    )
    loader = _make_loader_with_weights([])
    loader.quant_config = ComponentQuantizationConfig({"transformer": torchao_config})
    loader.load_config.pt_load_map_location = "cpu"
    files = ["/weights/model-10.bin", "/weights/model-2.bin"]
    mocker.patch.object(loader, "_prepare_weights", return_value=("/weights", files, False))
    pt_iterator = mocker.patch.object(
        loader_module,
        "pt_weights_iterator",
        return_value=iter([("block.weight", torch.ones(1))]),
    )
    safetensors_iterator = mocker.patch.object(loader_module, "safetensors_weights_iterator")
    multithread_iterator = mocker.patch.object(loader_module, "multi_thread_safetensors_weights_iterator")
    source = DiffusersPipelineLoader.ComponentSource(
        model_or_path="/weights",
        subfolder=None,
        revision=None,
        prefix="transformer.",
    )

    weights = list(loader._get_weights_iterator(source))

    pt_iterator.assert_called_once_with(
        ["/weights/model-2.bin", "/weights/model-10.bin"],
        loader.load_config.use_tqdm_on_load,
        "cpu",
    )
    safetensors_iterator.assert_not_called()
    multithread_iterator.assert_not_called()
    assert weights[0][0] == "transformer.block.weight"
    assert torch.equal(weights[0][1], torch.ones(1))


def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
    tmp_path: Path,
    monkeypatch,
):
    canonical_root = tmp_path / "canonical"
    canonical_root.mkdir()
    save_file(
        {"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)},
        str(canonical_root / "model.safetensors"),
    )
    store_root = tmp_path / "hwr-store"
    hash_calls = 0
    original_sha256 = source_identity_module._sha256_file

    def counted_sha256(path: Path, state: object) -> str:
        nonlocal hash_calls
        hash_calls += 1
        return original_sha256(path, state)  # type: ignore[arg-type]

    monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256)

    def make_loader() -> tuple[DiffusersPipelineLoader, _HWRPipeline]:
        loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(canonical_root, store_root))
        pipeline = _HWRPipeline(canonical_root)
        monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline)
        return loader, pipeline

    cold_loader, cold_model = make_loader()
    cold = cold_loader.load_model(load_device="cpu", device=torch.device("cpu"))
    assert cold is cold_model
    assert cold_model.load_count == 1
    assert cold_loader._hwr_state is not None

    warm_loader, warm_model = make_loader()
    monkeypatch.setattr(
        warm_loader,
        "_process_weights_after_loading",
        lambda *args, **kwargs: pytest.fail("warm HWR restore re-entered byte-changing finalization"),
    )
    warm = warm_loader.load_model(load_device="cpu", device=torch.device("cpu"))

    assert warm is warm_model
    assert warm_model.load_count == 0
    assert torch.equal(warm_model.transformer.weight, cold_model.transformer.weight)
    from vllm_omni.diffusion.offloader.startup import take_offload_startup_state

    startup_state = take_offload_startup_state(warm)
    assert startup_state is not None
    warm_plan = startup_state.host_weight_plan
    assert warm_plan is not None
    assert warm_plan.lease_carrier is not None
    warm_plan.lease_carrier.close()
    assert hash_calls == 1
    assert len(tuple((store_root / "source-digests-v1" / "entries").glob("*.json"))) == 1


def test_maybe_fuse_distilled_lora_skips_when_lora_path_unset():
    cfg = SimpleNamespace(
        lora_backend="distill",
        lora_path=None,
        lora_scale=1.0,
        dtype=torch.bfloat16,
        quantization_config=None,
        parallel_config=SimpleNamespace(use_hsdp=False),
    )
    loader = DiffusersPipelineLoader(LoadConfig(), cfg)

    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 False


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

    loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(tmp_path, tmp_path / "store"))
    models: list[_DummyPipelineModel] = []

    def init_model(*args, **kwargs):
        del args, kwargs
        model = _DummyPipelineModel(source_prefix="transformer.")
        models.append(model)
        return model

    def commit_error(*args, **kwargs):
        del args, kwargs
        raise loader_module._HWRCommitError("restore commit failed")

    monkeypatch.setattr(loader, "_init_from_load_format", init_model)
    monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: ())
    monkeypatch.setattr(loader, "_resolve_hwr", commit_error)
    monkeypatch.setattr(loader, "load_weights", lambda *args, **kwargs: None)
    monkeypatch.setattr(loader, "_process_weights_after_loading", lambda *args, **kwargs: None)
    monkeypatch.setattr(loader, "_apply_skip_softmax_calibration", lambda *args, **kwargs: None)

    recovered = loader.load_model(load_device="cpu", device=torch.device("cpu"))

    assert len(models) == 2
    assert recovered is models[1]
    assert loader.take_host_weight_plan() is None


def test_required_hwr_miss_fails_before_ordinary_loading_or_publication(
    tmp_path: Path,
    monkeypatch,
):
    canonical_root = tmp_path / "canonical"
    canonical_root.mkdir()
    save_file(
        {"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)},
        str(canonical_root / "model.safetensors"),
    )
    loader = DiffusersPipelineLoader(
        LoadConfig(),
        _hwr_config(canonical_root, tmp_path / "empty-store", mode="required"),
    )
    pipeline = _HWRPipeline(canonical_root)
    monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline)

    with pytest.raises(RuntimeError, match="Host Weight Runtime resolution failed"):
        loader.load_model(load_device="cpu", device=torch.device("cpu"))

    assert pipeline.load_count == 0
    assert loader.take_host_weight_plan() is None


def _make_dlo_online_quant_config(dp_size: int = 2) -> OmniDiffusionConfig:
    return OmniDiffusionConfig(
        model="",
        dtype=torch.float32,
        quantization_config="fp8",
        parallel_config=DiffusionParallelConfig(
            data_parallel_size=dp_size,
            sequence_parallel_size=1,
        ),
        enable_distributed_layerwise_offload=True,
        dlo_use_allgather=True,
    )


@pytest.mark.parametrize(
    ("dist_offload", "use_allgather", "mode"),
    [
        (False, False, "preferred"),
        (True, True, "preferred"),
        (True, False, "disabled"),
    ],
)
def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
    monkeypatch,
    tmp_path,
    dist_offload,
    use_allgather,
    mode,
):
    """Disabled and AllGather paths must never construct or probe HWR."""
    from vllm_omni.host_weight_runtime import HostWeightRuntime

    root = tmp_path / "must-not-be-touched"
    loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root, mode=mode))
    model = _DummyPipelineModel(source_prefix="transformer.")
    modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,))

    def unexpected_store_construction(*args, **kwargs):
        raise AssertionError(f"HWR store interaction was not eligible: {args}, {kwargs}")

    monkeypatch.setattr(HostWeightRuntime, "from_config", unexpected_store_construction)
    assert (
        loader._resolve_hwr(
            model,
            modules,
            dist_offload=dist_offload,
            use_allgather=use_allgather,
            load_format="default",
            sources=tuple(model.weights_sources),
        )
        is None
    )
    assert not root.exists()


def test_required_hwr_rejects_a_model_without_a_restore_contract(tmp_path):
    loader = DiffusersPipelineLoader(
        LoadConfig(),
        _hwr_config("dummy-model", tmp_path / "store", mode="required"),
    )
    model = _DummyPipelineModel(source_prefix="transformer.")
    modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,))

    with pytest.raises(ValueError, match="restore contract"):
        loader._resolve_hwr(
            model,
            modules,
            dist_offload=True,
            use_allgather=False,
            load_format="default",
            sources=tuple(model.weights_sources),
        )


@pytest.mark.parametrize("offline", [False, True])
def test_prepare_weights_honors_component_index_and_explicit_override(tmp_path, mocker, offline):
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod

    snapshot = tmp_path / "snapshot"
    transformer = snapshot / "transformer"
    transformer.mkdir(parents=True)
    indexed_files = [f"diffusion_pytorch_model-{index:05d}-of-00002.safetensors" for index in (1, 2)]
    stale_file = "diffusion_pytorch_model-00001-of-00008.safetensors"
    index_path = transformer / "diffusion_pytorch_model.safetensors.index.json"
    index_path.write_text(
        json.dumps({"weight_map": {f"weight.{index}": filename for index, filename in enumerate(indexed_files)}})
    )
    for filename in indexed_files + [stale_file]:
        (transformer / filename).touch()

    mocker.patch.object(loader_mod.huggingface_hub.constants, "HF_HUB_OFFLINE", offline)

    def download_index(*, filename, **_kwargs):
        if filename == "transformer/diffusion_pytorch_model.safetensors.index.json":
            return str(index_path)
        raise loader_mod.huggingface_hub.errors.EntryNotFoundError(filename)

    hub_api = mocker.Mock()
    hub_api.hf_hub_download.side_effect = download_index
    mocker.patch.object(loader_mod, "hf_api", return_value=hub_api)
    indexed_download = mocker.patch.object(loader_mod, "download_weights_from_hf_specific", return_value=str(snapshot))
    generic_download = mocker.patch.object(loader_mod, "download_weights_from_hf", return_value=str(snapshot))

    loader = _make_loader_with_weights([])
    cache_dir = str(tmp_path / "cache")
    loader.load_config.download_dir = cache_dir
    folder, files, use_safetensors = loader._prepare_weights(
        "org/model",
        subfolder="transformer",
        revision="revision",
        fall_back_to_pt=True,
        allow_patterns_overrides=None,
    )

    assert folder == str(transformer)
    assert [str(transformer / filename) for filename in indexed_files] == files
    assert use_safetensors
    assert hub_api.hf_hub_download.call_count == len(loader_mod.SAFETENSORS_INDEX_FILES)
    hub_api.hf_hub_download.assert_any_call(
        repo_id="org/model",
        filename="transformer/diffusion_pytorch_model.safetensors.index.json",
        cache_dir=cache_dir,
        revision="revision",
        local_files_only=offline,
    )
    indexed_download.assert_called_once_with(
        model_name_or_path="org/model",
        cache_dir=cache_dir,
        allow_patterns=[f"transformer/{filename}" for filename in indexed_files],
        revision="revision",
        ignore_patterns=loader.load_config.ignore_patterns,
        require_all=True,
    )

    _, override_files, _ = loader._prepare_weights(
        "org/model",
        subfolder="transformer",
        revision="revision",
        fall_back_to_pt=True,
        allow_patterns_overrides=[stale_file],
    )
    assert override_files == [str(transformer / stale_file)]
    generic_download.assert_called_once_with(
        "org/model",
        cache_dir,
        [stale_file],
        "revision",
        subfolder="transformer",
        ignore_patterns=loader.load_config.ignore_patterns,
    )


def test_prepare_local_weights_honors_component_index(tmp_path):
    transformer = tmp_path / "transformer"
    transformer.mkdir()
    indexed_file = "diffusion_pytorch_model-00001-of-00001.safetensors"
    stale_file = "diffusion_pytorch_model-00001-of-00008.safetensors"
    (transformer / "diffusion_pytorch_model.safetensors.index.json").write_text(
        json.dumps({"weight_map": {"weight": indexed_file}})
    )
    for filename in (indexed_file, stale_file):
        (transformer / filename).touch()

    folder, files, use_safetensors = _make_loader_with_weights([])._prepare_weights(
        tmp_path,
        subfolder="transformer",
        revision=None,
        fall_back_to_pt=True,
        allow_patterns_overrides=None,
    )

    assert folder == str(transformer)
    assert files == [str(transformer / indexed_file)]
    assert use_safetensors


def test_prepare_local_bin_index_is_authoritative_and_safetensors_remains_preferred(tmp_path):
    transformer = tmp_path / "transformer"
    transformer.mkdir()
    indexed_files = [f"diffusion_pytorch_model-{index:05d}-of-00002.bin" for index in (1, 2)]
    stale_file = "diffusion_pytorch_model-00001-of-00008.bin"
    (transformer / "diffusion_pytorch_model.bin.index.json").write_text(
        json.dumps({"weight_map": {f"weight.{index}": filename for index, filename in enumerate(indexed_files)}})
    )
    for filename in (*indexed_files, stale_file):
        (transformer / filename).touch()

    _, files, use_safetensors = _make_loader_with_weights([])._prepare_weights(
        tmp_path,
        subfolder="transformer",
        revision=None,
        fall_back_to_pt=True,
        allow_patterns_overrides=None,
    )

    assert files == [str(transformer / filename) for filename in indexed_files]
    assert not use_safetensors

    safetensors_file = "diffusion_pytorch_model.safetensors"
    (transformer / "diffusion_pytorch_model.safetensors.index.json").write_text(
        json.dumps({"weight_map": {"weight": safetensors_file}})
    )
    (transformer / safetensors_file).touch()
    _, files, use_safetensors = _make_loader_with_weights([])._prepare_weights(
        tmp_path,
        subfolder="transformer",
        revision=None,
        fall_back_to_pt=True,
        allow_patterns_overrides=None,
    )

    assert files == [str(transformer / safetensors_file)]
    assert use_safetensors


def test_prepare_weights_rejects_polluted_offline_cache_without_index(tmp_path, mocker):
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod

    snapshot = tmp_path / "snapshot"
    transformer = snapshot / "transformer"
    transformer.mkdir(parents=True)
    for shard_count in (4, 8):
        for shard in range(1, shard_count + 1):
            (transformer / f"diffusion_pytorch_model-{shard:05d}-of-{shard_count:05d}.safetensors").touch()

    mocker.patch.object(loader_mod.huggingface_hub.constants, "HF_HUB_OFFLINE", True)
    hub_api = mocker.Mock()
    hub_api.hf_hub_download.side_effect = loader_mod.huggingface_hub.errors.EntryNotFoundError("index is not cached")
    mocker.patch.object(loader_mod, "hf_api", return_value=hub_api)
    mocker.patch.object(loader_mod, "download_weights_from_hf", return_value=str(snapshot))

    loader = _make_loader_with_weights([])
    with pytest.raises(ValueError, match="conflicting shard totals"):
        loader._prepare_weights(
            "org/model",
            subfolder="transformer",
            revision="revision",
            fall_back_to_pt=True,
            allow_patterns_overrides=None,
        )

    assert hub_api.hf_hub_download.call_count == len(loader_mod.SAFETENSORS_INDEX_FILES) + len(
        loader_mod.PT_INDEX_FILES
    )
    assert all(call.kwargs["local_files_only"] for call in hub_api.hf_hub_download.call_args_list)


def test_strict_check_only_validates_source_prefix_parameters():
    model = _DummyPipelineModel(source_prefix="transformer.")
    loader = _make_loader_with_weights(["transformer.weight"])

    # Should not require VAE parameters because they are outside weights_sources.
    loader.load_weights(model)


def test_strict_check_raises_when_source_parameters_are_missing():
    model = _DummyPipelineModel(source_prefix="transformer.")
    loader = _make_loader_with_weights([])

    with pytest.raises(ValueError, match="transformer.weight"):
        loader.load_weights(model)


def test_empty_source_prefix_keeps_full_model_strict_check():
    model = _DummyPipelineModel(source_prefix="")
    loader = _make_loader_with_weights(["transformer.weight"])

    with pytest.raises(ValueError, match="vae.weight"):
        loader.load_weights(model)


def test_stream_online_quant_weights_offloads_layers_after_processing():
    from vllm.model_executor.model_loader.reload.layerwise import (
        get_layerwise_info,
    )

    events: list[str] = []

    class _OnlineQuantMethod:
        uses_meta_device = True

    class _TrackedLayer(nn.Linear):
        def __init__(self, name: str):
            super().__init__(2, 2, bias=False)
            self.name = name
            self.quant_method = _OnlineQuantMethod()
            get_layerwise_info(self).load_numel_total = self.weight.numel()

        def to(self, *args, **kwargs):
            events.append(self.name)
            return super().to(*args, **kwargs)

    class _StreamingModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.first = _TrackedLayer("first")
            self.second = _TrackedLayer("second")

    model = _StreamingModel()
    weights = iter(
        [
            ("first.weight", torch.zeros((2, 2))),
            ("second.weight", torch.zeros((2, 2))),
        ]
    )
    streamed = DiffusersPipelineLoader._stream_online_quant_weights_to_cpu(model, weights)

    assert next(streamed)[0] == "first.weight"
    get_layerwise_info(model.first).reset()
    assert next(streamed)[0] == "second.weight"
    assert events == ["first"]

    get_layerwise_info(model.second).reset()
    with pytest.raises(StopIteration):
        next(streamed)
    assert events == ["first", "second"]


def test_process_weights_skips_completed_online_quant_layer(monkeypatch):
    from unittest.mock import Mock

    from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
    from vllm.model_executor.model_loader.reload import layerwise

    class _TrackedLayer(nn.Linear):
        def __init__(self):
            super().__init__(2, 2, bias=False)
            self.to_calls: list[object] = []

        def to(self, *args, **kwargs):
            self.to_calls.append(args[0] if args else kwargs.get("device"))
            return super().to(*args, **kwargs)

    model = nn.Module()
    model.layer = _TrackedLayer()
    quant_method = Mock(spec=QuantizeMethodBase)
    quant_method.uses_meta_device = True
    model.layer.quant_method = quant_method
    model.layer._already_called_process_weights_after_loading = True
    finalize = Mock()
    monkeypatch.setattr(layerwise, "finalize_layerwise_processing", finalize)

    loader = _make_loader_with_weights([])
    loader._process_weights_after_loading(model, torch.device("cuda"))

    finalize.assert_called_once_with(model, model_config=None)
    quant_method.process_weights_after_loading.assert_not_called()
    assert model.layer.to_calls == []


class _ConfigAwareModel(nn.Module):
    def __init__(self, *, od_config):
        super().__init__()
        self.captured_config = get_current_diffusion_config()
        self.seen_config_during_init = get_current_diffusion_config_or_none()
        self.od_config = od_config


def test_initialize_model_sets_current_diffusion_config_during_model_construction(monkeypatch):
    import vllm_omni.diffusion.registry as registry_mod

    od_config = SimpleNamespace(
        model_class_name="DummyPipeline",
        parallel_config=DiffusionParallelConfig(vae_patch_parallel_size=1, sequence_parallel_size=1),
        vae_use_slicing=False,
        vae_use_tiling=False,
    )

    monkeypatch.setattr(
        registry_mod.DiffusionModelRegistry,
        "_try_load_model_cls",
        staticmethod(lambda _name: _ConfigAwareModel),
    )
    monkeypatch.setattr(registry_mod, "_apply_sequence_parallel_if_enabled", lambda *_args, **_kwargs: None)

    model = initialize_model(od_config)

    assert model.captured_config is od_config
    assert model.seen_config_during_init is od_config
    assert get_current_diffusion_config_or_none() is None


def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod

    class _DeviceContext:
        def __init__(self, device_type: str):
            self.type = device_type

        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc, tb):
            return False

    od_config = SimpleNamespace(
        dtype=torch.float32,
        parallel_config=SimpleNamespace(use_hsdp=False),
        quantization_config=None,
    )

    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
    loader.load_weights = lambda model: None  # type: ignore[assignment]
    loader._process_weights_after_loading = lambda model, target_device: None  # type: ignore[assignment]

    monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel)
    monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu"))

    model = loader.load_model(
        load_device="cpu",
        load_format="custom_pipeline",
        custom_pipeline_name="tests.dummy.ConfigAwarePipeline",
    )

    assert model.captured_config is od_config
    assert model.seen_config_during_init is od_config
    assert get_current_diffusion_config_or_none() is None


def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch):
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod

    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,
        model="unused",
    )
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
    model = nn.Module()
    model.transformer = nn.Linear(2, 2, bias=False)
    plan = HostWeightPlan(
        backing_kind="checkpoint_mmap",
        bindings={},
    )
    loaded_ordinary_weights = False

    def load_weights(_model):
        nonlocal loaded_ordinary_weights
        loaded_ordinary_weights = True

    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
    loader.load_weights = load_weights  # type: ignore[method-assign]
    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
    monkeypatch.setattr(
        loader_mod,
        "build_checkpoint_mmap_plan",
        lambda *_args, **_kwargs: HostWeightPlanResult(plan),
    )

    assert loader.load_model(load_device="cpu") is model
    assert not loaded_ordinary_weights
    assert loader.take_host_weight_plan() is plan
    assert loader.take_host_weight_plan() is None


def test_compact_rank_local_layer_offload_skips_dlo_mmap_planning(monkeypatch):
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
    from vllm_omni.diffusion.offloader.config import materialize_legacy_offload_flags

    od_config = SimpleNamespace(
        dtype=torch.float32,
        parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
        quantization_config=None,
        diffusion_offload_config={"mode": "layer", "components": ["dit"]},
        enable_cpu_offload=False,
        enable_layerwise_offload=False,
        enable_distributed_layerwise_offload=False,
        dlo_use_allgather=True,
        dlo_resident_layers=0,
        pin_cpu_memory=True,
        model="unused",
    )
    materialize_legacy_offload_flags(od_config)
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
    model = nn.Module()
    model.transformer = nn.Linear(2, 2, bias=False)
    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]
    monkeypatch.setattr(
        loader_mod,
        "build_checkpoint_mmap_plan",
        lambda *_args, **_kwargs: pytest.fail("ordinary rank-local offload must not plan DLO mmap"),
    )

    assert loader.load_model(load_device="cpu") is model
    assert calls == ["load", "process"]
    assert loader.take_host_weight_plan() is None


class _UnsupportedOnlineQuantMethod:
    uses_meta_device = True


def _compact_layer_offload_config(
    components: list[str],
    layer_options: dict[str, dict[str, str]],
) -> SimpleNamespace:
    return SimpleNamespace(
        dtype=torch.float32,
        parallel_config=SimpleNamespace(
            use_hsdp=False,
            tensor_parallel_size=1,
            data_parallel_size=2,
            sequence_parallel_size=1,
        ),
        quantization_config=None,
        diffusion_offload_config={
            "mode": "layer",
            "components": components,
            "layer_options": layer_options,
        },
        enable_cpu_offload=False,
        enable_layerwise_offload=False,
        enable_distributed_layerwise_offload=False,
        dlo_use_allgather=True,
        dlo_resident_layers=0,
        pin_cpu_memory=False,
        host_weight_runtime_mode="disabled",
        model="unused",
    )


def _stub_ordinary_loader(loader, model, monkeypatch, loader_mod) -> None:
    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
    loader.load_weights = lambda _model: None  # type: ignore[method-assign]
    loader._process_weights_after_loading = lambda *_args: None  # type: ignore[method-assign]
    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
    monkeypatch.setattr(
        loader_mod,
        "build_checkpoint_mmap_plan",
        lambda *_args, **_kwargs: HostWeightPlanResult(None, "test fallback"),
    )


def test_text_encoder_allgather_rejects_unsupported_online_quantization(monkeypatch):
    config = _compact_layer_offload_config(
        ["text_encoder"],
        {"text_encoder": {"weight_transfer": "allgather"}},
    )
    loader = DiffusersPipelineLoader(LoadConfig(), config)
    model = nn.Module()
    model.text_encoder = nn.Linear(2, 2, bias=False)
    model.text_encoder.quant_method = _UnsupportedOnlineQuantMethod()
    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]

    with pytest.raises(ValueError, match="unsupported online methods: _UnsupportedOnlineQuantMethod"):
        loader.load_model(load_device="cpu")


def test_dit_allgather_ignores_unsupported_quantization_on_unselected_encoder(monkeypatch):
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod

    config = _compact_layer_offload_config(
        ["dit"],
        {"dit": {"weight_transfer": "allgather"}},
    )
    loader = DiffusersPipelineLoader(LoadConfig(), config)
    model = nn.Module()
    model.transformer = nn.Linear(2, 2, bias=False)
    model.text_encoder = nn.Linear(2, 2, bias=False)
    model.text_encoder.quant_method = _UnsupportedOnlineQuantMethod()
    _stub_ordinary_loader(loader, model, monkeypatch, loader_mod)

    assert loader.load_model(load_device="cpu") is model


def test_rank_local_encoder_quantization_does_not_mark_dit_mmap_plan_online(monkeypatch):
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod

    config = _compact_layer_offload_config(
        ["dit", "text_encoder"],
        {
            "dit": {"weight_transfer": "allgather"},
            "text_encoder": {"weight_transfer": "rank-local"},
        },
    )
    loader = DiffusersPipelineLoader(LoadConfig(), config)
    model = nn.Module()
    model.transformer = nn.Linear(2, 2, bias=False)
    model.text_encoder = nn.Linear(2, 2, bias=False)
    model.text_encoder.quant_method = _UnsupportedOnlineQuantMethod()
    planned_online_quantization: list[bool] = []

    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
    loader.load_weights = lambda _model: None  # type: ignore[method-assign]
    loader._process_weights_after_loading = lambda *_args: None  # type: ignore[method-assign]
    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]

    def capture_plan(*_args, **kwargs):
        planned_online_quantization.append(kwargs["online_quantization"])
        return HostWeightPlanResult(None, "test fallback")

    monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", capture_plan)

    assert loader.load_model(load_device="cpu") is model
    assert planned_online_quantization == [False]


def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod

    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,
        model="unused",
    )
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)

    class MixedSourceModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.transformer = nn.Linear(2, 2, bias=False)
            self.text_encoder = nn.Linear(2, 2, bias=False)
            self.weights_sources = (
                DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="transformer."),
                DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="text_encoder."),
            )
            self.loaded_weight_names: list[str] = []

        def load_weights(self, weights):
            self.loaded_weight_names = [name for name, _ in weights]
            return set(self.loaded_weight_names)

    model = MixedSourceModel()
    plan = HostWeightPlan(
        backing_kind="checkpoint_mmap",
        bindings={
            "transformer.weight": TensorBinding(
                checkpoint_key="weight",
                file_path="unused",
            )
        },
        planned_source_prefixes=frozenset({"transformer."}),
    )
    requested_prefixes: list[str] = []

    def get_weights(source, model=None):
        del model
        requested_prefixes.append(source.prefix)
        yield source.prefix + "weight", torch.ones(2, 2)

    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
    loader._get_weights_iterator = get_weights  # type: ignore[method-assign]
    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
    monkeypatch.setattr(
        loader_mod,
        "build_checkpoint_mmap_plan",
        lambda *_args, **_kwargs: HostWeightPlanResult(plan),
    )

    assert loader.load_model(load_device="cpu") is model
    assert requested_prefixes == ["text_encoder."]
    assert model.loaded_weight_names == ["text_encoder.weight"]
    assert loader.take_host_weight_plan() is plan


def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod

    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,
        model="unused",
    )
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
    model = nn.Module()
    model.transformer = nn.Linear(2, 2, bias=False)
    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]
    monkeypatch.setattr(
        loader_mod,
        "build_checkpoint_mmap_plan",
        lambda *_args, **_kwargs: HostWeightPlanResult(None, "not direct-compatible"),
    )

    assert loader.load_model(load_device="cpu") is model
    assert calls == ["load", "process"]
    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,
    )

    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod

    od_config = _make_dlo_online_quant_config()
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
    model = nn.Module()
    model.transformer = nn.Linear(2, 2, bias=False)
    model.transformer.quant_method = object.__new__(Fp8PerTensorOnlineLinearMethod)
    model.transformer.quant_method.uses_meta_device = True
    calls: list[object] = []
    allowlist_models: list[nn.Module] = []

    original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods

    def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
        allowlist_models.append(candidate)
        return original_allowlist_check(candidate)

    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
    monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
    loader._request_offload_after_quant = lambda _model: 1  # type: ignore[method-assign]
    loader.load_weights = (  # type: ignore[method-assign]
        lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
    )
    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]
    monkeypatch.setattr(
        loader_mod,
        "build_checkpoint_mmap_plan",
        lambda *_args, **_kwargs: HostWeightPlanResult(
            None,
            "online quantization requires the ordinary loader",
        ),
    )

    assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
    assert allowlist_models == [model.transformer]
    assert calls == [("load", True), "process"]
    assert loader.take_host_weight_plan() is None


def test_dlo_allgather_online_int8_uses_ordinary_loader(monkeypatch):
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
    from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod

    od_config = _make_dlo_online_quant_config()
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
    model = nn.Module()
    model.transformer = nn.Linear(2, 2, bias=False)
    model.transformer.quant_method = object.__new__(NPUInt8OnlineLinearMethod)
    model.transformer.quant_method.uses_meta_device = True
    calls: list[object] = []
    allowlist_models: list[nn.Module] = []

    original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods

    def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
        allowlist_models.append(candidate)
        return original_allowlist_check(candidate)

    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
    monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
    loader._request_offload_after_quant = lambda _model: 1  # type: ignore[method-assign]
    loader.load_weights = (  # type: ignore[method-assign]
        lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
    )
    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]
    monkeypatch.setattr(
        loader_mod,
        "build_checkpoint_mmap_plan",
        lambda *_args, **_kwargs: HostWeightPlanResult(
            None,
            "online quantization requires the ordinary loader",
        ),
    )

    assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
    assert allowlist_models == [model.transformer]
    assert calls == [("load", True), "process"]
    assert loader.take_host_weight_plan() is None


def test_dlo_online_quant_group_size_one_skips_allgather_gate(monkeypatch):
    """A DLO group of one runs no weight collective, so an otherwise
    unvalidated online method must not be rejected by the AllGather gate."""

    class UnsupportedOnlineMethod:
        uses_meta_device = True

    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod

    od_config = _make_dlo_online_quant_config(dp_size=1)
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
    model = nn.Module()
    model.transformer = nn.Linear(2, 2, bias=False)
    model.transformer.quant_method = UnsupportedOnlineMethod()
    calls: list[object] = []
    allowlist_models: list[nn.Module] = []

    original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods

    def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
        allowlist_models.append(candidate)
        return original_allowlist_check(candidate)

    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
    monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
    loader._request_offload_after_quant = lambda _model: 1  # type: ignore[method-assign]
    loader.load_weights = (  # type: ignore[method-assign]
        lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
    )
    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]
    monkeypatch.setattr(
        loader_mod,
        "build_checkpoint_mmap_plan",
        lambda *_args, **_kwargs: HostWeightPlanResult(
            None,
            "online quantization requires the ordinary loader",
        ),
    )

    assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
    assert allowlist_models == []
    assert calls == [("load", True), "process"]


def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod

    class UnsupportedOnlineMethod:
        uses_meta_device = True

    od_config = _make_dlo_online_quant_config()
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
    model = nn.Module()
    model.transformer = nn.Linear(2, 2, bias=False)
    model.transformer.quant_method = UnsupportedOnlineMethod()

    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
    monkeypatch.setattr(
        loader_mod,
        "build_checkpoint_mmap_plan",
        lambda *_args, **_kwargs: HostWeightPlanResult(
            None,
            "online quantization requires the ordinary loader",
        ),
    )

    with pytest.raises(ValueError, match="per-tensor FP8, INT8, and MXFP8 linears"):
        loader.load_model(load_device="cpu")


def test_dlo_allgather_allows_unquantized_host_fallback():
    """Host-loaded unquantized fallback layers are plain contiguous bf16 —
    the same runtime layout DLO already shards on the ordinary path — so the
    allowlist must not reject them."""
    from vllm_omni.quantization.int8_config import UnquantizedHostLinearMethod

    model = nn.Module()
    model.transformer = nn.Linear(2, 2, bias=False)
    model.transformer.quant_method = object.__new__(UnquantizedHostLinearMethod)

    assert DiffusersPipelineLoader._unsupported_dlo_allgather_online_quant_methods(model) == ()


def test_dlo_load_model_keeps_host_fallback_on_cpu_through_post_load_sweep(monkeypatch, mocker):
    """Through load_model(): DLO + online quant must build the model inside
    load_unquantizable_fallback_on_cpu(), and the over-wide fallback weight
    must survive the post-load sweep (the _process_weights_after_loading pass
    plus model.to("cpu")) as the same host tensor it was loaded into — never
    bounced to the accelerator."""
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
    from vllm_omni.quantization import int8_config
    from vllm_omni.quantization.int8_config import (
        NPU_QUANT_MATMUL_MAX_OUT_FEATURES,
        DiffusionInt8Config,
        NPUInt8OnlineLinearMethod,
        UnquantizedHostLinearMethod,
    )

    # Same stand-in as TestHostFallbackLoading: the eager fallback delegates to
    # UnquantizedLinearMethod, which reads the TP group.
    mock_group = mocker.Mock()
    mock_group.rank_in_group = 0
    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
    mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)

    od_config = _make_dlo_online_quant_config()
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)

    out_features = NPU_QUANT_MATMUL_MAX_OUT_FEATURES + 1
    built_with_ctx: list[bool] = []

    def copy_loader(param, loaded_weight, *args, **kwargs):
        param.data.copy_(loaded_weight)

    def build_model(*_args, **_kwargs):
        built_with_ctx.append(int8_config._LOAD_UNQUANTIZABLE_FALLBACK_ON_CPU.get())
        quant_config = DiffusionInt8Config(is_checkpoint_int8_serialized=False, activation_scheme="dynamic")
        method = NPUInt8OnlineLinearMethod(quant_config)
        layer = nn.Module()
        layer.quant_method = method
        method.create_weights(
            layer,
            input_size_per_partition=8,
            output_partition_sizes=[out_features],
            input_size=8,
            output_size=out_features,
            params_dtype=torch.bfloat16,
            weight_loader=copy_loader,
        )
        model = nn.Module()
        model.transformer = layer
        return model

    def load_weights(_model, **_kwargs):
        layer = _model.transformer
        loaded = torch.arange(out_features * 8, dtype=torch.float32).reshape(out_features, 8).to(torch.bfloat16)
        layer.weight.weight_loader(layer.weight, loaded)

    loader._init_from_load_format = build_model  # type: ignore[method-assign]
    loader.load_weights = load_weights  # type: ignore[method-assign]
    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
    monkeypatch.setattr(
        loader_mod,
        "build_checkpoint_mmap_plan",
        lambda *_args, **_kwargs: HostWeightPlanResult(
            None,
            "online quantization requires the ordinary loader",
        ),
    )
    # Diffusion DiT models have no vLLM Attention layers, so the upstream
    # layerwise finalize is a no-op here; stub it to keep the real
    # _process_weights_after_loading sweep CPU-only.
    monkeypatch.setattr(
        "vllm.model_executor.model_loader.reload.layerwise.finalize_layerwise_processing",
        lambda *_args, **_kwargs: None,
    )

    model = loader.load_model(load_device="cpu", device=torch.device("cpu"))

    # Construction ran inside load_unquantizable_fallback_on_cpu(), so the
    # over-wide layer took the host-loading fallback...
    assert built_with_ctx == [True]
    assert type(model.transformer.quant_method) is UnquantizedHostLinearMethod
    # ...and the real post-load sweep skipped it via the fully-loaded flag and
    # model.to("cpu") left the loaded host tensor untouched.
    assert model.transformer._already_called_process_weights_after_loading
    assert model.transformer.weight.device.type == "cpu"
    expected = torch.arange(out_features * 8, dtype=torch.float32).reshape(out_features, 8).to(torch.bfloat16)
    assert torch.equal(model.transformer.weight, expected)


@pytest.mark.parametrize(("offload_after_quant", "ctx_expected"), [(True, True), (False, False)])
def test_hsdp_enters_host_fallback_context_only_when_offloading_after_quant(mocker, offload_after_quant, ctx_expected):
    """The HSDP path must apply the same host-fallback bound as the ordinary
    path: with a quant config it initializes on the accelerator, so the
    load_unquantizable_fallback_on_cpu() context is what keeps over-wide
    fallback weights off the device before sharding."""
    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
    from vllm_omni.quantization import int8_config

    od_config = SimpleNamespace(
        dtype=torch.float32,
        parallel_config=SimpleNamespace(
            use_hsdp=True,
            hsdp_replicate_size=1,
            hsdp_shard_size=2,
        ),
        quantization_config=None,
    )
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
    loader.quant_config = object()

    model = nn.Module()
    model.transformer = nn.Linear(2, 2, bias=False)
    ctx_seen: list[bool] = []

    def build_model(*_args, **_kwargs):
        ctx_seen.append(int8_config._LOAD_UNQUANTIZABLE_FALLBACK_ON_CPU.get())
        return model

    loader._init_from_load_format = build_model  # type: ignore[method-assign]
    loader.load_weights = lambda _model: None  # type: ignore[method-assign]
    loader._process_weights_after_loading = lambda *_args: None  # type: ignore[method-assign]
    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(
        "vllm_omni.diffusion.quantization.hsdp_fp8.prepare_fp8_layers_for_fsdp",
        side_effect=lambda _model: None,
    )
    mocker.patch.object(
        loader_mod,
        "apply_hsdp_to_model",
        side_effect=lambda *_args, **_kwargs: None,
    )

    loader._load_model_with_hsdp(torch.device("cpu"), offload_after_quant=offload_after_quant)

    assert ctx_seen == [ctx_expected]


def test_dlo_allgather_online_mxfp8_uses_ordinary_loader(monkeypatch):
    mxfp8_config = pytest.importorskip("vllm_omni.quantization.mxfp8_config")
    NPUMxfp8OnlineLinearMethod = mxfp8_config.NPUMxfp8OnlineLinearMethod

    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod

    od_config = _make_dlo_online_quant_config()
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
    model = nn.Module()
    model.transformer = nn.Linear(2, 2, bias=False)
    model.transformer.quant_method = object.__new__(NPUMxfp8OnlineLinearMethod)
    model.transformer.quant_method.uses_meta_device = True
    calls: list[object] = []
    allowlist_models: list[nn.Module] = []

    original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods

    def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
        allowlist_models.append(candidate)
        return original_allowlist_check(candidate)

    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
    monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
    loader._request_offload_after_quant = lambda _model: 1  # type: ignore[method-assign]
    loader.load_weights = (  # type: ignore[method-assign]
        lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
    )
    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]
    monkeypatch.setattr(
        loader_mod,
        "build_checkpoint_mmap_plan",
        lambda *_args, **_kwargs: HostWeightPlanResult(
            None,
            "online quantization requires the ordinary loader",
        ),
    )

    assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
    assert allowlist_models == [model.transformer]
    assert calls == [("load", True), "process"]
    assert loader.take_host_weight_plan() is None


def test_hsdp_processes_quantized_weights_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,
    )
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
    loader.quant_config = object()

    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._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(
        "vllm_omni.diffusion.quantization.hsdp_fp8.prepare_fp8_layers_for_fsdp",
        side_effect=lambda _model: events.append("prepare"),
    )
    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", "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(
        model_class_name="HeliosPipeline",
        model=model_path,
    )
    loader = DiffusersPipelineLoader(
        load_config=LoadConfig(),
        od_config=od_config,
    )
    pipeline = HeliosPipeline(od_config=od_config)

    weights = list(loader.get_all_weights(pipeline))
    assert len(weights) > 0


def test_load_model(prefetch_helios_model, mock_tp_group):
    """Ensure that load model creates an instance of the expected pipeline class."""
    od_config = OmniDiffusionConfig(
        model_class_name="HeliosPipeline",
        model=model_path,
    )
    loader = DiffusersPipelineLoader(
        load_config=LoadConfig(),
        od_config=od_config,
    )
    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"))


def test_pre_sharded_hsdp_rejects_quantization_before_loading():
    od_config = SimpleNamespace(
        dtype=torch.float32,
        hsdp_weight_load_strategy="pre_sharded",
        lora_path=None,
        quantization_config=object(),
        parallel_config=SimpleNamespace(
            use_hsdp=True,
            hsdp_replicate_size=1,
            hsdp_shard_size=2,
        ),
    )
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
    with pytest.raises(ValueError, match="does not support quantization"):
        loader._load_model_with_hsdp(torch.device("cpu"))


def test_pre_sharded_hsdp_strategy_dispatches_without_using_full_loader(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,
        hsdp_weight_load_strategy="pre_sharded",
        lora_path=None,
        quantization_config=None,
        parallel_config=SimpleNamespace(
            use_hsdp=True,
            hsdp_replicate_size=1,
            hsdp_shard_size=2,
        ),
    )
    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
    model = nn.Module()
    model.transformer = nn.Linear(2, 2, bias=False)
    discovered = PipelineModules(
        dits=[model.transformer],
        dit_names=["transformer"],
        vaes=[],
        encoders=[],
        encoder_names=[],
        resident_modules=[],
        resident_names=[],
    )
    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
    loader.load_weights = mocker.Mock(side_effect=AssertionError("full loader must not run"))  # type: ignore[method-assign]
    mocker.patch.object(loader_mod.ModuleDiscovery, "discover", return_value=discovered)
    pre_sharded_load = mocker.patch.object(
        loader,
        "_load_model_with_pre_sharded_hsdp",
        return_value=model,
    )

    assert loader._load_model_with_hsdp(torch.device("cpu")) is model
    pre_sharded_load.assert_called_once()


def test_pre_sharded_hsdp_rejects_tensor_transform():
    def transform(tensor):
        return tensor.transpose(-2, -1)

    load_plan = SimpleNamespace(
        bindings={
            "transformer.weight": TensorBinding(
                checkpoint_key="weight",
                file_path="model.safetensors",
                transform=transform,
            )
        }
    )

    with pytest.raises(ValueError, match="tensor transforms are unsupported"):
        DiffusersPipelineLoader._validate_pre_sharded_hsdp_bindings(load_plan)


def test_pre_sharded_hsdp_accepts_direct_safetensors_binding():
    load_plan = SimpleNamespace(
        bindings={
            "transformer.weight": TensorBinding(
                checkpoint_key="weight",
                file_path="model.safetensors",
            )
        }
    )

    DiffusersPipelineLoader._validate_pre_sharded_hsdp_bindings(load_plan)


def test_pre_sharded_hsdp_uses_fsdp_destination_directly():
    target = torch.empty(8, 4, 3)
    load_plan = SimpleNamespace(
        bindings={
            "transformer.weight": TensorBinding(
                checkpoint_key="source_weight",
                file_path="transformer/model.safetensors",
            )
        }
    )

    checkpoint_states, local_bytes = DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
        load_plan,
        {"transformer.weight": target},
    )
    checkpoint_target = checkpoint_states[Path("transformer")]["source_weight"]

    assert checkpoint_target is target
    assert checkpoint_target.shape == (8, 4, 3)
    assert checkpoint_target.device.type == "cpu"
    assert checkpoint_target.untyped_storage().data_ptr() == target.untyped_storage().data_ptr()
    assert local_bytes == target.numel() * target.element_size()


def test_pre_sharded_hsdp_checkpoint_writes_runtime_storage_directly():
    runtime_target = torch.empty(2, 3)
    load_plan = SimpleNamespace(
        bindings={
            "transformer.weight": TensorBinding(
                checkpoint_key="source_weight",
                file_path="transformer/model.safetensors",
            )
        }
    )
    checkpoint_tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3)
    checkpoint_states, _ = DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
        load_plan,
        {"transformer.weight": runtime_target},
    )
    checkpoint_states[Path("transformer")]["source_weight"].copy_(checkpoint_tensor)

    assert torch.equal(runtime_target, checkpoint_tensor)


def test_pre_sharded_hsdp_rejects_duplicate_checkpoint_destinations():
    binding = TensorBinding(checkpoint_key="weight", file_path="transformer/model.safetensors")
    load_plan = SimpleNamespace(
        bindings={
            "transformer.first": binding,
            "transformer.second": binding,
        }
    )

    with pytest.raises(ValueError, match="multiple runtime tensors"):
        DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
            load_plan,
            {"transformer.first": torch.empty(2, 3), "transformer.second": torch.empty(2, 3)},
        )


def test_pre_sharded_hsdp_managed_module_traversal_is_child_first_and_skips_nested_roots():
    root = nn.Module()
    root.left = nn.Sequential(nn.Linear(2, 2), nn.ReLU())
    root.right = nn.Sequential(nn.Linear(2, 2), nn.ReLU())

    names = {module: name or "root" for name, module in root.named_modules()}
    modules = DiffusersPipelineLoader._hsdp_managed_modules_post_order(
        root,
        nested_hsdp_roots={root.left},
    )

    assert [names[module] for module in modules] == ["right.0", "right.1", "right", "root"]


@pytest.mark.parametrize("checkpoint_key", ["renamed", "unexpected"])
def test_hsdp_checkpoint_plan_honors_remap_and_rejects_missing(tmp_path, checkpoint_key, monkeypatch):
    import vllm_omni.diffusion.model_loader.host_weight_plan as plan_mod

    monkeypatch.setattr(plan_mod, "get_direct_mmap_adapter", lambda _model: None)

    class Pipeline(nn.Module):
        def __init__(self):
            super().__init__()
            self.transformer = nn.Linear(2, 2, bias=False)

        @staticmethod
        def remap_checkpoint_key(name):
            return "transformer.weight" if name == "transformer.renamed" else name

    checkpoint = tmp_path / "model.safetensors"
    save_file({checkpoint_key: torch.ones(2, 2)}, str(checkpoint))
    source = SimpleNamespace(
        model_or_path=str(tmp_path),
        subfolder=None,
        revision=None,
        prefix="transformer.",
    )
    model = Pipeline()
    result = build_checkpoint_binding_plan(
        model,
        dit_modules=(("transformer", model.transformer),),
        sources=(source,),
        model_path=None,
        tensor_parallel_size=1,
        online_quantization=False,
    )

    if checkpoint_key == "renamed":
        assert result.plan is not None
        assert result.plan.bindings["transformer.weight"].checkpoint_key == "renamed"
    else:
        assert result.plan is None
        assert "no checkpoint binding" in result.fallback_reason