jl9876andChendi.Xue
0d5d1115d6
Signed-off-by: jl9876 <jia.li@intel.com> Co-authored-by: Chendi.Xue <chendi.xue@intel.com>
91 lines
3.0 KiB
Python
91 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Wan AutoRound-format MXFP4 configuration and layer-mapping tests."""
|
|
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from torch import nn
|
|
from vllm.model_executor.layers.linear import LinearBase
|
|
from vllm.model_executor.model_loader.utils import configure_quant_config
|
|
|
|
from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import (
|
|
WanTransformer3DModel,
|
|
)
|
|
from vllm_omni.quantization import build_quant_config
|
|
from vllm_omni.quantization.inc_config import OmniINCConfig
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
|
|
|
|
|
def _autoround_mxfp4_config() -> dict:
|
|
return {
|
|
"bits": 4,
|
|
"act_bits": 4,
|
|
"data_type": "mx_fp",
|
|
"act_data_type": "mx_fp",
|
|
"group_size": 32,
|
|
"act_group_size": 32,
|
|
"sym": True,
|
|
"act_sym": True,
|
|
"act_dynamic": True,
|
|
"enable_quanted_input": False,
|
|
"block_name_to_quantize": "blocks",
|
|
"quant_method": "auto-round",
|
|
"packing_format": "auto_round:llm_compressor",
|
|
"autoround_version": "0.14.2",
|
|
}
|
|
|
|
|
|
def test_build_wan_autoround_mxfp4_config():
|
|
config = build_quant_config(_autoround_mxfp4_config())
|
|
|
|
assert isinstance(config, OmniINCConfig)
|
|
assert config.get_name() == "inc"
|
|
assert config.weight_bits == 4
|
|
assert config.data_type == "mx_fp"
|
|
assert config.group_size == 32
|
|
assert config.packing_format == "auto_round:llm_compressor"
|
|
assert config.block_name_to_quantize == ["blocks"]
|
|
|
|
|
|
def test_wan_autoround_config_uses_runtime_layer_names():
|
|
config = build_quant_config(_autoround_mxfp4_config())
|
|
configure_quant_config(config, WanTransformer3DModel)
|
|
|
|
assert config.packed_modules_mapping is WanTransformer3DModel.packed_modules_mapping
|
|
assert config.block_name_to_quantize == ["blocks"]
|
|
|
|
layer = object.__new__(LinearBase)
|
|
nn.Module.__init__(layer)
|
|
for prefix in (
|
|
"blocks.0.ffn.net_0.proj",
|
|
"blocks.0.ffn.net_2",
|
|
"blocks.0.attn1.to_out",
|
|
):
|
|
layer_config = config.config_parser.resolve(layer, prefix)
|
|
assert layer_config.quantized
|
|
assert layer_config.is_mxfp4
|
|
|
|
|
|
def test_wan_autoround_scope_dispatches_only_blocks_to_mxfp4():
|
|
config = build_quant_config(_autoround_mxfp4_config())
|
|
configure_quant_config(config, WanTransformer3DModel)
|
|
|
|
layer = object.__new__(LinearBase)
|
|
nn.Module.__init__(layer)
|
|
layer_config = config.config_parser.resolve(layer, "blocks.0.ffn.net_0.proj")
|
|
outside_config = config.config_parser.resolve(layer, "condition_embedder.time_proj")
|
|
|
|
assert layer_config.quantized
|
|
assert layer_config.is_mxfp4
|
|
assert not outside_config.quantized
|
|
|
|
fake_method = object()
|
|
with patch(
|
|
"vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_scheme.INCMxfp4Scheme.get_linear_method",
|
|
return_value=fake_method,
|
|
):
|
|
method = config.get_quant_method(layer, "blocks.0.ffn.net_0.proj")
|
|
assert method is fake_method
|