diff --git a/tests/evals/gsm8k/configs/Qwen-1.5-MOE-W-MXFP4-A-MXFP6.yaml b/tests/evals/gsm8k/configs/Qwen-1.5-MOE-W-MXFP4-A-MXFP6.yaml new file mode 100644 index 0000000000..a7f2b23c09 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen-1.5-MOE-W-MXFP4-A-MXFP6.yaml @@ -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" diff --git a/tests/evals/gsm8k/configs/Qwen-1.5-MOE-W-MXFP6-A-MXFP6.yaml b/tests/evals/gsm8k/configs/Qwen-1.5-MOE-W-MXFP6-A-MXFP6.yaml new file mode 100644 index 0000000000..d1671ab85a --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen-1.5-MOE-W-MXFP6-A-MXFP6.yaml @@ -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" diff --git a/tests/evals/gsm8k/configs/Qwen3-30B-A3B-NVFP4-quark.yaml b/tests/evals/gsm8k/configs/Qwen3-30B-A3B-NVFP4-quark.yaml new file mode 100644 index 0000000000..0ff5e68d28 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3-30B-A3B-NVFP4-quark.yaml @@ -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" diff --git a/tests/evals/gsm8k/configs/models-gfx950-large.txt b/tests/evals/gsm8k/configs/models-gfx950-large.txt index aa7d00bebc..9027afd638 100644 --- a/tests/evals/gsm8k/configs/models-gfx950-large.txt +++ b/tests/evals/gsm8k/configs/models-gfx950-large.txt @@ -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 \ No newline at end of file diff --git a/tests/evals/gsm8k/configs/models-mi3xx.txt b/tests/evals/gsm8k/configs/models-mi3xx.txt index e8759d7d02..71fb98e213 100644 --- a/tests/evals/gsm8k/configs/models-mi3xx.txt +++ b/tests/evals/gsm8k/configs/models-mi3xx.txt @@ -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 \ No newline at end of file +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 \ No newline at end of file diff --git a/tests/quantization/test_auto_awq.py b/tests/quantization/test_auto_awq.py index dcb2b11c8f..e1a0572f76 100644 --- a/tests/quantization/test_auto_awq.py +++ b/tests/quantization/test_auto_awq.py @@ -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(): diff --git a/tests/quantization/test_auto_gptq.py b/tests/quantization/test_auto_gptq.py index b8507022bd..e91093737e 100644 --- a/tests/quantization/test_auto_gptq.py +++ b/tests/quantization/test_auto_gptq.py @@ -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(): diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py index 304ef1efb0..004983cb19 100644 --- a/tests/quantization/test_compressed_tensors.py +++ b/tests/quantization/test_compressed_tensors.py @@ -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( diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py index 0a89d2a035..1c73573d8d 100644 --- a/tests/quantization/test_fp8.py +++ b/tests/quantization/test_fp8.py @@ -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 diff --git a/tests/quantization/test_fp8_per_channel.py b/tests/quantization/test_fp8_per_channel.py index b8ec3998f4..2c3fddcb0f 100644 --- a/tests/quantization/test_fp8_per_channel.py +++ b/tests/quantization/test_fp8_per_channel.py @@ -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]) diff --git a/tests/quantization/test_gptq_dynamic.py b/tests/quantization/test_gptq_dynamic.py index 809f31aec2..d4da7c4fd9 100644 --- a/tests/quantization/test_gptq_dynamic.py +++ b/tests/quantization/test_gptq_dynamic.py @@ -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) diff --git a/tests/quantization/test_lm_head.py b/tests/quantization/test_lm_head.py index 196949b808..a8fa8fd0d2 100644 --- a/tests/quantization/test_lm_head.py +++ b/tests/quantization/test_lm_head.py @@ -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) diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index c4b549a91a..3fc63a26ad 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -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(): diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index b918a162b7..3ba7b70458 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -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( diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index 36c1c9411c..1c1537a122 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -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( diff --git a/tests/quantization/test_register_quantization_config.py b/tests/quantization/test_register_quantization_config.py index ac7746935a..5529a65b77 100644 --- a/tests/quantization/test_register_quantization_config.py +++ b/tests/quantization/test_register_quantization_config.py @@ -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() diff --git a/tests/quantization/test_torchao.py b/tests/quantization/test_torchao.py index a724803b9a..50817ce330 100644 --- a/tests/quantization/test_torchao.py +++ b/tests/quantization/test_torchao.py @@ -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") diff --git a/tests/quantization/utils.py b/tests/quantization/utils.py index 5e34199f94..4b1b0343a8 100644 --- a/tests/quantization/utils.py +++ b/tests/quantization/utils.py @@ -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(): diff --git a/tests/utils_/test_torch_utils.py b/tests/utils_/test_torch_utils.py index 6d63fc5db4..6901a162a4 100644 --- a/tests/utils_/test_torch_utils.py +++ b/tests/utils_/test_torch_utils.py @@ -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"), [ diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 2d428e46c8..4f6be20fb4 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -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: