From 963ba1ab7a02b2c1f687285fcd2072cd57cd0b34 Mon Sep 17 00:00:00 2001 From: Gao Han Date: Sat, 6 Jun 2026 08:39:10 +0800 Subject: [PATCH] [CI] Revert "[Feature] gen-only FP8 Quantization Support for SenseNova-U1" (#4196) --- docs/user_guide/quantization/fp8.md | 11 -- .../quantization/test_quantization_fp8.py | 91 --------- .../quantization/test_sensenova_u1_fp8.py | 175 ------------------ .../sensenova_u1/pipeline_sensenova_u1.py | 42 +---- 4 files changed, 1 insertion(+), 318 deletions(-) delete mode 100644 tests/diffusion/quantization/test_sensenova_u1_fp8.py diff --git a/docs/user_guide/quantization/fp8.md b/docs/user_guide/quantization/fp8.md index 69e161a60..c8f14548c 100644 --- a/docs/user_guide/quantization/fp8.md +++ b/docs/user_guide/quantization/fp8.md @@ -52,17 +52,6 @@ guide. FP8 on Ampere may use a weight-only path where available. Audio encoder, vision encoder, talker, and code2wav stay in BF16 unless a model-specific guide says otherwise. -### MoT Diffusion Model (SenseNova-U1) - -| Model | HF models | Scope | Status | Notes | -|-------|-----------|-------|--------|-------| -| SenseNova-U1 | `SenseNova/SenseNova-U1-8B-MoT` | Gen-path only (`*_mot_gen` layers) | Validated | Understanding-path layers stay in BF16; no `ignored_layers` needed | - -SenseNova-U1 uses Mixture-of-Tokenizers (MoT) with duplicate weights per -layer: understanding path (`qkv_proj`, `mlp`) and generation path -(`qkv_proj_mot_gen`, `mlp_mot_gen`). When `quantization="fp8"` is set, -the pipeline automatically applies FP8 only to gen-path GEMMs. - ### Multi-Stage Diffusion Model (BAGEL, GLM-Image) | Model | Scope | Status | Notes | diff --git a/tests/diffusion/quantization/test_quantization_fp8.py b/tests/diffusion/quantization/test_quantization_fp8.py index 8ff16e1f1..58f91742d 100644 --- a/tests/diffusion/quantization/test_quantization_fp8.py +++ b/tests/diffusion/quantization/test_quantization_fp8.py @@ -6,13 +6,11 @@ End-to-end tests for the unified quantization framework (PR #1764). Validates FP8 quantization works correctly for all supported model types: - Single-stage diffusion models (FLUX.1-dev, Qwen-Image, Z-Image-Turbo) - Multi-stage models (BAGEL: LLM + Diffusion) - - MoT gen-only models (SenseNova-U1: gen-path FP8, und-path BF16) Tests verify: 1. FP8 quantization produces valid images 2. Memory usage is lower than BF16 baseline 3. Multi-stage models only quantize the diffusion stage (not the LLM stage) - 4. MoT models only quantize gen-path layers (mot_gen) Usage: # Run all FP8 quantization tests @@ -438,92 +436,3 @@ def test_single_stage_quantization_config_key(): quantization_config="fp8", ) assert len(images) >= 1 - - -# ─── SenseNova-U1 gen-only FP8 tests ───────────────────────────────────────── - -_SENSENOVA_MODEL = "SenseNova/SenseNova-U1-8B-MoT" - - -def _generate_sensenova_u1_image( - quantization_config: str | None = None, - num_inference_steps: int = 20, -) -> tuple[Any, float]: - """Generate an image with SenseNova-U1 (single-stage MoT, gen-only FP8). - - Returns (generated_image, peak_memory_gib). - """ - omni_kwargs: dict[str, Any] = {} - if quantization_config: - omni_kwargs["quantization_config"] = quantization_config - - with OmniRunner(_SENSENOVA_MODEL, **omni_kwargs) as runner: - torch.accelerator.reset_peak_memory_stats() - - sampling_params = OmniDiffusionSamplingParams( - height=1024, - width=1024, - seed=42, - num_inference_steps=num_inference_steps, - extra_args={ - "cfg_scale": 4.0, - "cfg_norm": "none", - "timestep_shift": 3.0, - }, - ) - omni_outputs = list( - runner.omni.generate( - prompts={"prompt": "A cat sitting on a windowsill", "modalities": ["image"]}, - sampling_params_list=sampling_params, - ) - ) - - peak_mem = torch.accelerator.max_memory_allocated() / (1024**3) - - generated_image = None - for req_output in omni_outputs: - if images := getattr(req_output, "images", None): - generated_image = images[0] - break - - assert generated_image is not None, "No images generated from SenseNova-U1" - assert generated_image.size == (1024, 1024), f"Expected 1024x1024, got {generated_image.size}" - - return generated_image, peak_mem - - -@hardware_test(res={"cuda": "H100"}) -def test_sensenova_u1_fp8_generates_image(): - """SenseNova-U1 with gen-only FP8 generates a valid image. - - FP8 is applied only to gen-path (mot_gen) layers; understanding-path - layers stay in BF16. - """ - image, _ = _generate_sensenova_u1_image(quantization_config="fp8") - image.save("test_sensenova_u1_fp8.png") - - -@hardware_test(res={"cuda": "H100"}) -def test_sensenova_u1_bf16_generates_image(): - """SenseNova-U1 without quantization generates a valid image (baseline).""" - image, _ = _generate_sensenova_u1_image(quantization_config=None) - image.save("test_sensenova_u1_bf16.png") - - -@hardware_test(res={"cuda": "H100"}) -def test_sensenova_u1_fp8_uses_less_memory(): - """Gen-only FP8 should use less peak memory than BF16 for SenseNova-U1.""" - _, mem_bf16 = _generate_sensenova_u1_image( - quantization_config=None, - num_inference_steps=10, - ) - torch.accelerator.empty_cache() - - _, mem_fp8 = _generate_sensenova_u1_image( - quantization_config="fp8", - num_inference_steps=10, - ) - - print(f"SenseNova-U1 BF16 peak memory: {mem_bf16:.2f} GiB") - print(f"SenseNova-U1 FP8 peak memory: {mem_fp8:.2f} GiB") - assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)" diff --git a/tests/diffusion/quantization/test_sensenova_u1_fp8.py b/tests/diffusion/quantization/test_sensenova_u1_fp8.py deleted file mode 100644 index 4ed30e9a6..000000000 --- a/tests/diffusion/quantization/test_sensenova_u1_fp8.py +++ /dev/null @@ -1,175 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""CPU tests for SenseNova-U1 gen-only FP8 quantization. - -Validates that the ``_GenOnlyQuantConfig`` wrapper correctly routes FP8 -quantization to only the gen-path (``*_mot_gen``) layers while leaving -understanding-path layers in BF16. -""" - -from types import SimpleNamespace -from unittest.mock import MagicMock - -import pytest -import torch -from vllm.distributed import parallel_state -from vllm.model_executor.layers.linear import LinearBase - -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] - - -@pytest.fixture(autouse=True) -def _fake_tp_group(): - """Provide a fake TP group so vLLM parallel layers can be instantiated. - - Note: a shared ``init_fake_tp_group`` fixture exists in - ``tests/helpers/fixtures/env.py`` but requires pytest-mock. - """ - mock_tp = MagicMock() - mock_tp.world_size = 1 - mock_tp.rank_in_group = 0 - old = parallel_state._TP - parallel_state._TP = mock_tp - yield - parallel_state._TP = old - - -def _make_minimal_llm_config(num_layers: int = 2) -> SimpleNamespace: - """Create a minimal SenseNova-U1 LLM config for unit testing.""" - return SimpleNamespace( - hidden_size=256, - intermediate_size=512, - num_hidden_layers=num_layers, - num_attention_heads=4, - num_key_value_heads=4, - head_dim=64, - hidden_act="silu", - rms_norm_eps=1e-6, - vocab_size=1024, - rope_theta=1000000.0, - rope_theta_hw=10000.0, - max_position_embeddings=4096, - max_position_embeddings_hw=1024, - attention_bias=True, - layer_types=["full_attention"] * num_layers, - tie_word_embeddings=False, - ) - - -def _build_model_with_gen_only_quant(quant_method: str | None = "fp8"): - """Construct SenseNovaU1ForCausalLM with gen-only FP8 wrapping.""" - from unittest.mock import MagicMock, patch - - from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import ( - _GenOnlyQuantConfig, - ) - from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import ( - SenseNovaU1ForCausalLM, - ) - from vllm_omni.quantization import build_quant_config - - inner_config = build_quant_config(quant_method) - if inner_config is not None: - quant_config = _GenOnlyQuantConfig(inner_config) - else: - quant_config = None - - llm_cfg = _make_minimal_llm_config(num_layers=2) - - mock_vllm_config = MagicMock() - mock_vllm_config.model_config.dtype = torch.bfloat16 - mock_vllm_config.model_config.quantization = quant_method - - with ( - torch.device("meta"), - patch("vllm.model_executor.layers.quantization.fp8.get_current_vllm_config", return_value=mock_vllm_config), - ): - model = SenseNovaU1ForCausalLM( - llm_cfg, - quant_config=quant_config, - prefix="language_model", - ) - return model - - -def test_mot_gen_layers_have_fp8_quant_method(): - """Gen-path (mot_gen) Linear layers should have a real FP8 quant_method.""" - from vllm.model_executor.layers.linear import UnquantizedLinearMethod - - model = _build_model_with_gen_only_quant("fp8") - - mot_gen_linears = [ - (name, mod) for name, mod in model.named_modules() if isinstance(mod, LinearBase) and "mot_gen" in name - ] - - assert len(mot_gen_linears) > 0, "No mot_gen LinearBase modules found" - - for name, mod in mot_gen_linears: - assert mod.quant_method is not None, f"mot_gen layer {name!r} should have quant_method with FP8" - assert not isinstance(mod.quant_method, UnquantizedLinearMethod), ( - f"mot_gen layer {name!r} should NOT use UnquantizedLinearMethod" - ) - - -def test_und_layers_have_no_quant_method(): - """Understanding-path Linear layers should use UnquantizedLinearMethod (stay BF16).""" - from vllm.model_executor.layers.linear import UnquantizedLinearMethod - - model = _build_model_with_gen_only_quant("fp8") - - und_linears = [ - (name, mod) for name, mod in model.named_modules() if isinstance(mod, LinearBase) and "mot_gen" not in name - ] - - assert len(und_linears) > 0, "No und LinearBase modules found" - - for name, mod in und_linears: - assert isinstance(mod.quant_method, UnquantizedLinearMethod), ( - f"und layer {name!r} should use UnquantizedLinearMethod, got {type(mod.quant_method).__name__}" - ) - - -def test_no_quant_method_without_quantization(): - """Without quantization, all layers should use UnquantizedLinearMethod.""" - from vllm.model_executor.layers.linear import UnquantizedLinearMethod - - model = _build_model_with_gen_only_quant(None) - - for name, mod in model.named_modules(): - if isinstance(mod, LinearBase): - assert isinstance(mod.quant_method, UnquantizedLinearMethod), ( - f"LinearBase module {name!r} has quant_method=" - f"{type(mod.quant_method).__name__} " - "but no quantization was configured" - ) - - -def test_und_and_gen_layers_both_present(): - """Sanity check: the model has both und and gen layers.""" - model = _build_model_with_gen_only_quant(None) - - layer_names = [name for name, _ in model.named_modules()] - - und_qkv = [n for n in layer_names if "qkv_proj" in n and "mot_gen" not in n] - gen_qkv = [n for n in layer_names if "qkv_proj_mot_gen" in n] - und_mlp = [n for n in layer_names if ".mlp." in n and "mot_gen" not in n] - gen_mlp = [n for n in layer_names if "mlp_mot_gen" in n] - - assert len(und_qkv) > 0, "No und qkv_proj layers found" - assert len(gen_qkv) > 0, "No gen qkv_proj_mot_gen layers found" - assert len(und_mlp) > 0, "No und mlp layers found" - assert len(gen_mlp) > 0, "No gen mlp_mot_gen layers found" - - -def test_gen_only_wrapper_delegates_attributes(): - """_GenOnlyQuantConfig should delegate attribute access to the inner config.""" - from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import ( - _GenOnlyQuantConfig, - ) - from vllm_omni.quantization import build_quant_config - - inner = build_quant_config("fp8") - wrapper = _GenOnlyQuantConfig(inner) - - assert wrapper.get_name() == "fp8" - assert wrapper.activation_scheme == inner.activation_scheme diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py index 917cea13e..fe8695b07 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py +++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py @@ -31,7 +31,6 @@ import torchvision.transforms as T from PIL import Image from transformers import AutoTokenizer from vllm.logger import init_logger -from vllm.model_executor.layers.quantization import QuantizationConfig from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device @@ -484,39 +483,6 @@ def _optimized_scale(positive_flat, negative_flat): return dot / sq_norm -# --------------------------------------------------------------------------- -# Gen-only FP8 wrapper -# --------------------------------------------------------------------------- - - -class _GenOnlyQuantConfig: - """Wraps an inner QuantizationConfig so only ``*_mot_gen`` layers are quantized. - - SenseNova-U1 uses Mixture-of-Tokenizers with duplicate weights per layer: - understanding path (``qkv_proj``, ``mlp``) and generation path - (``qkv_proj_mot_gen``, ``mlp_mot_gen``). The RFC requires FP8 only on the - gen-path GEMMs during denoising; und-path layers stay in BF16. - - This is *not* a full ``QuantizationConfig`` subclass — it only needs to - satisfy the duck-typed interface that ``LinearBase.__init__`` uses: - ``get_quant_method(layer, prefix)``. The outer registry and weight-loader - still operate on the original config stored in ``od_config.quantization_config``. - """ - - def __init__(self, inner: QuantizationConfig): - self._inner = inner - - def get_quant_method(self, layer, prefix: str = ""): - from vllm.model_executor.layers.linear import UnquantizedLinearMethod - - if "mot_gen" in prefix: - return self._inner.get_quant_method(layer, prefix) - return UnquantizedLinearMethod() - - def __getattr__(self, name): - return getattr(self._inner, name) - - # --------------------------------------------------------------------------- # Pipeline # --------------------------------------------------------------------------- @@ -583,15 +549,9 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli self.img_context_token_id = self.tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN) self.img_start_token_id = self.tokenizer.convert_tokens_to_ids(IMG_START_TOKEN) - # Language model (TP-aware). - # Wrap quant_config so only gen-path (mot_gen) layers are quantized; - # understanding-path layers stay in BF16. - quant_config = od_config.quantization_config - if quant_config is not None: - quant_config = _GenOnlyQuantConfig(quant_config) + # Language model (TP-aware) self.language_model = SenseNovaU1ForCausalLM( self.llm_cfg, - quant_config=quant_config, prefix="language_model", )