anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Quantization] Migrate GGUF diffusion model support to OOT plugin (#4769)

Browse Source 
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
This commit is contained in:
Isotr0py authored and GitHub committed 2026-07-11 18:48:43 +00:00
1 parent ff6f0da271
commit b95905deb1
35 files changed
+272 -1614
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py
-112
View File
@@ -1,112 +0,0 @@
# 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]))
tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py
-184
View File
@@ -1,184 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from types import SimpleNamespace
import pytest
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
from vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image import (
QwenImageGGUFAdapter,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def _make_od_config(
model_class_name: str = "QwenImagePipeline",
model_type: str = "qwen_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)
class _FakeTransformer:
def __init__(self) -> None:
self._params = {
"img_in.qweight": 1,
"img_in.qweight_type": 1,
"txt_in.qweight": 1,
"txt_in.qweight_type": 1,
"proj_out.qweight": 1,
"proj_out.qweight_type": 1,
"time_text_embed.timestep_embedder.linear_1.qweight": 1,
"time_text_embed.timestep_embedder.linear_1.qweight_type": 1,
"time_text_embed.timestep_embedder.linear_2.qweight": 1,
"time_text_embed.timestep_embedder.linear_2.qweight_type": 1,
"transformer_blocks.0.img_mod.1.qweight": 1,
"transformer_blocks.0.img_mod.1.qweight_type": 1,
"transformer_blocks.0.txt_mod.1.qweight": 1,
"transformer_blocks.0.txt_mod.1.qweight_type": 1,
"transformer_blocks.0.attn.to_qkv.qweight": 1,
"transformer_blocks.0.attn.to_qkv.qweight_type": 1,
"transformer_blocks.0.attn.add_kv_proj.qweight": 1,
"transformer_blocks.0.attn.add_kv_proj.qweight_type": 1,
"transformer_blocks.0.attn.to_out.weight": 1,
"transformer_blocks.0.img_mlp.net.0.proj.qweight": 1,
"transformer_blocks.0.img_mlp.net.0.proj.qweight_type": 1,
}
def named_parameters(self):
return list(self._params.items())
def named_buffers(self):
return []
def get_submodule(self, _prefix: str):
return self
def test_qwen_adapter_selected_for_qwen_image_family():
adapter = get_gguf_adapter(
"dummy.gguf",
_FakeTransformer(),
_make_source(),
_make_od_config(),
)
assert isinstance(adapter, QwenImageGGUFAdapter)
def test_qwen_adapter_matches_multiple_pipeline_variants():
for model_class_name in (
"QwenImagePipeline",
"QwenImageEditPipeline",
"QwenImageEditPlusPipeline",
"QwenImageLayeredPipeline",
):
assert QwenImageGGUFAdapter.is_compatible(
_make_od_config(model_class_name=model_class_name),
_FakeTransformer(),
_make_source(),
)
def test_qwen_adapter_preserves_split_projection_names(monkeypatch: pytest.MonkeyPatch):
import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
monkeypatch.setattr(
qwen_image_module,
"gguf_quant_weights_iterator",
lambda _path: iter(
[
("transformer_blocks.0.attn.to_q.qweight_type", 1),
("transformer_blocks.0.attn.to_q.qweight", 2),
("transformer_blocks.0.attn.to_k.qweight_type", 3),
("transformer_blocks.0.attn.to_k.qweight", 4),
("transformer_blocks.0.attn.to_out.0.qweight_type", 5),
("transformer_blocks.0.attn.to_out.0.qweight", 6),
]
),
)
adapter = QwenImageGGUFAdapter(
"dummy.gguf",
_FakeTransformer(),
_make_source(),
_make_od_config(),
)
weights = list(adapter.weights_iterator())
assert ("transformer_blocks.0.attn.to_q.qweight_type", 1) in weights
assert ("transformer_blocks.0.attn.to_q.qweight", 2) in weights
assert ("transformer_blocks.0.attn.to_k.qweight_type", 3) in weights
assert ("transformer_blocks.0.attn.to_k.qweight", 4) in weights
assert ("transformer_blocks.0.attn.to_out.0.qweight_type", 5) in weights
assert ("transformer_blocks.0.attn.to_out.0.qweight", 6) in weights
def test_qwen_adapter_keeps_top_level_quantized_weights(monkeypatch: pytest.MonkeyPatch):
import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
monkeypatch.setattr(
qwen_image_module,
"gguf_quant_weights_iterator",
lambda _path: iter(
[
("img_in.qweight_type", 1),
("img_in.qweight", 2),
("time_text_embed.timestep_embedder.linear_1.qweight_type", 3),
("time_text_embed.timestep_embedder.linear_1.qweight", 4),
]
),
)
adapter = QwenImageGGUFAdapter(
"dummy.gguf",
_FakeTransformer(),
_make_source(),
_make_od_config(),
)
weights = list(adapter.weights_iterator())
assert ("img_in.qweight_type", 1) in weights
assert ("img_in.qweight", 2) in weights
assert ("time_text_embed.timestep_embedder.linear_1.qweight_type", 3) in weights
assert ("time_text_embed.timestep_embedder.linear_1.qweight", 4) in weights
def test_qwen_adapter_keeps_modulation_quantized_weights(monkeypatch: pytest.MonkeyPatch):
import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
monkeypatch.setattr(
qwen_image_module,
"gguf_quant_weights_iterator",
lambda _path: iter(
[
("transformer_blocks.0.img_mod.1.qweight_type", 1),
("transformer_blocks.0.img_mod.1.qweight", 2),
("transformer_blocks.0.txt_mod.1.qweight_type", 3),
("transformer_blocks.0.txt_mod.1.qweight", 4),
]
),
)
adapter = QwenImageGGUFAdapter(
"dummy.gguf",
_FakeTransformer(),
_make_source(),
_make_od_config(),
)
weights = list(adapter.weights_iterator())
assert ("transformer_blocks.0.img_mod.1.qweight_type", 1) in weights
assert ("transformer_blocks.0.img_mod.1.qweight", 2) in weights
assert ("transformer_blocks.0.txt_mod.1.qweight_type", 3) in weights
assert ("transformer_blocks.0.txt_mod.1.qweight", 4) in weights
tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py
-83
View File
@@ -1,83 +0,0 @@
# 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
tests/diffusion/model_loader/test_diffusers_loader.py
-23
View File
@@ -16,7 +16,6 @@ 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 OmniDiffusionConfig
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
from vllm_omni.diffusion.models.helios import HeliosPipeline
from vllm_omni.diffusion.registry import initialize_model
@@ -111,27 +110,6 @@ def test_empty_source_prefix_keeps_full_model_strict_check():
loader.load_weights(model)
def test_qwen_model_class_selects_qwen_gguf_adapter():
od_config = type(
"Config",
(),
{
"model_class_name": "QwenImagePipeline",
"tf_model_config": {"model_type": "qwen_image"},
},
)()
source = DiffusersPipelineLoader.ComponentSource(
model_or_path="dummy",
subfolder="transformer",
revision=None,
prefix="transformer.",
)
adapter = get_gguf_adapter("dummy.gguf", object(), source, od_config)
assert adapter.__class__.__name__ == "QwenImageGGUFAdapter"
class _ConfigAwareModel(nn.Module):
def __init__(self, *, od_config):
super().__init__()
@@ -186,7 +164,6 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
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]
loader._is_gguf_quantization = lambda: False # type: ignore[assignment]
monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel)
monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu"))
tests/diffusion/model_loader/test_diffusers_loader_gguf.py
-259
View File
@@ -1,259 +0,0 @@
# 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.quant_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()
loader.quant_config = {
"method": "gguf",
"gguf_model": "weights.gguf",
}
assert loader._is_gguf_quantization() is True
def test_is_gguf_quantization_rejects_non_gguf_dict_config():
loader = _make_loader()
loader.quant_config = {
"method": "fp8",
"gguf_model": "weights.gguf",
}
assert loader._is_gguf_quantization() is False
def test_is_gguf_quantization_requires_gguf_model_for_dict_config():
loader = _make_loader()
loader.quant_config = {"method": "gguf"}
with pytest.raises(ValueError, match="gguf_model"):
loader._is_gguf_quantization()
def test_is_gguf_quantization_accepts_object_config():
loader = _make_loader()
loader.quant_config = SimpleNamespace(
get_name=lambda: "gguf",
gguf_model="weights.gguf",
)
assert loader._is_gguf_quantization() is True
def test_is_gguf_quantization_uses_fallback_object_without_get_name():
loader = _make_loader()
loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
assert loader._is_gguf_quantization() 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()
loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
loader.od_config = SimpleNamespace(revision="main")
source = DiffusersPipelineLoader.ComponentSource(
model_or_path="dummy",
subfolder="transformer",
revision=None,
prefix="transformer.",
)
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()))
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()
monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
monkeypatch.setattr(
loader,
"_get_gguf_weights_iterator",
lambda source, model: 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)
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()
monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
monkeypatch.setattr(
loader,
"_get_gguf_weights_iterator",
lambda source, model: 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)
assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
assert hf_calls == ["vae"]
tests/diffusion/quantization/test_fp8_config.py
+1 -16
View File
@@ -154,21 +154,6 @@ def test_component_config_with_default():
assert config.resolve("vae.encoder") is None
def test_gguf_config():
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig
config = build_quant_config(
{
"method": "gguf",
"gguf_model": "path/to/model.gguf",
}
)
assert isinstance(config, DiffusionGGUFConfig)
assert config.gguf_model == "path/to/model.gguf"
assert config.get_name() == "gguf"
def test_integration_string():
from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -236,7 +221,7 @@ def test_transformer_config_auto_detects_modelopt_fp8():
def test_supported_methods_includes_vllm():
from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
for method in ["fp8", "gguf", "awq", "gptq", "bitsandbytes", "modelopt"]:
for method in ["fp8", "awq", "gptq", "bitsandbytes", "modelopt"]:
assert method in SUPPORTED_QUANTIZATION_METHODS, f"{method} missing"
tests/diffusion/quantization/test_gguf_config.py
-128
View File
@@ -1,128 +0,0 @@
# 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, UnquantizedLinearMethod
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.gguf_config import (
UNQUANTIZED_TYPES,
DiffusionGGUFConfig,
DiffusionGGUFLinearMethod,
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"
def test_build_quant_config_builds_gguf_config():
config = build_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 = DiffusionGGUFConfig(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 = DiffusionGGUFConfig(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 = DiffusionGGUFConfig(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(DiffusionGGUFConfig(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(DiffusionGGUFConfig(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))
tests/diffusion/quantization/test_trajectory_similarity_tool.py
-1
View File
@@ -99,7 +99,6 @@ def test_candidate_model_can_point_to_offline_checkpoint_without_online_quantiza
candidate_quantization_config_json=None,
candidate_ignored_layers=None,
ignored_layers=None,
candidate_gguf_model=None,
candidate_load_format="default",
)