diff --git a/tests/diffusion/quantization/artifacts/fp8_z_image_quality/README.md b/tests/diffusion/quantization/artifacts/fp8_z_image_quality/README.md index 1e9c3a069..8f935dd7b 100644 --- a/tests/diffusion/quantization/artifacts/fp8_z_image_quality/README.md +++ b/tests/diffusion/quantization/artifacts/fp8_z_image_quality/README.md @@ -13,16 +13,28 @@ VLLM_OMNI_QUALITY_OUTPUT_DIR=$PWD/tests/diffusion/quantization/artifacts/fp8_z_i Quantized config: ```python -{"method": "fp8", "ignored_layers": ["img_mlp"]} +{ + "transformer": { + "method": "fp8", + "ignored_layers": [ + "img_mlp", + "layers.15..29.{attention.to_qkv,attention.to_out.0,feed_forward.w13,feed_forward.w2}", + ], + }, + "default": None, +} ``` -Result: failed with the current `max_lpips=0.15` threshold. +Result: passed with `max_lpips=0.10`. | Metric | Value | |--------|-------| -| LPIPS | 0.250578 | -| PSNR | 17.042153 dB | -| MAE | 0.076899 | +| LPIPS | 0.096935 | +| PSNR | 21.893253 dB | +| MAE | 0.037334 | + +Timing from the scan run for this config: BF16 `2.297244s`, FP8 `2.062006s` +(`0.8976x` of BF16 generation time). ## Comparison diff --git a/tests/diffusion/quantization/artifacts/fp8_z_image_quality/fp8_z_image_comparison.png b/tests/diffusion/quantization/artifacts/fp8_z_image_quality/fp8_z_image_comparison.png index 9d12c278f..844cb0726 100644 Binary files a/tests/diffusion/quantization/artifacts/fp8_z_image_quality/fp8_z_image_comparison.png and b/tests/diffusion/quantization/artifacts/fp8_z_image_quality/fp8_z_image_comparison.png differ diff --git a/tests/diffusion/quantization/artifacts/fp8_z_image_quality/fp8_z_image_quantized.png b/tests/diffusion/quantization/artifacts/fp8_z_image_quality/fp8_z_image_quantized.png index 4029bc9b0..4266b5e9b 100644 Binary files a/tests/diffusion/quantization/artifacts/fp8_z_image_quality/fp8_z_image_quantized.png and b/tests/diffusion/quantization/artifacts/fp8_z_image_quality/fp8_z_image_quantized.png differ diff --git a/tests/diffusion/quantization/test_fp8_config.py b/tests/diffusion/quantization/test_fp8_config.py index a1f79030e..64d645d2f 100644 --- a/tests/diffusion/quantization/test_fp8_config.py +++ b/tests/diffusion/quantization/test_fp8_config.py @@ -213,6 +213,35 @@ def test_integration_per_component(): assert config.quantization_config.component_configs["vae"] is None +def test_recursive_replace_linear_resolves_component_quantization(monkeypatch): + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.models import utils + + model = nn.Module() + model.text_encoder = nn.Sequential(nn.Linear(2, 2)) + model.transformer = nn.Sequential(nn.Linear(2, 2)) + config = OmniDiffusionConfig( + model="test", + quantization_config={ + "transformer": {"method": "fp8"}, + "default": None, + }, + ) + + calls = {} + + def fake_replace(linear, style, quant_config, *, prefix): + calls[prefix] = quant_config + return nn.Identity() + + monkeypatch.setattr(utils, "replace_linear_class", fake_replace) + + utils.recursive_replace_linear(model, config) + + assert calls["text_encoder.0"] is None + assert calls["transformer.0"].get_name() == "fp8" + + def test_transformer_config_auto_detects_modelopt_fp8(): from vllm.model_executor.layers.quantization.modelopt import ModelOptFp8Config diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py index 33b0dd9c6..cd9665307 100644 --- a/tests/diffusion/quantization/test_quantization_quality.py +++ b/tests/diffusion/quantization/test_quantization_quality.py @@ -118,7 +118,25 @@ QUALITY_CONFIGS = [ QualityTestConfig( id="fp8_z_image", model="Tongyi-MAI/Z-Image-Turbo", - quantization={"method": "fp8", "ignored_layers": ["img_mlp"]}, + quantization={ + "transformer": { + "method": "fp8", + "ignored_layers": [ + "img_mlp", + *[ + f"layers.{layer_id}.{suffix}" + for layer_id in range(15, 30) + for suffix in ( + "attention.to_qkv", + "attention.to_out.0", + "feed_forward.w13", + "feed_forward.w2", + ) + ], + ], + }, + "default": None, + }, task="t2i", prompt=( "A breathtaking twilight scene atop a floating archipelago of crystalline islands suspended in an " @@ -141,7 +159,7 @@ QUALITY_CONFIGS = [ "starlight waterfalls. Ethereal, melancholic, and transcendent mood—like a moment of quiet revelation " "at the edge of existence." ), - max_lpips=0.15, + max_lpips=0.10, num_inference_steps=20, ), QualityTestConfig( diff --git a/vllm_omni/diffusion/models/utils.py b/vllm_omni/diffusion/models/utils.py index 122646219..f4b7ae83d 100644 --- a/vllm_omni/diffusion/models/utils.py +++ b/vllm_omni/diffusion/models/utils.py @@ -19,6 +19,7 @@ from vllm.model_executor.models.utils import maybe_prefix from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.quantization import build_quant_config +from vllm_omni.quantization.component_config import ComponentQuantizationConfig if TYPE_CHECKING: from transformers import PretrainedConfig, PreTrainedModel @@ -87,7 +88,17 @@ def recursive_replace_linear(model: nn.Module, od_config: OmniDiffusionConfig): # Replace modules as needed if isinstance(child_module, nn.Linear): style = "replicate" - new_module = replace_linear_class(child_module, style, quant_config, prefix=qual_name) + child_quant_config = ( + quant_config.resolve(qual_name) + if isinstance(quant_config, ComponentQuantizationConfig) + else quant_config + ) + new_module = replace_linear_class( + child_module, + style, + child_quant_config, + prefix=qual_name, + ) else: _recursive_replace(child_module, prefix=qual_name) if new_module is not child_module: diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py index 388b6316c..bf690d3a2 100644 --- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py +++ b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py @@ -45,6 +45,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) +from vllm_omni.quantization.component_config import ComponentQuantizationConfig logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -218,7 +219,10 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): vae_config = DistributedAutoencoderKL.load_config(model, subfolder="vae", local_files_only=local_files_only) self.vae = DistributedAutoencoderKL.from_config(vae_config).to(self._execution_device) - self.transformer = ZImageTransformer2DModel(quant_config=od_config.quantization_config) + quant_config = od_config.quantization_config + if isinstance(quant_config, ComponentQuantizationConfig): + quant_config = quant_config.resolve("transformer") + self.transformer = ZImageTransformer2DModel(quant_config=quant_config) self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) # Note: Context parallelism is applied centrally in registry.initialize_model()