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

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

Browse Source 
Signed-off-by: CAICAIIs <3360776475@qq.com>
This commit is contained in:
CAICAII authored and GitHub committed 2026-07-03 03:06:44 +00:00
1 parent 06329c6450
commit b443504d61
7 files changed
+816 -9
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py
+58 -4
View File
@@ -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
vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py
+266
View File
@@ -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 "<none>"
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,
)