# 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 from unittest.mock import MagicMock 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 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, ) 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 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_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_hwr_warm_snapshot_present(): cfg = SimpleNamespace( lora_backend="distill", lora_path="/path/to/lora.safetensors", lora_scale=1.0, dtype=torch.bfloat16, quantization_config=None, parallel_config=SimpleNamespace(use_hsdp=False), ) loader = DiffusersPipelineLoader(LoadConfig(), cfg) loader._hwr_state = {"warm_snapshot": {"transformer.weight": (1234, "hash")}} model = nn.Module() model.load_lora_weights = MagicMock() loader._maybe_fuse_distilled_lora(model) model.load_lora_weights.assert_not_called() assert getattr(model, "lora_is_fused", False) is True def test_maybe_fuse_distilled_lora_fuses_when_not_warm_snapshot(): cfg = SimpleNamespace( lora_backend=LoRABackend.DISTILL, lora_path="/path/to/lora.safetensors", lora_scale=1.0, dtype=torch.bfloat16, quantization_config=None, parallel_config=SimpleNamespace(use_hsdp=False), ) loader = DiffusersPipelineLoader(LoadConfig(), cfg) loader._hwr_state = None model = nn.Module() model.load_lora_weights = MagicMock() loader._maybe_fuse_distilled_lora(model) model.load_lora_weights.assert_called_once_with("/path/to/lora.safetensors") assert getattr(model, "lora_is_fused", False) is True def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_path: Path, monkeypatch): from vllm_omni.diffusion.model_loader import diffusers_loader as loader_module 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=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_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_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_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)