diff --git a/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md b/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md index 1adc61193..c6be3c014 100644 --- a/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md +++ b/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md @@ -9,7 +9,8 @@ vLLM-Omni supports running diffusion models with the diffusers backend, directly The diffusers backend is a black-box adapter. Its primary focus is to serve diffusion models online. Currently, the following features are NOT yet supported. -It is not guaranteed whether they will be supported in the future. +Community contributions are welcome to add these features when they can be +implemented with clear Diffusers or vLLM-Omni behavior. - CFG parallel execution - Sequence parallel execution @@ -66,6 +67,35 @@ This is suitable for sampling parameters not available through the vLLM-Omni int When a parameter is available in the vLLM-Omni interface, it will be adapted here. But if that parameter is simultaneously set in both the vLLM-Omni interface and `diffusers_call_kwargs`, the **former** will take precedence (because it is set at request time). +### Quantization + +Use `diffusers_load_kwargs` for Diffusers-native quantization options. If +`diffusers_load_kwargs["quantization_config"]` is provided as a dictionary, the +diffusers backend builds a Diffusers quantization config and lets Diffusers +validate it before calling `DiffusionPipeline.from_pretrained()`. + +The diffusers backend also provides a small compatibility shortcut for +vLLM-Omni quantization configs when `diffusers_load_kwargs` does not already +contain `quantization_config`. Currently this maps online/dynamic `fp8` and +online/dynamic `int8` to Diffusers/TorchAO dynamic quantization for transformer +components such as `transformer` or `transformer_2`. This path requires +`torchao` to be installed. + +For example, the CLI can request dynamic FP8 through the vLLM-Omni interface: + +```bash +vllm serve "Qwen/Qwen-Image" \ + --omni \ + --diffusion-load-format diffusers \ + --quantization-config '{"method": "fp8"}' +``` + +Other vLLM-Omni quantization methods, such as `gguf`, `modelopt`, `mxfp4`, +`mxfp8`, serialized checkpoints, static FP8 configs, and layer-name skip lists +such as `ignored_layers`, are intentionally not translated by this compatibility +shortcut. Use Diffusers-native configuration through `diffusers_load_kwargs` or +a native vLLM-Omni pipeline for those cases. + ### Attention Backends The diffusers backend converts diff --git a/examples/online_serving/diffusers_pipeline_adapter/README.md b/examples/online_serving/diffusers_pipeline_adapter/README.md index d6f41f632..37bf48aa0 100644 --- a/examples/online_serving/diffusers_pipeline_adapter/README.md +++ b/examples/online_serving/diffusers_pipeline_adapter/README.md @@ -6,7 +6,8 @@ vLLM-Omni supports running diffusion models with the diffusers backend, directly The diffusers backend is a black-box adapter. Its primary focus is to serve diffusion models online. Currently, the following features are NOT yet supported. -It is not guaranteed whether they will be supported in the future. +Community contributions are welcome to add these features when they can be +implemented with clear Diffusers or vLLM-Omni behavior. - CFG parallel execution - Sequence parallel execution @@ -68,6 +69,35 @@ But if that parameter is simultaneously set in both the vLLM-Omni interface and Consider referring to [`OmniDiffusionSamplingParams`](../../../vllm_omni/inputs/data.py) for the default sampling parameters in the vLLM-Omni interface, and the corresponding [diffusers pipeline](https://huggingface.co/docs/diffusers/main/en/api/pipelines/overview)'s `__call__` function documentation. +### Quantization + +Use `diffusers_load_kwargs` for Diffusers-native quantization options. If +`diffusers_load_kwargs["quantization_config"]` is provided as a dictionary, the +diffusers backend builds a Diffusers quantization config and lets Diffusers +validate it before calling `DiffusionPipeline.from_pretrained()`. + +The diffusers backend also provides a small compatibility shortcut for +vLLM-Omni quantization configs when `diffusers_load_kwargs` does not already +contain `quantization_config`. Currently this maps online/dynamic `fp8` and +online/dynamic `int8` to Diffusers/TorchAO dynamic quantization for transformer +components such as `transformer` or `transformer_2`. This path requires +`torchao` to be installed. + +For example, the CLI can request dynamic FP8 through the vLLM-Omni interface: + +```bash +vllm serve "Qwen/Qwen-Image" \ + --omni \ + --diffusion-load-format diffusers \ + --quantization-config '{"method": "fp8"}' +``` + +Other vLLM-Omni quantization methods, such as `gguf`, `modelopt`, `mxfp4`, +`mxfp8`, serialized checkpoints, static FP8 configs, and layer-name skip lists +such as `ignored_layers`, are intentionally not translated by this compatibility +shortcut. Use Diffusers-native configuration through `diffusers_load_kwargs` or +a native vLLM-Omni pipeline for those cases. + ### Attention Backends The diffusers backend converts diff --git a/tests/diffusion/diffusion_backend/test_diffusers_backend.py b/tests/diffusion/diffusion_backend/test_diffusers_backend.py index b9b4a72b7..babc23b4a 100644 --- a/tests/diffusion/diffusion_backend/test_diffusers_backend.py +++ b/tests/diffusion/diffusion_backend/test_diffusers_backend.py @@ -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.""" diff --git a/tests/diffusion/diffusion_backend/test_diffusers_quantization_real.py b/tests/diffusion/diffusion_backend/test_diffusers_quantization_real.py new file mode 100644 index 000000000..5faa42eaf --- /dev/null +++ b/tests/diffusion/diffusion_backend/test_diffusers_quantization_real.py @@ -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, + ) diff --git a/tests/diffusion/diffusion_backend/test_diffusers_quantization_utils.py b/tests/diffusion/diffusion_backend/test_diffusers_quantization_utils.py new file mode 100644 index 000000000..6def7ebad --- /dev/null +++ b/tests/diffusion/diffusion_backend/test_diffusers_quantization_utils.py @@ -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" diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py index 4d48053a6..94fcb6699 100644 --- a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py +++ b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py @@ -24,6 +24,11 @@ from torch import nn from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import BasePipelineUtils, get_pipeline_utils +from vllm_omni.diffusion.models.diffusers_adapter.quantization_utils import ( + apply_diffusers_quantization_config, + convert_diffusers_quantization_config, + ensure_supported_diffusers_quantization, +) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniPromptType, OmniTextPrompt @@ -31,6 +36,20 @@ from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) +_DIFFUSERS_CONFIG_LOAD_KWARGS = { + "cache_dir", + "dduf_entries", + "force_download", + "local_dir", + "local_dir_use_symlinks", + "local_files_only", + "proxies", + "revision", + "subfolder", + "token", + "user_agent", +} + class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): """Black-box adapter that delegates full pipeline execution to a diffusers pipeline. @@ -80,12 +99,19 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): "torch_dtype": dtype, **self.od_config.diffusers_load_kwargs, } - logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}") + convert_diffusers_quantization_config(load_kwargs) pipeline_class = self.od_config.diffusers_pipeline_cls pipeline_class_name = pipeline_class.__name__ if pipeline_class is not None else None self._pipeline_utils = get_pipeline_utils(pipeline_class_name) self._pipeline_utils.update_load_kwargs(self.od_config, load_kwargs) + component_names = ( + self._load_diffusers_component_names(model_id, load_kwargs) + if self.od_config.quantization_config is not None and "quantization_config" not in load_kwargs + else {} + ) + apply_diffusers_quantization_config(self.od_config, load_kwargs, component_names) + logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}") self._pipeline = DiffusionPipeline.from_pretrained(model_id, **load_kwargs) self._pipeline_utils.apply_post_load_updates(self._pipeline, self.od_config) @@ -190,10 +216,38 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): "Eager execution is not supported with the diffusers backend. " "Use a native pipeline for continuous batching mode." ) - if self.od_config.quantization_config is not None: - raise NotImplementedError( - "Quantization is not supported with the diffusers backend. Use a native pipeline for quantization." + if ( + self.od_config.quantization_config is not None + and "quantization_config" not in self.od_config.diffusers_load_kwargs + ): + ensure_supported_diffusers_quantization(self.od_config.quantization_config) + if self.od_config.diffusers_load_kwargs.get("dduf_file"): + raise NotImplementedError( + "Diffusers backend quantization conversion does not support " + "diffusers_load_kwargs.dduf_file yet. The preflight component " + "discovery would need to mirror Diffusers' DDUF config loading. " + "Use diffusers_load_kwargs.quantization_config for a native " + "Diffusers quantization config, or omit dduf_file for vLLM-Omni " + "quantization conversion." + ) + + def _load_diffusers_component_names( + self, + model_id: str, + load_kwargs: dict[str, Any], + ) -> dict[str, Any]: + config_load_kwargs = {k: load_kwargs[k] for k in _DIFFUSERS_CONFIG_LOAD_KWARGS if k in load_kwargs} + pipeline_config = DiffusionPipeline.load_config(model_id, **config_load_kwargs) + return { + name: value + for name, value in pipeline_config.items() + if ( + isinstance(value, list) + and len(value) > 0 + and value[0] is not None + and (name not in load_kwargs or load_kwargs[name] is not None) ) + } # ------------------------------------------------------------------ # Wrap settings, inputs, and outputs diff --git a/vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py b/vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py new file mode 100644 index 000000000..a7228d21e --- /dev/null +++ b/vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py @@ -0,0 +1,266 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Quantization helpers for the Diffusers backend.""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from typing import Any + +from diffusers import PipelineQuantizationConfig, TorchAoConfig + +logger = logging.getLogger(__name__) + +_DIFFUSERS_TRANSFORMER_COMPONENTS = ("transformer", "transformer_2") +_QuantizationConfigSpec = tuple[Callable[[Any], None], str] + + +def ensure_supported_diffusers_quantization(quant_config: Any) -> None: + """Validate that a vLLM-Omni quantization config has a Diffusers mapping.""" + + method = _get_quant_method_name(quant_config) + validator, _ = _get_quantization_config_spec(method) + validator(quant_config) + + +def convert_diffusers_quantization_config(load_kwargs: dict[str, Any]) -> None: + """Build Diffusers-native quantization config objects from dict kwargs.""" + + quant_config = load_kwargs.get("quantization_config") + if isinstance(quant_config, dict): + load_kwargs["quantization_config"] = _build_diffusers_native_quantization_config(quant_config) + + +def apply_diffusers_quantization_config( + od_config: Any, + load_kwargs: dict[str, Any], + component_names: dict[str, Any], +) -> bool: + """Inject a courtesy-converted quantization_config into load kwargs. + + ``diffusers_load_kwargs`` is the canonical Diffusers backend configuration + path, so an explicit ``quantization_config`` already present in ``load_kwargs`` + is never replaced. Returns whether a config was injected by this helper. + """ + + quant_config = getattr(od_config, "quantization_config", None) + if quant_config is None: + return False + + if "quantization_config" in load_kwargs: + logger.warning( + "Both vLLM-Omni quantization_config and diffusers_load_kwargs.quantization_config " + "were provided for the diffusers backend. Using the Diffusers-native " + "quantization_config from diffusers_load_kwargs." + ) + return False + + transformer_components = _get_diffusers_transformer_components(component_names) + load_kwargs["quantization_config"] = _build_diffusers_quantization_config( + quant_config, + transformer_components, + ) + return True + + +def _normalize_method_name(method: Any) -> str: + method = getattr(method, "value", method) + return str(method).lower().replace("-", "_") + + +def _get_quant_method_name(quant_config: Any) -> str: + get_name = getattr(quant_config, "get_name", None) + if get_name is not None: + return _normalize_method_name(get_name() if callable(get_name) else get_name) + + method = getattr(quant_config, "quant_method", None) + if method is None: + method = getattr(quant_config, "method", None) + if method is None: + raise NotImplementedError( + "Diffusers backend quantization conversion requires a quantization " + "config with get_name(), quant_method, or method." + ) + return _normalize_method_name(method) + + +def _ensure_no_ignored_layers(quant_config: Any) -> None: + ignored_layers = getattr(quant_config, "ignored_layers", None) + if not ignored_layers: + ignored_layers = getattr(quant_config, "modules_to_not_convert", None) + if ignored_layers: + raise NotImplementedError( + "Diffusers backend quantization conversion does not map vLLM " + "ignored_layers/modules_to_not_convert names to Diffusers module " + "names. Use diffusers_load_kwargs for a native Diffusers config." + ) + + +def _get_torchao_quant_type_cls(class_name: str) -> type[Any]: + try: + import torchao.quantization as torchao_quantization + except ImportError as exc: + raise ImportError( + "Diffusers backend quantization conversion for fp8/int8 requires " + "torchao. Install torchao or pass a Diffusers-native " + "quantization_config through diffusers_load_kwargs." + ) from exc + + try: + return getattr(torchao_quantization, class_name) + except AttributeError as exc: + raise ImportError(f"torchao.quantization.{class_name} is required for this quantization mapping.") from exc + + +def _build_diffusers_native_quantization_config(config: dict[str, Any]) -> Any: + config = config.copy() + quant_mapping = config.get("quant_mapping") + if isinstance(quant_mapping, dict): + config["quant_mapping"] = { + component_name: _build_diffusers_component_quant_config(component_config) + for component_name, component_config in quant_mapping.items() + } + return PipelineQuantizationConfig(**config) + + +def _build_diffusers_component_quant_config(config: Any) -> Any: + if not isinstance(config, dict): + return config + + quant_method = config.get("quant_method") + if quant_method is None: + raise ValueError("Diffusers quant_mapping entries must provide quant_method.") + + quant_method = _normalize_method_name(quant_method) + config_cls = _get_diffusers_quantization_config_cls(quant_method) + from_dict = getattr(config_cls, "from_dict", None) + if from_dict is not None: + return from_dict(config) + + init_kwargs = config.copy() + init_kwargs.pop("quant_method", None) + return config_cls(**init_kwargs) + + +def _get_diffusers_quantization_config_cls(quant_method: str) -> type[Any]: + from diffusers.quantizers.auto import AUTO_QUANTIZATION_CONFIG_MAPPING as DIFFUSERS_QUANT_CONFIG_MAPPING + + config_cls = DIFFUSERS_QUANT_CONFIG_MAPPING.get(quant_method) + if config_cls is not None: + return config_cls + + try: + from transformers.quantizers.auto import ( + AUTO_QUANTIZATION_CONFIG_MAPPING as TRANSFORMERS_QUANT_CONFIG_MAPPING, + ) + except ImportError: + TRANSFORMERS_QUANT_CONFIG_MAPPING = {} + + config_cls = TRANSFORMERS_QUANT_CONFIG_MAPPING.get(quant_method) + if config_cls is None: + raise ValueError(f"Diffusers quant_mapping entry uses unknown quant_method={quant_method!r}.") + return config_cls + + +def _build_torchao_pipeline_quant_config( + torchao_quant_type_name: str, + transformer_components: list[str], +) -> Any: + """Build a Diffusers config after the caller has validated the vLLM config.""" + + quant_type_cls = _get_torchao_quant_type_cls(torchao_quant_type_name) + return PipelineQuantizationConfig( + quant_mapping={ + component_name: TorchAoConfig( + quant_type=quant_type_cls(), + ) + for component_name in transformer_components + } + ) + + +def _get_diffusers_transformer_components(component_names: dict[str, Any]) -> list[str]: + transformer_components = [ + component_name for component_name in component_names if component_name in _DIFFUSERS_TRANSFORMER_COMPONENTS + ] + if transformer_components: + return transformer_components + + available_components = ", ".join(sorted(component_names.keys())) or "" + raise NotImplementedError( + "Diffusers backend quantization conversion currently only supports " + "pipelines with transformer or transformer_2 components. " + f"Found pipeline components: {available_components}. Use " + "diffusers_load_kwargs for a native Diffusers quantization config." + ) + + +def _validate_fp8_quant_config(quant_config: Any) -> None: + if getattr(quant_config, "is_checkpoint_fp8_serialized", False): + raise NotImplementedError( + "Diffusers backend fp8 conversion only supports online/dynamic " + "TorchAO quantization; serialized vLLM fp8 checkpoints are not mapped." + ) + + activation_scheme = getattr(quant_config, "activation_scheme", "dynamic") + if activation_scheme != "dynamic": + raise NotImplementedError( + f"Diffusers backend fp8 conversion only supports activation_scheme='dynamic'. Got {activation_scheme!r}." + ) + + weight_block_size = getattr(quant_config, "weight_block_size", None) + if weight_block_size is not None: + raise NotImplementedError( + "Diffusers backend fp8 conversion does not map vLLM weight_block_size " + "to TorchAO. Use diffusers_load_kwargs for a native Diffusers config." + ) + + _ensure_no_ignored_layers(quant_config) + + +def _validate_int8_quant_config(quant_config: Any) -> None: + if getattr(quant_config, "is_checkpoint_int8_serialized", False): + raise NotImplementedError( + "Diffusers backend int8 conversion only supports online/dynamic " + "TorchAO quantization; serialized vLLM int8 checkpoints are not mapped." + ) + + activation_scheme = getattr(quant_config, "activation_scheme", "dynamic") + if activation_scheme != "dynamic": + raise NotImplementedError( + f"Diffusers backend int8 conversion only supports activation_scheme='dynamic'. Got {activation_scheme!r}." + ) + + _ensure_no_ignored_layers(quant_config) + + +_QUANTIZATION_CONFIG_SPECS: dict[str, _QuantizationConfigSpec] = { + "fp8": (_validate_fp8_quant_config, "Float8DynamicActivationFloat8WeightConfig"), + "int8": (_validate_int8_quant_config, "Int8DynamicActivationInt8WeightConfig"), +} + + +def _get_quantization_config_spec(method: str) -> _QuantizationConfigSpec: + spec = _QUANTIZATION_CONFIG_SPECS.get(method) + if spec is None: + raise NotImplementedError( + f"Diffusers backend quantization conversion does not support {method!r}. " + "Use diffusers_load_kwargs for a native Diffusers quantization config, " + "or use a native vLLM-Omni pipeline for this quantization method." + ) + return spec + + +def _build_diffusers_quantization_config( + quant_config: Any, + transformer_components: list[str], +) -> Any: + """Build a Diffusers PipelineQuantizationConfig from a supported config.""" + + method = _get_quant_method_name(quant_config) + _, torchao_quant_type_name = _get_quantization_config_spec(method) + return _build_torchao_pipeline_quant_config( + torchao_quant_type_name, + transformer_components, + )