0 Fork 0

Files

T

SYLAR 6b93459826 [Core] Unified quantization framework (#1764 )

2026-03-24 18:56:55 +08:00

17 KiB

Raw Blame History

Support FP8 Quantization

This section describes how to add FP8 quantization support to a diffusion model. We use the Qwen-Image pipeline as the reference implementation.

Overview
Step-by-Step Implementation
Testing
Troubleshooting
Reference Implementations
Summary

Overview

What is FP8 Quantization?

FP8 quantization reduces model memory footprint by storing weights in 8-bit floating point (float8_e4m3fn) instead of 16-bit (bfloat16/float16). This saves ~50% of weight memory.

A diffusion pipeline typically has three quantizable components:

Component	Layer Types	Quantization Approach
DiT (transformer)	`nn.Linear`	Replace with vLLM quantized linear layers
Text encoder	`nn.Linear`	FP8 weight storage via hooks
VAE	`nn.Conv2d`, `nn.Conv3d`	FP8 weight storage via hooks

Architecture

The quantization system has two mechanisms:

1. vLLM Quantized Linear Layers (for DiT)

vLLM provides drop-in replacements for nn.Linear that perform FP8 W8A8 computation on supported hardware. These are configured via quant_config passed to layer constructors.

Key classes:

Class	Location	Purpose
`build_quant_config`	`vllm_omni/quantization/__init__.py`	Build quantization config from string/dict/per-component spec
`ComponentQuantizationConfig`	`vllm_omni/quantization/component_config.py`	Per-component routing by layer prefix

2. FP8 Weight Storage via Hooks (for text encoder and VAE)

For models loaded via from_pretrained() (outside vLLM's weight pipeline), we use forward hooks to store weights in FP8 and dequantize on-the-fly:

At rest:     weight.data = fp8_tensor (half memory)
Pre-hook:    weight.data = fp8_weight.to(bf16) * scale
Forward:     Normal computation in BF16
Post-hook:   weight.data = fp8_weight (back to FP8)

Key function:

Function	Location	Purpose
`apply_fp8_weight_storage`	`vllm_omni/diffusion/models/utils.py`	Apply FP8 hooks to Linear/Conv layers (introduced in #1412)

Step-by-Step Implementation

Step 1: Add `quant_config` to the Transformer

Pass quant_config to vLLM parallel linear layers in the transformer. This enables FP8 W8A8 for DiT linear layers.

1.1: Accept `quant_config` in Constructor

from vllm.model_executor.layers.quantization.base_config import QuantizationConfig

class YourTransformer2DModel(nn.Module):
    def __init__(
        self,
        *,
        od_config: OmniDiffusionConfig,
        quant_config: QuantizationConfig | None = None,
        # ... other params
    ):
        self.quant_config = quant_config
        # Pass to sub-modules
        self.blocks = nn.ModuleList([
            YourTransformerBlock(quant_config=quant_config, ...)
            for _ in range(num_layers)
        ])

1.2: Pass `quant_config` to vLLM Linear Layers

Replace nn.Linear with vLLM parallel linear layers that accept quant_config:

from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear,
)

class YourAttentionBlock(nn.Module):
    def __init__(self, hidden_size, num_heads, quant_config=None, prefix=""):
        # QKV projection
        self.to_qkv = QKVParallelLinear(
            hidden_size=hidden_size,
            head_size=hidden_size // num_heads,
            total_num_heads=num_heads,
            quant_config=quant_config,
            prefix=f"{prefix}.to_qkv",
        )
        # Output projection
        self.to_out = RowParallelLinear(
            input_size=hidden_size,
            output_size=hidden_size,
            quant_config=quant_config,
            prefix=f"{prefix}.to_out",
        )

class YourFeedForward(nn.Module):
    def __init__(self, hidden_size, intermediate_size, quant_config=None, prefix=""):
        self.gate_up_proj = ColumnParallelLinear(
            input_size=hidden_size,
            output_size=intermediate_size * 2,
            quant_config=quant_config,
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            input_size=intermediate_size,
            output_size=hidden_size,
            quant_config=quant_config,
            prefix=f"{prefix}.down_proj",
        )

!!! important Always pass the prefix parameter. It is used by ignored_layers matching to skip sensitive layers.

1.3: Identify Sensitive Layers

Some layers degrade output quality when quantized. Common sensitive layers:

Layer	Why Sensitive	Typical Models	Quant Methods Affected
`img_mlp`	Processes denoising latents with shifting dynamic range	Qwen-Image	FP8, Int8
`feed_forward`	FFN layers in DiT blocks, large dynamic range	Z-Image	Int8
`proj_out`	Final output projection, small quantization errors amplified	Various	FP8, Int8
`lm_head`	Output vocabulary projection, precision-critical	Omni (thinker)	NVFP4
`mlp.gate`	MoE router gates, routing decisions are precision-critical	MoE models	NVFP4

Quantization quality benchmarks (LPIPS: lower = better, <0.01 imperceptible, >0.1 noticeable):

Model	Method	All Layers	With `ignored_layers`	Recommendation
Qwen-Image-2512	Int8	LPIPS 0.0197	0.0027 (skip `img_mlp`)	Skip `img_mlp`
Z-Image-Turbo	Int8	LPIPS 0.1597	0.0290 (skip `feed_forward`)	Skip `feed_forward`
Z-Image-Turbo	FP8	LPIPS ~0.005	—	All layers OK

To identify sensitive layers for a new model:

Run inference with all layers quantized
Compare output quality with BF16 baseline using LPIPS (see benchmarks/diffusion/quantization_quality.py)
Selectively disable quantization per layer (via ignored_layers) until quality matches
Document the recommended ignored_layers in the model's supported models table

Step 2: Wire Up Quantization in the Pipeline

2.1: Create `quant_config` and Pass to Transformer

In the pipeline's __init__, extract the vLLM quant config and pass it to the transformer:

from vllm_omni.quantization import build_quant_config

class YourPipeline(nn.Module):
    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
        super().__init__()
        # ... load scheduler, text_encoder, vae ...

        # Build quant config from OmniDiffusionConfig
        quant_config = build_quant_config(od_config.quantization_config)

        # Pass to transformer
        self.transformer = YourTransformer2DModel(
            od_config=od_config,
            quant_config=quant_config,
            # ... other kwargs
        )

2.2: Apply FP8 Weight Storage to Text Encoder and VAE

For models loaded via from_pretrained(), apply hook-based FP8 weight storage.

!!! note The apply_fp8_weight_storage function is provided by vllm_omni/diffusion/models/utils.py, introduced in PR #1412.

from vllm_omni.diffusion.models.utils import apply_fp8_weight_storage

class YourPipeline(nn.Module):
    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
        super().__init__()

        self.text_encoder = SomeTextEncoder.from_pretrained(...)
        self.vae = SomeVAE.from_pretrained(...).to(self.device)

        # Apply FP8 weight storage to text encoder and VAE
        if (
            od_config.quantization_config is not None
            and getattr(od_config.quantization_config, "quant_method", None) == "fp8"
        ):
            apply_fp8_weight_storage(self.vae)
            apply_fp8_weight_storage(self.text_encoder)

        # ... rest of init

apply_fp8_weight_storage walks all nn.Linear, nn.Conv2d, and nn.Conv3d modules in the model and applies FP8 hooks to each.

2.3: Fix `load_weights` to Include Pre-loaded Components

Since text encoders and VAEs are loaded via from_pretrained() (not through vLLM's weight pipeline), their parameters won't appear in the loaded_weights set. Mark them explicitly:

def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
    loader = AutoWeightsLoader(self)
    loaded_weights = loader.load_weights(weights)
    # Mark pre-loaded component weights as loaded
    loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
    loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
    return loaded_weights

Step 3: Add to Supported Models Documentation

Update the FP8 supported models table in docs/user_guide/diffusion/quantization/fp8.md:

| Your-Model | `org/Your-Model` | ✅ | ✅ | ✅ | `sensitive_layer` |

Testing

Quick Verification

from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

# Test with FP8 enabled
omni = Omni(model="your-model", quantization="fp8")

outputs = omni.generate(
    "A cat sitting on a windowsill",
    OmniDiffusionSamplingParams(num_inference_steps=50),
)

Quality Gate Test (LPIPS)

We provide a pytest-based quality gate at tests/diffusion/quantization/test_quantization_quality.py. It generates outputs with both BF16 and your quantized method using the same seed, computes LPIPS perceptual distance, and fails if it exceeds a threshold. This is the recommended way to validate that a quantization method does not unacceptably degrade output quality.

Step 1: Add your test config — edit QUALITY_CONFIGS in the test file:

QualityTestConfig(
    id="int8_z_image",              # pytest ID
    model="Tongyi-MAI/Z-Image-Turbo",
    quantization="int8",            # your quantization method
    task="t2i",                     # "t2i" or "t2v"
    prompt="a cup of coffee on a wooden table, morning light",
    max_lpips=0.15,                 # fail threshold (see guide below)
    num_inference_steps=20,
),

Step 2: Run the test:

# Install dependency
pip install lpips

# Run all quality tests
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m ""

# Run only your method
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "int8"

# Run a specific model
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "z_image"

Step 3: Read the output:

Quantization Quality: int8_z_image
  Model:         Tongyi-MAI/Z-Image-Turbo
  Method:        int8
  LPIPS:         0.0290  (threshold: 0.15)
  BF16 memory:   19.15 GiB
  Quant memory:  14.80 GiB  (23% reduction)
  Result:        PASS

LPIPS threshold guide:

LPIPS Range	Meaning	Recommendation
< 0.01	Imperceptible	Ideal — no visible difference
0.01 – 0.05	Minor differences	Good — acceptable for most use cases
0.05 – 0.15	Noticeable on close inspection	Acceptable with `ignored_layers` tuning
0.15 – 0.30	Clearly visible	Investigate — likely needs `ignored_layers`
> 0.30	Significant degradation	Quantization method may not suit this model

!!! tip If your test fails, don't raise the threshold — instead, identify sensitive layers and add them to ignored_layers (see next section). Then re-run the test with the lower LPIPS.

Quality Benchmark (Detailed)

For a full multi-prompt benchmark with Markdown output (suitable for PR descriptions), use the standalone benchmark script:

python benchmarks/diffusion/quantization_quality.py \
    --model Tongyi-MAI/Z-Image-Turbo \
    --quantization fp8 \
    --prompts \
        "a cup of coffee on the table" \
        "an aerial view of a coral reef" \
    --height 1024 --width 1024 \
    --num-inference-steps 50 --seed 42

This generates BF16 and quantized outputs, computes per-prompt LPIPS, and prints a Markdown summary table ready to paste into your PR.

Identifying Sensitive Layers with `ignored_layers`

Not all layers tolerate quantization equally. ignored_layers lets you keep specific layers in BF16 while quantizing everything else. This is critical for achieving good quality with aggressive quantization methods (Int8, NVFP4).

Why some layers are sensitive:

Layer Type	Why Sensitive	Typical Impact
`img_mlp`	Processes denoising latents with shifting dynamic range across timesteps. No built-in normalization to absorb quantization error.	Color shifts, blurriness
`feed_forward`	FFN layers with large dynamic range in DiT blocks.	Artifacts, loss of detail
`proj_out`	Final output projection — small errors get amplified since there are no subsequent layers to correct them.	Overall quality loss
`lm_head`	Vocabulary projection in LLMs — precision-critical for token selection.	Garbage text output
`mlp.gate`	MoE router gates — routing decisions are binary and precision-critical.	Wrong expert selection, quality collapse

How to find sensitive layers for your model:

Run the quality gate test with all layers quantized:

pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "your_model"

If LPIPS exceeds your threshold, try skipping common sensitive layers one at a time:

QualityTestConfig(
    id="fp8_your_model_skip_mlp",
    model="your-org/Your-Model",
    quantization="fp8",
    task="t2i",
    prompt="a cup of coffee on a wooden table",
    max_lpips=0.05,
),

And in the test, modify the Omni call:

omni_qt = Omni(
    model=config.model,
    quantization_config={
        "method": config.quantization,
        "ignored_layers": ["img_mlp"],
    },
)

Compare LPIPS with and without ignored_layers to isolate the problematic layers.
Document the recommended ignored_layers in the supported models table.

Real-world examples:

Model	Method	All Layers (LPIPS)	With `ignored_layers`	Which Layers to Skip
Qwen-Image-2512	Int8	0.0197	0.0027	`img_mlp`
Z-Image-Turbo	Int8	0.1597	0.0290	`feed_forward`
Z-Image-Turbo	FP8	~0.005	—	None needed

What to Verify (Checklist)

Model loads without errors — Check logs for quantization initialization messages
Quality gate passes — LPIPS within threshold vs BF16 baseline
Memory reduction — Expect ~30-50% savings on quantized components
Sensitive layers identified — Document which layers (if any) need ignored_layers
ignored_layers pattern works — Verify that the layer name prefix matches your transformer's module names

Troubleshooting

Issue: Output Quality Degradation

Symptoms: Generated images have visible artifacts, color shifts, or blurriness compared to BF16 baseline.

Causes & Solutions:

Sensitive layers being quantized:

Solution: Identify and skip sensitive layers:

omni = Omni(
    model="your-model",
    quantization_config={
        "method": "fp8",
        "ignored_layers": ["img_mlp"],
    },
)

All layers sensitive (rare):

Solution: The model may not be suitable for FP8 quantization. Consider weight-only quantization or other methods.

Issue: Missing Weights Warning

Symptoms: Warnings about unloaded weights for vae.* or text_encoder.* parameters.

Solution: Ensure load_weights marks pre-loaded component weights. See Step 2.3.

Issue: CUDA Out of Memory

Symptoms: OOM errors during model loading despite FP8 being enabled.