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 8f935dd7b..17edd31c1 100644 --- a/tests/diffusion/quantization/artifacts/fp8_z_image_quality/README.md +++ b/tests/diffusion/quantization/artifacts/fp8_z_image_quality/README.md @@ -14,27 +14,26 @@ Quantized config: ```python { - "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, + "method": "fp8", + "ignored_layers": [ + "img_mlp", + "layers.15..29.{attention.to_qkv,attention.to_out.0,feed_forward.w13,feed_forward.w2}", + "model.layers.28..35.{self_attn.q_proj,self_attn.k_proj,self_attn.v_proj,self_attn.o_proj," + "mlp.gate_proj,mlp.up_proj,mlp.down_proj}", + ], } ``` -Result: passed with `max_lpips=0.10`. +Result: passed with `max_lpips=0.15`. | Metric | Value | |--------|-------| -| LPIPS | 0.096935 | -| PSNR | 21.893253 dB | -| MAE | 0.037334 | +| LPIPS | 0.128773 | +| PSNR | 20.125577 dB | +| MAE | 0.047712 | -Timing from the scan run for this config: BF16 `2.297244s`, FP8 `2.062006s` -(`0.8976x` of BF16 generation time). +This uses regular online FP8 quantization for the text encoder, with only the final +8 text-encoder blocks routed through `ignored_layers` for BF16 fallback. ## 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 844cb0726..99c486c7e 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 4266b5e9b..758957115 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 64d645d2f..a1f79030e 100644 --- a/tests/diffusion/quantization/test_fp8_config.py +++ b/tests/diffusion/quantization/test_fp8_config.py @@ -213,35 +213,6 @@ 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 cd9665307..5f0dc4d62 100644 --- a/tests/diffusion/quantization/test_quantization_quality.py +++ b/tests/diffusion/quantization/test_quantization_quality.py @@ -119,23 +119,33 @@ QUALITY_CONFIGS = [ id="fp8_z_image", model="Tongyi-MAI/Z-Image-Turbo", 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", - ) - ], + "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, + *[ + f"model.layers.{layer_id}.{suffix}" + for layer_id in range(28, 36) + for suffix in ( + "self_attn.q_proj", + "self_attn.k_proj", + "self_attn.v_proj", + "self_attn.o_proj", + "mlp.gate_proj", + "mlp.up_proj", + "mlp.down_proj", + ) + ], + ], }, task="t2i", prompt=( @@ -159,7 +169,7 @@ QUALITY_CONFIGS = [ "starlight waterfalls. Ethereal, melancholic, and transcendent mood—like a moment of quiet revelation " "at the edge of existence." ), - max_lpips=0.10, + max_lpips=0.15, num_inference_steps=20, ), QualityTestConfig( diff --git a/vllm_omni/diffusion/models/utils.py b/vllm_omni/diffusion/models/utils.py index f4b7ae83d..122646219 100644 --- a/vllm_omni/diffusion/models/utils.py +++ b/vllm_omni/diffusion/models/utils.py @@ -19,7 +19,6 @@ 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 @@ -88,17 +87,7 @@ def recursive_replace_linear(model: nn.Module, od_config: OmniDiffusionConfig): # Replace modules as needed if isinstance(child_module, nn.Linear): style = "replicate" - 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, - ) + new_module = replace_linear_class(child_module, style, 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 bf690d3a2..388b6316c 100644 --- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py +++ b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py @@ -45,7 +45,6 @@ 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 @@ -219,10 +218,7 @@ 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) - 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.transformer = ZImageTransformer2DModel(quant_config=od_config.quantization_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()