From 602f831068afba3c2561c857df0aebccaceef210 Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Mon, 23 Mar 2026 14:56:15 +0800 Subject: [PATCH] [Bugfix] fix gguf TypeError: GGUFConfig.get_name() missing 1 required positional argument: 'self', test: add diffusion gguf unit coverage (#1865) Signed-off-by: David Chen <530634352@qq.com> --- .../test_flux2_klein_gguf_adapter.py | 112 ++++++++ .../test_z_image_gguf_adapter.py | 83 ++++++ .../test_diffusers_loader_gguf.py | 268 ++++++++++++++++++ .../quantization/test_gguf_config.py | 133 +++++++++ vllm_omni/diffusion/quantization/base.py | 12 +- 5 files changed, 607 insertions(+), 1 deletion(-) create mode 100644 tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py create mode 100644 tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py create mode 100644 tests/diffusion/model_loader/test_diffusers_loader_gguf.py create mode 100644 tests/diffusion/quantization/test_gguf_config.py diff --git a/tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py b/tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py new file mode 100644 index 000000000..0c81dc801 --- /dev/null +++ b/tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the Flux2-Klein GGUF adapter.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter +from vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein import ( + Flux2KleinGGUFAdapter, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _make_od_config( + model_class_name: str = "Flux2Pipeline", + model_type: str = "flux", +): + return SimpleNamespace( + model_class_name=model_class_name, + tf_model_config={"model_type": model_type}, + ) + + +def _make_source(prefix: str = "", subfolder: str = "transformer"): + return SimpleNamespace(prefix=prefix, subfolder=subfolder) + + +def test_flux2_klein_adapter_selected_for_flux_family(): + adapter = get_gguf_adapter( + "dummy.gguf", + object(), + _make_source(), + _make_od_config(), + ) + + assert isinstance(adapter, Flux2KleinGGUFAdapter) + + +def test_flux2_klein_adapter_matches_flux_model_type(): + assert Flux2KleinGGUFAdapter.is_compatible( + _make_od_config(model_class_name="OtherPipeline", model_type="flux-dev"), + object(), + _make_source(), + ) + + +def test_flux2_klein_adapter_renames_core_projection_paths(monkeypatch: pytest.MonkeyPatch): + import vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein as flux_module + + monkeypatch.setattr( + flux_module, + "gguf_quant_weights_iterator", + lambda _path: iter( + [ + ("img_in.qweight", torch.ones((1, 1))), + ("time_in.in_layer.qweight_type", torch.tensor(1)), + ("double_blocks.0.img_attn.norm.query_norm.weight", torch.full((1, 1), 2.0)), + ("final_layer.linear.qweight", torch.full((1, 1), 3.0)), + ] + ), + ) + + adapter = Flux2KleinGGUFAdapter( + "dummy.gguf", + object(), + _make_source(), + _make_od_config(), + ) + + weights = list(adapter.weights_iterator()) + names = [name for name, _ in weights] + + assert "x_embedder.qweight" in names + assert "time_guidance_embed.timestep_embedder.linear_1.qweight_type" in names + assert "transformer_blocks.0.attn.norm_q.weight" in names + assert "proj_out.qweight" in names + + +def test_flux2_klein_adapter_swaps_final_adaln_shift_and_scale(monkeypatch: pytest.MonkeyPatch): + import vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein as flux_module + + monkeypatch.setattr( + flux_module, + "gguf_quant_weights_iterator", + lambda _path: iter( + [ + ( + "final_layer.adaLN_modulation.1.scale", + torch.tensor([1.0, 2.0, 3.0, 4.0]), + ) + ] + ), + ) + + adapter = Flux2KleinGGUFAdapter( + "dummy.gguf", + object(), + _make_source(), + _make_od_config(), + ) + + weights = list(adapter.weights_iterator()) + + assert len(weights) == 1 + assert weights[0][0] == "norm_out.linear.weight" + assert torch.equal(weights[0][1], torch.tensor([3.0, 4.0, 1.0, 2.0])) diff --git a/tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py b/tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py new file mode 100644 index 000000000..cccea56fb --- /dev/null +++ b/tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the Z-Image GGUF adapter.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter +from vllm_omni.diffusion.model_loader.gguf_adapters.z_image import ( + ZImageGGUFAdapter, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _make_od_config( + model_class_name: str = "ZImagePipeline", + model_type: str = "z_image", +): + return SimpleNamespace( + model_class_name=model_class_name, + tf_model_config={"model_type": model_type}, + ) + + +def _make_source(prefix: str = "", subfolder: str = "transformer"): + return SimpleNamespace(prefix=prefix, subfolder=subfolder) + + +def test_z_image_adapter_selected_for_z_image_family(): + adapter = get_gguf_adapter( + "dummy.gguf", + object(), + _make_source(), + _make_od_config(), + ) + + assert isinstance(adapter, ZImageGGUFAdapter) + + +def test_z_image_adapter_matches_model_type_variants(): + for model_type in ("z_image", "zimage", "z-image"): + assert ZImageGGUFAdapter.is_compatible( + _make_od_config(model_class_name="OtherPipeline", model_type=model_type), + object(), + _make_source(), + ) + + +def test_z_image_adapter_renames_known_gguf_tensor_paths(monkeypatch: pytest.MonkeyPatch): + import vllm_omni.diffusion.model_loader.gguf_adapters.z_image as z_image_module + + monkeypatch.setattr( + z_image_module, + "gguf_quant_weights_iterator", + lambda _path: iter( + [ + ("model.diffusion_model.final_layer.qweight", torch.ones((1, 1))), + ("model.diffusion_model.x_embedder.qweight_type", torch.tensor(1)), + ("transformer_blocks.0.attention.out.weight", torch.full((1, 1), 2.0)), + ("transformer_blocks.0.attention.qkv.qweight", torch.full((1, 1), 3.0)), + ] + ), + ) + + adapter = ZImageGGUFAdapter( + "dummy.gguf", + object(), + _make_source(), + _make_od_config(), + ) + + weights = list(adapter.weights_iterator()) + names = [name for name, _ in weights] + + assert "all_final_layer.2-1.qweight" in names + assert "all_x_embedder.2-1.qweight_type" in names + assert "transformer_blocks.0.attention.to_out.0.weight" in names + assert "transformer_blocks.0.attention.to_qkv.qweight" in names diff --git a/tests/diffusion/model_loader/test_diffusers_loader_gguf.py b/tests/diffusion/model_loader/test_diffusers_loader_gguf.py new file mode 100644 index 000000000..9aae9135e --- /dev/null +++ b/tests/diffusion/model_loader/test_diffusers_loader_gguf.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for GGUF-specific DiffusersPipelineLoader behavior.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +import vllm_omni.diffusion.model_loader.diffusers_loader as loader_module +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +class _DummyModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.transformer = nn.Linear(2, 2, bias=True) + self.vae = nn.Linear(2, 2, bias=False) + self.register_buffer("transformer_buffer", torch.ones(1)) + self.calls: list[list[str]] = [] + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path="dummy", + subfolder="transformer", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ), + DiffusersPipelineLoader.ComponentSource( + model_or_path="dummy", + subfolder="vae", + revision=None, + prefix="vae.", + fall_back_to_pt=True, + ), + ] + + def load_weights(self, weights): + loadable = dict(self.named_parameters()) + loadable.update(dict(self.named_buffers())) + seen: list[str] = [] + loaded: set[str] = set() + for name, tensor in weights: + seen.append(name) + if name in loadable: + target = loadable[name] + target.data.copy_(tensor.to(dtype=target.dtype)) + loaded.add(name) + self.calls.append(seen) + return loaded + + +def _make_loader() -> DiffusersPipelineLoader: + loader = object.__new__(DiffusersPipelineLoader) + loader.load_config = SimpleNamespace( + download_dir="cache-dir", + ignore_patterns=["*.tmp"], + ) + loader.od_config = None + loader.counter_before_loading_weights = 0.0 + loader.counter_after_loading_weights = 0.0 + return loader + + +def test_is_gguf_quantization_accepts_dict_config(): + loader = _make_loader() + od_config = SimpleNamespace( + quantization_config={ + "method": "gguf", + "gguf_model": "weights.gguf", + } + ) + + assert loader._is_gguf_quantization(od_config) is True + + +def test_is_gguf_quantization_rejects_non_gguf_dict_config(): + loader = _make_loader() + od_config = SimpleNamespace( + quantization_config={ + "method": "fp8", + "gguf_model": "weights.gguf", + } + ) + + assert loader._is_gguf_quantization(od_config) is False + + +def test_is_gguf_quantization_requires_gguf_model_for_dict_config(): + loader = _make_loader() + od_config = SimpleNamespace(quantization_config={"method": "gguf"}) + + with pytest.raises(ValueError, match="gguf_model"): + loader._is_gguf_quantization(od_config) + + +def test_is_gguf_quantization_accepts_object_config(): + loader = _make_loader() + od_config = SimpleNamespace( + quantization_config=SimpleNamespace( + get_name=lambda: "gguf", + gguf_model="weights.gguf", + ) + ) + + assert loader._is_gguf_quantization(od_config) is True + + +def test_is_gguf_quantization_uses_fallback_object_without_get_name(): + loader = _make_loader() + od_config = SimpleNamespace(quantization_config=SimpleNamespace(gguf_model="weights.gguf")) + + assert loader._is_gguf_quantization(od_config) is True + + +def test_get_model_loadable_names_collects_parameters_and_buffers(): + loader = _make_loader() + model = _DummyModel() + + names = loader._get_model_loadable_names(model) + + assert "transformer.weight" in names + assert "transformer.bias" in names + assert "transformer_buffer" in names + + +def test_resolve_gguf_model_path_returns_local_file(tmp_path): + loader = _make_loader() + gguf_file = tmp_path / "model.gguf" + gguf_file.write_bytes(b"gguf") + + resolved = loader._resolve_gguf_model_path(str(gguf_file), revision=None) + + assert resolved == str(gguf_file) + + +def test_resolve_gguf_model_path_downloads_explicit_gguf_filename(monkeypatch: pytest.MonkeyPatch): + loader = _make_loader() + + monkeypatch.setattr(loader_module.os.path, "isfile", lambda _path: False) + monkeypatch.setattr( + loader_module, + "hf_hub_download", + lambda repo_id, filename, revision, cache_dir: f"{cache_dir}/{repo_id}/{filename}", + ) + + resolved = loader._resolve_gguf_model_path("owner/repo/model-Q4_0.gguf", revision="main") + + assert resolved == "cache-dir/owner/repo/model-Q4_0.gguf" + + +def test_resolve_gguf_model_path_downloads_by_quant_type(monkeypatch: pytest.MonkeyPatch): + loader = _make_loader() + + monkeypatch.setattr(loader_module.os.path, "isfile", lambda _path: False) + monkeypatch.setattr( + loader_module, + "download_gguf", + lambda repo_id, quant_type, cache_dir, revision, ignore_patterns: f"{cache_dir}/{repo_id}/{quant_type}.gguf", + ) + + resolved = loader._resolve_gguf_model_path("owner/repo:Q4_0", revision="main") + + assert resolved == "cache-dir/owner/repo/Q4_0.gguf" + + +def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.MonkeyPatch): + loader = _make_loader() + source = DiffusersPipelineLoader.ComponentSource( + model_or_path="dummy", + subfolder="transformer", + revision=None, + prefix="transformer.", + ) + od_config = SimpleNamespace( + revision="main", + quantization_config=SimpleNamespace(gguf_model="weights.gguf"), + ) + + class _Adapter: + def weights_iterator(self): + yield "weight", torch.ones((2, 2)) + yield "bias", torch.zeros(2) + + monkeypatch.setattr(loader, "_resolve_gguf_model_path", lambda gguf_model, revision: "resolved.gguf") + monkeypatch.setattr(loader_module, "get_gguf_adapter", lambda gguf_file, model, source, od_config: _Adapter()) + + weights = list(loader._get_gguf_weights_iterator(source, object(), od_config)) + + assert weights[0][0] == "transformer.weight" + assert torch.equal(weights[0][1], torch.ones((2, 2))) + assert weights[1][0] == "transformer.bias" + assert torch.equal(weights[1][1], torch.zeros(2)) + + +def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(monkeypatch: pytest.MonkeyPatch): + loader = _make_loader() + model = _DummyModel() + od_config = SimpleNamespace() + + monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources)) + monkeypatch.setattr( + loader, + "_get_gguf_weights_iterator", + lambda source, model, od_config: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]), + ) + + def _hf_iter(source): + if source.subfolder == "transformer": + return iter([("transformer.bias", torch.zeros(2, dtype=torch.float32))]) + return iter([("vae.weight", torch.full((2, 2), 2.0, dtype=torch.float32))]) + + monkeypatch.setattr(loader, "_get_weights_iterator", _hf_iter) + monkeypatch.setattr( + loader, + "_get_expected_parameter_names", + lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"}, + ) + + loaded = loader._load_weights_with_gguf(model, od_config) + + assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"} + assert model.calls[0] == ["transformer.weight"] + assert model.calls[1] == ["transformer.bias"] + assert model.calls[2] == ["vae.weight"] + + +def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete(monkeypatch: pytest.MonkeyPatch): + loader = _make_loader() + model = _DummyModel() + od_config = SimpleNamespace() + + monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources)) + monkeypatch.setattr( + loader, + "_get_gguf_weights_iterator", + lambda source, model, od_config: iter( + [ + ("transformer.weight", torch.ones((2, 2), dtype=torch.float32)), + ("transformer.bias", torch.zeros(2, dtype=torch.float32)), + ] + ), + ) + + hf_calls: list[str] = [] + + def _hf_iter(source): + hf_calls.append(source.subfolder or "") + if source.subfolder == "transformer": + return iter([("transformer.bias", torch.zeros(2, dtype=torch.float32))]) + return iter([("vae.weight", torch.full((2, 2), 2.0, dtype=torch.float32))]) + + monkeypatch.setattr(loader, "_get_weights_iterator", _hf_iter) + monkeypatch.setattr( + loader, + "_get_expected_parameter_names", + lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"}, + ) + + loaded = loader._load_weights_with_gguf(model, od_config) + + assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"} + assert hf_calls == ["vae"] diff --git a/tests/diffusion/quantization/test_gguf_config.py b/tests/diffusion/quantization/test_gguf_config.py new file mode 100644 index 000000000..bd3d1fcb2 --- /dev/null +++ b/tests/diffusion/quantization/test_gguf_config.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for diffusion GGUF quantization config and linear method.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from vllm.model_executor.layers.linear import LinearBase +from vllm.model_executor.layers.quantization.gguf import UNQUANTIZED_TYPES + +from vllm_omni.diffusion.quantization import ( + DiffusionGgufConfig, + get_diffusion_quant_config, +) +from vllm_omni.diffusion.quantization.gguf import ( + DiffusionGGUFLinearMethod, + UnquantizedLinearMethod, + _GGUFConfig, + dequant_gemm_gguf, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def test_gguf_config_creation_and_delegation(): + config = DiffusionGgufConfig( + gguf_model="weights.gguf", + unquantized_modules=["proj_out"], + ) + + assert config.gguf_model == "weights.gguf" + assert config.unquantized_modules == ["proj_out"] + assert config.get_name() == "gguf" + assert isinstance(config.get_vllm_quant_config(), _GGUFConfig) + + +def test_get_diffusion_quant_config_builds_gguf_config(): + config = get_diffusion_quant_config( + "gguf", + gguf_model="repo/model-Q4_0.gguf", + unquantized_modules=["norm_out"], + ) + + assert isinstance(config, DiffusionGgufConfig) + assert config.gguf_model == "repo/model-Q4_0.gguf" + assert config.unquantized_modules == ["norm_out"] + + +def test_gguf_vllm_config_returns_diffusion_linear_method_for_linear_layers(): + linear = object.__new__(LinearBase) + method = _GGUFConfig(unquantized_modules=[]).get_quant_method(linear, "transformer.img_in") + + assert isinstance(method, DiffusionGGUFLinearMethod) + + +def test_gguf_vllm_config_respects_unquantized_modules(): + linear = object.__new__(LinearBase) + method = _GGUFConfig(unquantized_modules=["proj_out"]).get_quant_method(linear, "transformer.proj_out") + + assert isinstance(method, UnquantizedLinearMethod) + + +def test_gguf_vllm_config_returns_none_for_non_linear_layers(): + method = _GGUFConfig(unquantized_modules=[]).get_quant_method(torch.nn.LayerNorm(4), "norm") + + assert method is None + + +def test_dequant_gemm_gguf_uses_plain_matmul_for_unquantized_types(): + qweight_type = next(iter(UNQUANTIZED_TYPES)) + x = torch.tensor([[1.0, 2.0]], dtype=torch.float32) + qweight = torch.tensor([[3.0, 4.0], [5.0, 6.0]], dtype=torch.float32) + + out = dequant_gemm_gguf(x, qweight, qweight_type) + + assert torch.allclose(out, x @ qweight.T) + + +def test_diffusion_gguf_linear_method_applies_bias_on_unquantized_weight(): + qweight_type = next(iter(UNQUANTIZED_TYPES)) + method = DiffusionGGUFLinearMethod(_GGUFConfig(unquantized_modules=[])) + layer = SimpleNamespace( + qweight=torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32), + qweight_type=SimpleNamespace(weight_type=qweight_type), + ) + x = torch.tensor([[[1.0, 0.5]]], dtype=torch.float32) + bias = torch.tensor([0.25, -0.5], dtype=torch.float32) + + out = method.apply(layer, x, bias) + + expected = x @ layer.qweight.T + bias + assert torch.allclose(out, expected) + + +def test_diffusion_gguf_linear_method_concatenates_sharded_outputs(): + qweight_type = next(iter(UNQUANTIZED_TYPES)) + method = DiffusionGGUFLinearMethod(_GGUFConfig(unquantized_modules=[])) + qweight = torch.nn.Parameter( + torch.tensor( + [ + [1.0, 0.0], + [0.0, 1.0], + [2.0, 0.0], + [0.0, 2.0], + ], + dtype=torch.float32, + ), + requires_grad=False, + ) + qweight.shard_id = ["left", "right"] + qweight.shard_offset_map = { + "left": (0, 2, 2), + "right": (2, 4, 2), + } + layer = SimpleNamespace( + qweight=qweight, + qweight_type=SimpleNamespace( + shard_weight_type={ + "left": qweight_type, + "right": qweight_type, + } + ), + ) + x = torch.tensor([[1.5, 2.0]], dtype=torch.float32) + + out = method.apply(layer, x) + + expected_q = x @ qweight[:2].T + expected_k = x @ qweight[2:4].T + assert torch.allclose(out, torch.cat([expected_q, expected_k], dim=-1)) diff --git a/vllm_omni/diffusion/quantization/base.py b/vllm_omni/diffusion/quantization/base.py index 8a35547f9..f008eca0b 100644 --- a/vllm_omni/diffusion/quantization/base.py +++ b/vllm_omni/diffusion/quantization/base.py @@ -38,7 +38,17 @@ class DiffusionQuantizationConfig(ABC): Delegates to the underlying vLLM config class's get_name(). """ if cls.quant_config_cls is not None: - return cls.quant_config_cls.get_name() + try: + return cls.quant_config_cls.get_name() + except TypeError as exc: + # Upstream quantization configs are not fully consistent: + # some expose get_name as a classmethod (e.g. FP8), while + # others expose it as an instance method (e.g. GGUF in + # vLLM 0.18.0). Fall back to a default-constructed instance + # so diffusion wrappers work with both styles. + if "required positional argument" not in str(exc): + raise + return cls.quant_config_cls().get_name() raise NotImplementedError("Subclass must set quant_config_cls or override get_name().") def get_vllm_quant_config(self) -> "QuantizationConfig | None":