iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[CI] Speed up quantization test group (#53291)

Browse Source 
Signed-off-by: Felix Marty <Felix.Marty@amd.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
This commit is contained in:
fxmarty-amd
2026-09-01 13:15:24 -05:00
committed by GitHub
co-authored by Andreas Karatzas
parent e7cf4730d6
commit d1c15e589d
20 changed files with 824 additions and 725 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/evals/gsm8k/configs/Qwen-1.5-MOE-W-MXFP4-A-MXFP6.yaml
+5
View File
@@ -0,0 +1,5 @@
model_name: "fxmarty/qwen1.5_moe_a2.7b_chat_w_fp4_a_fp6_e2m3"
accuracy_threshold: 0.4405
num_questions: 50
num_fewshot: 5
server_args: "--enforce-eager --max-model-len 4096"
tests/evals/gsm8k/configs/Qwen-1.5-MOE-W-MXFP6-A-MXFP6.yaml
+5
View File
@@ -0,0 +1,5 @@
model_name: "fxmarty/qwen1.5_moe_a2.7b_chat_w_fp6_e3m2_a_fp6_e3m2"
accuracy_threshold: 0.5049
num_questions: 50
num_fewshot: 5
server_args: "--enforce-eager --max-model-len 4096"
tests/evals/gsm8k/configs/Qwen3-30B-A3B-NVFP4-quark.yaml
+5
View File
@@ -0,0 +1,5 @@
model_name: "amd-quark/Qwen3-30B-A3B-nvfp4-quark"
accuracy_threshold: 0.89
num_questions: 1319
num_fewshot: 5
server_args: "--enforce-eager --max-model-len 4096"
tests/evals/gsm8k/configs/models-gfx950-large.txt
+3
View File
@@ -6,3 +6,6 @@
# startup timeout fix has landed and is verified green.
# DeepSeek-V4-Pro-NVFP4.yaml
DeepSeek-V4-Flash-NVFP4.yaml
Qwen3-30B-A3B-NVFP4-quark.yaml
Qwen-1.5-MOE-W-MXFP4-A-MXFP6.yaml
Qwen-1.5-MOE-W-MXFP6-A-MXFP6.yaml
tests/evals/gsm8k/configs/models-mi3xx.txt
+4 -1
View File
@@ -3,5 +3,8 @@ DeepSeek-R1-DP_MI325.yaml
DeepSeek-V3.2-TP_MI325.yaml
DeepSeek-V3.2-DP_MI325.yaml
Qwen3-30B-A3B-NVFP4.yaml
Qwen3-30B-A3B-NVFP4-quark.yaml
Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml
Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml
Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml
Qwen-1.5-MOE-W-MXFP4-A-MXFP6.yaml
Qwen-1.5-MOE-W-MXFP6-A-MXFP6.yaml
tests/quantization/test_auto_awq.py
+37 -29
View File
@@ -18,7 +18,14 @@ from __future__ import annotations
import pytest
import torch
from tests.quantization.utils import is_quant_method_supported
from tests.quantization.utils import (
is_quant_method_supported,
load_model_without_vllm_runner,
)
from vllm.config import set_current_vllm_config
from vllm.forward_context import set_forward_context
from vllm.model_executor.layers.attention import Attention
from vllm.platforms import current_platform
def _get_auto_awq_config_source() -> str:
@@ -186,42 +193,43 @@ AWQ_MODELS = [
reason="auto_awq is not supported on this GPU type.",
)
@pytest.mark.parametrize("model_id", AWQ_MODELS)
def test_auto_awq_quantization_method(vllm_runner, model_id: str, monkeypatch):
def test_auto_awq_quantization_method(
model_id: str, monkeypatch, dist_init, workspace_init
):
"""Test that quantization='auto_awq' loads and runs correctly."""
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
with vllm_runner(
model, vllm_config = load_model_without_vllm_runner(
model_id,
dtype=torch.float16,
quantization="auto_awq",
max_model_len=2048,
enforce_eager=True,
) as llm:
model_config_kwargs={"max_model_len": 2048},
)
target_device = torch.device(current_platform.device_type)
def check_model(model):
from vllm.model_executor.layers.quantization.auto_awq import (
AutoAWQLinearMethod,
AutoAWQMarlinLinearMethod,
)
from vllm.model_executor.layers.quantization.auto_awq import (
AutoAWQLinearMethod,
AutoAWQMarlinLinearMethod,
)
for name, submodule in model.named_modules():
if name == "model.layers.0.self_attn.qkv_proj":
# Should use either AutoAWQLinearMethod (Triton) or
# AutoAWQMarlinLinearMethod (Marlin) depending on hardware
assert isinstance(
submodule.quant_method,
(AutoAWQLinearMethod, AutoAWQMarlinLinearMethod),
), (
f"Expected AutoAWQLinearMethod or AutoAWQMarlinLinearMethod "
f"for {name}, got {type(submodule.quant_method)}"
)
break
qkv_proj = model.model.layers[0].self_attn.qkv_proj
assert isinstance(
qkv_proj.quant_method,
(AutoAWQLinearMethod, AutoAWQMarlinLinearMethod),
), (
"Expected AutoAWQLinearMethod or AutoAWQMarlinLinearMethod, "
f"got {type(qkv_proj.quant_method)}"
)
llm.apply_model(check_model)
outputs = llm.generate_greedy([PROMPT], max_tokens=8)
assert outputs
assert len(outputs[0][1]) > 0
monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
input_ids = torch.tensor([1, 2, 3, 4], device=target_device)
positions = torch.arange(input_ids.numel(), device=target_device)
with (
set_current_vllm_config(vllm_config),
set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
):
hidden_states = model(input_ids, positions, None)
logits = model.compute_logits(hidden_states)
assert torch.isfinite(logits).all()
def test_auto_awq_config_get_name():
tests/quantization/test_auto_gptq.py
+27 -17
View File
@@ -10,13 +10,20 @@ from types import SimpleNamespace
import pytest
import torch
from tests.quantization.utils import is_quant_method_supported
from tests.quantization.utils import (
is_quant_method_supported,
load_model_without_vllm_runner,
)
from vllm.config import set_current_vllm_config
from vllm.forward_context import set_forward_context
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.fused_moe import RoutedExperts
from vllm.model_executor.layers.quantization.auto_gptq import (
AutoGPTQConfig,
AutoGPTQLinearMethod,
AutoGPTQMoEMethod,
)
from vllm.platforms import current_platform
PROMPT = "On the surface of Mars, we found"
@@ -30,29 +37,32 @@ MODELS = [
reason="auto_gptq is not supported on this GPU type.",
)
@pytest.mark.parametrize("model_id", MODELS)
def test_auto_gptq_quantization_method(vllm_runner, model_id: str, monkeypatch):
def test_auto_gptq_quantization_method(
model_id: str, monkeypatch, dist_init, workspace_init
):
"""Test that quantization='auto_gptq' loads and runs correctly."""
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
with vllm_runner(
model, vllm_config = load_model_without_vllm_runner(
model_id,
dtype=torch.float16,
quantization="auto_gptq",
max_model_len=2048,
enforce_eager=True,
) as llm:
model_config_kwargs={"max_model_len": 2048},
)
def check_model(model):
for name, submodule in model.named_modules():
if name == "model.layers.0.self_attn.qkv_proj":
assert isinstance(submodule.quant_method, AutoGPTQLinearMethod)
break
qkv_proj = model.model.layers[0].self_attn.qkv_proj
assert isinstance(qkv_proj.quant_method, AutoGPTQLinearMethod)
llm.apply_model(check_model)
outputs = llm.generate_greedy([PROMPT], max_tokens=8)
assert outputs
assert len(outputs[0][1]) > 0
target_device = torch.device(current_platform.device_type)
monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
input_ids = torch.tensor([1, 2, 3, 4], device=target_device)
positions = torch.arange(input_ids.numel(), device=target_device)
with (
set_current_vllm_config(vllm_config),
set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
):
hidden_states = model(input_ids, positions, None)
logits = model.compute_logits(hidden_states)
assert torch.isfinite(logits).all()
def test_auto_gptq_config_get_name():
tests/quantization/test_compressed_tensors.py
+125 -164
View File
@@ -5,7 +5,6 @@
Run `pytest tests/quantization/test_compressed_tensors.py`.
"""
from contextlib import contextmanager
from unittest.mock import Mock
import pytest
@@ -18,12 +17,17 @@ from compressed_tensors.quantization import (
)
from tests.models.utils import check_logprobs_close
from tests.quantization.utils import load_model_without_vllm_runner
from vllm.config import set_current_vllm_config
from vllm.config.kernel import KernelConfig
from vllm.forward_context import set_forward_context
from vllm.model_executor.kernels.linear import (
Fp8BlockScaledMMLinearKernel,
)
from vllm.model_executor.kernels.linear.scaled_mm import (
MarlinFP8ScaledMMLinearKernel,
)
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.fused_moe import UnquantizedFusedMoEMethod
from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
@@ -99,7 +103,7 @@ def enable_pickle(monkeypatch):
),
],
)
def test_compressed_tensors_w8a8_static_setup(vllm_runner, model_args):
def test_compressed_tensors_w8a8_static_setup(model_args, dist_init, workspace_init):
model_path, strategy, quant_type, shape_0, is_symmetric = model_args
if (
@@ -108,55 +112,45 @@ def test_compressed_tensors_w8a8_static_setup(vllm_runner, model_args):
):
pytest.skip(f"Skip model {model_path} as it is not supported on ROCm.")
with vllm_runner(model_path, enforce_eager=True) as llm:
model, _ = load_model_without_vllm_runner(model_path)
layer = model.model.layers[0]
def check_model(model):
layer = model.model.layers[0]
qkv_proj = layer.self_attn.qkv_proj
o_proj = layer.self_attn.o_proj
gate_up_proj = layer.mlp.gate_up_proj
down_proj = layer.mlp.down_proj
qkv_proj = layer.self_attn.qkv_proj
o_proj = layer.self_attn.o_proj
gate_up_proj = layer.mlp.gate_up_proj
down_proj = layer.mlp.down_proj
def zp_valid(zp: torch.Tensor | None):
if is_symmetric:
return zp is None
# assert zp for symmetric and asymmetric cases
def zp_valid(zp: torch.Tensor | None):
if is_symmetric:
return zp is None
return zp is not None and zp.dtype is torch.int32
return zp is not None and zp.dtype is torch.int32
assert zp_valid(qkv_proj.input_zero_point)
assert zp_valid(o_proj.input_zero_point)
assert zp_valid(gate_up_proj.input_zero_point)
assert zp_valid(down_proj.input_zero_point)
assert zp_valid(qkv_proj.input_zero_point)
assert zp_valid(o_proj.input_zero_point)
assert zp_valid(gate_up_proj.input_zero_point)
assert zp_valid(down_proj.input_zero_point)
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(o_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(gate_up_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(down_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Int8)
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(o_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(gate_up_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(down_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Int8)
assert qkv_proj.scheme.strategy == strategy
assert qkv_proj.scheme.is_static_input_scheme
expected_type = torch.int8
assert qkv_proj.scheme.strategy == strategy
assert qkv_proj.scheme.is_static_input_scheme
expected_type = torch.int8
assert qkv_proj.weight.dtype is expected_type
assert o_proj.weight.dtype is expected_type
assert gate_up_proj.weight.dtype is expected_type
assert qkv_proj.weight.dtype is expected_type
assert o_proj.weight.dtype is expected_type
assert gate_up_proj.weight.dtype is expected_type
if qkv_proj.scheme.strategy == "tensor":
# Make sure it is a channelwise buffer
# After running process_weights_after_loading
assert len(qkv_proj.weight_scale.shape) == 2
assert qkv_proj.weight_scale.shape[0] == shape_0
assert qkv_proj.weight_scale.shape[1] == 1
assert qkv_proj.weight_scale.dtype is torch.float32
assert qkv_proj.input_scale.dtype is torch.float32
llm.apply_model(check_model)
output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
assert output
if qkv_proj.scheme.strategy == "tensor":
assert len(qkv_proj.weight_scale.shape) == 2
assert qkv_proj.weight_scale.shape[0] == shape_0
assert qkv_proj.weight_scale.shape[1] == 1
assert qkv_proj.weight_scale.dtype is torch.float32
assert qkv_proj.input_scale.dtype is torch.float32
@pytest.mark.parametrize(
@@ -243,10 +237,7 @@ def test_compressed_tensors_no_enforce_eager(vllm_runner):
"use_aiter", [True, False] if current_platform.is_rocm() else [False]
)
def test_compressed_tensors_w8a8_dynamic_per_token(
vllm_runner,
model_args,
use_aiter,
monkeypatch,
model_args, use_aiter, monkeypatch, dist_init, workspace_init
):
model_path, strategy = model_args
@@ -262,23 +253,13 @@ def test_compressed_tensors_w8a8_dynamic_per_token(
# this will enable VLLM_ROCM_USE_AITER_LINEAR
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
with vllm_runner(model_path, enforce_eager=True, dtype=torch.float16) as llm:
def check_model(model):
layer = model.model.layers[0]
qkv_proj = layer.self_attn.qkv_proj
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Int8)
assert not qkv_proj.scheme.is_static_input_scheme
assert qkv_proj.scheme.strategy == strategy
assert qkv_proj.weight.dtype is torch.int8
llm.apply_model(check_model)
output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
assert output
model, _ = load_model_without_vllm_runner(model_path, dtype=torch.float16)
qkv_proj = model.model.layers[0].self_attn.qkv_proj
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Int8)
assert not qkv_proj.scheme.is_static_input_scheme
assert qkv_proj.scheme.strategy == strategy
assert qkv_proj.weight.dtype is torch.int8
@pytest.mark.parametrize(
@@ -329,33 +310,36 @@ def test_compressed_tensors_wNa16(vllm_runner, wNa16_args):
assert output
def test_compressed_tensors_fp8(vllm_runner):
def test_compressed_tensors_fp8(monkeypatch, dist_init, workspace_init):
model_path = "nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test"
with vllm_runner(model_path, enforce_eager=True) as llm:
model, vllm_config = load_model_without_vllm_runner(model_path)
target_device = torch.device(current_platform.device_type)
def check_model(model):
layer = model.model.layers[0]
layer = model.model.layers[0]
qkv_proj = layer.self_attn.qkv_proj
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(
qkv_proj.scheme,
(CompressedTensorsW8A8Fp8, CompressedTensorsW8A16Fp8),
)
assert qkv_proj.input_scale.dtype is torch.float32
qkv_proj = layer.self_attn.qkv_proj
if isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8):
assert len(qkv_proj.input_scale.shape) == 0
assert qkv_proj.weight.dtype is current_platform.fp8_dtype()
assert qkv_proj.weight_scale.dtype is torch.float32
assert len(qkv_proj.weight_scale.shape) == 0
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(
qkv_proj.scheme,
(CompressedTensorsW8A8Fp8, CompressedTensorsW8A16Fp8),
)
assert qkv_proj.input_scale.dtype is torch.float32
if isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8):
assert len(qkv_proj.input_scale.shape) == 0
assert qkv_proj.weight.dtype is current_platform.fp8_dtype()
assert qkv_proj.weight_scale.dtype is torch.float32
assert len(qkv_proj.weight_scale.shape) == 0
llm.apply_model(check_model)
output = llm.generate_greedy("Hello my name is", max_tokens=4)
assert output
monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
input_ids = torch.tensor([1, 2, 3, 4], device=target_device)
positions = torch.arange(input_ids.numel(), device=target_device)
with (
set_current_vllm_config(vllm_config),
set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
):
hidden_states = model(input_ids, positions, None)
logits = model.compute_logits(hidden_states)
assert torch.isfinite(logits).all()
def test_compressed_tensors_w8a8_fp8_moe_forwards_swiglu_params():
@@ -413,27 +397,6 @@ def test_compressed_tensors_kv_cache_fp8_per_attn_head(vllm_runner):
assert output
@contextmanager
def _nvfp4_marlin_error_context(model, capfd):
is_rocm_and_unsupported = (
model == "nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4A16"
and current_platform.is_rocm()
)
if is_rocm_and_unsupported:
expected_error = (
"ValueError: Forced NVFP4 kernel MarlinNvFp4LinearKernel is not "
"supported: Marlin FP4 not available"
)
with pytest.raises(RuntimeError, match="Engine core initialization failed"):
yield
captured = capfd.readouterr()
assert expected_error in captured.out + captured.err
else:
yield
@pytest.mark.parametrize(
"args",
[
@@ -441,26 +404,30 @@ def _nvfp4_marlin_error_context(model, capfd):
("nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4", False),
],
)
def test_compressed_tensors_nvfp4(vllm_runner, args, capfd):
model, use_a16 = args
with (
_nvfp4_marlin_error_context(model, capfd),
vllm_runner(model, enforce_eager=True) as llm,
):
def test_compressed_tensors_nvfp4(args, dist_init, workspace_init):
model_path, use_a16 = args
is_rocm_and_unsupported = (
model_path == "nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4A16"
and current_platform.is_rocm()
)
if is_rocm_and_unsupported:
with pytest.raises(
ValueError,
match="Forced NVFP4 kernel MarlinNvFp4LinearKernel is not supported",
):
load_model_without_vllm_runner(model_path)
return
def check_model(model):
layer = model.model.layers[0]
try:
model, _ = load_model_without_vllm_runner(model_path)
except ValueError as exc:
pytest.fail(f"Unexpected NVFP4 loading failure: {exc}")
qkv_proj = layer.self_attn.qkv_proj
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(qkv_proj.scheme, CompressedTensorsW4A4Fp4)
assert qkv_proj.scheme.use_a16 == use_a16
assert qkv_proj.scheme.group_size == 16
llm.apply_model(check_model)
output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
print(output)
assert output
qkv_proj = model.model.layers[0].self_attn.qkv_proj
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(qkv_proj.scheme, CompressedTensorsW4A4Fp4)
assert qkv_proj.scheme.use_a16 == use_a16
assert qkv_proj.scheme.group_size == 16
@pytest.mark.skipif(
@@ -533,47 +500,41 @@ def test_compressed_tensors_transforms_perplexity(
@pytest.mark.parametrize(
"linear_backend", ["auto", "marlin"] if current_platform.is_cuda() else ["auto"]
)
def test_compressed_tensors_fp8_block_enabled(vllm_runner, linear_backend):
def test_compressed_tensors_fp8_block_enabled(
linear_backend, dist_init, workspace_init
):
model_path = "RedHatAI/Qwen3-0.6B-FP8-BLOCK"
with vllm_runner(
model_path, enforce_eager=True, linear_backend=linear_backend
) as llm:
fp8_dtype = current_platform.fp8_dtype()
model, _ = load_model_without_vllm_runner(
model_path,
vllm_config_kwargs={
"kernel_config": KernelConfig(linear_backend=linear_backend)
},
)
fp8_dtype = current_platform.fp8_dtype()
layer = model.model.layers[0]
def check_model(model):
layer = model.model.layers[0]
qkv_proj = layer.self_attn.qkv_proj
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8)
if linear_backend == "marlin":
assert isinstance(qkv_proj.scheme.fp8_linear, MarlinFP8ScaledMMLinearKernel)
assert qkv_proj.weight.dtype is torch.int32
assert qkv_proj.weight_scale.dtype is qkv_proj.orig_dtype
else:
assert isinstance(qkv_proj.scheme.fp8_linear, Fp8BlockScaledMMLinearKernel)
assert qkv_proj.weight.dtype is fp8_dtype
assert qkv_proj.weight_scale.dtype is torch.float32
assert len(qkv_proj.weight.shape) == 2
assert len(qkv_proj.weight_scale.shape) == 2
qkv_proj = layer.self_attn.qkv_proj
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8)
if linear_backend == "marlin":
assert isinstance(
qkv_proj.scheme.fp8_linear, MarlinFP8ScaledMMLinearKernel
)
assert qkv_proj.weight.dtype is torch.int32
assert qkv_proj.weight_scale.dtype is qkv_proj.orig_dtype
else:
assert isinstance(
qkv_proj.scheme.fp8_linear, Fp8BlockScaledMMLinearKernel
)
assert qkv_proj.weight.dtype is fp8_dtype
assert qkv_proj.weight_scale.dtype is torch.float32
assert len(qkv_proj.weight.shape) == 2
assert len(qkv_proj.weight_scale.shape) == 2
if linear_backend == "auto":
input_quant_op = qkv_proj.scheme.fp8_linear.quant_fp8
assert isinstance(input_quant_op, QuantFP8)
assert input_quant_op._forward_method in (
input_quant_op.forward_cuda,
input_quant_op.forward_hip,
input_quant_op.forward_xpu,
)
llm.apply_model(check_model)
output = llm.generate_greedy("Hello my name is", max_tokens=4)
assert output
if linear_backend == "auto":
input_quant_op = qkv_proj.scheme.fp8_linear.quant_fp8
assert isinstance(input_quant_op, QuantFP8)
assert input_quant_op._forward_method in (
input_quant_op.forward_cuda,
input_quant_op.forward_hip,
input_quant_op.forward_xpu,
)
@pytest.mark.skipif(
tests/quantization/test_fp8.py
+43 -25
View File
@@ -12,12 +12,20 @@ import pytest
import regex as re
import torch
from tests.quantization.utils import is_quant_method_supported
from tests.quantization.utils import (
is_quant_method_supported,
load_model_without_vllm_runner,
)
from vllm import _custom_ops as ops
from vllm.config import set_current_vllm_config
from vllm.config.cache import CacheConfig
from vllm.config.kernel import KernelConfig
from vllm.config.model import ModelConfig
from vllm.forward_context import set_forward_context
from vllm.model_executor.kernels.linear.scaled_mm import (
MarlinFP8ScaledMMLinearKernel,
)
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.attention.attention import (
set_default_quant_scales,
)
@@ -116,21 +124,33 @@ def test_static_fp8_moe_input_scales_remain_scalar() -> None:
"use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
)
def test_model_load_and_run(
vllm_runner, model_id: str, force_marlin: bool, use_rocm_aiter: bool, monkeypatch
model_id: str,
force_marlin: bool,
use_rocm_aiter: bool,
monkeypatch,
dist_init,
workspace_init,
) -> None:
if use_rocm_aiter:
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
kwargs = {}
if force_marlin:
kwargs["linear_backend"] = "marlin"
kwargs["moe_backend"] = "marlin"
with vllm_runner(model_id, enforce_eager=True, **kwargs) as llm:
# note: this does not test accuracy, just that we can run through
# see lm-eval tests for accuracy
outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4)
print(outputs[0][1])
kernel_config = KernelConfig(
linear_backend="marlin" if force_marlin else "auto",
moe_backend="marlin" if force_marlin else "auto",
)
model, vllm_config = load_model_without_vllm_runner(
model_id,
model_config_kwargs={"hf_overrides": {"num_hidden_layers": 3}},
vllm_config_kwargs={"kernel_config": kernel_config},
)
monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
with (
set_current_vllm_config(vllm_config),
set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
):
model(input_ids, positions, None)
@pytest.mark.skipif(
@@ -525,20 +545,18 @@ def test_kv_cache_scale_sync_to_host_copies():
not is_quant_method_supported("fp8"),
reason="FP8 is not supported on this GPU type.",
)
def test_kv_cache_dtype_skip_layers(vllm_runner, monkeypatch):
def test_kv_cache_dtype_skip_layers(monkeypatch, dist_init, workspace_init):
"""Test that kv_cache_dtype_skip_layers skips quantization for specified layers."""
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
with vllm_runner(
model, _ = load_model_without_vllm_runner(
"facebook/opt-125m",
kv_cache_dtype="fp8",
kv_cache_dtype_skip_layers=["0", "2"],
enforce_eager=True,
) as llm:
def check_layers(model):
for i, layer in enumerate(model.model.decoder.layers):
expected = "auto" if str(i) in ["0", "2"] else "fp8"
assert layer.self_attn.attn.kv_cache_dtype == expected
llm.apply_model(check_layers)
vllm_config_kwargs={
"cache_config": CacheConfig(
cache_dtype="fp8", kv_cache_dtype_skip_layers=["0", "2"]
)
},
)
for i, layer in enumerate(model.model.decoder.layers):
expected = "auto" if str(i) in ["0", "2"] else "fp8"
assert layer.self_attn.attn.kv_cache_dtype == expected
tests/quantization/test_fp8_per_channel.py
-34
View File
@@ -67,37 +67,3 @@ def test_scaled_fp8_quant_per_channel_shape() -> None:
assert y.dtype == current_platform.fp8_dtype()
assert s.shape == (96, 1)
assert s.dtype == torch.float32
@pytest.mark.skipif(
not is_quant_method_supported("fp8"),
reason="FP8 is not supported on this GPU type.",
)
def test_fp8_per_channel_online_quantization(
vllm_runner,
monkeypatch,
) -> None:
"""End-to-end smoke: load `facebook/opt-125m` bf16 with
`quantization='fp8_per_channel'`, check a dense Linear is wrapped by
`Fp8PtpcOnlineLinearMethod`, its weights are fp8 with per-channel
scales (shape `[N, 1]`), and a short greedy generation works.
"""
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
with vllm_runner(
"facebook/opt-125m",
quantization="fp8_per_channel",
enforce_eager=True,
) as llm:
def check_model(model):
fc1 = model.model.decoder.layers[0].fc1
assert isinstance(fc1.quant_method, Fp8PtpcOnlineLinearMethod)
assert fc1.weight.dtype == current_platform.fp8_dtype()
assert fc1.weight_scale.ndim == 2
assert fc1.weight_scale.shape[-1] == 1
assert fc1.input_scale is None
llm.apply_model(check_model)
outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4)
print(outputs[0][1])
tests/quantization/test_gptq_dynamic.py
+30 -39
View File
@@ -10,6 +10,7 @@ Note: Only symmetric GPTQ models are supported after consolidation to Marlin.
import pytest
import torch
from tests.quantization.utils import load_model_without_vllm_runner
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQLinearMethod
from vllm.model_executor.layers.quantization.utils.gptq_utils import (
@@ -28,48 +29,38 @@ MODELS = [
@pytest.mark.parametrize("model_id", MODELS)
def test_gptq_with_dynamic(vllm_runner, model_id: str, monkeypatch):
def test_gptq_with_dynamic(model_id: str, monkeypatch, dist_init, workspace_init):
# `LLM.apply_model` requires pickling a function.
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
linear_method_cls = AutoGPTQLinearMethod
with vllm_runner(
model_id, dtype=torch.float16, max_model_len=2048, enforce_eager=True
) as llm:
model, _ = load_model_without_vllm_runner(
model_id,
dtype=torch.float16,
model_config_kwargs={
"max_model_len": 2048,
"hf_overrides": {"num_hidden_layers": 3},
},
)
def check_model(model):
for name, submodule in model.named_modules():
if name == "lm_head":
assert isinstance(submodule.quant_method, linear_method_cls)
elif name == "model.layers.0.self_attn.qkv_proj":
# The first layer is quantized using bits=4, group_size=128
# desc_act=True
assert isinstance(submodule.quant_method, linear_method_cls)
config = submodule.quant_method.quant_config
assert config.weight_bits == 4
assert config.group_size == 128
assert config.desc_act
elif name == "model.layers.1.self_attn.qkv_proj":
# The second layer is quantized using bits=8, group_size=32
# desc_act=False
assert isinstance(submodule.quant_method, linear_method_cls)
config = submodule.quant_method.quant_config
assert (
get_dynamic_override(config, layer_name=name, key="bits") == 8
)
assert (
get_dynamic_override(config, layer_name=name, key="group_size")
== 32
)
assert not get_dynamic_override(
config, layer_name=name, key="desc_act"
)
elif (
name == "model.layers.2.self_attn.qkv_proj"
or name == "model.layers.2.mlp.gate_up_proj"
):
# All other layers (layer index >= 2) are not quantized
assert isinstance(submodule.quant_method, UnquantizedLinearMethod)
llm.apply_model(check_model)
for name, submodule in model.named_modules():
if name == "lm_head":
assert isinstance(submodule.quant_method, linear_method_cls)
elif name == "model.layers.0.self_attn.qkv_proj":
assert isinstance(submodule.quant_method, linear_method_cls)
config = submodule.quant_method.quant_config
assert config.weight_bits == 4
assert config.group_size == 128
assert config.desc_act
elif name == "model.layers.1.self_attn.qkv_proj":
assert isinstance(submodule.quant_method, linear_method_cls)
config = submodule.quant_method.quant_config
assert get_dynamic_override(config, layer_name=name, key="bits") == 8
assert get_dynamic_override(config, layer_name=name, key="group_size") == 32
assert not get_dynamic_override(config, layer_name=name, key="desc_act")
elif name in (
"model.layers.2.self_attn.qkv_proj",
"model.layers.2.mlp.gate_up_proj",
):
assert isinstance(submodule.quant_method, UnquantizedLinearMethod)
tests/quantization/test_lm_head.py
+16 -20
View File
@@ -8,6 +8,7 @@ Run `pytest tests/quantization/test_quant_lm_head_true.py --forked`.
import pytest
import torch
from tests.quantization.utils import load_model_without_vllm_runner
from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQLinearMethod
from vllm.model_executor.layers.vocab_parallel_embedding import (
UnquantizedEmbeddingMethod,
@@ -23,29 +24,24 @@ MODELS_QUANT = [
@pytest.mark.parametrize("model_id, lm_head_quantized", MODELS_QUANT)
def test_lm_head(
vllm_runner,
model_id: str,
lm_head_quantized: bool,
monkeypatch,
dist_init,
workspace_init,
) -> None:
# `LLM.apply_model` requires pickling a function.
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
with vllm_runner(
model_id, dtype=torch.float16, max_model_len=2048, enforce_eager=True
) as vllm_model:
def check_model(model):
lm_head_layer = model.lm_head
if lm_head_quantized:
assert isinstance(
lm_head_layer.quant_method,
AutoGPTQLinearMethod,
)
else:
assert isinstance(
lm_head_layer.quant_method, UnquantizedEmbeddingMethod
)
vllm_model.apply_model(check_model)
print(vllm_model.generate_greedy(["Hello my name is"], max_tokens=4)[0][1])
model, _ = load_model_without_vllm_runner(
model_id,
dtype=torch.float16,
model_config_kwargs={
"max_model_len": 2048,
"hf_overrides": {"num_hidden_layers": 3},
},
)
lm_head_layer = model.lm_head
if lm_head_quantized:
assert isinstance(lm_head_layer.quant_method, AutoGPTQLinearMethod)
else:
assert isinstance(lm_head_layer.quant_method, UnquantizedEmbeddingMethod)
tests/quantization/test_modelopt.py
+81 -93
View File
@@ -12,14 +12,19 @@ from unittest.mock import MagicMock, Mock, patch
import pytest
import torch
from tests.quantization.utils import is_quant_method_supported
from tests.quantization.utils import (
is_quant_method_supported,
load_model_without_vllm_runner,
)
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.config.model import ModelConfig
from vllm.forward_context import set_forward_context
from vllm.model_executor.kernels.linear import (
FlashInferCuteDslNvFp4W4A16LinearKernel,
HummingNvFp4LinearKernel,
MarlinNvFp4LinearKernel,
)
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
from vllm.model_executor.layers.quantization.modelopt import (
ModelOptFp8Config,
@@ -419,124 +424,107 @@ def test_modelopt_fp8_checkpoint_setup(default_vllm_config, vllm_runner):
not is_quant_method_supported("modelopt"),
reason="ModelOpt FP8 is not supported on this GPU type.",
)
def test_modelopt_fp8_pc_pt_checkpoint_setup(default_vllm_config, vllm_runner):
def test_modelopt_fp8_pc_pt_checkpoint_setup(monkeypatch, dist_init, workspace_init):
"""Test ModelOpt FP8_PER_CHANNEL_PER_TOKEN checkpoint setup."""
model_id = "CedricHwang/qwen2.5-0.5b-modelopt-fp8-pc-pt"
model_path = _snapshot_download_or_skip(model_id)
# Set model config as model_config.dtype is required in ModelOptFp8LinearMethod.
default_vllm_config.model_config = ModelConfig()
with vllm_runner(model_path, quantization="modelopt", enforce_eager=True) as llm:
model, vllm_config = load_model_without_vllm_runner(
model_path,
quantization="modelopt",
)
layer = model.model.layers[0]
def check_model(model):
layer = model.model.layers[0]
qkv_proj = layer.self_attn.qkv_proj
o_proj = layer.self_attn.o_proj
gate_up_proj = layer.mlp.gate_up_proj
down_proj = layer.mlp.down_proj
qkv_proj = layer.self_attn.qkv_proj
o_proj = layer.self_attn.o_proj
gate_up_proj = layer.mlp.gate_up_proj
down_proj = layer.mlp.down_proj
from vllm.model_executor.layers.quantization.modelopt import (
ModelOptFp8PcPtLinearMethod,
)
from vllm.model_executor.layers.quantization.modelopt import (
ModelOptFp8PcPtLinearMethod,
)
assert isinstance(qkv_proj.quant_method, ModelOptFp8PcPtLinearMethod)
assert isinstance(o_proj.quant_method, ModelOptFp8PcPtLinearMethod)
assert isinstance(gate_up_proj.quant_method, ModelOptFp8PcPtLinearMethod)
assert isinstance(down_proj.quant_method, ModelOptFp8PcPtLinearMethod)
assert isinstance(qkv_proj.quant_method, ModelOptFp8PcPtLinearMethod)
assert isinstance(o_proj.quant_method, ModelOptFp8PcPtLinearMethod)
assert isinstance(gate_up_proj.quant_method, ModelOptFp8PcPtLinearMethod)
assert isinstance(down_proj.quant_method, ModelOptFp8PcPtLinearMethod)
fp8_dtype = current_platform.fp8_dtype()
assert qkv_proj.weight.dtype == fp8_dtype
assert o_proj.weight.dtype == fp8_dtype
assert gate_up_proj.weight.dtype == fp8_dtype
assert down_proj.weight.dtype == fp8_dtype
fp8_dtype = current_platform.fp8_dtype()
assert qkv_proj.weight.dtype == fp8_dtype
assert o_proj.weight.dtype == fp8_dtype
assert gate_up_proj.weight.dtype == fp8_dtype
assert down_proj.weight.dtype == fp8_dtype
# Per-channel scales; activations are dynamically scaled per token.
for projection in (qkv_proj, o_proj, gate_up_proj, down_proj):
assert hasattr(projection, "weight_scale")
assert projection.weight_scale.dtype == torch.float32
assert projection.weight_scale.dim() == 1
assert not hasattr(projection, "input_scale")
# Per-channel scales; activations are dynamically scaled per token.
assert hasattr(qkv_proj, "weight_scale")
assert qkv_proj.weight_scale.dtype == torch.float32
assert qkv_proj.weight_scale.dim() == 1
assert not hasattr(qkv_proj, "input_scale")
assert hasattr(o_proj, "weight_scale")
assert o_proj.weight_scale.dtype == torch.float32
assert o_proj.weight_scale.dim() == 1
assert not hasattr(o_proj, "input_scale")
assert hasattr(gate_up_proj, "weight_scale")
assert gate_up_proj.weight_scale.dtype == torch.float32
assert gate_up_proj.weight_scale.dim() == 1
assert not hasattr(gate_up_proj, "input_scale")
assert hasattr(down_proj, "weight_scale")
assert down_proj.weight_scale.dtype == torch.float32
assert down_proj.weight_scale.dim() == 1
assert not hasattr(down_proj, "input_scale")
llm.apply_model(check_model)
output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
assert output
print(f"ModelOpt FP8_PER_CHANNEL_PER_TOKEN output: {output}")
monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
input_ids = torch.tensor([1, 2, 3, 4], device=current_platform.device_type)
positions = torch.arange(input_ids.numel(), device=current_platform.device_type)
with (
set_current_vllm_config(vllm_config),
set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
):
hidden_states = model(input_ids, positions, None)
logits = model.compute_logits(hidden_states)
assert torch.isfinite(logits).all()
@pytest.mark.skipif(
not is_quant_method_supported("modelopt"),
reason="ModelOpt FP8 is not supported on this GPU type.",
)
def test_modelopt_fp8_pb_wo_checkpoint_setup(default_vllm_config, vllm_runner):
def test_modelopt_fp8_pb_wo_checkpoint_setup(monkeypatch, dist_init, workspace_init):
"""Test ModelOpt FP8_PB_WO checkpoint setup."""
model_id = "CedricHwang/qwen2.5-0.5b-modelopt-fp8-pb-wo"
model_path = _snapshot_download_or_skip(model_id)
# Set model config as model_config.dtype is required in ModelOptFp8LinearMethod.
default_vllm_config.model_config = ModelConfig()
with vllm_runner(model_path, quantization="modelopt", enforce_eager=True) as llm:
model, vllm_config = load_model_without_vllm_runner(
model_path,
quantization="modelopt",
)
layer = model.model.layers[0]
def check_model(model):
layer = model.model.layers[0]
qkv_proj = layer.self_attn.qkv_proj
o_proj = layer.self_attn.o_proj
gate_up_proj = layer.mlp.gate_up_proj
down_proj = layer.mlp.down_proj
qkv_proj = layer.self_attn.qkv_proj
o_proj = layer.self_attn.o_proj
gate_up_proj = layer.mlp.gate_up_proj
down_proj = layer.mlp.down_proj
from vllm.model_executor.layers.quantization.modelopt import (
ModelOptFp8PbWoLinearMethod,
)
from vllm.model_executor.layers.quantization.modelopt import (
ModelOptFp8PbWoLinearMethod,
)
assert isinstance(qkv_proj.quant_method, ModelOptFp8PbWoLinearMethod)
assert isinstance(o_proj.quant_method, ModelOptFp8PbWoLinearMethod)
assert isinstance(gate_up_proj.quant_method, ModelOptFp8PbWoLinearMethod)
assert isinstance(down_proj.quant_method, ModelOptFp8PbWoLinearMethod)
assert isinstance(qkv_proj.quant_method, ModelOptFp8PbWoLinearMethod)
assert isinstance(o_proj.quant_method, ModelOptFp8PbWoLinearMethod)
assert isinstance(gate_up_proj.quant_method, ModelOptFp8PbWoLinearMethod)
assert isinstance(down_proj.quant_method, ModelOptFp8PbWoLinearMethod)
fp8_dtype = current_platform.fp8_dtype()
assert qkv_proj.weight.dtype == fp8_dtype
assert o_proj.weight.dtype == fp8_dtype
assert gate_up_proj.weight.dtype == fp8_dtype
assert down_proj.weight.dtype == fp8_dtype
fp8_dtype = current_platform.fp8_dtype()
assert qkv_proj.weight.dtype == fp8_dtype
assert o_proj.weight.dtype == fp8_dtype
assert gate_up_proj.weight.dtype == fp8_dtype
assert down_proj.weight.dtype == fp8_dtype
# Block scales are materialized as a 2D [out_blk, in_blk] tensor.
for projection in (qkv_proj, o_proj, gate_up_proj, down_proj):
assert hasattr(projection, "weight_scale")
assert projection.weight_scale.dtype == torch.float32
assert projection.weight_scale.dim() == 2
# Block scales; should be materialized as a 2D [out_blk, in_blk] tensor.
assert hasattr(qkv_proj, "weight_scale")
assert qkv_proj.weight_scale.dtype == torch.float32
assert qkv_proj.weight_scale.dim() == 2
assert hasattr(o_proj, "weight_scale")
assert o_proj.weight_scale.dtype == torch.float32
assert o_proj.weight_scale.dim() == 2
assert hasattr(gate_up_proj, "weight_scale")
assert gate_up_proj.weight_scale.dtype == torch.float32
assert gate_up_proj.weight_scale.dim() == 2
assert hasattr(down_proj, "weight_scale")
assert down_proj.weight_scale.dtype == torch.float32
assert down_proj.weight_scale.dim() == 2
llm.apply_model(check_model)
output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
assert output
print(f"ModelOpt FP8_PB_WO output: {output}")
monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
input_ids = torch.tensor([1, 2, 3, 4], device=current_platform.device_type)
positions = torch.arange(input_ids.numel(), device=current_platform.device_type)
with (
set_current_vllm_config(vllm_config),
set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
):
hidden_states = model(input_ids, positions, None)
logits = model.compute_logits(hidden_states)
assert torch.isfinite(logits).all()
def test_modelopt_nvfp4_config_dispatches_w4a4_method():
tests/quantization/test_online.py
+119 -52
View File
@@ -13,16 +13,22 @@ from torch.distributed import ProcessGroup
from tests.quantization.utils import (
_test_online_quant_peak_mem_impl,
is_quant_method_supported,
load_model_without_vllm_runner,
)
from vllm import _custom_ops as ops
from vllm._aiter_ops import rocm_aiter_ops
from vllm._custom_ops import scaled_fp4_quant
from vllm.config.quantization import resolve_quantization_config
from vllm.forward_context import set_forward_context
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
from vllm.model_executor.layers.quantization.online.fp8 import (
Fp8PerBlockOnlineLinearMethod,
Fp8PerBlockOnlineMoEMethod,
Fp8PerTensorOnlineLinearMethod,
Fp8PerTensorOnlineMoEMethod,
Fp8PtpcOnlineLinearMethod,
Fp8PtpcOnlineMoEMethod,
_fp8_channel_scale,
_fp8_quant_per_channel,
_fp8_scale,
@@ -43,6 +49,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
amax_for_tp_weight_quant,
weight_amax,
)
from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader
from vllm.model_executor.models.granitemoe import (
GraniteMoeModel,
)
from vllm.platforms import current_platform
from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
@@ -128,6 +138,12 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
Fp8PerBlockOnlineLinearMethod,
Fp8PerBlockOnlineMoEMethod,
),
(
"fp8_per_channel",
None,
Fp8PtpcOnlineLinearMethod,
Fp8PtpcOnlineMoEMethod,
),
# quantization='online' with per-layer-kind overrides
(
"online",
@@ -160,13 +176,14 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
"use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
)
def test_online_quantization(
vllm_runner,
quant_scheme: str,
online_quant_args: dict | None,
expected_linear_cls,
expected_moe_cls,
use_rocm_aiter: bool,
monkeypatch,
dist_init,
workspace_init,
) -> None:
"""
Tests that online quantization frontend configuration works -
@@ -188,67 +205,117 @@ def test_online_quantization(
if current_platform.is_xpu() and quant_scheme == "fp8_per_block":
pytest.skip("Skip test for online fp8_per_block on XPU platform.")
# `LLM.apply_model` requires pickling a function.
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
# a tiny model with both dense and MoE layers
model_name = "ibm-granite/granite-3.0-1b-a400m-base"
runner_kwargs = dict(
model, vllm_config = load_model_without_vllm_runner(
model_name,
dtype="bfloat16",
quantization=quant_scheme,
enforce_eager=True,
model_config_kwargs={
"quantization_config": resolve_quantization_config(
quant_scheme, online_quant_args
),
"hf_overrides": {
"num_hidden_layers": 3,
"vocab_size": 256,
"hidden_size": 256,
"intermediate_size": 512,
"num_attention_heads": 4,
"num_key_value_heads": 4,
"max_position_embeddings": 64,
"num_local_experts": 4,
"num_experts_per_tok": 2,
},
},
model_loader_cls=DummyModelLoader,
)
if online_quant_args is not None:
runner_kwargs["quantization_config"] = online_quant_args
monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
o_proj = model.model.layers[0].self_attn.o_proj
moe = model.model.layers[0].block_sparse_moe.experts
assert isinstance(o_proj.quant_method, expected_linear_cls)
assert isinstance(moe._quant_method, expected_moe_cls)
if quant_scheme == "mxfp4":
assert o_proj.weight.dtype == torch.uint8
elif current_platform.is_cuda() or current_platform.is_xpu():
assert o_proj.weight.dtype == torch.float8_e4m3fn
elif current_platform.is_rocm():
assert o_proj.weight.dtype == current_platform.fp8_dtype()
else:
pytest.skip("Only runs on CUDA and ROCm.")
if quant_scheme == "fp8_per_channel":
assert o_proj.weight_scale.ndim == 2
assert o_proj.weight_scale.shape[-1] == 1
assert o_proj.input_scale is None
if isinstance(online_quant_args, dict) and "ignore" in online_quant_args:
for layer_idx in range(len(model.model.layers)):
o_proj = model.model.layers[layer_idx].self_attn.o_proj
if layer_idx == 1:
assert isinstance(o_proj.quant_method, UnquantizedLinearMethod)
else:
assert isinstance(o_proj.quant_method, expected_linear_cls)
for layer in model.model.layers:
assert isinstance(
layer.self_attn.qkv_proj.quant_method, UnquantizedLinearMethod
)
input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE)
positions = torch.arange(input_ids.numel(), device=DEVICE)
with set_forward_context(None, vllm_config, num_tokens=input_ids.numel()):
hidden_states = model(input_ids, positions, None)
logits = model.compute_logits(hidden_states)
assert torch.isfinite(logits).all()
@pytest.mark.skipif(
not is_quant_method_supported("fp8"),
reason="FP8 is not supported on this GPU type.",
)
def test_online_quantization_loads_real_weights(vllm_runner, monkeypatch) -> None:
"""Verify online quantization loads a Granite-MoE checkpoint end to end."""
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
original_load_weights = GraniteMoeModel.load_weights
def load_weights(self, weights):
weights = (
(name, weight)
for name, weight in weights
if not name.startswith("layers.") or int(name.split(".")[1]) < 3
)
return original_load_weights(self, weights)
monkeypatch.setattr(GraniteMoeModel, "load_weights", load_weights)
with vllm_runner(
model_name,
**runner_kwargs,
"ibm-granite/granite-3.0-1b-a400m-base",
quantization="fp8_per_tensor",
dtype="bfloat16",
enforce_eager=True,
hf_overrides={"num_hidden_layers": 3},
max_model_len=16,
max_num_seqs=1,
) as llm:
def check_model(model):
# checks further down in the test case are hardcoded for this
# model
assert model_name == "ibm-granite/granite-3.0-1b-a400m-base"
o_proj = model.model.layers[0].self_attn.o_proj
moe = model.model.layers[0].block_sparse_moe.experts
# o_proj and moe in layer 0 are always quantized (never ignored)
# because of how we craft the test case inputs
assert isinstance(o_proj.quant_method, expected_linear_cls)
if moe is not None:
assert isinstance(moe._quant_method, expected_moe_cls)
if quant_scheme == "mxfp4":
# Packed e2m1 values, two per byte.
assert o_proj.weight.dtype == torch.uint8
elif current_platform.is_cuda() or current_platform.is_xpu():
assert o_proj.weight.dtype == torch.float8_e4m3fn
elif current_platform.is_rocm():
assert o_proj.weight.dtype == current_platform.fp8_dtype()
else:
pytest.skip("Only runs on CUDA and ROCm.")
# Verify ignored layers are unquantized.
if isinstance(online_quant_args, dict) and "ignore" in online_quant_args:
# only .*1.self_attn_o_proj is skipped
for layer_idx in range(len(model.model.layers)):
o_proj = model.model.layers[layer_idx].self_attn.o_proj
if layer_idx == 1:
assert isinstance(o_proj.quant_method, UnquantizedLinearMethod)
else:
assert isinstance(o_proj.quant_method, expected_linear_cls)
# every .*self_attn.qkv_proj is skipped
for layer_idx in range(len(model.model.layers)):
qkv_proj = model.model.layers[layer_idx].self_attn.qkv_proj
assert isinstance(qkv_proj.quant_method, UnquantizedLinearMethod)
layer = model.model.layers[0]
assert isinstance(
layer.self_attn.o_proj.quant_method,
Fp8PerTensorOnlineLinearMethod,
)
assert isinstance(
layer.block_sparse_moe.experts._quant_method,
Fp8PerTensorOnlineMoEMethod,
)
llm.apply_model(check_model)
outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4)
print(outputs[0][1])
outputs = llm.generate_greedy(["Hello my name is"], max_tokens=1)
assert outputs
@pytest.mark.skipif(
tests/quantization/test_quark.py
+143 -208
View File
@@ -18,7 +18,12 @@ import pytest
import torch
from packaging import version
from vllm._aiter_ops import is_aiter_found_and_supported
from tests.quantization.utils import load_model_without_vllm_runner
from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
from vllm.config import set_current_vllm_config
from vllm.config.cache import CacheConfig
from vllm.forward_context import set_forward_context
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.fused_moe import (
RoutedExperts,
UnquantizedFusedMoEMethod,
@@ -76,6 +81,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
kNvfp4Dynamic,
kNvfp4Static,
)
from vllm.model_executor.models.llama import LlamaForCausalLM
from vllm.platforms import current_platform
from vllm.transformers_utils.repo_utils import hf_api
@@ -903,157 +909,184 @@ def test_quant_method_dispatch_instantiation(case, monkeypatch, default_vllm_con
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
@pytest.mark.parametrize("tp", [1])
def test_quark_fp8_w_per_tensor_a_per_tensor(vllm_runner, kv_cache_dtype, tp):
def test_quark_fp8_w_per_tensor_a_per_tensor(
kv_cache_dtype: str, monkeypatch, dist_init, workspace_init
):
model_path = "amd/Llama-3.1-8B-Instruct-FP8-KV-Quark-test"
with vllm_runner(
checkpoint_scales = {}
scale_names = {
"model.layers.0.self_attn.k_proj.output_scale",
"model.layers.0.self_attn.v_proj.output_scale",
}
original_load_weights = LlamaForCausalLM.load_weights
def load_weights(self, weights):
def capture_scales():
for name, weight in weights:
if name in scale_names:
checkpoint_scales[name] = weight.detach().cpu()
yield name, weight
return original_load_weights(self, capture_scales())
monkeypatch.setattr(LlamaForCausalLM, "load_weights", load_weights)
model, vllm_config = load_model_without_vllm_runner(
model_path,
enforce_eager=True,
kv_cache_dtype=kv_cache_dtype,
tensor_parallel_size=tp,
) as llm:
model_config_kwargs={"hf_overrides": {"num_hidden_layers": 3}},
vllm_config_kwargs={"cache_config": CacheConfig(cache_dtype=kv_cache_dtype)},
)
def check_model(model):
layer = model.model.layers[0]
qkv_proj = model.model.layers[0].self_attn.qkv_proj
assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
assert isinstance(qkv_proj.scheme, QuarkW8A8Fp8)
assert len(qkv_proj.input_scale.shape) == 0
assert qkv_proj.weight.dtype is current_platform.fp8_dtype()
assert len(qkv_proj.weight_scale.shape) == 0
qkv_proj = layer.self_attn.qkv_proj
attn = model.model.layers[0].self_attn.attn
if kv_cache_dtype == "fp8":
assert checkpoint_scales.keys() == scale_names
scale_multiplier = 2 if current_platform.is_fp8_fnuz() else 1
assert attn._k_scale_float == (
checkpoint_scales["model.layers.0.self_attn.k_proj.output_scale"].item()
* scale_multiplier
)
assert attn._v_scale_float == (
checkpoint_scales["model.layers.0.self_attn.v_proj.output_scale"].item()
* scale_multiplier
)
else:
assert attn._k_scale_float == 1.0
assert attn._v_scale_float == 1.0
assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
assert isinstance(qkv_proj.scheme, QuarkW8A8Fp8)
if isinstance(qkv_proj.scheme, QuarkW8A8Fp8):
assert len(qkv_proj.input_scale.shape) == 0
assert qkv_proj.weight.dtype is current_platform.fp8_dtype()
assert len(qkv_proj.weight_scale.shape) == 0
llm.apply_model(check_model)
output = llm.generate_greedy("Hello my name is", max_tokens=4)
assert output
monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
with (
set_current_vllm_config(vllm_config),
set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
):
hidden_states = model(input_ids, positions, None)
logits = model.compute_logits(hidden_states)
assert torch.isfinite(logits).all()
@pytest.mark.parametrize("tp", [1])
def test_quark_fp8_w_per_channel_a_per_token(vllm_runner, tp):
def test_quark_fp8_w_per_channel_a_per_token(monkeypatch, dist_init, workspace_init):
model_path = "amd/Qwen2.5-1.5B-Instruct-ptpc-Quark-ts"
with vllm_runner(model_path, enforce_eager=True, tensor_parallel_size=tp) as llm:
model, vllm_config = load_model_without_vllm_runner(
model_path,
model_config_kwargs={"hf_overrides": {"num_hidden_layers": 3}},
)
def check_model(model):
layer = model.model.layers[0]
qkv_proj = model.model.layers[0].self_attn.qkv_proj
assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
assert isinstance(qkv_proj.scheme, QuarkW8A8Fp8)
assert qkv_proj.weight.dtype is current_platform.fp8_dtype()
assert qkv_proj.weight_scale.shape[0] == qkv_proj.weight.shape[1]
assert qkv_proj.weight_scale.shape[1] == 1
qkv_proj = layer.self_attn.qkv_proj
assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
assert isinstance(qkv_proj.scheme, QuarkW8A8Fp8)
if isinstance(qkv_proj.scheme, QuarkW8A8Fp8):
assert qkv_proj.weight.dtype is current_platform.fp8_dtype()
assert qkv_proj.weight_scale.shape[0] == qkv_proj.weight.shape[1]
assert qkv_proj.weight_scale.shape[1] == 1
llm.apply_model(check_model)
output = llm.generate_greedy("Hello my name is", max_tokens=4)
assert output
monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
with (
set_current_vllm_config(vllm_config),
set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
):
hidden_states = model(input_ids, positions, None)
logits = model.compute_logits(hidden_states)
assert torch.isfinite(logits).all()
@pytest.mark.parametrize("tp", [1])
def test_quark_int8_w_per_tensor_a_per_tensor(vllm_runner, tp):
def test_quark_int8_w_per_tensor_a_per_tensor(monkeypatch, dist_init, workspace_init):
model_path = "amd/Llama-3.1-8B-Instruct-w-int8-a-int8-sym-test"
with vllm_runner(model_path, enforce_eager=True, tensor_parallel_size=tp) as llm:
model, vllm_config = load_model_without_vllm_runner(
model_path,
model_config_kwargs={"hf_overrides": {"num_hidden_layers": 3}},
)
with set_current_vllm_config(vllm_config):
qkv_proj = model.model.layers[0].self_attn.qkv_proj
assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
assert isinstance(qkv_proj.scheme, QuarkW8A8Int8)
def check_model(model):
layer = model.model.layers[0]
qkv_proj = layer.self_attn.qkv_proj
assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
assert isinstance(qkv_proj.scheme, QuarkW8A8Int8)
llm.apply_model(check_model)
output = llm.generate_greedy("Hello my name is", max_tokens=4)
assert output
monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
with set_forward_context(None, vllm_config, num_tokens=input_ids.numel()):
hidden_states = model(input_ids, positions, None)
logits = model.compute_logits(hidden_states)
assert torch.isfinite(logits).all()
@pytest.mark.parametrize("tp", [1])
def test_quark_int8_w8a8_moe(vllm_runner, tp):
def test_quark_int8_w8a8_moe(monkeypatch, dist_init, workspace_init):
"""Test W8A8 INT8 MoE quantization with a tiny Qwen3 MoE model."""
model_path = "amd/tiny-qwen3-moe-w8a8-int8"
with vllm_runner(
model, vllm_config = load_model_without_vllm_runner(
model_path,
enforce_eager=True,
tensor_parallel_size=tp,
gpu_memory_utilization=0.1,
) as llm:
model_config_kwargs={"hf_overrides": {"num_hidden_layers": 3}},
)
def check_model(model):
layer = model.model.layers[0]
# MoE experts should use QuarkW8A8Int8MoEMethod
moe = layer.mlp.experts
assert isinstance(moe._quant_method, QuarkW8A8Int8MoEMethod), (
f"Expected QuarkW8A8Int8MoEMethod, got {type(moe._quant_method)}"
)
# Non-MoE linear layers should use QuarkW8A8Int8
qkv_proj = layer.self_attn.qkv_proj
assert isinstance(qkv_proj.scheme, QuarkW8A8Int8)
layer = model.model.layers[0]
moe = layer.mlp.experts
assert isinstance(moe._quant_method, QuarkW8A8Int8MoEMethod), (
f"Expected QuarkW8A8Int8MoEMethod, got {type(moe._quant_method)}"
)
qkv_proj = layer.self_attn.qkv_proj
assert isinstance(qkv_proj.scheme, QuarkW8A8Int8)
llm.apply_model(check_model)
output = llm.generate_greedy("Hello", max_tokens=4)
assert output
monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
with (
set_current_vllm_config(vllm_config),
set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
):
hidden_states = model(input_ids, positions, None)
logits = model.compute_logits(hidden_states)
assert torch.isfinite(logits).all()
@pytest.mark.skipif(
not (on_gfx950() or on_gfx942()),
reason="Quark W4A8 (INT4-FP8) MoE requires the AITER kernel on gfx942/gfx950",
)
@pytest.mark.parametrize("tp", [1])
def test_quark_w4a8_fp8_moe(vllm_runner, monkeypatch, tp):
def test_quark_w4a8_fp8_moe(monkeypatch, dist_init, workspace_init):
"""Test W4A8 (INT4 weight + FP8 activation) MoE with a tiny Qwen3 MoE model.
W4A8 dispatches through the AITER fused MoE kernel, so AITER must be on.
"""
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
rocm_aiter_ops.refresh_env_variables()
model_path = "amd/tiny-qwen3-moe-w4a8"
with vllm_runner(
model, vllm_config = load_model_without_vllm_runner(
model_path,
enforce_eager=True,
tensor_parallel_size=tp,
gpu_memory_utilization=0.1,
) as llm:
)
with set_current_vllm_config(vllm_config):
moe = model.model.layers[0].mlp.experts
assert isinstance(moe._quant_method, QuarkW4A8Fp8MoEMethod), (
f"Expected QuarkW4A8Fp8MoEMethod, got {type(moe._quant_method)}"
)
def check_model(model):
moe = model.model.layers[0].mlp.experts
assert isinstance(moe._quant_method, QuarkW4A8Fp8MoEMethod), (
f"Expected QuarkW4A8Fp8MoEMethod, got {type(moe._quant_method)}"
)
llm.apply_model(check_model)
output = llm.generate_greedy("Hello", max_tokens=4)
assert output
monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
with set_forward_context(None, vllm_config, num_tokens=input_ids.numel()):
hidden_states = model(input_ids, positions, None)
logits = model.compute_logits(hidden_states)
assert torch.isfinite(logits).all()
def test_quark_fp8_parity(vllm_runner):
def test_quark_fp8_parity(dist_init, workspace_init):
quark_model_id = "amd-quark/llama-tiny-fp8-quark-quant-method"
fp8_model_id = "amd-quark/llama-tiny-fp8-quant-method"
llm_kwargs = {
"tensor_parallel_size": 1,
"enforce_eager": True,
"gpu_memory_utilization": 0.1,
}
with (
vllm_runner(quark_model_id, **llm_kwargs) as quark_handle,
vllm_runner(fp8_model_id, **llm_kwargs) as fp8_handle,
):
def load_state_dict(model_id: str) -> dict[str, torch.Tensor]:
model, _ = load_model_without_vllm_runner(model_id)
return {k: v.cpu() for k, v in model.state_dict().items()}
def get_state_dict(model):
return {k: v.cpu() for k, v in model.state_dict().items()}
(quark_state_dict,) = quark_handle.apply_model(get_state_dict)
(fp8_state_dict,) = fp8_handle.apply_model(get_state_dict)
quark_state_dict = load_state_dict(quark_model_id)
fp8_state_dict = load_state_dict(fp8_model_id)
assert fp8_state_dict.keys() == quark_state_dict.keys()
@@ -1097,104 +1130,6 @@ GSM8K_ACCURACY_CONFIGS = [
),
]
WIKITEXT_ACCURACY_CONFIGS = [
AccuracyTestConfig(
model_name="fxmarty/qwen1.5_moe_a2.7b_chat_w_fp4_a_fp6_e2m3",
excepted_value=11.3,
),
AccuracyTestConfig(
model_name="fxmarty/qwen1.5_moe_a2.7b_chat_w_fp6_e3m2_a_fp6_e3m2",
excepted_value=10.6,
),
AccuracyTestConfig(
model_name="fxmarty/qwen_1.5-moe-a2.7b-mxfp4", excepted_value=12.45
),
]
@pytest.mark.skipif(
not QUARK_MXFP4_AVAILABLE,
reason=f"amd-quark>={QUARK_MXFP4_MIN_VERSION} is not available",
)
@pytest.mark.parametrize(
"config",
[pytest.param(val, id=f"config:{val}") for val in WIKITEXT_ACCURACY_CONFIGS],
)
@pytest.mark.parametrize(
"tp_size", [pytest.param(val, id=f"tp_size:{val}") for val in [1, 2]]
)
def test_ocp_mx_wikitext_correctness(config: AccuracyTestConfig, tp_size: int):
device_count = torch.accelerator.device_count()
if device_count < tp_size:
pytest.skip(f"This test requires >={tp_size} gpus, got only {device_count}")
task = "wikitext"
rtol = 0.1
# Smaller cudagraph_capture_sizes to speed up the test.
results = lm_eval.simple_evaluate(
model="vllm",
model_args=config.get_model_args(
tp_size=tp_size, kwargs={"cudagraph_capture_sizes": [16]}
),
tasks=task,
batch_size=64,
)
EXPECTED_VALUE = config.excepted_value
measured_value = results["results"][task]["word_perplexity,none"]
assert (
measured_value < EXPECTED_VALUE + rtol
and measured_value > EXPECTED_VALUE - rtol
), f"Expected: {EXPECTED_VALUE} | Measured: {measured_value}"
@pytest.mark.skipif(
not QUARK_MXFP4_AVAILABLE,
reason=f"amd-quark>={QUARK_MXFP4_MIN_VERSION} is not available",
)
@pytest.mark.parametrize("tp_size", [1, 2])
def test_nvfp4_wikitext_correctness(tp_size: int):
device_count = torch.accelerator.device_count()
if device_count < tp_size:
pytest.skip(f"This test requires >={tp_size} gpus, got only {device_count}")
# NOTE: expected_value from nvidia/Qwen3-30B-A3B-NVFP4
expected_value = 11.2391
model_name = "amd-quark/Qwen3-30B-A3B-nvfp4-quark"
task = "wikitext"
rtol = 0.25
config = AccuracyTestConfig(
model_name=model_name,
excepted_value=expected_value,
)
model_args = config.get_model_args(
tp_size=tp_size,
kwargs={
"cudagraph_capture_sizes": [16],
},
)
model_args.pop("add_bos_token")
# Smaller cudagraph_capture_sizes to speed up the test.
results = lm_eval.simple_evaluate(
model="vllm",
model_args=model_args,
tasks=task,
batch_size=64,
)
EXPECTED_VALUE = config.excepted_value
measured_value = results["results"][task]["word_perplexity,none"]
assert (
measured_value < EXPECTED_VALUE + rtol
and measured_value > EXPECTED_VALUE - rtol
), f"Expected: {EXPECTED_VALUE} | Measured: {measured_value}"
@pytest.mark.parametrize("config", GSM8K_ACCURACY_CONFIGS)
@pytest.mark.skipif(
tests/quantization/test_register_quantization_config.py
+25 -17
View File
@@ -14,6 +14,10 @@ import pytest
import torch
import torch.nn.functional as F
from tests.quantization.utils import load_model_without_vllm_runner
from vllm.config import set_current_vllm_config
from vllm.forward_context import set_forward_context
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.linear import (
LinearBase, # noqa: E501
UnquantizedLinearMethod,
@@ -26,6 +30,7 @@ from vllm.model_executor.layers.quantization import (
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig, # noqa: E501
)
from vllm.platforms import current_platform
class FakeQuantLinearMethod(UnquantizedLinearMethod):
@@ -124,23 +129,26 @@ def test_register_quantization_config(caplog_vllm):
"meta-llama/Llama-3.2-1B-Instruct",
],
)
def test_custom_quant(vllm_runner, model, monkeypatch):
def test_custom_quant(model, monkeypatch, dist_init, workspace_init):
"""Test infer with the custom quantization method."""
# `LLM.apply_model` requires pickling a function.
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
vllm_model, vllm_config = load_model_without_vllm_runner(
model,
quantization="custom_quant",
)
layer = vllm_model.model.layers[0]
qkv_proj = layer.self_attn.qkv_proj
with vllm_runner(
model_name=model, quantization="custom_quant", enforce_eager=True
) as llm:
# Check the quantization method is FakeQuantLinearMethod.
assert isinstance(qkv_proj.quant_method, FakeQuantLinearMethod)
def check_model(model):
layer = model.model.layers[0]
qkv_proj = layer.self_attn.qkv_proj
# Check the quantization method is FakeQuantLinearMethod
assert isinstance(qkv_proj.quant_method, FakeQuantLinearMethod)
llm.apply_model(check_model)
output = llm.generate_greedy("Hello my name is", max_tokens=1)
assert output
monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
device_type = current_platform.device_type
input_ids = torch.tensor([1, 2, 3, 4], device=device_type)
positions = torch.arange(input_ids.numel(), device=device_type)
with (
set_current_vllm_config(vllm_config),
set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
):
hidden_states = vllm_model(input_ids, positions, None)
logits = vllm_model.compute_logits(hidden_states)
assert torch.isfinite(logits).all()
tests/quantization/test_torchao.py
+68 -24
View File
@@ -7,6 +7,11 @@ import pytest
import torch
from packaging import version
from tests.quantization.utils import load_model_without_vllm_runner
from vllm.config import set_current_vllm_config
from vllm.config.load import LoadConfig
from vllm.forward_context import set_forward_context
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.model_loader import get_model_loader
from vllm.platforms import current_platform
@@ -27,20 +32,39 @@ TORCHAO_VERSION_0_18_AVAILABLE = TORCHAO_AVAILABLE and version.parse(
) >= version.parse("0.18.0")
# Fix for leaking `torch.set_float32_matmul_precision("high")`
# https://github.com/pytorch/ao/blob/v0.17.0/torchao/quantization/utils.py#L696
@pytest.fixture(autouse=True)
def restore_float32_matmul_precision():
original_precision = torch.get_float32_matmul_precision()
try:
yield
finally:
torch.set_float32_matmul_precision(original_precision)
@pytest.mark.skipif(
current_platform.is_rocm() and current_platform.is_fp8_fnuz(),
reason="Only fp8_fnuz supported on CDNA3 architecture",
)
@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
def test_pre_quantized_model(vllm_runner):
with vllm_runner(
def test_pre_quantized_model(monkeypatch, dist_init, workspace_init):
model, vllm_config = load_model_without_vllm_runner(
"torchao-testing/opt-125m-Float8WeightOnlyConfig-v2-0.15.0",
quantization="torchao",
dtype="bfloat16",
enforce_eager=True,
) as llm:
output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
assert output
)
monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
with (
set_current_vllm_config(vllm_config),
set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
):
hidden_states = model(input_ids, positions, None)
logits = model.compute_logits(hidden_states)
assert torch.isfinite(logits).all()
@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
@@ -51,35 +75,55 @@ def test_pre_quantized_model(vllm_runner):
# {"": "cuda"},
],
)
def test_opt_125m_int8wo_model_loading_with_params(vllm_runner, pt_load_map_location):
def test_opt_125m_int8wo_model_loading_with_params(
pt_load_map_location, monkeypatch, dist_init, workspace_init
):
torch._dynamo.reset()
model_name = "jerryzh168/opt-125m-int8wo-partial-quant"
with vllm_runner(
model_name=model_name,
model, vllm_config = load_model_without_vllm_runner(
model_name,
quantization="torchao",
dtype="bfloat16",
pt_load_map_location=pt_load_map_location,
enforce_eager=True,
) as llm:
output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
vllm_config_kwargs={
"load_config": LoadConfig(pt_load_map_location=pt_load_map_location)
},
)
assert output
monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
with (
set_current_vllm_config(vllm_config),
set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
):
hidden_states = model(input_ids, positions, None)
logits = model.compute_logits(hidden_states)
assert torch.isfinite(logits).all()
@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
def test_qwenvl_int8wo_model_loading_with_params(vllm_runner):
def test_qwenvl_int8wo_model_loading_with_params(
monkeypatch, dist_init, workspace_init
):
torch._dynamo.reset()
model_name = "mobicham/Qwen2.5-VL-3B-Instruct_int8wo_ao"
with vllm_runner(
model_name=model_name,
model, vllm_config = load_model_without_vllm_runner(
model_name,
quantization="torchao",
dtype="bfloat16",
pt_load_map_location=f"{DEVICE_TYPE}:0",
enforce_eager=True,
) as llm:
output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
vllm_config_kwargs={
"load_config": LoadConfig(pt_load_map_location=f"{DEVICE_TYPE}:0")
},
)
assert output
monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE).expand(3, -1)
with (
set_current_vllm_config(vllm_config),
set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
):
hidden_states = model(input_ids, positions, None)
logits = model.compute_logits(hidden_states)
assert torch.isfinite(logits).all()
@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
tests/quantization/utils.py
+70
View File
@@ -2,13 +2,83 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import logging
from typing import Any
import regex as re
import torch
from vllm.config import (
CompilationConfig,
ModelConfig,
VllmConfig,
set_current_vllm_config,
)
from vllm.model_executor.layers.quantization import get_quantization_config
from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
from vllm.platforms import current_platform
def _limit_num_hidden_layers(
model: torch.nn.Module, num_hidden_layers: int | None
) -> None:
if num_hidden_layers is None:
return
original_load_weights = model.load_weights
def should_load_weight(name: str) -> bool:
for prefix in ("model.layers.", "layers."):
if name.startswith(prefix):
layer_idx = int(name.removeprefix(prefix).split(".", 1)[0])
return layer_idx < num_hidden_layers
return True
def load_weights(weights):
weights = (
(name, weight) for name, weight in weights if should_load_weight(name)
)
return original_load_weights(weights)
model.load_weights = load_weights
def load_model_without_vllm_runner(
model_path: str,
*,
dtype: str | torch.dtype = "bfloat16",
quantization: str | None = None,
model_config_kwargs: dict[str, Any] | None = None,
vllm_config_kwargs: dict[str, Any] | None = None,
model_loader_cls: type = DefaultModelLoader,
) -> tuple[torch.nn.Module, VllmConfig]:
"""Instantiate a model, load weights, and process them for inference."""
model_config = ModelConfig(
model=model_path,
dtype=dtype,
quantization=quantization,
**(model_config_kwargs or {}),
)
vllm_config_args = dict(vllm_config_kwargs or {})
vllm_config_args.setdefault("compilation_config", CompilationConfig(mode=0))
vllm_config = VllmConfig(model_config=model_config, **vllm_config_args)
hf_overrides = (model_config_kwargs or {}).get("hf_overrides") or {}
num_hidden_layers = hf_overrides.get("num_hidden_layers")
with set_current_vllm_config(vllm_config):
model_loader = model_loader_cls(vllm_config.load_config)
if num_hidden_layers is not None:
original_load_weights = model_loader.load_weights
def load_weights(model, model_config):
_limit_num_hidden_layers(model, num_hidden_layers)
original_load_weights(model, model_config)
model_loader.load_weights = load_weights
model = model_loader.load_model(vllm_config, model_config)
return model, vllm_config
def is_quant_method_supported(quant_method: str) -> bool:
# Currently, quantization tests only run GPUs
if current_platform.is_cpu():
tests/utils_/test_torch_utils.py
+14
View File
@@ -11,6 +11,7 @@ from vllm.utils.torch_utils import (
get_kv_cache_torch_dtype,
is_lossless_cast,
is_quantized_kv_cache,
set_default_torch_dtype,
set_torch_threads_for_runtime,
startup_omp_num_threads,
)
@@ -28,6 +29,19 @@ def test_nvfp4_4over6_cache_dtype() -> None:
assert get_kv_quant_mode(cache_config.cache_dtype) == KVQuantMode.NVFP4
def test_set_default_torch_dtype_restores_dtype_after_exception() -> None:
original_dtype = torch.get_default_dtype()
with (
pytest.raises(RuntimeError, match="expected failure"),
set_default_torch_dtype(torch.bfloat16),
):
assert torch.get_default_dtype() == torch.bfloat16
raise RuntimeError("expected failure")
assert torch.get_default_dtype() == original_dtype
@pytest.mark.parametrize(
("src_dtype", "tgt_dtype", "expected_result"),
[
vllm/utils/torch_utils.py
+4 -2
View File
@@ -164,8 +164,10 @@ def set_default_torch_dtype(dtype: torch.dtype):
"""Sets the default torch dtype to the given dtype."""
old_dtype = torch.get_default_dtype()
torch.set_default_dtype(dtype)
yield
torch.set_default_dtype(old_dtype)
try:
yield
finally:
torch.set_default_dtype(old_dtype)
def _cgroup_cpu_limit() -> float | None: