[Diffusion] Add Diffusers quantization config bridge (#3837)

Signed-off-by: CAICAIIs <3360776475@qq.com>
2026-07-03 03:06:44 +00:00
1 parent 06329c6450
commit b443504d61
@@ -67,6 +67,50 @@ def _make_request(**overrides) -> OmniDiffusionRequest:
    return OmniDiffusionRequest(**defaults)


+def _patch_fake_diffusers_quantization_backends(mocker):
+    class FakePipelineQuantizationConfig:
+        def __init__(self, *, quant_mapping=None, quant_backend=None, quant_kwargs=None, components_to_quantize=None):
+            self.quant_mapping = quant_mapping
+            self.quant_backend = quant_backend
+            self.quant_kwargs = quant_kwargs
+            self.components_to_quantize = components_to_quantize
+
+    class FakeTorchAoConfig:
+        def __init__(self, quant_type, modules_to_not_convert=None):
+            self.quant_type = quant_type
+            self.modules_to_not_convert = modules_to_not_convert
+
+        @classmethod
+        def from_dict(cls, config):
+            return cls(quant_type=config["quant_type"])
+
+    class FakeInt8DynamicActivationInt8WeightConfig:
+        pass
+
+    mocker.patch(
+        "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils.PipelineQuantizationConfig",
+        FakePipelineQuantizationConfig,
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils.TorchAoConfig",
+        FakeTorchAoConfig,
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils._get_torchao_quant_type_cls",
+        return_value=FakeInt8DynamicActivationInt8WeightConfig,
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils._get_diffusers_quantization_config_cls",
+        return_value=FakeTorchAoConfig,
+    )
+    return FakePipelineQuantizationConfig, FakeTorchAoConfig, FakeInt8DynamicActivationInt8WeightConfig
+
+
+class _TransformerBackedPipeline(DiffusionPipeline):
+    def __init__(self, transformer):
+        pass
+
+
 def _make_batch(**overrides) -> DiffusionRequestBatch:
    """Wrap a single request in a DiffusionRequestBatch, matching the forward contract."""
    return DiffusionRequestBatch(requests=[_make_request(**overrides)])
@@ -98,7 +142,7 @@ class TestPipelineArgumentsHandling:

    @pytest.mark.parametrize(
        "feature_id",
-        ["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "quantization"],
+        ["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "unsupported_quantization"],
    )
    def test_adapter_guard_unsupported_feature(self, feature_id):
        if feature_id == "cfg_parallel":
@@ -128,14 +172,188 @@ class TestPipelineArgumentsHandling:
            )
        elif feature_id == "enforce_eager":
            od_config = _make_od_config(enforce_eager=True)
-        elif feature_id == "quantization":
-            od_config = _make_od_config(quantization_config=SimpleNamespace(quant_method="fp8"))
+        elif feature_id == "unsupported_quantization":
+            od_config = _make_od_config(quantization_config=SimpleNamespace(quant_method="gguf"))
        else:
            raise ValueError(f"Unknown feature ID: {feature_id}")

        with pytest.raises(NotImplementedError):
            DiffusersAdapterPipeline(od_config=od_config)

+    def test_adapter_load_weights_injects_supported_quantization_config(self, mocker):
+        (
+            FakePipelineQuantizationConfig,
+            FakeTorchAoConfig,
+            FakeInt8DynamicActivationInt8WeightConfig,
+        ) = _patch_fake_diffusers_quantization_backends(mocker)
+
+        class MockPipeline:
+            def __call__(self, prompt=None):
+                return None
+
+            def to(self, device):
+                return self
+
+        mock_from_pretrained = mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained",
+            return_value=MockPipeline(),
+        )
+        mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.load_config",
+            return_value={
+                "_class_name": "TransformerBackedPipeline",
+                "scheduler": ["diffusers", "SchedulerMixin"],
+                "transformer": ["diffusers", "Transformer2DModel"],
+                "transformer_2": ["diffusers", "Transformer2DModel"],
+                "vae": ["diffusers", "AutoencoderKL"],
+            },
+        )
+
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                quantization_config=SimpleNamespace(quant_method="int8"),
+                diffusers_pipeline_cls=_TransformerBackedPipeline,
+            )
+        )
+        adapter.load_weights()
+
+        pipeline_quant_config = mock_from_pretrained.call_args.kwargs["quantization_config"]
+        torchao_config = pipeline_quant_config.quant_mapping["transformer"]
+        assert isinstance(pipeline_quant_config, FakePipelineQuantizationConfig)
+        assert isinstance(torchao_config, FakeTorchAoConfig)
+        assert isinstance(torchao_config.quant_type, FakeInt8DynamicActivationInt8WeightConfig)
+        assert isinstance(
+            pipeline_quant_config.quant_mapping["transformer_2"].quant_type,
+            FakeInt8DynamicActivationInt8WeightConfig,
+        )
+
+    @pytest.mark.parametrize(
+        ("pipeline_config", "diffusers_load_kwargs", "diffusers_pipeline_cls"),
+        [
+            (
+                {
+                    "_class_name": "NoTransformerPipeline",
+                    "scheduler": ["diffusers", "SchedulerMixin"],
+                    "unet": ["diffusers", "UNet2DConditionModel"],
+                    "vae": ["diffusers", "AutoencoderKL"],
+                },
+                {},
+                _TransformerBackedPipeline,
+            ),
+            (
+                {
+                    "_class_name": "TransformerPipeline",
+                    "scheduler": ["diffusers", "SchedulerMixin"],
+                    "transformer": ["diffusers", "Transformer2DModel"],
+                    "vae": ["diffusers", "AutoencoderKL"],
+                },
+                {"transformer": None},
+                None,
+            ),
+        ],
+    )
+    def test_adapter_load_weights_rejects_converted_quantization_without_transformer_component(
+        self,
+        pipeline_config,
+        diffusers_load_kwargs,
+        diffusers_pipeline_cls,
+        mocker,
+    ):
+        _patch_fake_diffusers_quantization_backends(mocker)
+        mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.load_config",
+            return_value=pipeline_config,
+        )
+        mock_from_pretrained = mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained"
+        )
+
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                quantization_config=SimpleNamespace(quant_method="int8"),
+                diffusers_load_kwargs=diffusers_load_kwargs,
+                diffusers_pipeline_cls=diffusers_pipeline_cls,
+            )
+        )
+
+        with pytest.raises(NotImplementedError, match="transformer"):
+            adapter.load_weights()
+        mock_from_pretrained.assert_not_called()
+
+    def test_adapter_rejects_dduf_file_with_converted_quantization(self):
+        with pytest.raises(NotImplementedError, match="dduf_file"):
+            DiffusersAdapterPipeline(
+                od_config=_make_od_config(
+                    quantization_config=SimpleNamespace(quant_method="int8"),
+                    diffusers_load_kwargs={"dduf_file": "model.dduf"},
+                )
+            )
+
+    def test_adapter_load_weights_preserves_diffusers_native_quantization_config(self, mocker):
+        class MockPipeline:
+            def __call__(self, prompt=None):
+                return None
+
+            def to(self, device):
+                return self
+
+        diffusers_native_quantization_config = object()
+        mock_from_pretrained = mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained",
+            return_value=MockPipeline(),
+        )
+
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                quantization_config=SimpleNamespace(quant_method="gguf"),
+                diffusers_load_kwargs={
+                    "dduf_file": "model.dduf",
+                    "quantization_config": diffusers_native_quantization_config,
+                },
+            )
+        )
+        adapter.load_weights()
+
+        kwargs = mock_from_pretrained.call_args.kwargs
+        assert kwargs["quantization_config"] is diffusers_native_quantization_config
+        assert kwargs["dduf_file"] == "model.dduf"
+
+    def test_adapter_load_weights_builds_diffusers_native_quantization_config_from_dict(self, mocker):
+        FakePipelineQuantizationConfig, FakeTorchAoConfig, _ = _patch_fake_diffusers_quantization_backends(mocker)
+
+        class MockPipeline:
+            def __call__(self, prompt=None):
+                return None
+
+            def to(self, device):
+                return self
+
+        mock_from_pretrained = mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained",
+            return_value=MockPipeline(),
+        )
+
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                diffusers_load_kwargs={
+                    "quantization_config": {
+                        "quant_mapping": {
+                            "transformer": {
+                                "quant_method": "torchao",
+                                "quant_type": "fake-int8-config",
+                            },
+                        },
+                    },
+                },
+            )
+        )
+        adapter.load_weights()
+
+        quant_config = mock_from_pretrained.call_args.kwargs["quantization_config"]
+        assert isinstance(quant_config, FakePipelineQuantizationConfig)
+        assert isinstance(quant_config.quant_mapping["transformer"], FakeTorchAoConfig)
+        assert quant_config.quant_mapping["transformer"].quant_type == "fake-int8-config"
+
    def test_adapter_guard_unknown_output_type(self, mocker):
        """Test that the adapter wraps an unknown output type as-is.
        This is useful when `return_dict=True` and the diffusers pipeline returns an OrderedDict subclass."""
@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Optional real-backend checks for Diffusers quantization config conversion."""
+
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _load_quantization_utils():
+    module_path = (
+        Path(__file__).resolve().parents[3] / "vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py"
+    )
+    spec = importlib.util.spec_from_file_location("diffusers_quantization_utils_real_test", module_path)
+    assert spec is not None
+    assert spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_convert_diffusers_quantization_config_from_real_torchao_dict():
+    diffusers = pytest.importorskip("diffusers")
+    torchao_quantization = pytest.importorskip("torchao.quantization")
+    quant_type_cls = getattr(torchao_quantization, "Int8DynamicActivationInt8WeightConfig", None)
+    if quant_type_cls is None:
+        pytest.skip("torchao does not expose Int8DynamicActivationInt8WeightConfig")
+
+    quantization_utils = _load_quantization_utils()
+
+    real_torchao_config = diffusers.TorchAoConfig(quant_type=quant_type_cls())
+    serialized_torchao_config = real_torchao_config.to_dict()
+    assert serialized_torchao_config["quant_type"].keys() == {"default"}
+
+    load_kwargs = {
+        "quantization_config": {
+            "quant_mapping": {
+                "transformer": serialized_torchao_config,
+            },
+        },
+    }
+
+    quantization_utils.convert_diffusers_quantization_config(load_kwargs)
+
+    quant_config = load_kwargs["quantization_config"]
+    component_config = quant_config.quant_mapping["transformer"]
+    assert isinstance(quant_config, diffusers.PipelineQuantizationConfig)
+    assert isinstance(component_config, diffusers.TorchAoConfig)
+    assert isinstance(
+        component_config.quant_type,
+        quant_type_cls,
+    )
@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Focused tests for Diffusers backend quantization conversion helpers."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.models.diffusers_adapter import quantization_utils
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _FakePipelineQuantizationConfig:
+    def __init__(self, *, quant_mapping=None, quant_backend=None, quant_kwargs=None, components_to_quantize=None):
+        self.quant_mapping = quant_mapping
+        self.quant_backend = quant_backend
+        self.quant_kwargs = quant_kwargs
+        self.components_to_quantize = components_to_quantize
+
+
+class _FakeTorchAoConfig:
+    def __init__(self, quant_type):
+        self.quant_type = quant_type
+
+    @classmethod
+    def from_dict(cls, config):
+        return cls(quant_type=config["quant_type"])
+
+
+class _FakeFloat8DynamicActivationFloat8WeightConfig:
+    pass
+
+
+class _FakeInt8DynamicActivationInt8WeightConfig:
+    pass
+
+
+@pytest.fixture(autouse=True)
+def patch_quantization_backends(monkeypatch):
+    monkeypatch.setattr(
+        quantization_utils,
+        "PipelineQuantizationConfig",
+        _FakePipelineQuantizationConfig,
+    )
+    monkeypatch.setattr(quantization_utils, "TorchAoConfig", _FakeTorchAoConfig)
+
+    def fake_get_torchao_quant_type_cls(class_name: str):
+        return {
+            "Float8DynamicActivationFloat8WeightConfig": _FakeFloat8DynamicActivationFloat8WeightConfig,
+            "Int8DynamicActivationInt8WeightConfig": _FakeInt8DynamicActivationInt8WeightConfig,
+        }[class_name]
+
+    monkeypatch.setattr(quantization_utils, "_get_torchao_quant_type_cls", fake_get_torchao_quant_type_cls)
+    monkeypatch.setattr(quantization_utils, "_get_diffusers_quantization_config_cls", lambda _: _FakeTorchAoConfig)
+
+
+def _quant_config(method: str, **kwargs):
+    return SimpleNamespace(get_name=lambda: method, **kwargs)
+
+
+@pytest.mark.parametrize(
+    ("method", "expected_type"),
+    [
+        ("fp8", _FakeFloat8DynamicActivationFloat8WeightConfig),
+        ("int8", _FakeInt8DynamicActivationInt8WeightConfig),
+    ],
+)
+def test_apply_injects_converted_quantization_config(method, expected_type):
+    od_config = SimpleNamespace(quantization_config=_quant_config(method))
+    load_kwargs = {}
+
+    injected = quantization_utils.apply_diffusers_quantization_config(
+        od_config,
+        load_kwargs,
+        {"transformer": ["diffusers", "Transformer2DModel"]},
+    )
+
+    assert injected is True
+    torchao_config = load_kwargs["quantization_config"].quant_mapping["transformer"]
+    assert isinstance(torchao_config.quant_type, expected_type)
+
+
+def test_apply_injects_converted_quantization_config_for_transformer_2():
+    od_config = SimpleNamespace(quantization_config=_quant_config("int8"))
+    load_kwargs = {}
+
+    quantization_utils.apply_diffusers_quantization_config(
+        od_config,
+        load_kwargs,
+        {
+            "transformer": ["diffusers", "Transformer2DModel"],
+            "transformer_2": ["diffusers", "Transformer2DModel"],
+        },
+    )
+
+    assert sorted(load_kwargs["quantization_config"].quant_mapping) == ["transformer", "transformer_2"]
+
+
+@pytest.mark.parametrize("method", ["gguf", "modelopt", "mxfp4", "mxfp8", "inc"])
+def test_unsupported_methods_fail_explicitly(method):
+    with pytest.raises(NotImplementedError, match=method):
+        quantization_utils.ensure_supported_diffusers_quantization(_quant_config(method))
+
+
+@pytest.mark.parametrize(
+    "quant_config",
+    [
+        _quant_config("fp8", activation_scheme="static"),
+        _quant_config("int8", is_checkpoint_int8_serialized=True),
+        _quant_config("fp8", weight_block_size=[128, 128]),
+        _quant_config("int8", ignored_layers=["transformer.proj_out"]),
+    ],
+)
+def test_ambiguous_mappings_fail_explicitly(quant_config):
+    with pytest.raises(NotImplementedError):
+        quantization_utils.ensure_supported_diffusers_quantization(quant_config)
+
+
+def test_apply_preserves_diffusers_load_kwargs_quantization_config(mocker):
+    od_config = SimpleNamespace(quantization_config=_quant_config("fp8"))
+    existing = object()
+    load_kwargs = {"quantization_config": existing}
+    mock_warning = mocker.patch.object(quantization_utils.logger, "warning")
+
+    injected = quantization_utils.apply_diffusers_quantization_config(od_config, load_kwargs, {})
+
+    assert injected is False
+    assert load_kwargs["quantization_config"] is existing
+    assert "Using the Diffusers-native quantization_config" in mock_warning.call_args.args[0]
+
+
+def test_convert_diffusers_quantization_config_from_dict():
+    load_kwargs = {
+        "quantization_config": {
+            "quant_mapping": {
+                "transformer": {
+                    "quant_method": "torchao",
+                    "quant_type": "fake-int8-config",
+                }
+            }
+        }
+    }
+
+    quantization_utils.convert_diffusers_quantization_config(load_kwargs)
+
+    quant_config = load_kwargs["quantization_config"]
+    assert isinstance(quant_config, _FakePipelineQuantizationConfig)
+    assert isinstance(quant_config.quant_mapping["transformer"], _FakeTorchAoConfig)
+    assert quant_config.quant_mapping["transformer"].quant_type == "fake-int8-config"