diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index 7e0942bf6..c9b60fe2c 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -684,3 +684,24 @@ steps: hostPath: path: /mnt/hf-cache type: DirectoryOrCreate + + - label: "Diffusion ยท GGUF Plugin Test" + soft_fail: true + source_file_dependencies: + - tests/plugin/test_gguf_diffusion.py + - vllm_omni/diffusion/model_loader + commands: + - pip install "vllm-gguf-plugin>=0.0.3" + - timeout 30m pytest -sv tests/plugin/test_gguf_diffusion.py + agents: + queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + environment: + - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" diff --git a/docs/user_guide/quantization/gguf.md b/docs/user_guide/quantization/gguf.md index e70d2d26e..206e25319 100644 --- a/docs/user_guide/quantization/gguf.md +++ b/docs/user_guide/quantization/gguf.md @@ -9,6 +9,15 @@ for the transformer. GGUF is static quantization: the quantized weights are produced before serving. +!!! note + GGUF support has migrated to OOT [vllm-gguf-plugin](https://github.com/vllm-project/vllm-gguf-plugin). Make sure you have GGUF plugin installed before serving a GGUF model. + +Before serving a GGUF model, make sure to install the [vllm-gguf-plugin](https://github.com/vllm-project/vllm-gguf-plugin): + +```bash +uv pip install vllm-gguf-plugin>=0.0.3 +``` + ## Hardware Support | Device | Support | diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 1a0691ebb..2d1e70697 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -220,7 +220,7 @@ def parse_args() -> argparse.Namespace: "--quantization", type=str, default=None, - choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"], + choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"], help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).", ) diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index a8d980c6f..374356ef4 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -202,17 +202,11 @@ def parse_args() -> argparse.Namespace: "--quantization", type=str, default=None, - choices=["fp8", "int8", "gguf"], + choices=["fp8", "int8"], help="Quantization method for the transformer. " - "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8), 'gguf' (GGUF quantized weights). " + "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8). " "Default: None (no quantization, uses BF16).", ) - parser.add_argument( - "--gguf-model", - type=str, - default=None, - help=("GGUF file path or HF reference for transformer weights. Required when --quantization gguf is set."), - ) parser.add_argument( "--ignored-layers", type=str, @@ -402,14 +396,7 @@ def main(): # ignored_layers is specified so the list flows through OmniDiffusionConfig quant_kwargs: dict[str, Any] = {} ignored_layers = [s.strip() for s in args.ignored_layers.split(",") if s.strip()] if args.ignored_layers else None - if args.quantization == "gguf": - if not args.gguf_model: - raise ValueError("--gguf-model is required when --quantization gguf is set.") - quant_kwargs["quantization_config"] = { - "method": "gguf", - "gguf_model": args.gguf_model, - } - elif args.quantization and ignored_layers: + if args.quantization and ignored_layers: quant_kwargs["quantization_config"] = { "method": args.quantization, "ignored_layers": ignored_layers, diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md index 4c2b5a83a..edc75bb26 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/examples/offline_inference/text_to_video/text_to_video.md @@ -231,7 +231,7 @@ python text_to_video.py \ - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2). - `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio; LTX-2.3 outputs 48kHz audio). -- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF). +- `--quantization`: quantization method (such as `fp8` for FP8). - `--flow-shift`: scheduler flow_shift parameter. - `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above). diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 200515a83..0f9bb8c8e 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -227,7 +227,7 @@ def parse_args() -> argparse.Namespace: "--quantization", type=str, default=None, - choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"], + choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"], help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).", ) diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 8f15347d1..bf2dba701 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -751,3 +751,14 @@ def test_diffusion_config_field_classification_covers_current_fields(): "distributed_executor_backend", } <= omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS assert "prompt_file_path" in omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS + + +def test_diffusion_config_projection_keeps_mapping_quantization_config_serializable(): + quantization_config = { + "method": "example_quant", + "weights": "weights.bin", + } + + cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(quantization_config=quantization_config) + + assert cfg.quantization_config == quantization_config 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 deleted file mode 100644 index 0c81dc801..000000000 --- a/tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py +++ /dev/null @@ -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])) diff --git a/tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py b/tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py deleted file mode 100644 index 6b7392c80..000000000 --- a/tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py +++ /dev/null @@ -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 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 deleted file mode 100644 index cccea56fb..000000000 --- a/tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py +++ /dev/null @@ -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 diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 376dbe685..f884f3667 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -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")) diff --git a/tests/diffusion/model_loader/test_diffusers_loader_gguf.py b/tests/diffusion/model_loader/test_diffusers_loader_gguf.py deleted file mode 100644 index f747dca81..000000000 --- a/tests/diffusion/model_loader/test_diffusers_loader_gguf.py +++ /dev/null @@ -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"] diff --git a/tests/diffusion/quantization/test_fp8_config.py b/tests/diffusion/quantization/test_fp8_config.py index a1f79030e..9bad16792 100644 --- a/tests/diffusion/quantization/test_fp8_config.py +++ b/tests/diffusion/quantization/test_fp8_config.py @@ -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" diff --git a/tests/diffusion/quantization/test_gguf_config.py b/tests/diffusion/quantization/test_gguf_config.py deleted file mode 100644 index 201bcdaba..000000000 --- a/tests/diffusion/quantization/test_gguf_config.py +++ /dev/null @@ -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)) diff --git a/tests/diffusion/quantization/test_trajectory_similarity_tool.py b/tests/diffusion/quantization/test_trajectory_similarity_tool.py index 45d51d356..252e6f1c1 100644 --- a/tests/diffusion/quantization/test_trajectory_similarity_tool.py +++ b/tests/diffusion/quantization/test_trajectory_similarity_tool.py @@ -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", ) diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 1f18d975b..3d29a4332 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -359,3 +359,15 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo assert not hasattr(stage_configs[0].engine_args, "additional_config") assert stage_configs[1].engine_args.additional_config == {"torchair_graph_config": {"enabled": True}} + + +def test_default_stage_config_includes_quantization_config(): + """Ensure structured quantization_config survives default diffusion-stage creation.""" + quantization_config = { + "method": "example_quant", + "weights": "weights.bin", + } + + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"quantization_config": quantization_config})[0] + + assert stage_cfg["engine_args"]["quantization_config"] == quantization_config diff --git a/tests/plugin/__init__.py b/tests/plugin/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/tests/plugin/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/plugin/test_gguf_diffusion.py b/tests/plugin/test_gguf_diffusion.py new file mode 100644 index 000000000..44de4a1fb --- /dev/null +++ b/tests/plugin/test_gguf_diffusion.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end tests for GGUF quantization on diffusion models. + +Validates that GGUF-quantized diffusion models generate valid images and +use less peak GPU memory than BF16 baseline. + +Requires vllm-omni to be installed alongside the plugin. + +Usage: + pytest tests/diffusion/test_gguf_memory.py -v +""" + +from __future__ import annotations + +from typing import NamedTuple + +import numpy as np +import pytest +import torch + +from tests.helpers.env import DeviceMemoryMonitor +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.platforms import current_omni_platform + + +class DiffusionGGUFTestConfig(NamedTuple): + artifact_prefix: str + gguf_model: str + hf_model: str + min_cosine_similarity: float = 0.95 + + +Z_IMAGE_CONFIG = DiffusionGGUFTestConfig( + artifact_prefix="zimage", + hf_model="Tongyi-MAI/Z-Image-Turbo", + gguf_model="unsloth/Z-Image-Turbo-GGUF:Q4_0", +) + + +FLUX_CONFIG = DiffusionGGUFTestConfig( + artifact_prefix="flux2_klein", + hf_model="black-forest-labs/FLUX.2-klein-4B", + gguf_model="unsloth/FLUX.2-klein-4B-GGUF:Q8_0", +) + + +def _image_cosine_similarity(hf_image, gguf_image) -> float: + hf_tensor = torch.as_tensor(np.array(hf_image.convert("RGB"), copy=True), dtype=torch.float32).flatten() + gguf_tensor = torch.as_tensor(np.array(gguf_image.convert("RGB"), copy=True), dtype=torch.float32).flatten() + return torch.nn.functional.cosine_similarity(hf_tensor, gguf_tensor, dim=0).item() + + +def _generate_single_stage_image( + model: str, + height: int = 256, + width: int = 256, + num_inference_steps: int = 20, + seed: int = 42, + **extra_kwargs, +) -> tuple[list, float]: + """Generate an image with a single-stage diffusion model. + + Returns (images, peak_memory_gib). + """ + omni_kwargs = dict(extra_kwargs) + + memory_monitor = DeviceMemoryMonitor(device_index=0) + memory_monitor.start() + try: + with OmniRunner(model, **omni_kwargs) as runner: + generator = torch.Generator( + device=current_omni_platform.device_type, + ).manual_seed(seed) + outputs = runner.omni.generate( + "a photo of a cat sitting on a laptop keyboard", + OmniDiffusionSamplingParams( + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=0.0, + generator=generator, + ), + ) + finally: + memory_monitor.stop() + peak_mem = memory_monitor.peak_used_mb / 1024 + + first_output = outputs[0] + assert first_output.final_output_type == "image" + if hasattr(first_output, "images") and first_output.images: + images = first_output.images + else: + assert hasattr(first_output, "request_output") and first_output.request_output + request_output = first_output.request_output + if isinstance(request_output, list): + req_out = request_output[0] + else: + req_out = request_output + assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images") + images = req_out.images + assert len(images) >= 1 + assert images[0].width == width + assert images[0].height == height + + return images, peak_mem + + +@hardware_test(res={"cuda": "L4"}) +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.slow +@pytest.mark.parametrize("model", [Z_IMAGE_CONFIG, FLUX_CONFIG], ids=["Z-Image-Turbo", "FLUX.2-klein"]) +def test_single_stage_diffusion_gguf(model: DiffusionGGUFTestConfig, monkeypatch: pytest.MonkeyPatch) -> None: + """Z-Image-Turbo GGUF generates valid images and uses less memory than BF16.""" + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + + # BF16 baseline + hf_images, mem_bf16 = _generate_single_stage_image( + model=model.hf_model, + ) + + # GGUF + images, mem_gguf = _generate_single_stage_image( + model=model.hf_model, + diffusion_quantization_config={ + "method": "gguf", + "gguf_model": model.gguf_model, + }, + ) + + hf_image_path = f"test_{model.artifact_prefix}_hf.png" + gguf_image_path = f"test_{model.artifact_prefix}_gguf.png" + assert len(hf_images) >= 1 + hf_images[0].save(hf_image_path) + assert len(images) >= 1 + images[0].save(gguf_image_path) + print(f"Saved HF image: {hf_image_path}") + print(f"Saved GGUF image: {gguf_image_path}") + + image_cosine_similarity = _image_cosine_similarity(hf_images[0], images[0]) + print(f"{model.artifact_prefix} image cosine similarity: {image_cosine_similarity:.4f}") + assert image_cosine_similarity >= model.min_cosine_similarity, ( + f"GGUF image cosine similarity ({image_cosine_similarity:.4f}) should be >= {model.min_cosine_similarity:.4f}" + ) + + print(f"{model.artifact_prefix} BF16 peak VRAM delta: {mem_bf16:.2f} GiB") + print(f"{model.artifact_prefix} GGUF peak VRAM delta: {mem_gguf:.2f} GiB") + reduction = (mem_bf16 - mem_gguf) / mem_bf16 * 100 + print(f"VRAM reduction: {reduction:.1f}%") + assert mem_gguf < mem_bf16, f"GGUF ({mem_gguf:.2f} GiB) should use less VRAM than BF16 ({mem_bf16:.2f} GiB)" diff --git a/tests/test_version.py b/tests/test_version.py index 07e622a7d..911a79611 100644 --- a/tests/test_version.py +++ b/tests/test_version.py @@ -56,3 +56,26 @@ def test_warning_contains_version_strings(): msg = str(record[0].message) assert "0.19.0rc2.dev21" in msg assert "0.18.0" in msg + + +def test_omni_plugin_loader_invokes_vllm_general_plugins_first(monkeypatch): + import vllm.plugins as vllm_plugins + + import vllm_omni.plugins as omni_plugins + + calls = [] + monkeypatch.setattr(omni_plugins, "omni_plugins_loaded", False) + monkeypatch.setattr(vllm_plugins, "load_general_plugins", lambda: calls.append("vllm")) + monkeypatch.setattr( + omni_plugins, + "load_omni_plugins_by_group", + lambda group: {"omni": lambda: calls.append(group)}, + ) + + omni_plugins.load_omni_general_plugins() + + assert calls == ["vllm", omni_plugins.OMNI_DEFAULT_PLUGINS_GROUP] + + omni_plugins.load_omni_general_plugins() + + assert calls == ["vllm", omni_plugins.OMNI_DEFAULT_PLUGINS_GROUP] diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index c44bd3ea9..dc5f4ef9c 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -551,7 +551,7 @@ class _DiffusionConfigProjection: elif isinstance(self.quantization_config, str): self.quantization_config = build_quant_config(self.quantization_config) elif isinstance(self.quantization_config, Mapping): - self.quantization_config = build_quant_config(dict(self.quantization_config)) + self.quantization_config = dict(self.quantization_config) else: raise TypeError( "quantization_config must be str, dict, QuantizationConfig, or None, " diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 7ee2b4395..943c55a52 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -11,7 +11,6 @@ from pathlib import Path from typing import cast import torch -from huggingface_hub import hf_hub_download from torch import nn from vllm.config import ModelConfig from vllm.config.load import LoadConfig @@ -36,7 +35,6 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( get_checkpoint_adapter, ) -from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery from vllm_omni.diffusion.registry import initialize_model @@ -381,10 +379,7 @@ class DiffusersPipelineLoader: # DiffusionPipeline.from_pretrained() internally; it does # not use our native customized pipeline classes. cast(DiffusersAdapterPipeline, model).load_weights() - elif self._is_gguf_quantization(): - self._load_weights_with_gguf(model) else: - # Quantization does not happen in `load_weights` but after it self.load_weights(model) # Process weights after loading for quantization (e.g., FP8 online quantization) @@ -454,7 +449,7 @@ class DiffusersPipelineLoader: def _is_expected_quantized_weight(name: str) -> bool: """Return True if *name* is a quantization-specific parameter. - Quantization methods (GPTQ, AWQ, FP8, GGUF, Autoround, etc.) create extra + Quantization methods (GPTQ, AWQ, FP8, Autoround, etc.) create extra parameters that have no counterpart in an unquantized checkpoint. These are expected to be absent and should not trigger a load error. """ @@ -467,8 +462,6 @@ class DiffusersPipelineLoader: ".weight_scale", ".weight_scale_inv", ".input_scale", - # GGUF - ".qweight_type", # INT8 (weight_scale already covered above) ) return name.endswith(_QUANTIZED_WEIGHT_SUFFIXES) @@ -500,116 +493,6 @@ class DiffusersPipelineLoader: if unexpected_missing: raise ValueError(f"Following weights were not initialized from checkpoint: {unexpected_missing}") - def _get_gguf_model_from_config(self) -> str | None: - """Extract gguf_model from quant_config (handles both dict and object).""" - if self.quant_config is None: - return None - if isinstance(self.quant_config, dict): - return self.quant_config.get("gguf_model") - return getattr(self.quant_config, "gguf_model", None) - - def _is_gguf_quantization(self) -> bool: - """Check whether or not this pipeline loader is pointing at a GGUF config.""" - from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig - - maybe_gguf_model = self._get_gguf_model_from_config() - is_gguf = ( - isinstance(self.quant_config, DiffusionGGUFConfig) - or (isinstance(self.quant_config, dict) and self.quant_config.get("method") == "gguf") - or (hasattr(self.quant_config, "get_name") and self.quant_config.get_name() == "gguf") - ) - if is_gguf: - if not maybe_gguf_model: - raise ValueError("GGUF quantization requires gguf_model") - return True - - if isinstance(self.quant_config, dict): - return False - - return bool(maybe_gguf_model) - - def _is_transformer_source(self, source: "ComponentSource") -> bool: - if source.subfolder == "transformer": - return True - return source.prefix.startswith("transformer.") - - def _get_model_loadable_names(self, model: nn.Module) -> set[str]: - # Avoid model.state_dict() here because GGUF uses UninitializedParameter - # which raises during detach(). Collect names directly. - return {name for name, _ in model.named_parameters()} | {name for name, _ in model.named_buffers()} - - def _resolve_gguf_model_path(self, gguf_model: str, revision: str | None) -> str: - if os.path.isfile(gguf_model): - return gguf_model - # repo_id/filename.gguf - if "/" in gguf_model and gguf_model.endswith(".gguf"): - repo_id, filename = gguf_model.rsplit("/", 1) - return hf_hub_download( - repo_id=repo_id, - filename=filename, - revision=revision, - cache_dir=self.load_config.download_dir, - ) - # repo_id:quant_type - if "/" in gguf_model and ":" in gguf_model: - repo_id, quant_type = gguf_model.rsplit(":", 1) - return download_gguf( - repo_id, - quant_type, - cache_dir=self.load_config.download_dir, - revision=revision, - ignore_patterns=self.load_config.ignore_patterns, - ) - raise ValueError( - f"Unrecognized GGUF reference: {gguf_model!r} (expected local file, " - "/.gguf, or :)" - ) - - def _get_gguf_weights_iterator( - self, - source: "ComponentSource", - model: nn.Module, - ) -> Generator[tuple[str, torch.Tensor], None, None]: - gguf_model = self._get_gguf_model_from_config() - if gguf_model is None: - raise ValueError("GGUF quantization requires quantization_config.gguf_model") - gguf_file = self._resolve_gguf_model_path(gguf_model, self.od_config.revision) - adapter = get_gguf_adapter(gguf_file, model, source, self.od_config) - weights_iter = adapter.weights_iterator() - return ((source.prefix + name, tensor) for (name, tensor) in weights_iter) - - def _load_weights_with_gguf(self, model: nn.Module) -> set[str]: - sources = self._get_weight_sources(model) - loaded: set[str] = set() - loadable_names: set[str] | None = None - - for source in sources: - if self._is_transformer_source(source): - loaded |= model.load_weights(self._get_gguf_weights_iterator(source, model)) - - # GGUF checkpoints can be transformer-only or partially quantized. - # Only fall back to HF if this source still has missing loadable weights. - loadable_names = loadable_names or self._get_model_loadable_names(model) - has_missing_for_source = any( - name.startswith(source.prefix) and name not in loaded for name in loadable_names - ) - if not has_missing_for_source: - continue - - hf_iter = self._get_weights_iterator(source) - hf_iter = ( - (name, tensor) for (name, tensor) in hf_iter if name in loadable_names and name not in loaded - ) - loaded |= model.load_weights(hf_iter) - else: - loaded |= model.load_weights(self._get_weights_iterator(source)) - - weights_to_load = self._get_expected_parameter_names(model) - weights_not_loaded = weights_to_load - loaded - if weights_not_loaded: - raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}") - return loaded - def _init_from_load_format( self, load_format: str, diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/__init__.py b/vllm_omni/diffusion/model_loader/gguf_adapters/__init__.py deleted file mode 100644 index 3266ba06c..000000000 --- a/vllm_omni/diffusion/model_loader/gguf_adapters/__init__.py +++ /dev/null @@ -1,47 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -from typing import TYPE_CHECKING - -import torch - -from .base import GGUFAdapter -from .flux2_klein import Flux2KleinGGUFAdapter -from .qwen_image import QwenImageGGUFAdapter -from .z_image import ZImageGGUFAdapter - -if TYPE_CHECKING: - from vllm_omni.diffusion.data import OmniDiffusionConfig - from vllm_omni.diffusion.model_loader.diffusers_loader import ( - DiffusersPipelineLoader, - ) - - -def get_gguf_adapter( - gguf_file: str, - model: torch.nn.Module, - source: DiffusersPipelineLoader.ComponentSource, - od_config: OmniDiffusionConfig, -) -> GGUFAdapter: - adapter_classes = (QwenImageGGUFAdapter, ZImageGGUFAdapter, Flux2KleinGGUFAdapter) - for adapter_cls in adapter_classes: - if adapter_cls.is_compatible(od_config, model, source): - return adapter_cls(gguf_file, model, source, od_config) - model_type = None - if od_config.tf_model_config is not None: - model_type = od_config.tf_model_config.get("model_type") - supported = ", ".join(cls.__name__ for cls in adapter_classes) - raise ValueError( - "No GGUF adapter matched diffusion model " - f"(model_class_name={od_config.model_class_name!r}, model_type={model_type!r}). " - f"Supported adapters: {supported}." - ) - - -__all__ = [ - "GGUFAdapter", - "Flux2KleinGGUFAdapter", - "QwenImageGGUFAdapter", - "ZImageGGUFAdapter", - "get_gguf_adapter", -] diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/base.py b/vllm_omni/diffusion/model_loader/gguf_adapters/base.py deleted file mode 100644 index f3f4d28a0..000000000 --- a/vllm_omni/diffusion/model_loader/gguf_adapters/base.py +++ /dev/null @@ -1,109 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -from abc import ABC, abstractmethod -from collections.abc import Generator -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any - -import numpy as np -import torch - -try: - import gguf -except ImportError: - # Optional dependency: only required for GGUF-quantized diffusion models. - # Keep this import lazy so a missing ``gguf`` does not break importing the - # diffusion model loader (and thus inspection/loading of every other - # diffusion model that imports through this package). - gguf = None - -if TYPE_CHECKING: - from vllm_omni.diffusion.data import OmniDiffusionConfig - from vllm_omni.diffusion.model_loader.diffusers_loader import ( - DiffusersPipelineLoader, - ) - - -@dataclass -class MappedTensor: - name: str - tensor: Any - tensor_type: Any - row_slice: slice | None = None - swap_scale_shift: bool = False - - -class GGUFAdapter(ABC): - """Base class for model-specific GGUF adapters.""" - - def __init__( - self, - gguf_file: str, - model: torch.nn.Module, - source: DiffusersPipelineLoader.ComponentSource, - od_config: OmniDiffusionConfig, - ) -> None: - self.gguf_file = gguf_file - self.model = model - self.source = source - self.od_config = od_config - - @staticmethod - def is_compatible( - od_config: OmniDiffusionConfig, - model: torch.nn.Module, - source: DiffusersPipelineLoader.ComponentSource, - ) -> bool: - return False - - @abstractmethod - def weights_iterator(self) -> Generator[tuple[str, torch.Tensor], None, None]: - raise NotImplementedError - - -# FIXME(Isotr0py): Sync implemnentation with upstream vLLM? -def gguf_quant_weights_iterator(gguf_file: str) -> Generator[tuple[str, torch.Tensor]]: - """ - Iterate over the quant weights in the model gguf files and convert - them to torch tensors. - Be careful of the order of yielding weight types and weights data, - we have to yield all weight types first before yielding any weights. - Otherwise it would cause issue when loading weights with for packed - layer with different quant types. - """ - - if gguf is None: - raise ImportError( - "The 'gguf' package is required to load GGUF-quantized diffusion models. " - "Install it with `pip install gguf`." - ) - reader = gguf.GGUFReader(gguf_file) - - for tensor in reader.tensors: - weight_type = tensor.tensor_type - name = tensor.name - - if weight_type.name not in ("F32", "F16"): - weight_type_name = name.replace("weight", "qweight_type") - weight_type = torch.tensor(weight_type) - yield weight_type_name, weight_type - - for tensor in reader.tensors: - weight = tensor.data - weight_type = tensor.tensor_type - name = tensor.name - if weight_type.name not in ("F32", "F16"): - name = name.replace("weight", "qweight") - if weight_type.name == "BF16" and tensor.data.dtype == np.uint8: - # BF16 is currently the only "quantization" type that isn't - # actually quantized but is read as a raw byte tensor. - # Reinterpret as `torch.bfloat16` tensor. - weight = weight.view(np.uint16) - if reader.byte_order == "S": - # GGUF endianness != system endianness - weight = weight.byteswap() - param = torch.tensor(weight).view(torch.bfloat16) - else: - param = torch.tensor(weight) - yield name, param diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py b/vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py deleted file mode 100644 index 4ee6439b2..000000000 --- a/vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py +++ /dev/null @@ -1,98 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -from collections.abc import Iterable - -import torch -from vllm.model_executor.models.utils import WeightsMapper - -from .base import GGUFAdapter, gguf_quant_weights_iterator - -FLUX2_TRANSFORMER_KEYS_RENAME_DICT = { - "single_blocks.": "single_transformer_blocks.", - # Image and text input projections - "img_in": "x_embedder", - "txt_in": "context_embedder", - # Timestep and guidance embeddings - "time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1", - "time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2", - "guidance_in.in_layer": "time_guidance_embed.guidance_embedder.linear_1", - "guidance_in.out_layer": "time_guidance_embed.guidance_embedder.linear_2", - # Modulation parameters - "double_stream_modulation_img.lin": "double_stream_modulation_img.linear", - "double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear", - "single_stream_modulation.lin": "single_stream_modulation.linear", - # Final output layer - # "final_layer.adaLN_modulation.1": "norm_out.linear", # Handle separately since we need to swap mod params - "final_layer.linear": "proj_out", -} - -FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP = { - "final_layer.adaLN_modulation.1": "norm_out.linear", -} - -FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP = { - "double_blocks.": "transformer_blocks.", - # Handle fused QKV projections separately as we need to break into Q, K, V projections - "img_attn.norm.query_norm": "attn.norm_q", - "img_attn.norm.key_norm": "attn.norm_k", - "img_attn.proj": "attn.to_out.0", - "img_mlp.0": "ff.linear_in", - "img_mlp.2": "ff.linear_out", - "txt_attn.norm.query_norm": "attn.norm_added_q", - "txt_attn.norm.key_norm": "attn.norm_added_k", - "txt_attn.proj": "attn.to_add_out", - "txt_mlp.0": "ff_context.linear_in", - "txt_mlp.2": "ff_context.linear_out", - # Additional for fuse qkv - "img_attn.qkv": "attn.to_qkv", - "txt_attn.qkv": "attn.add_kv_proj", -} - -FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP = { - "linear1": "attn.to_qkv_mlp_proj", - "linear2": "attn.to_out", - "norm.query_norm": "attn.norm_q", - "norm.key_norm": "attn.norm_k", -} - - -class Flux2KleinGGUFAdapter(GGUFAdapter): - """GGUF adapter for Flux2-Klein models with qkv splitting and adaLN swap.""" - - @staticmethod - def is_compatible(od_config, model: torch.nn.Module, source) -> bool: - model_class = od_config.model_class_name or "" - if model_class.startswith("Flux2"): - return True - cfg = od_config.tf_model_config - if cfg is not None: - model_type = str(cfg.get("model_type", "")).lower() - if model_type.startswith("flux"): - return True - return False - - gguf_to_hf_mapper = WeightsMapper( - # double_stream_modulation - orig_to_new_prefix=FLUX2_TRANSFORMER_KEYS_RENAME_DICT | FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP, - orig_to_new_substr=FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP | FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP, - ) - - def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]: - def custom_weights_adapter( - weights: Iterable[tuple[str, torch.Tensor]], - ) -> Iterable[tuple[str, torch.Tensor]]: - for name, weight in weights: - # Handle the special case for adaLN modulation parameters that require swapping shift and scale - if name.endswith(".scale"): - name = name.replace(".scale", ".weight") - if name == "norm_out.linear.weight": - shift, scale = weight.chunk(2, dim=0) - weight = torch.cat([scale, shift], dim=0) - yield name, weight - else: - yield name, weight - - weights = gguf_quant_weights_iterator(self.gguf_file) - weights = self.gguf_to_hf_mapper.apply(weights) - yield from custom_weights_adapter(weights) diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py b/vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py deleted file mode 100644 index cc8ccea5c..000000000 --- a/vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py +++ /dev/null @@ -1,27 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -from collections.abc import Iterable - -import torch - -from .base import GGUFAdapter, gguf_quant_weights_iterator - - -class QwenImageGGUFAdapter(GGUFAdapter): - """GGUF adapter for the Qwen-Image transformer family.""" - - @staticmethod - def is_compatible(od_config, model: torch.nn.Module, source) -> bool: - model_class = od_config.model_class_name or "" - if model_class.startswith("QwenImage"): - return True - cfg = od_config.tf_model_config - if cfg is not None: - model_type = str(cfg.get("model_type", "")).lower() - if model_type.startswith("qwen_image"): - return True - return False - - def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]: - yield from gguf_quant_weights_iterator(self.gguf_file) diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py b/vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py deleted file mode 100644 index 7d8963355..000000000 --- a/vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py +++ /dev/null @@ -1,43 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -from collections.abc import Iterable - -import torch -from vllm.model_executor.models.utils import WeightsMapper - -from .base import GGUFAdapter, gguf_quant_weights_iterator - -Z_IMAGE_KEYS_RENAME_DICT = { - "final_layer.": "all_final_layer.2-1.", - "x_embedder.": "all_x_embedder.2-1.", - ".attention.qkv": ".attention.to_qkv", - ".attention.k_norm": ".attention.norm_k", - ".attention.q_norm": ".attention.norm_q", - ".attention.out": ".attention.to_out.0", - "model.diffusion_model.": "", -} - - -class ZImageGGUFAdapter(GGUFAdapter): - """GGUF adapter for Z-Image models with QKV/FFN shard support.""" - - @staticmethod - def is_compatible(od_config, model: torch.nn.Module, source) -> bool: - model_class = od_config.model_class_name or "" - if model_class.startswith("ZImage"): - return True - cfg = od_config.tf_model_config - if cfg is not None: - model_type = str(cfg.get("model_type", "")).lower() - if model_type in {"z_image", "zimage", "z-image"}: - return True - return False - - gguf_to_hf_mapper = WeightsMapper( - orig_to_new_substr=Z_IMAGE_KEYS_RENAME_DICT, - ) - - def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]: - weights = gguf_quant_weights_iterator(self.gguf_file) - yield from self.gguf_to_hf_mapper.apply(weights) diff --git a/vllm_omni/diffusion/models/flux2/flux2_transformer.py b/vllm_omni/diffusion/models/flux2/flux2_transformer.py index 5a65db257..52e0a4430 100644 --- a/vllm_omni/diffusion/models/flux2/flux2_transformer.py +++ b/vllm_omni/diffusion/models/flux2/flux2_transformer.py @@ -1034,7 +1034,7 @@ class Flux2Transformer2DModel(nn.Module): name = original_name if name not in params_dict and ".to_out.0." in name: name = name.replace(".to_out.0.", ".to_out.") - # Some GGUF checkpoints include quantized tensors for modules that + # Some quantized checkpoints include tensors for modules that # are intentionally left unquantized in this model. param = params_dict.get(name) if param is None: diff --git a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py index cf2c9f78a..e62fccba0 100644 --- a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py +++ b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py @@ -1044,7 +1044,7 @@ class Flux2Transformer2DModel(nn.Module): name = original_name if name not in params_dict and ".to_out.0." in name: name = name.replace(".to_out.0.", ".to_out.") - # Some GGUF checkpoints include quantized tensors for modules that + # Some quantized checkpoints include tensors for modules that # are intentionally left unquantized in this model. param = params_dict.get(name) if param is None: diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index e71af65d3..ec9a9698c 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1043,6 +1043,7 @@ class AsyncOmniEngine: "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4), "quantization": kwargs.get("quantization", None), + "quantization_config": kwargs.get("quantization_config", None), "diffusion_kv_cache_dtype": kwargs.get("diffusion_kv_cache_dtype", None), "diffusion_kv_cache_skip_steps": kwargs.get("diffusion_kv_cache_skip_steps", None), "diffusion_kv_cache_skip_layers": kwargs.get("diffusion_kv_cache_skip_layers", None), @@ -1233,7 +1234,7 @@ class AsyncOmniEngine: kwargs.get("diffusion_attention_config"), attention_backend=kwargs.get("diffusion_attention_backend"), ) - quantization_config = kwargs.get("diffusion_quantization_config") + quantization_config = kwargs.get("diffusion_quantization_config") or kwargs.get("quantization_config") if quantization_config is not None: if ( not hasattr(cfg.engine_args, "quantization_config") diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 86576ea55..5d435bc2f 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -478,7 +478,7 @@ class OmniServeCommand(CLISubcommand): default=None, help=( "JSON string for diffusion quantization_config. " - 'Example: \'{"method":"gguf","gguf_model":"/path/to/model.gguf"}\'.' + 'Example: \'{"method":"fp8","activation_scheme":"dynamic"}\'.' ), ) omni_config_group.add_argument( diff --git a/vllm_omni/plugins/__init__.py b/vllm_omni/plugins/__init__.py index 41bf398ca..240394b2c 100644 --- a/vllm_omni/plugins/__init__.py +++ b/vllm_omni/plugins/__init__.py @@ -68,6 +68,10 @@ def load_omni_general_plugins() -> None: return omni_plugins_loaded = True + from vllm.plugins import load_general_plugins + + load_general_plugins() + plugins = load_omni_plugins_by_group(group=OMNI_DEFAULT_PLUGINS_GROUP) # General plugins: we only need to execute the loaded functions. for func in plugins.values(): diff --git a/vllm_omni/quantization/__init__.py b/vllm_omni/quantization/__init__.py index d8cad3a98..8896903e8 100644 --- a/vllm_omni/quantization/__init__.py +++ b/vllm_omni/quantization/__init__.py @@ -12,10 +12,10 @@ Adds per-component quantization for multi-stage models. """ from .component_config import ComponentQuantizationConfig -from .factory import SUPPORTED_QUANTIZATION_METHODS, build_quant_config +from .factory import SUPPORTED_QUANTIZATION_METHODS, build_quant_config, register_quantization_override from .inc_config import OmniINCConfig -# Heavy configs (GGUF, MXFP8) are NOT imported here to avoid pulling in +# Heavy configs are NOT imported here to avoid pulling in # optional dependencies (pynvml, torch_npu) at module load time. # Import them directly when needed: # from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config @@ -25,4 +25,5 @@ __all__ = [ "ComponentQuantizationConfig", "OmniINCConfig", "SUPPORTED_QUANTIZATION_METHODS", + "register_quantization_override", ] diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index d9a7a31bf..2129d8f04 100644 --- a/vllm_omni/quantization/factory.py +++ b/vllm_omni/quantization/factory.py @@ -3,8 +3,8 @@ """Factory for building quantization configs. build_quant_config() delegates to vLLM's quantization registry. -The only extension point is _OVERRIDES for methods that need a -different QuantizationConfig subclass in the OMNI context (e.g. GGUF). +Out-of-tree integrations can register Omni-specific builders with +register_quantization_override(). """ from __future__ import annotations @@ -81,13 +81,6 @@ from .component_config import ComponentQuantizationConfig # noqa: E402 logger = init_logger(__name__) -def _build_gguf(**kw: Any) -> QuantizationConfig: - """Lazy import to avoid pulling in CUDA/pynvml at module load time.""" - from .gguf_config import DiffusionGGUFConfig - - return DiffusionGGUFConfig(**kw) - - def _build_int8(**kw: Any) -> QuantizationConfig: """Lazy import for Int8 diffusion config (supports CUDA + NPU).""" from .int8_config import DiffusionInt8Config @@ -140,7 +133,6 @@ def _build_inc(**kw: Any) -> QuantizationConfig: _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = { - "gguf": _build_gguf, "int8": _build_int8, "mxfp8": _build_mxfp8, "mxfp4": _build_mxfp4, @@ -150,7 +142,24 @@ _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = { "auto_round": _build_inc, } -SUPPORTED_QUANTIZATION_METHODS: list[str] = list(dict.fromkeys(QUANTIZATION_METHODS + list(_OVERRIDES.keys()))) + +def _compute_supported_quantization_methods() -> list[str]: + return list(dict.fromkeys(QUANTIZATION_METHODS + list(_OVERRIDES.keys()))) + + +SUPPORTED_QUANTIZATION_METHODS: list[str] = _compute_supported_quantization_methods() + + +def _refresh_registered_methods() -> None: + SUPPORTED_QUANTIZATION_METHODS[:] = _compute_supported_quantization_methods() + global _CACHED_ALIAS_MAP + _CACHED_ALIAS_MAP = None + + +def register_quantization_override(method: str, builder: Callable[..., QuantizationConfig]) -> None: + """Register an Omni-specific quantization config builder.""" + _OVERRIDES[_normalize_method_name(method)] = builder + _refresh_registered_methods() def _build_reverse_alias_map() -> dict[str, str]: diff --git a/vllm_omni/quantization/gguf_config.py b/vllm_omni/quantization/gguf_config.py deleted file mode 100644 index 186ac4ee7..000000000 --- a/vllm_omni/quantization/gguf_config.py +++ /dev/null @@ -1,323 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""GGUF quantization config for diffusion transformers. - -Uses dequant+GEMM instead of the fused kernel path (which expects 2D inputs). - -The GGUF quantization module was migrated to an external plugin in upstream -vLLM (commit 6635279d8). This file inlines the base classes that -DiffusionGGUFConfig / DiffusionGGUFLinearMethod depend on, reproduced from -the last upstream version before the migration. -""" - -from __future__ import annotations - -from collections.abc import Mapping -from types import MappingProxyType -from typing import TYPE_CHECKING, Any - -import gguf -import torch -from gguf import GGMLQuantizationType as WeightType -from torch.nn.parameter import Parameter, UninitializedParameter -from vllm import _custom_ops as ops -from vllm.logger import init_logger -from vllm.model_executor.layers.linear import ( - LinearBase, - LinearMethodBase, - UnquantizedLinearMethod, -) -from vllm.model_executor.layers.quantization.base_config import ( - QuantizationConfig, - QuantizeMethodBase, -) -from vllm.model_executor.utils import set_weight_attrs - -if TYPE_CHECKING: - from vllm.model_executor.layers.quantization import QuantizationMethods - -logger = init_logger(__name__) - -# --------------------------------------------------------------------------- -# GGUF base types (inlined from upstream vllm before migration to plugin) -# --------------------------------------------------------------------------- - -UNQUANTIZED_TYPES: set[WeightType] = {WeightType.F32, WeightType.F16, WeightType.BF16} - -STANDARD_QUANT_TYPES: set[WeightType] = { - WeightType.Q4_0, - WeightType.Q4_1, - WeightType.Q5_0, - WeightType.Q5_1, - WeightType.Q8_0, - WeightType.Q8_1, -} -KQUANT_TYPES: set[WeightType] = { - WeightType.Q2_K, - WeightType.Q3_K, - WeightType.Q4_K, - WeightType.Q5_K, - WeightType.Q6_K, -} -IMATRIX_QUANT_TYPES: set[WeightType] = { - WeightType.IQ1_M, - WeightType.IQ1_S, - WeightType.IQ2_XXS, - WeightType.IQ2_XS, - WeightType.IQ2_S, - WeightType.IQ3_XXS, - WeightType.IQ3_S, - WeightType.IQ4_XS, - WeightType.IQ4_NL, -} -DEQUANT_TYPES: set[WeightType] = STANDARD_QUANT_TYPES | KQUANT_TYPES | IMATRIX_QUANT_TYPES - - -def is_layer_skipped_gguf( - prefix: str, - unquantized_modules: list[str], - fused_mapping: Mapping[str, list[str]] = MappingProxyType({}), -) -> bool: - # Fused layers like gate_up_proj or qkv_proj will not be fused - # in the safetensors checkpoint. So, we convert the name - # from the fused version to unfused + check to make sure that - # each shard of the fused layer has the same scheme. - proj_name = prefix.split(".")[-1] - if proj_name in fused_mapping: - shard_prefixes = [prefix.replace(proj_name, shard_proj_name) for shard_proj_name in fused_mapping[proj_name]] - - is_skipped = None - for shard_prefix in shard_prefixes: - is_shard_skipped = any(shard_prefix in module_name for module_name in unquantized_modules) - - if is_skipped is None: - is_skipped = is_shard_skipped - elif is_shard_skipped != is_skipped: - raise ValueError( - f"Detected some but not all shards of {prefix} " - "are quantized. All shards of fused layers " - "to have the same precision." - ) - else: - is_skipped = any(module_name in prefix for module_name in unquantized_modules) - - assert is_skipped is not None - return is_skipped - - -class _GGUFUninitializedParameter(UninitializedParameter): - """UninitializedParameter subclass that stores GGUF weight data in a list. - - The actual tensor is materialized later via _create_padded_weight_param. - """ - - cls_to_become = Parameter - data_container: list[torch.Tensor] - - -class GGUFConfig(QuantizationConfig): - """Config class for GGUF (inlined from upstream vllm).""" - - def __init__(self, unquantized_modules: list[str] | None = None) -> None: - super().__init__() - self.unquantized_modules = unquantized_modules or [] - - def __repr__(self) -> str: - return "GGUFConfig()" - - def get_name(self) -> QuantizationMethods: - return "gguf" - - def get_supported_act_dtypes(self) -> list[torch.dtype]: - return [torch.half, torch.bfloat16, torch.float32] - - @classmethod - def get_min_capability(cls) -> int: - return 60 - - @classmethod - def get_config_filenames(cls) -> list[str]: - return [] # no extra configs. - - @classmethod - def from_config(cls, config: dict[str, Any]) -> GGUFConfig: - return cls() - - @classmethod - def override_quantization_method( - cls, hf_quant_cfg: dict[str, Any], user_quant: str | None, hf_config=None - ) -> QuantizationMethods | None: - if user_quant == "gguf": - return "gguf" - return None - - def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None: - if isinstance(layer, LinearBase): - if is_layer_skipped_gguf(prefix, self.unquantized_modules, self.packed_modules_mapping): - return UnquantizedLinearMethod() - return GGUFLinearMethod(self) - return None - - -class GGUFLinearMethod(LinearMethodBase): - """Linear method for GGUF (inlined from upstream vllm).""" - - def __init__(self, quant_config: GGUFConfig): - self.quant_config = quant_config - - def create_weights( - self, - layer: torch.nn.Module, - input_size_per_partition: int, - output_partition_sizes: list[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - self.params_dtype = params_dtype - output_size_per_partition = sum(output_partition_sizes) - - tensor_shape = (output_size_per_partition, input_size_per_partition) - qweight = _GGUFUninitializedParameter(requires_grad=False) - set_weight_attrs( - qweight, - { - "input_dim": 1, - "output_dim": 0, - "tensor_shape": tensor_shape, - "is_gguf_weight": True, - "data_container": [], - "shard_id": [], - "shard_id_map": {}, - }, - ) - set_weight_attrs(qweight, extra_weight_attrs) - layer.register_parameter("qweight", qweight) - - qweight_type = Parameter( - torch.empty(len(output_partition_sizes), dtype=torch.uint8), - requires_grad=False, - ) - set_weight_attrs( - qweight_type, - { - "is_gguf_weight_type": True, - "weight_type": 0, - "shard_weight_type": {}, - "ignore_warning": True, - }, - ) - set_weight_attrs(qweight_type, extra_weight_attrs) - layer.register_parameter("qweight_type", qweight_type) - - def process_weights_after_loading(self, layer: torch.nn.Module): - qweight_type = layer.qweight_type.weight_type - if not (qweight_type in UNQUANTIZED_TYPES or qweight_type in DEQUANT_TYPES): - qweight_type = WeightType(qweight_type) - raise ValueError(f"Unsupported GGUF quantization type {qweight_type} in layer {layer}.") - self._create_padded_weight_param(layer) - - def _create_padded_weight_param(self, layer: torch.nn.Module): - """Create padded weight parameter for GGUF MergedLinear layer.""" - qweight = layer.qweight - shard_id_map = qweight.shard_id_map - shard_id = qweight.shard_id - if len(data_container := qweight.data_container) > 1: - dtype = {data.dtype for data in data_container} - assert len(dtype) == 1, ValueError(f"Data container has mixed dtypes: {dtype}") - dtype = next(iter(dtype)) - padded_side = max(x.size(1) for x in data_container) - concat_side = sum(x.size(0) for x in data_container) - padded_data = torch.zeros((concat_side, padded_side), dtype=dtype, device=qweight.device) - shard_offset_map = dict[str, tuple[int, int, int]]() - for idx in shard_id: - id_in_container = shard_id_map[idx] - start = sum(x.size(0) for x in data_container[:id_in_container]) - end = start + data_container[id_in_container].size(0) - size = data_container[id_in_container].size(1) - padded_data[start:end, :size] = data_container[id_in_container] - shard_offset_map[idx] = (start, end, size) - qweight.data_container.clear() - padded_param = Parameter(padded_data, requires_grad=False) - set_weight_attrs(padded_param, vars(qweight)) - set_weight_attrs(padded_param, {"shard_offset_map": shard_offset_map}) - layer.register_parameter("qweight", padded_param) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - # Base apply: overridden by DiffusionGGUFLinearMethod. - raise NotImplementedError("Use DiffusionGGUFLinearMethod for diffusion GGUF") - - -# --------------------------------------------------------------------------- -# Diffusion-specific GGUF (dequant+GEMM, replaces fused kernel path) -# --------------------------------------------------------------------------- - - -def dequant_gemm_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int) -> torch.Tensor: - if qweight_type in UNQUANTIZED_TYPES: - return x @ qweight.T - block_size, type_size = gguf.GGML_QUANT_SIZES[qweight_type] - shape = (qweight.shape[0], qweight.shape[1] // type_size * block_size) - weight = ops.ggml_dequantize(qweight, qweight_type, *shape, x.dtype) - return x @ weight.T - - -class DiffusionGGUFLinearMethod(GGUFLinearMethod): - """GGUF linear method using dequant+GEMM for N-D diffusion tensors.""" - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - shard_id = getattr(layer.qweight, "shard_id", None) - - if shard_id: - shard_id = ["q", "k", "v"] if "q" in shard_id else shard_id - qweight = layer.qweight - result = [] - for idx in shard_id: - start, end, offset = layer.qweight.shard_offset_map[idx] - qweight_type = layer.qweight_type.shard_weight_type[idx] - result.append( - dequant_gemm_gguf( - x, - qweight[start:end, :offset].contiguous(), - qweight_type, - ) - ) - out = torch.cat(result, axis=-1) - else: - qweight = layer.qweight - qweight_type = layer.qweight_type.weight_type - out = dequant_gemm_gguf(x, qweight, qweight_type) - - if bias is not None: - out.add_(bias) - return out - - -class DiffusionGGUFConfig(GGUFConfig): - """GGUF config that carries gguf_model path and uses dequant+GEMM.""" - - def __init__( - self, - gguf_model: str | None = None, - unquantized_modules: list[str] | None = None, - ) -> None: - super().__init__(unquantized_modules=unquantized_modules or []) - self.gguf_model = gguf_model - - def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None: - if isinstance(layer, LinearBase): - if is_layer_skipped_gguf(prefix, self.unquantized_modules, self.packed_modules_mapping): - return UnquantizedLinearMethod() - return DiffusionGGUFLinearMethod(self) - return None diff --git a/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py b/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py index 3d4f68d23..77b191738 100644 --- a/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py +++ b/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py @@ -280,7 +280,6 @@ def _build_quantization_config( quantization: str | None, quantization_config_json: str | None, ignored_layers: str | None, - gguf_model: str | None, ) -> dict[str, Any] | None: explicit = _load_json_arg(quantization_config_json) if explicit is not None: @@ -292,8 +291,6 @@ def _build_quantization_config( ignored = _split_csv(ignored_layers) if ignored: config["ignored_layers"] = ignored - if gguf_model is not None: - config["gguf_model"] = gguf_model return config @@ -305,7 +302,6 @@ def _build_variant_config(args: argparse.Namespace, variant: str) -> VariantConf quantization=quantization, quantization_config_json=getattr(args, f"{variant}_quantization_config_json"), ignored_layers=ignored_layers, - gguf_model=getattr(args, f"{variant}_gguf_model"), ) return VariantConfig( label=variant, @@ -554,8 +550,6 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--reference-ignored-layers", help="Comma-separated reference ignored layer patterns.") parser.add_argument("--candidate-ignored-layers", help="Comma-separated candidate ignored layer patterns.") parser.add_argument("--ignored-layers", help="Alias for --candidate-ignored-layers.") - parser.add_argument("--reference-gguf-model", help="Reference GGUF file or HF reference when using gguf.") - parser.add_argument("--candidate-gguf-model", help="Candidate GGUF file or HF reference when using gguf.") parser.add_argument("--reference-load-format", default="default", help="Reference diffusion load format.") parser.add_argument("--candidate-load-format", default="default", help="Candidate diffusion load format.") parser.add_argument("--prompt", default="a cup of coffee on the table")