# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Tests for the DiffusersPipelineLoader. """ import json from pathlib import Path from types import SimpleNamespace import pytest import torch import torch.nn as nn from huggingface_hub import snapshot_download from safetensors.torch import save_file from vllm.config.load import LoadConfig from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlan, HostWeightPlanResult, TensorBinding, ) 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 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.""" 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_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_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 == 2 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_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 == 2 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=SimpleNamespace(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_configure_vae_stack_tiling_enables_tiling_and_native_capability(): from vllm_omni.diffusion.registry import _configure_vae_stack_tiling calls = [] vae = SimpleNamespace( supports_stack_tiling=True, set_stack_tiling=lambda mode: calls.append(mode), ) model = SimpleNamespace(vae=vae) od_config = SimpleNamespace( model_class_name="MiniMaxH3Pipeline", vae_stack_tiling="auto", vae_use_tiling=False, ) _configure_vae_stack_tiling(model, od_config) assert od_config.vae_use_tiling is True assert calls == ["auto"] def test_configure_vae_stack_tiling_rejects_unsupported_vae(): from vllm_omni.diffusion.registry import _configure_vae_stack_tiling od_config = SimpleNamespace( model_class_name="UnsupportedPipeline", vae_stack_tiling="true", vae_use_tiling=False, ) with pytest.raises(ValueError, match="does not declare stacked-tiling support"): _configure_vae_stack_tiling(SimpleNamespace(vae=SimpleNamespace()), od_config) def test_configure_vae_stack_tiling_auto_falls_back_for_unsupported_vae(): from vllm_omni.diffusion.registry import _configure_vae_stack_tiling od_config = SimpleNamespace( model_class_name="UnsupportedPipeline", vae_stack_tiling="auto", vae_use_tiling=False, ) _configure_vae_stack_tiling(SimpleNamespace(vae=SimpleNamespace()), od_config) assert od_config.vae_use_tiling is False 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_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_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] 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] 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_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] 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_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)