| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests whether gptq models with quantized lm_head can be loaded. |
| 4 | |
| 5 | Run `pytest tests/quantization/test_quant_lm_head_true.py --forked`. |
| 6 | """ |
| 7 | |
| 8 | from types import SimpleNamespace |
| 9 | |
| 10 | import pytest |
| 11 | import torch |
| 12 | from compressed_tensors.quantization import preset_name_to_scheme |
| 13 | |
| 14 | from tests.quantization.utils import load_model_without_vllm_runner |
| 15 | from vllm.model_executor.layers.linear import ColumnParallelLinear |
| 16 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 17 | from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQLinearMethod |
| 18 | from vllm.model_executor.layers.quantization.base_config import QuantizationConfig |
| 19 | from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 |
| 20 | CompressedTensorsConfig, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config |
| 23 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 24 | ParallelLMHead, |
| 25 | UnquantizedEmbeddingMethod, |
| 26 | ) |
| 27 | from vllm.platforms import current_platform |
| 28 | from vllm.utils.import_utils import has_humming |
| 29 | |
| 30 | PROMPT = "On the surface of Mars, we found" |
| 31 | |
| 32 | MODELS_QUANT = [ |
| 33 | ("LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit", False), |
| 34 | ] |
| 35 | |
| 36 | |
| 37 | @pytest.mark.parametrize("model_id, lm_head_quantized", MODELS_QUANT) |
| 38 | def test_lm_head( |
| 39 | model_id: str, |
| 40 | lm_head_quantized: bool, |
| 41 | monkeypatch, |
| 42 | dist_init, |
| 43 | workspace_init, |
| 44 | ) -> None: |
| 45 | # `LLM.apply_model` requires pickling a function. |
| 46 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 47 | model, _ = load_model_without_vllm_runner( |
| 48 | model_id, |
| 49 | dtype=torch.float16, |
| 50 | model_config_kwargs={ |
| 51 | "max_model_len": 2048, |
| 52 | "hf_overrides": {"num_hidden_layers": 3}, |
| 53 | }, |
| 54 | ) |
| 55 | lm_head_layer = model.lm_head |
| 56 | if lm_head_quantized: |
| 57 | assert isinstance(lm_head_layer.quant_method, AutoGPTQLinearMethod) |
| 58 | else: |
| 59 | assert isinstance(lm_head_layer.quant_method, UnquantizedEmbeddingMethod) |
| 60 | |
| 61 | |
| 62 | @pytest.mark.skipif( |
| 63 | not current_platform.is_cuda() or not has_humming(), reason="requires Humming/CUDA" |
| 64 | ) |
| 65 | @pytest.mark.parametrize( |
| 66 | "preset,quant_format", |
| 67 | [ |
| 68 | ("FP8_DYNAMIC", "float-quantized"), |
| 69 | ("NVFP4A16", "nvfp4-pack-quantized"), |
| 70 | ("W4A16", "pack-quantized"), |
| 71 | ("W4A16_NVFP4", None), |
| 72 | ], |
| 73 | ) |
| 74 | @pytest.mark.parametrize("bias", [False, True]) |
| 75 | @torch.inference_mode() |
| 76 | def test_quantized_lm_head_matches_linear( |
| 77 | dist_init, default_vllm_config, preset, quant_format, bias |
| 78 | ): |
| 79 | """An explicitly quantized head must produce the same logits as a linear layer.""" |
| 80 | default_vllm_config.model_config = SimpleNamespace( |
| 81 | dtype=torch.bfloat16, head_dtype=None |
| 82 | ) |
| 83 | default_vllm_config.kernel_config.linear_backend = "humming" |
| 84 | quant_config: QuantizationConfig |
| 85 | if quant_format is None: |
| 86 | quant_config = ModelOptNvFp4Config( |
| 87 | quant_method=preset, |
| 88 | is_checkpoint_nvfp4_serialized=True, |
| 89 | kv_cache_quant_algo=None, |
| 90 | exclude_modules=[], |
| 91 | ) |
| 92 | else: |
| 93 | scheme = preset_name_to_scheme(preset, targets=["Linear", "lm_head"]) |
| 94 | quant_config = CompressedTensorsConfig.from_config( |
| 95 | {"config_groups": {"group_0": scheme.model_dump()}, "format": quant_format} |
| 96 | ) |
| 97 | with torch.device("cuda"): |
| 98 | kwargs = dict( |
| 99 | bias=bias, |
| 100 | params_dtype=torch.bfloat16, |
| 101 | quant_config=quant_config, |
| 102 | disable_tp=True, |
| 103 | ) |
| 104 | head = ParallelLMHead(500, 256, prefix="lm_head", **kwargs) |
| 105 | linear = ColumnParallelLinear(256, 512, prefix="proj", **kwargs) |
| 106 | for name, param in linear.named_parameters(): |
| 107 | if "scale" in name: |
| 108 | param.fill_(1.0) |
| 109 | elif name == "weight_shape": |
| 110 | param.copy_(torch.tensor([512, 256])) |
| 111 | elif param.is_floating_point(): |
| 112 | param.copy_(torch.randn(param.shape, dtype=torch.float32)) |
| 113 | else: |
| 114 | param.random_(0, 127) |
| 115 | head.load_state_dict(linear.state_dict()) |
| 116 | for layer in (head, linear): |
| 117 | layer.quant_method.process_weights_after_loading(layer) |
| 118 | x = torch.randn(8, 256, dtype=torch.bfloat16) |
| 119 | expected, _ = linear(x) |
| 120 | actual = LogitsProcessor(500)(head, x, head.bias) |
| 121 | assert torch.isfinite(actual).all() |
| 122 | torch.testing.assert_close(actual, expected[:, :500], rtol=0, atol=0) |