[AutoRound] Add offline quantized W4A16 model support (#1777)

Signed-off-by: yiliu30 <yi4.liu@intel.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/layers/test_adalayernorm.py

@@ -0,0 +1,237 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for shared AdaLayerNorm layers used by FLUX and other models."""
+
+import os
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.fixture(autouse=True)
+def _init_distributed():
+    """Initialize the minimal distributed environment required by
+    ReplicatedLinear (tensor-parallel group must exist)."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29501")
+    init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method="env://",
+    )
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
+@pytest.fixture(autouse=True)
+def _force_default_gemm(monkeypatch):
+    """Force CPU-compatible GEMM dispatch for tests using CPU tensors.
+
+    vLLM's dispatch_unquantized_gemm() selects the backend by platform
+    (e.g. rocm_unquantized_gemm on AMD machines), not by tensor device.
+    CPU test tensors crash with NotImplementedError on ROCm.  Monkeypatch
+    the dispatcher to always return the default (torch.nn.functional.linear)
+    implementation which works on any device."""
+    from vllm.model_executor.layers.utils import default_unquantized_gemm
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
+        lambda: default_unquantized_gemm,
+    )
+
+
+def test_adalayernorm_import_from_shared_module():
+    """Verify imports work from the shared adalayernorm module."""
+    from vllm_omni.diffusion.layers.adalayernorm import (  # noqa: F401
+        AdaLayerNormContinuous,
+        AdaLayerNormZero,
+        AdaLayerNormZeroSingle,
+    )
+
+
+def test_adalayernorm_zero_forward_shape():
+    """AdaLayerNormZero produces correct output shapes (x, gate, shift, scale, gate)."""
+    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZero
+
+    dim = 64
+    batch = 2
+    seq_len = 4
+    norm = AdaLayerNormZero(dim)
+
+    x = torch.randn(batch, seq_len, dim)
+    emb = torch.randn(batch, dim)
+
+    out_x, gate_msa, shift_mlp, scale_mlp, gate_mlp = norm(x, emb)
+
+    assert out_x.shape == (batch, seq_len, dim)
+    assert gate_msa.shape == (batch, dim)
+    assert shift_mlp.shape == (batch, dim)
+    assert scale_mlp.shape == (batch, dim)
+    assert gate_mlp.shape == (batch, dim)
+
+
+def test_adalayernorm_zero_single_forward_shape():
+    """AdaLayerNormZeroSingle produces (x, gate) with correct shapes."""
+    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZeroSingle
+
+    dim = 64
+    batch = 2
+    seq_len = 4
+    norm = AdaLayerNormZeroSingle(dim)
+
+    x = torch.randn(batch, seq_len, dim)
+    emb = torch.randn(batch, dim)
+
+    out_x, gate = norm(x, emb)
+
+    assert out_x.shape == (batch, seq_len, dim)
+    assert gate.shape == (batch, dim)
+
+
+def test_adalayernorm_continuous_forward_shape():
+    """AdaLayerNormContinuous produces correct output shape."""
+    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormContinuous
+
+    dim = 64
+    cond_dim = 64
+    batch = 2
+    seq_len = 4
+    norm = AdaLayerNormContinuous(dim, cond_dim)
+
+    x = torch.randn(batch, seq_len, dim)
+    conditioning = torch.randn(batch, cond_dim)
+
+    out = norm(x, conditioning)
+
+    assert out.shape == (batch, seq_len, dim)
+
+
+def test_adalayernorm_zero_accepts_quant_config():
+    """Constructor accepts quant_config=None and prefix='test' without error."""
+    from vllm_omni.diffusion.layers.adalayernorm import (
+        AdaLayerNormContinuous,
+        AdaLayerNormZero,
+        AdaLayerNormZeroSingle,
+    )
+
+    # Should not raise with quant_config=None and prefix
+    AdaLayerNormZero(64, quant_config=None, prefix="test.norm1")
+    AdaLayerNormZeroSingle(64, quant_config=None, prefix="test.norm")
+    AdaLayerNormContinuous(64, 64, quant_config=None, prefix="test.norm_out")
+
+
+def test_adalayernorm_uses_replicated_linear():
+    """Verify .linear is a ReplicatedLinear instance (not nn.Linear)."""
+    from vllm.model_executor.layers.linear import ReplicatedLinear
+
+    from vllm_omni.diffusion.layers.adalayernorm import (
+        AdaLayerNormContinuous,
+        AdaLayerNormZero,
+        AdaLayerNormZeroSingle,
+    )
+
+    norm_zero = AdaLayerNormZero(64)
+    assert isinstance(norm_zero.linear, ReplicatedLinear)
+
+    norm_zero_single = AdaLayerNormZeroSingle(64)
+    assert isinstance(norm_zero_single.linear, ReplicatedLinear)
+
+    norm_continuous = AdaLayerNormContinuous(64, 64)
+    assert isinstance(norm_continuous.linear, ReplicatedLinear)
+
+
+# ── Numerical equivalence tests against diffusers originals ──
+
+
+def _copy_weights(src_linear, dst_replicated_linear):
+    """Copy weights from nn.Linear to ReplicatedLinear for comparison."""
+    dst_replicated_linear.weight.data.copy_(src_linear.weight.data)
+    if src_linear.bias is not None and dst_replicated_linear.bias is not None:
+        dst_replicated_linear.bias.data.copy_(src_linear.bias.data)
+
+
+def test_adalayernorm_zero_matches_diffusers():
+    """Verify AdaLayerNormZero produces identical output to diffusers original."""
+    from diffusers.models.normalization import (
+        AdaLayerNormZero as DiffusersAdaLayerNormZero,
+    )
+
+    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZero
+
+    dim = 64
+    torch.manual_seed(42)
+    ours = AdaLayerNormZero(dim)
+    ref = DiffusersAdaLayerNormZero(dim)
+
+    # Copy weights: nn.Linear -> ReplicatedLinear
+    _copy_weights(ref.linear, ours.linear)
+
+    x = torch.randn(2, 4, dim)
+    emb = torch.randn(2, dim)
+
+    out_ours = ours(x, emb)
+    out_ref = ref(x, emb=emb)
+
+    for o, r in zip(out_ours, out_ref):
+        torch.testing.assert_close(o, r, atol=1e-5, rtol=1e-5)
+
+
+def test_adalayernorm_zero_single_matches_diffusers():
+    """Verify AdaLayerNormZeroSingle produces identical output to diffusers original."""
+    from diffusers.models.normalization import (
+        AdaLayerNormZeroSingle as DiffusersAdaLayerNormZeroSingle,
+    )
+
+    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZeroSingle
+
+    dim = 64
+    torch.manual_seed(42)
+    ours = AdaLayerNormZeroSingle(dim)
+    ref = DiffusersAdaLayerNormZeroSingle(dim)
+
+    _copy_weights(ref.linear, ours.linear)
+
+    x = torch.randn(2, 4, dim)
+    emb = torch.randn(2, dim)
+
+    out_ours = ours(x, emb)
+    out_ref = ref(x, emb=emb)
+
+    for o, r in zip(out_ours, out_ref):
+        torch.testing.assert_close(o, r, atol=1e-5, rtol=1e-5)
+
+
+def test_adalayernorm_continuous_matches_diffusers():
+    """Verify AdaLayerNormContinuous produces identical output to diffusers original."""
+    from diffusers.models.normalization import (
+        AdaLayerNormContinuous as DiffusersAdaLayerNormContinuous,
+    )
+
+    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormContinuous
+
+    dim = 64
+    cond_dim = 64
+    torch.manual_seed(42)
+    # Match constructor args: diffusers defaults elementwise_affine=True, eps=1e-5
+    ours = AdaLayerNormContinuous(dim, cond_dim, elementwise_affine=False, eps=1e-6)
+    ref = DiffusersAdaLayerNormContinuous(dim, cond_dim, elementwise_affine=False, eps=1e-6)
+
+    _copy_weights(ref.linear, ours.linear)
+
+    x = torch.randn(2, 4, dim)
+    cond = torch.randn(2, cond_dim)
+
+    out_ours = ours(x, cond)
+    out_ref = ref(x, cond)
+
+    torch.testing.assert_close(out_ours, out_ref, atol=1e-5, rtol=1e-5)

tests/diffusion/models/flux/test_flux_prefix_propagation.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests that FLUX transformer blocks correctly propagate `quant_config` and
+`prefix` through all sub-layers.
+
+The tests instantiate blocks with a known prefix and verify that all quantization-
+aware sub-layers (AdaLayerNorm, FeedForward, Attention projections) receive the
+prefix rooted at the block prefix.  This is critical for quantized weight loading
+to match checkpoint keys to the correct model parameters.
+"""
+
+import os
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+# Standard dimensions for a minimal FLUX block
+_DIM = 64
+_HEADS = 2
+_HEAD_DIM = _DIM // _HEADS
+
+
+@pytest.fixture(autouse=True)
+def _init_distributed():
+    """Initialize the minimal distributed environment required by
+    vLLM parallel linear layers (tensor-parallel group must exist)."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29502")
+    init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method="env://",
+    )
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
+def _param_names(module) -> set[str]:
+    """Return the set of all parameter names in a module."""
+    return {name for name, _ in module.named_parameters()}
+
+
+def test_flux_transformer_block_passes_prefix():
+    """FluxTransformerBlock propagates prefix to norm1, norm1_context, attn, ff, ff_context."""
+    from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformerBlock
+
+    prefix = "transformer_blocks.0"
+    block = FluxTransformerBlock(
+        dim=_DIM,
+        num_attention_heads=_HEADS,
+        attention_head_dim=_HEAD_DIM,
+        quant_config=None,
+        prefix=prefix,
+    )
+
+    params = _param_names(block)
+
+    # norm1 and norm1_context (AdaLayerNormZero) should have linear weights
+    assert any(name.startswith("norm1.linear.") for name in params), (
+        f"norm1.linear.* not found in params: {sorted(params)}"
+    )
+    assert any(name.startswith("norm1_context.linear.") for name in params), (
+        f"norm1_context.linear.* not found in params: {sorted(params)}"
+    )
+
+    # attn should have QKV projections
+    assert any(name.startswith("attn.to_qkv.") for name in params), (
+        f"attn.to_qkv.* not found in params: {sorted(params)}"
+    )
+
+    # ff and ff_context should have net layers
+    assert any(name.startswith("ff.net.") for name in params), f"ff.net.* not found in params: {sorted(params)}"
+    assert any(name.startswith("ff_context.net.") for name in params), (
+        f"ff_context.net.* not found in params: {sorted(params)}"
+    )
+
+
+def test_flux_single_transformer_block_passes_prefix():
+    """FluxSingleTransformerBlock propagates prefix to norm, proj_mlp, attn."""
+    from vllm_omni.diffusion.models.flux.flux_transformer import FluxSingleTransformerBlock
+
+    prefix = "single_transformer_blocks.0"
+    block = FluxSingleTransformerBlock(
+        dim=_DIM,
+        num_attention_heads=_HEADS,
+        attention_head_dim=_HEAD_DIM,
+        quant_config=None,
+        prefix=prefix,
+    )
+
+    params = _param_names(block)
+
+    # norm (AdaLayerNormZeroSingle) should have linear weights
+    assert any(name.startswith("norm.linear.") for name in params), (
+        f"norm.linear.* not found in params: {sorted(params)}"
+    )
+
+    # proj_mlp (ReplicatedLinear) should have weight
+    assert any(name.startswith("proj_mlp.") for name in params), f"proj_mlp.* not found in params: {sorted(params)}"
+
+    # attn should have QKV projection
+    assert any(name.startswith("attn.to_qkv.") for name in params), (
+        f"attn.to_qkv.* not found in params: {sorted(params)}"
+    )
+
+
+def test_flux_feedforward_passes_prefix():
+    """FeedForward propagates prefix to net.0 (GELU proj) and net.2 (output proj)."""
+    from vllm_omni.diffusion.models.flux.flux_transformer import FeedForward
+
+    prefix = "transformer_blocks.0.ff"
+    ff = FeedForward(
+        dim=_DIM,
+        dim_out=_DIM,
+        quant_config=None,
+        prefix=prefix,
+    )
+
+    params = _param_names(ff)
+
+    # net.0 is ColumnParallelApproxGELU which wraps a ColumnParallelLinear
+    assert any("net.0" in name for name in params), f"net.0 not found in params: {sorted(params)}"
+
+    # net.2 is RowParallelLinear
+    assert any("net.2" in name for name in params), f"net.2 not found in params: {sorted(params)}"

tests/diffusion/quantization/test_inc_config.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for INC/AutoRound quantization via the unified framework."""
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+def test_build_quant_config_autoround():
+    """build_quant_config("auto-round", ...) should produce an INCConfig."""
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+
+    from vllm_omni.quantization import build_quant_config
+
+    config = build_quant_config(
+        "auto-round",
+        bits=4,
+        group_size=128,
+        sym=True,
+        packing_format="auto_round:auto_gptq",
+    )
+    assert config is not None
+    assert isinstance(config, INCConfig)
+    assert config.weight_bits == 4
+    assert config.group_size == 128
+
+
+def test_build_quant_config_inc():
+    """build_quant_config("inc", ...) should also produce an INCConfig."""
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+
+    from vllm_omni.quantization import build_quant_config
+
+    config = build_quant_config("inc", bits=4, group_size=128)
+    assert isinstance(config, INCConfig)
+    assert config.weight_bits == 4
+
+
+def test_build_quant_config_autoround_dict():
+    """Dict-style config with method=auto-round should work."""
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+
+    from vllm_omni.quantization import build_quant_config
+
+    config = build_quant_config(
+        {
+            "method": "auto-round",
+            "bits": 4,
+            "group_size": 128,
+            "sym": True,
+            "packing_format": "auto_round:auto_gptq",
+        }
+    )
+    assert isinstance(config, INCConfig)
+    assert config.weight_bits == 4
+
+
+def test_build_quant_config_autoround_filters_metadata():
+    """Checkpoint metadata keys (autoround_version, batch_size, iters)
+    should be silently filtered out instead of causing TypeError."""
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+
+    from vllm_omni.quantization import build_quant_config
+
+    config = build_quant_config(
+        "auto-round",
+        bits=4,
+        group_size=128,
+        sym=True,
+        packing_format="auto_round:auto_gptq",
+        block_name_to_quantize="transformer_blocks,single_transformer_blocks",
+        autoround_version="0.12.0",  # metadata — must be filtered
+        batch_size=1,  # metadata — must be filtered
+        iters=0,  # metadata — must be filtered
+    )
+    assert isinstance(config, INCConfig)
+    assert config.weight_bits == 4
+    assert config.group_size == 128
+
+
+def test_build_quant_config_bits_to_weight_bits_mapping():
+    """The 'bits' key from checkpoints should be mapped to 'weight_bits'."""
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+
+    from vllm_omni.quantization import build_quant_config
+
+    # If weight_bits is already provided, bits should be ignored
+    config = build_quant_config("auto-round", weight_bits=4, group_size=128)
+    assert isinstance(config, INCConfig)
+    assert config.weight_bits == 4
+
+
+def test_autoround_in_supported_methods():
+    """auto-round and inc should appear in SUPPORTED_QUANTIZATION_METHODS."""
+    from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
+
+    assert "auto-round" in SUPPORTED_QUANTIZATION_METHODS
+    assert "inc" in SUPPORTED_QUANTIZATION_METHODS
+
+
+def test_integration_autoround_via_omni_diffusion_config():
+    """OmniDiffusionConfig with auto-round quantization dict should resolve."""
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    config = OmniDiffusionConfig(
+        model="test",
+        quantization_config={
+            "method": "auto-round",
+            "bits": 4,
+            "group_size": 128,
+            "sym": True,
+        },
+    )
+    assert isinstance(config.quantization_config, INCConfig)
+    assert config.quantization_config.weight_bits == 4
+
+
+def test_integration_autodetect_from_transformer_config():
+    """When TransformerConfig has quant_config, OmniDiffusionConfig should
+    auto-detect it even without explicit quantization_config."""
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+
+    from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
+
+    tf_config = TransformerConfig.from_dict(
+        {
+            "quantization_config": {
+                "quant_method": "auto-round",
+                "bits": 4,
+                "group_size": 128,
+                "sym": True,
+                "packing_format": "auto_round:auto_gptq",
+                "autoround_version": "0.12.0",
+                "batch_size": 1,
+                "iters": 0,
+            }
+        }
+    )
+    assert tf_config.quant_method == "auto-round"
+    assert isinstance(tf_config.quant_config, INCConfig)
+
+    od_config = OmniDiffusionConfig(model="test", tf_model_config=tf_config)
+    assert isinstance(od_config.quantization_config, INCConfig)
+    assert od_config.quantization_config.weight_bits == 4