| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Test model set-up and weight loading for llmcompressor-quantized models. |
| 4 | |
| 5 | Run `pytest tests/quantization/test_compressed_tensors.py`. |
| 6 | """ |
| 7 | |
| 8 | from unittest.mock import Mock |
| 9 | |
| 10 | import pytest |
| 11 | import torch |
| 12 | from compressed_tensors.quantization import ( |
| 13 | ActivationOrdering, |
| 14 | QuantizationArgs, |
| 15 | QuantizationStrategy, |
| 16 | QuantizationType, |
| 17 | ) |
| 18 | |
| 19 | from tests.models.utils import check_logprobs_close |
| 20 | from tests.quantization.utils import load_model_without_vllm_runner |
| 21 | from vllm.config import set_current_vllm_config |
| 22 | from vllm.config.kernel import KernelConfig |
| 23 | from vllm.forward_context import set_forward_context |
| 24 | from vllm.model_executor.kernels.linear import ( |
| 25 | Fp8BlockScaledMMLinearKernel, |
| 26 | ) |
| 27 | from vllm.model_executor.kernels.linear.scaled_mm import ( |
| 28 | MarlinFP8ScaledMMLinearKernel, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.attention import Attention |
| 31 | from vllm.model_executor.layers.fused_moe import UnquantizedFusedMoEMethod |
| 32 | from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend |
| 33 | from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 |
| 34 | CompressedTensorsConfig, |
| 35 | CompressedTensorsLinearMethod, |
| 36 | CompressedTensorsW4A4Fp4, |
| 37 | CompressedTensorsW4A4Mxfp4, |
| 38 | CompressedTensorsW4A8Fp8, |
| 39 | CompressedTensorsW8A8Fp8, |
| 40 | CompressedTensorsW8A8Int8, |
| 41 | CompressedTensorsW8A8Mxfp8, |
| 42 | CompressedTensorsW8A16Fp8, |
| 43 | CompressedTensorsWNA8O8Int, |
| 44 | CompressedTensorsWNA16, |
| 45 | ) |
| 46 | from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_fp8 import ( # noqa: E501 |
| 47 | CompressedTensorsW8A8Fp8MoEMethod, |
| 48 | ) |
| 49 | from vllm.model_executor.layers.quantization.compressed_tensors.utils import ( |
| 50 | find_matched_target, |
| 51 | ) |
| 52 | from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 |
| 53 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 54 | GroupShape, |
| 55 | QuantKey, |
| 56 | ScaleDesc, |
| 57 | ) |
| 58 | from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead |
| 59 | from vllm.platforms import current_platform |
| 60 | from vllm.scalar_type import scalar_types |
| 61 | from vllm.v1.attention.backends.fa_utils import get_flash_attn_version |
| 62 | |
| 63 | # AITER only supports per-channel-per-channel INT8 gemm |
| 64 | # and per-tensor-per-tensor INT8 GEMM. |
| 65 | # It does not support mix precision MM and mix quantization scheme. |
| 66 | ROCM_AITER_SUPPORTED_INT8_MODEL = [ |
| 67 | "neuralmagic/Llama-3.2-1B-quantized.w8a8", |
| 68 | "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2", |
| 69 | ] |
| 70 | |
| 71 | # TritonInt8ScaledMMLinearKernel only supports symmetric quantization. |
| 72 | ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL = [ |
| 73 | "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change", |
| 74 | "nm-testing/tinyllama-oneshot-w8-channel-a8-tensor", |
| 75 | "neuralmagic/Llama-3.2-1B-quantized.w8a8", |
| 76 | "nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2", |
| 77 | "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2", |
| 78 | ] |
| 79 | |
| 80 | |
| 81 | @pytest.fixture(scope="function", autouse=True) |
| 82 | def enable_pickle(monkeypatch): |
| 83 | """`LLM.apply_model` requires pickling a function.""" |
| 84 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 85 | |
| 86 | |
| 87 | def _compressed_tensors_actorder_config(actorder): |
| 88 | return { |
| 89 | "format": "pack-quantized", |
| 90 | "config_groups": { |
| 91 | "group_0": { |
| 92 | "targets": ["Linear"], |
| 93 | "weights": { |
| 94 | "num_bits": 4, |
| 95 | "type": "int", |
| 96 | "strategy": "group", |
| 97 | "group_size": 128, |
| 98 | "symmetric": True, |
| 99 | "dynamic": False, |
| 100 | "actorder": actorder, |
| 101 | }, |
| 102 | } |
| 103 | }, |
| 104 | } |
| 105 | |
| 106 | |
| 107 | @pytest.mark.parametrize( |
| 108 | "actorder", [ActivationOrdering.GROUP, ActivationOrdering.DYNAMIC, True] |
| 109 | ) |
| 110 | def test_compressed_tensors_rejects_runtime_activation_order(actorder): |
| 111 | with pytest.raises(ValueError, match="group/dynamic activation ordering"): |
| 112 | CompressedTensorsConfig.from_config( |
| 113 | _compressed_tensors_actorder_config(actorder) |
| 114 | ) |
| 115 | |
| 116 | |
| 117 | @pytest.mark.parametrize( |
| 118 | "actorder", [ActivationOrdering.STATIC, ActivationOrdering.WEIGHT, False, None] |
| 119 | ) |
| 120 | def test_compressed_tensors_accepts_preordered_weights(actorder): |
| 121 | config = CompressedTensorsConfig.from_config( |
| 122 | _compressed_tensors_actorder_config(actorder) |
| 123 | ) |
| 124 | assert config.target_scheme_map["Linear"]["weights"].actorder == ( |
| 125 | actorder.value if isinstance(actorder, ActivationOrdering) else None |
| 126 | ) |
| 127 | |
| 128 | |
| 129 | @pytest.mark.parametrize( |
| 130 | "model_args", |
| 131 | [ |
| 132 | ( |
| 133 | "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change", |
| 134 | "tensor", |
| 135 | QuantizationType.INT, |
| 136 | 2560, |
| 137 | True, |
| 138 | ), |
| 139 | ( |
| 140 | "nm-testing/asym-w8w8-int8-static-per-tensor-tiny-llama", |
| 141 | "tensor", |
| 142 | QuantizationType.INT, |
| 143 | 2560, |
| 144 | False, |
| 145 | ), |
| 146 | ], |
| 147 | ) |
| 148 | def test_compressed_tensors_w8a8_static_setup(model_args, dist_init, workspace_init): |
| 149 | model_path, strategy, quant_type, shape_0, is_symmetric = model_args |
| 150 | |
| 151 | if ( |
| 152 | current_platform.is_rocm() |
| 153 | and model_path not in ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL |
| 154 | ): |
| 155 | pytest.skip(f"Skip model {model_path} as it is not supported on ROCm.") |
| 156 | |
| 157 | model, _ = load_model_without_vllm_runner(model_path) |
| 158 | layer = model.model.layers[0] |
| 159 | |
| 160 | qkv_proj = layer.self_attn.qkv_proj |
| 161 | o_proj = layer.self_attn.o_proj |
| 162 | gate_up_proj = layer.mlp.gate_up_proj |
| 163 | down_proj = layer.mlp.down_proj |
| 164 | |
| 165 | def zp_valid(zp: torch.Tensor | None): |
| 166 | if is_symmetric: |
| 167 | return zp is None |
| 168 | |
| 169 | return zp is not None and zp.dtype is torch.int32 |
| 170 | |
| 171 | assert zp_valid(qkv_proj.input_zero_point) |
| 172 | assert zp_valid(o_proj.input_zero_point) |
| 173 | assert zp_valid(gate_up_proj.input_zero_point) |
| 174 | assert zp_valid(down_proj.input_zero_point) |
| 175 | |
| 176 | assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) |
| 177 | assert isinstance(o_proj.quant_method, CompressedTensorsLinearMethod) |
| 178 | assert isinstance(gate_up_proj.quant_method, CompressedTensorsLinearMethod) |
| 179 | assert isinstance(down_proj.quant_method, CompressedTensorsLinearMethod) |
| 180 | assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Int8) |
| 181 | |
| 182 | assert qkv_proj.scheme.strategy == strategy |
| 183 | assert qkv_proj.scheme.is_static_input_scheme |
| 184 | expected_type = torch.int8 |
| 185 | |
| 186 | assert qkv_proj.weight.dtype is expected_type |
| 187 | assert o_proj.weight.dtype is expected_type |
| 188 | assert gate_up_proj.weight.dtype is expected_type |
| 189 | |
| 190 | if qkv_proj.scheme.strategy == "tensor": |
| 191 | assert len(qkv_proj.weight_scale.shape) == 2 |
| 192 | assert qkv_proj.weight_scale.shape[0] == shape_0 |
| 193 | assert qkv_proj.weight_scale.shape[1] == 1 |
| 194 | assert qkv_proj.weight_scale.dtype is torch.float32 |
| 195 | assert qkv_proj.input_scale.dtype is torch.float32 |
| 196 | |
| 197 | |
| 198 | @pytest.mark.parametrize( |
| 199 | "model_path", |
| 200 | [ |
| 201 | "neuralmagic/Llama-3.2-1B-quantized.w8a8", |
| 202 | ], |
| 203 | ) |
| 204 | @pytest.mark.parametrize("max_tokens", [4]) |
| 205 | @pytest.mark.parametrize("num_logprobs", [10]) |
| 206 | @pytest.mark.parametrize( |
| 207 | "use_aiter", [True, False] if current_platform.is_rocm() else [False] |
| 208 | ) |
| 209 | def test_compressed_tensors_w8a8_logprobs( |
| 210 | hf_runner, |
| 211 | vllm_runner, |
| 212 | example_prompts, |
| 213 | model_path, |
| 214 | max_tokens, |
| 215 | num_logprobs, |
| 216 | use_aiter, |
| 217 | monkeypatch, |
| 218 | ): |
| 219 | if ( |
| 220 | current_platform.is_rocm() |
| 221 | and model_path not in ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL |
| 222 | ): |
| 223 | pytest.skip(f"Skip model {model_path} as it is not supported on ROCm.") |
| 224 | |
| 225 | if use_aiter: |
| 226 | if model_path not in ROCM_AITER_SUPPORTED_INT8_MODEL: |
| 227 | pytest.skip(f"Skip model {model_path} as it is not support by aiter.") |
| 228 | # this will enable VLLM_ROCM_USE_AITER_LINEAR |
| 229 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") |
| 230 | |
| 231 | dtype = "bfloat16" |
| 232 | |
| 233 | # skip language translation prompt for the static per tensor models |
| 234 | if model_path in ( |
| 235 | "nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Static-Per-Tensor-Sym", |
| 236 | "nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Static-Per-Tensor-Asym", |
| 237 | ): |
| 238 | example_prompts = example_prompts[0:-1] |
| 239 | |
| 240 | with hf_runner(model_path, dtype=dtype) as hf_model: |
| 241 | hf_outputs = hf_model.generate_greedy_logprobs_limit( |
| 242 | example_prompts, max_tokens, num_logprobs |
| 243 | ) |
| 244 | |
| 245 | with vllm_runner(model_path, dtype=dtype, enforce_eager=True) as vllm_model: |
| 246 | vllm_outputs = vllm_model.generate_greedy_logprobs( |
| 247 | example_prompts, max_tokens, num_logprobs |
| 248 | ) |
| 249 | |
| 250 | check_logprobs_close( |
| 251 | outputs_0_lst=hf_outputs, |
| 252 | outputs_1_lst=vllm_outputs, |
| 253 | name_0="hf", |
| 254 | name_1="vllm", |
| 255 | ) |
| 256 | |
| 257 | if current_platform.is_rocm(): |
| 258 | torch.accelerator.synchronize() |
| 259 | |
| 260 | |
| 261 | def test_compressed_tensors_no_enforce_eager(vllm_runner): |
| 262 | model_path = "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change" |
| 263 | with vllm_runner(model_path) as llm: |
| 264 | output = llm.generate_greedy("Hello my name is", max_tokens=4) |
| 265 | assert output |
| 266 | |
| 267 | |
| 268 | @pytest.mark.parametrize( |
| 269 | "model_args", |
| 270 | [ |
| 271 | ("nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2", "tensor"), |
| 272 | ( |
| 273 | "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2", |
| 274 | "channel", |
| 275 | ), |
| 276 | ], |
| 277 | ) |
| 278 | @pytest.mark.parametrize( |
| 279 | "use_aiter", [True, False] if current_platform.is_rocm() else [False] |
| 280 | ) |
| 281 | def test_compressed_tensors_w8a8_dynamic_per_token( |
| 282 | model_args, use_aiter, monkeypatch, dist_init, workspace_init |
| 283 | ): |
| 284 | model_path, strategy = model_args |
| 285 | |
| 286 | if ( |
| 287 | current_platform.is_rocm() |
| 288 | and model_path not in ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL |
| 289 | ): |
| 290 | pytest.skip(f"Skip model {model_path} as it is not supported on ROCm.") |
| 291 | |
| 292 | if use_aiter: |
| 293 | if model_path not in ROCM_AITER_SUPPORTED_INT8_MODEL: |
| 294 | pytest.skip(f"Skip model {model_path} as it is not support by aiter.") |
| 295 | # this will enable VLLM_ROCM_USE_AITER_LINEAR |
| 296 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") |
| 297 | |
| 298 | model, _ = load_model_without_vllm_runner(model_path, dtype=torch.float16) |
| 299 | qkv_proj = model.model.layers[0].self_attn.qkv_proj |
| 300 | assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) |
| 301 | assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Int8) |
| 302 | assert not qkv_proj.scheme.is_static_input_scheme |
| 303 | assert qkv_proj.scheme.strategy == strategy |
| 304 | assert qkv_proj.weight.dtype is torch.int8 |
| 305 | |
| 306 | |
| 307 | @pytest.mark.parametrize( |
| 308 | "wNa16_args", |
| 309 | [ |
| 310 | ( |
| 311 | "nm-testing/tinyllama-oneshot-w4a16-channel-v2", |
| 312 | "channel", |
| 313 | None, |
| 314 | 8, |
| 315 | True, |
| 316 | ), |
| 317 | ], |
| 318 | ) |
| 319 | @pytest.mark.skipif( |
| 320 | not current_platform.is_cuda(), reason="The tests are skipped on non-CUDA platform." |
| 321 | ) |
| 322 | def test_compressed_tensors_wNa16(vllm_runner, wNa16_args): |
| 323 | model, strategy, group, pack_factor, symmetric = wNa16_args |
| 324 | with vllm_runner(model, enforce_eager=True) as llm: |
| 325 | |
| 326 | def check_model(model): |
| 327 | layer = model.model.layers[0] |
| 328 | |
| 329 | qkv_proj = layer.self_attn.qkv_proj |
| 330 | assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) |
| 331 | assert isinstance(qkv_proj.scheme, CompressedTensorsWNA16) |
| 332 | |
| 333 | assert qkv_proj.scheme.strategy == strategy |
| 334 | assert qkv_proj.scheme.group_size == (-1 if group is None else group) |
| 335 | |
| 336 | assert qkv_proj.scheme.pack_factor == pack_factor |
| 337 | assert qkv_proj.scheme.symmetric == symmetric |
| 338 | |
| 339 | llm.apply_model(check_model) |
| 340 | |
| 341 | output = llm.generate_greedy("Hello my name is", max_tokens=4) |
| 342 | assert output |
| 343 | |
| 344 | |
| 345 | def test_compressed_tensors_fp8(monkeypatch, dist_init, workspace_init): |
| 346 | model_path = "nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test" |
| 347 | model, vllm_config = load_model_without_vllm_runner(model_path) |
| 348 | target_device = torch.device(current_platform.device_type) |
| 349 | |
| 350 | layer = model.model.layers[0] |
| 351 | qkv_proj = layer.self_attn.qkv_proj |
| 352 | assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) |
| 353 | assert isinstance( |
| 354 | qkv_proj.scheme, |
| 355 | (CompressedTensorsW8A8Fp8, CompressedTensorsW8A16Fp8), |
| 356 | ) |
| 357 | assert qkv_proj.input_scale.dtype is torch.float32 |
| 358 | |
| 359 | if isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8): |
| 360 | assert len(qkv_proj.input_scale.shape) == 0 |
| 361 | assert qkv_proj.weight.dtype is current_platform.fp8_dtype() |
| 362 | assert qkv_proj.weight_scale.dtype is torch.float32 |
| 363 | assert len(qkv_proj.weight_scale.shape) == 0 |
| 364 | |
| 365 | monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous()) |
| 366 | input_ids = torch.tensor([1, 2, 3, 4], device=target_device) |
| 367 | positions = torch.arange(input_ids.numel(), device=target_device) |
| 368 | with ( |
| 369 | set_current_vllm_config(vllm_config), |
| 370 | set_forward_context(None, vllm_config, num_tokens=input_ids.numel()), |
| 371 | ): |
| 372 | hidden_states = model(input_ids, positions, None) |
| 373 | logits = model.compute_logits(hidden_states) |
| 374 | assert torch.isfinite(logits).all() |
| 375 | |
| 376 | |
| 377 | def test_compressed_tensors_w8a8_fp8_moe_forwards_swiglu_params(): |
| 378 | quant_method = object.__new__(CompressedTensorsW8A8Fp8MoEMethod) |
| 379 | quant_method.input_quant = QuantizationArgs( |
| 380 | num_bits=8, |
| 381 | type=QuantizationType.FLOAT, |
| 382 | strategy=QuantizationStrategy.TOKEN, |
| 383 | dynamic=True, |
| 384 | symmetric=True, |
| 385 | ) |
| 386 | quant_method.weight_block_size = None |
| 387 | quant_method.fp8_backend = Fp8MoeBackend.TRITON |
| 388 | |
| 389 | layer = Mock() |
| 390 | layer.w13_weight_scale = torch.ones(2, 4, 1) |
| 391 | layer.w2_weight_scale = torch.ones(2, 8, 1) |
| 392 | layer.w13_input_scale = None |
| 393 | layer.w2_input_scale = None |
| 394 | layer.swiglu_alpha = 1.702 |
| 395 | layer.swiglu_beta = None |
| 396 | layer.swiglu_limit = 7.0 |
| 397 | |
| 398 | quant_config = quant_method.get_fused_moe_quant_config(layer) |
| 399 | |
| 400 | assert quant_config.gemm1_alpha == 1.702 |
| 401 | assert quant_config.gemm1_beta is None |
| 402 | assert quant_config.gemm1_clamp_limit == 7.0 |
| 403 | |
| 404 | |
| 405 | @pytest.mark.skipif( |
| 406 | not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." |
| 407 | ) |
| 408 | def test_compressed_tensors_kv_cache_fp8_per_tensor(vllm_runner): |
| 409 | model_path = "nm-testing/TinyLlama-1.1B-Chat-v1.0-kvcache-fp8-tensor" |
| 410 | with vllm_runner(model_path) as llm: |
| 411 | output = llm.generate_greedy("Hello world!", max_tokens=4) |
| 412 | assert output |
| 413 | |
| 414 | |
| 415 | @pytest.mark.skipif( |
| 416 | not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." |
| 417 | ) |
| 418 | def test_compressed_tensors_kv_cache_fp8_per_attn_head(vllm_runner): |
| 419 | model_path = "nm-testing/TinyLlama-1.1B-Chat-v1.0-kvcache-fp8-attn_head" |
| 420 | try: |
| 421 | fa_version = get_flash_attn_version() |
| 422 | except Exception: |
| 423 | pytest.skip("This test requires FlashAttention backend.") |
| 424 | if fa_version is None or fa_version < 3: |
| 425 | pytest.skip("This test requires FlashAttention version >= 3.") |
| 426 | |
| 427 | with vllm_runner(model_path, attention_config={"backend": "FLASH_ATTN"}) as llm: |
| 428 | output = llm.generate_greedy("Hello world!", max_tokens=4) |
| 429 | assert output |
| 430 | |
| 431 | |
| 432 | @pytest.mark.parametrize( |
| 433 | "args", |
| 434 | [ |
| 435 | ("nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4A16", True), |
| 436 | ("nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4", False), |
| 437 | ], |
| 438 | ) |
| 439 | def test_compressed_tensors_nvfp4(args, dist_init, workspace_init): |
| 440 | model_path, use_a16 = args |
| 441 | is_rocm_and_unsupported = ( |
| 442 | model_path == "nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4A16" |
| 443 | and current_platform.is_rocm() |
| 444 | ) |
| 445 | if is_rocm_and_unsupported: |
| 446 | with pytest.raises( |
| 447 | ValueError, |
| 448 | match="Forced NVFP4 kernel MarlinNvFp4LinearKernel is not supported", |
| 449 | ): |
| 450 | load_model_without_vllm_runner(model_path) |
| 451 | return |
| 452 | |
| 453 | try: |
| 454 | model, _ = load_model_without_vllm_runner(model_path) |
| 455 | except ValueError as exc: |
| 456 | pytest.fail(f"Unexpected NVFP4 loading failure: {exc}") |
| 457 | |
| 458 | qkv_proj = model.model.layers[0].self_attn.qkv_proj |
| 459 | assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) |
| 460 | assert isinstance(qkv_proj.scheme, CompressedTensorsW4A4Fp4) |
| 461 | assert qkv_proj.scheme.use_a16 == use_a16 |
| 462 | assert qkv_proj.scheme.group_size == 16 |
| 463 | |
| 464 | |
| 465 | @pytest.mark.skipif( |
| 466 | not current_platform.is_cuda() or not current_platform.has_device_capability(90), |
| 467 | reason="W4A8 FP8 is not yet supported on this GPU type.", |
| 468 | ) |
| 469 | @pytest.mark.parametrize( |
| 470 | "args", |
| 471 | [("czhu-cohere/TinyLlama-1.1B-Chat-v1.0-W4A8-e2e", CompressedTensorsW4A8Fp8)], |
| 472 | ) |
| 473 | def test_compressed_tensors_w4a8_fp8(vllm_runner, args): |
| 474 | model, scheme = args |
| 475 | with vllm_runner(model, enforce_eager=True) as llm: |
| 476 | |
| 477 | def check_model(model): |
| 478 | layer = model.model.layers[0] |
| 479 | |
| 480 | qkv_proj = layer.self_attn.qkv_proj |
| 481 | o_proj = layer.self_attn.o_proj |
| 482 | gate_up_proj = layer.mlp.gate_up_proj |
| 483 | down_proj = layer.mlp.down_proj |
| 484 | |
| 485 | for proj in (qkv_proj, o_proj, gate_up_proj, down_proj): |
| 486 | assert isinstance(proj.quant_method, CompressedTensorsLinearMethod) |
| 487 | assert isinstance(proj.scheme, scheme) |
| 488 | |
| 489 | assert proj.weight_packed.dtype is torch.int32 |
| 490 | assert proj.weight_scale.dtype is torch.float8_e4m3fn |
| 491 | assert proj.weight_chan_scale.dtype is torch.float32 |
| 492 | assert proj.scheme.group_size == 128 |
| 493 | |
| 494 | llm.apply_model(check_model) |
| 495 | output = llm.generate_greedy("Hello my name is", max_tokens=4) |
| 496 | print(output) |
| 497 | assert output |
| 498 | |
| 499 | |
| 500 | @pytest.mark.skipif( |
| 501 | not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." |
| 502 | ) |
| 503 | @pytest.mark.parametrize( |
| 504 | "model,prompt,exp_perplexity", |
| 505 | [ |
| 506 | ( |
| 507 | "nm-testing/Llama-3.2-1B-Instruct-spinquantR1R2R4-w4a16", |
| 508 | "Flat is better than nested.\nSparse is better than dense.", |
| 509 | 150.0, |
| 510 | ), |
| 511 | ( |
| 512 | "nm-testing/Llama-3.2-1B-Instruct-quip-w4a16", |
| 513 | "Flat is better than nested.\nSparse is better than dense.", |
| 514 | 150.0, |
| 515 | ), |
| 516 | ( |
| 517 | "nm-testing/Llama-3.2-1B-Instruct-quipv16-nvfp4", |
| 518 | "Flat is better than nested.\nSparse is better than dense.", |
| 519 | 150.0, |
| 520 | ), |
| 521 | ], |
| 522 | ) |
| 523 | def test_compressed_tensors_transforms_perplexity( |
| 524 | vllm_runner, model, prompt, exp_perplexity |
| 525 | ): |
| 526 | with vllm_runner(model, enforce_eager=True) as llm: |
| 527 | perplexity = llm.generate_prompt_perplexity([prompt])[0] |
| 528 | print(perplexity) |
| 529 | assert perplexity <= exp_perplexity |
| 530 | |
| 531 | |
| 532 | @pytest.mark.parametrize( |
| 533 | "linear_backend", ["auto", "marlin"] if current_platform.is_cuda() else ["auto"] |
| 534 | ) |
| 535 | def test_compressed_tensors_fp8_block_enabled( |
| 536 | linear_backend, dist_init, workspace_init |
| 537 | ): |
| 538 | model_path = "RedHatAI/Qwen3-0.6B-FP8-BLOCK" |
| 539 | model, _ = load_model_without_vllm_runner( |
| 540 | model_path, |
| 541 | vllm_config_kwargs={ |
| 542 | "kernel_config": KernelConfig(linear_backend=linear_backend) |
| 543 | }, |
| 544 | ) |
| 545 | fp8_dtype = current_platform.fp8_dtype() |
| 546 | layer = model.model.layers[0] |
| 547 | |
| 548 | qkv_proj = layer.self_attn.qkv_proj |
| 549 | assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) |
| 550 | assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8) |
| 551 | if linear_backend == "marlin": |
| 552 | assert isinstance(qkv_proj.scheme.fp8_linear, MarlinFP8ScaledMMLinearKernel) |
| 553 | assert qkv_proj.weight.dtype is torch.int32 |
| 554 | assert qkv_proj.weight_scale.dtype is qkv_proj.orig_dtype |
| 555 | else: |
| 556 | assert isinstance(qkv_proj.scheme.fp8_linear, Fp8BlockScaledMMLinearKernel) |
| 557 | assert qkv_proj.weight.dtype is fp8_dtype |
| 558 | assert qkv_proj.weight_scale.dtype is torch.float32 |
| 559 | assert len(qkv_proj.weight.shape) == 2 |
| 560 | assert len(qkv_proj.weight_scale.shape) == 2 |
| 561 | |
| 562 | if linear_backend == "auto": |
| 563 | input_quant_op = qkv_proj.scheme.fp8_linear.quant_fp8 |
| 564 | assert isinstance(input_quant_op, QuantFP8) |
| 565 | assert input_quant_op._forward_method in ( |
| 566 | input_quant_op.forward_cuda, |
| 567 | input_quant_op.forward_hip, |
| 568 | input_quant_op.forward_xpu, |
| 569 | ) |
| 570 | |
| 571 | |
| 572 | @pytest.mark.skipif( |
| 573 | not current_platform.is_cuda(), |
| 574 | reason="This test is not for non-CUDA platforms", |
| 575 | ) |
| 576 | def test_compressed_tensors_moe_ignore_with_model(vllm_runner): |
| 577 | """Integration test for MoE layer ignore functionality with a real model. |
| 578 | |
| 579 | This test would verify that when loading a compressed-tensors quantized |
| 580 | MoE model where some MoE layers are in the ignore list, those layers |
| 581 | use UnquantizedFusedMoEMethod while non-ignored layers use the |
| 582 | quantized method. |
| 583 | |
| 584 | Expected model structure: |
| 585 | - Compressed-tensors quantized MoE model (e.g., Mixtral-based) |
| 586 | - Config with ignore list containing specific MoE layers |
| 587 | - Multiple MoE layers where some are quantized and some are not |
| 588 | """ |
| 589 | # model_path = "nm-testing/tinysmokeqwen3moe-W4A16-first-only" # CT 12.3 |
| 590 | model_path = "nm-testing/tinysmokeqwen3moe-W4A16-first-only-CTstable" # CT 12.2 |
| 591 | |
| 592 | with vllm_runner(model_path, enforce_eager=True) as llm: |
| 593 | |
| 594 | def check_model(model): |
| 595 | from vllm.model_executor.layers.fused_moe import MoERunner |
| 596 | from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 |
| 597 | CompressedTensorsMoEMethod, |
| 598 | ) |
| 599 | |
| 600 | # Check layer 0 MoE (should be quantized) |
| 601 | layer_quantized = model.model.layers[0].mlp.experts |
| 602 | assert isinstance(layer_quantized, MoERunner) |
| 603 | assert isinstance(layer_quantized._quant_method, CompressedTensorsMoEMethod) |
| 604 | |
| 605 | # Check layer 10 MoE (should be unquantized + ignored) |
| 606 | layer_unquantized = model.model.layers[3].mlp.experts |
| 607 | assert isinstance(layer_unquantized, MoERunner) |
| 608 | assert isinstance( |
| 609 | layer_unquantized._quant_method, UnquantizedFusedMoEMethod |
| 610 | ) |
| 611 | |
| 612 | llm.apply_model(check_model) |
| 613 | |
| 614 | # Verify the model can generate output |
| 615 | output = llm.generate_greedy("Hello, my name is", max_tokens=4) |
| 616 | assert output |
| 617 | |
| 618 | |
| 619 | def test_w4a16_moe_torch_compile(vllm_runner): |
| 620 | """Regression test: MoE quant_config must be initialized inside the |
| 621 | moe_forward custom op, not just in forward_native which is compiled by |
| 622 | Dynamo (attribute mutations are not replayed at runtime). |
| 623 | |
| 624 | Without the fix in _moe_forward/_moe_forward_shared, this hits: |
| 625 | AssertionError: Hidden size mismatch 2048 != 1024 |
| 626 | because use_int4_w4a16 is False (moe_quant_config stays None). |
| 627 | """ |
| 628 | model_path = "nm-testing/tinysmokeqwen3moe-W4A16-first-only-CTstable" |
| 629 | |
| 630 | with vllm_runner( |
| 631 | model_path, |
| 632 | enforce_eager=False, |
| 633 | max_model_len=256, |
| 634 | compilation_config={ |
| 635 | "cudagraph_mode": "NONE", |
| 636 | }, |
| 637 | ) as llm: |
| 638 | output = llm.generate_greedy("Hi", max_tokens=1) |
| 639 | assert output |
| 640 | |
| 641 | |
| 642 | def _make_ct_config(*, target: str = "Linear") -> CompressedTensorsConfig: |
| 643 | """Build a minimal CompressedTensorsConfig with INT8 channel quant.""" |
| 644 | weight_quant = QuantizationArgs( |
| 645 | num_bits=8, |
| 646 | type=QuantizationType.INT, |
| 647 | strategy=QuantizationStrategy.CHANNEL, |
| 648 | symmetric=True, |
| 649 | dynamic=False, |
| 650 | ) |
| 651 | return CompressedTensorsConfig( |
| 652 | target_scheme_map={ |
| 653 | target: { |
| 654 | "weights": weight_quant, |
| 655 | "input_activations": None, |
| 656 | "format": "pack-quantized", |
| 657 | } |
| 658 | }, |
| 659 | ignore=[], |
| 660 | quant_format="pack-quantized", |
| 661 | ) |
| 662 | |
| 663 | |
| 664 | def test_get_quant_method_returns_linear_method_for_parallel_lm_head(): |
| 665 | """ParallelLMHead whose name matches a target must get a quantised method.""" |
| 666 | config = _make_ct_config(target="re:.*lm_head") |
| 667 | mock_lm_head = Mock(spec=ParallelLMHead) |
| 668 | mock_lm_head.__class__ = ParallelLMHead |
| 669 | |
| 670 | method = config.get_quant_method(mock_lm_head, prefix="model.lm_head") |
| 671 | |
| 672 | assert isinstance(method, CompressedTensorsLinearMethod), ( |
| 673 | f"Expected CompressedTensorsLinearMethod, got {type(method).__name__}" |
| 674 | ) |
| 675 | |
| 676 | |
| 677 | def test_get_quant_method_returns_none_for_ignored_parallel_lm_head(): |
| 678 | """ParallelLMHead on the ignore list should be left unquantized (None).""" |
| 679 | config = _make_ct_config(target="re:.*lm_head") |
| 680 | config.ignore = ["re:.*lm_head"] |
| 681 | mock_lm_head = Mock(spec=ParallelLMHead) |
| 682 | mock_lm_head.__class__ = ParallelLMHead |
| 683 | |
| 684 | method = config.get_quant_method(mock_lm_head, prefix="model.lm_head") |
| 685 | |
| 686 | assert method is None, ( |
| 687 | f"Expected None for ignored ParallelLMHead, got {type(method).__name__}" |
| 688 | ) |
| 689 | |
| 690 | |
| 691 | def test_get_quant_method_returns_none_for_unmatched_parallel_lm_head(): |
| 692 | """ParallelLMHead with target='Linear' (typical real model) must not crash. |
| 693 | |
| 694 | Most compressed-tensors models only target 'Linear'. ParallelLMHead does |
| 695 | not match that target, so get_quant_method should return None (unquantized) |
| 696 | instead of raising ValueError. |
| 697 | """ |
| 698 | config = _make_ct_config(target="Linear") |
| 699 | mock_lm_head = Mock(spec=ParallelLMHead) |
| 700 | mock_lm_head.__class__ = ParallelLMHead |
| 701 | |
| 702 | method = config.get_quant_method(mock_lm_head, prefix="model.lm_head") |
| 703 | |
| 704 | assert method is None, ( |
| 705 | f"Expected None for unmatched ParallelLMHead, got {type(method).__name__}" |
| 706 | ) |
| 707 | |
| 708 | |
| 709 | def test_find_matched_target_returns_none_on_no_match(): |
| 710 | result = find_matched_target( |
| 711 | layer_name="model.layers.0.self_attn.qkv_proj", |
| 712 | module=Mock(spec=torch.nn.Linear), |
| 713 | targets=["no_match_target"], |
| 714 | ) |
| 715 | assert result is None |
| 716 | |
| 717 | |
| 718 | def test_get_scheme_dict_returns_none_on_no_match(): |
| 719 | config = _make_ct_config(target="matched_layer") |
| 720 | result = config.get_scheme_dict( |
| 721 | layer=Mock(spec=torch.nn.Linear), |
| 722 | layer_name="model.layers.0.unmatched_layer", |
| 723 | ) |
| 724 | assert result is None |
| 725 | |
| 726 | |
| 727 | # Test constants for activation quantization |
| 728 | _STATIC_SYM_INT8_ACT = QuantizationArgs( |
| 729 | num_bits=8, |
| 730 | type=QuantizationType.INT, |
| 731 | strategy=QuantizationStrategy.TENSOR.value, |
| 732 | symmetric=True, |
| 733 | dynamic=False, |
| 734 | ) |
| 735 | |
| 736 | _STATIC_ASYM_INT8_ACT = QuantizationArgs( |
| 737 | num_bits=8, |
| 738 | type=QuantizationType.INT, |
| 739 | strategy=QuantizationStrategy.TENSOR.value, |
| 740 | symmetric=False, |
| 741 | dynamic=False, |
| 742 | ) |
| 743 | |
| 744 | _DYNAMIC_INT8_ACT = QuantizationArgs( |
| 745 | num_bits=8, |
| 746 | type=QuantizationType.INT, |
| 747 | strategy=QuantizationStrategy.TOKEN.value, |
| 748 | symmetric=True, |
| 749 | dynamic=True, |
| 750 | ) |
| 751 | |
| 752 | |
| 753 | @pytest.mark.parametrize( |
| 754 | "weight_bits,weight_strategy,input_act,output_act,format,expected_scheme", |
| 755 | [ |
| 756 | # W8A8 int-quantized -> W8A8Int8 (regression test for #46389) |
| 757 | pytest.param( |
| 758 | 8, |
| 759 | QuantizationStrategy.CHANNEL.value, |
| 760 | _STATIC_SYM_INT8_ACT, |
| 761 | None, |
| 762 | "int-quantized", |
| 763 | CompressedTensorsW8A8Int8, |
| 764 | id="w8a8_channel_static_sym", |
| 765 | ), |
| 766 | pytest.param( |
| 767 | 8, |
| 768 | QuantizationStrategy.CHANNEL.value, |
| 769 | _STATIC_ASYM_INT8_ACT, |
| 770 | None, |
| 771 | "int-quantized", |
| 772 | CompressedTensorsW8A8Int8, |
| 773 | id="w8a8_channel_static_asym", |
| 774 | ), |
| 775 | pytest.param( |
| 776 | 8, |
| 777 | QuantizationStrategy.TENSOR.value, |
| 778 | _STATIC_SYM_INT8_ACT, |
| 779 | None, |
| 780 | "int-quantized", |
| 781 | CompressedTensorsW8A8Int8, |
| 782 | id="w8a8_tensor_static", |
| 783 | ), |
| 784 | pytest.param( |
| 785 | 8, |
| 786 | QuantizationStrategy.CHANNEL.value, |
| 787 | _DYNAMIC_INT8_ACT, |
| 788 | None, |
| 789 | "int-quantized", |
| 790 | CompressedTensorsW8A8Int8, |
| 791 | id="w8a8_channel_dynamic", |
| 792 | ), |
| 793 | # W8A8O8 int-quantized -> WNA8O8Int (both input and output) |
| 794 | pytest.param( |
| 795 | 8, |
| 796 | QuantizationStrategy.CHANNEL.value, |
| 797 | _STATIC_SYM_INT8_ACT, |
| 798 | _STATIC_SYM_INT8_ACT, |
| 799 | "int-quantized", |
| 800 | CompressedTensorsWNA8O8Int, |
| 801 | id="w8a8o8_channel", |
| 802 | ), |
| 803 | pytest.param( |
| 804 | 4, |
| 805 | QuantizationStrategy.GROUP.value, |
| 806 | _STATIC_SYM_INT8_ACT, |
| 807 | _STATIC_SYM_INT8_ACT, |
| 808 | "int-quantized", |
| 809 | CompressedTensorsWNA8O8Int, |
| 810 | id="w4a8o8_group", |
| 811 | ), |
| 812 | # Weight-only pack-quantized -> WNA16 |
| 813 | pytest.param( |
| 814 | 8, |
| 815 | QuantizationStrategy.CHANNEL.value, |
| 816 | None, |
| 817 | None, |
| 818 | "pack-quantized", |
| 819 | CompressedTensorsWNA16, |
| 820 | id="w8_pack", |
| 821 | ), |
| 822 | pytest.param( |
| 823 | 4, |
| 824 | QuantizationStrategy.GROUP.value, |
| 825 | None, |
| 826 | None, |
| 827 | "pack-quantized", |
| 828 | CompressedTensorsWNA16, |
| 829 | id="w4_pack", |
| 830 | ), |
| 831 | pytest.param( |
| 832 | 2, |
| 833 | QuantizationStrategy.GROUP.value, |
| 834 | None, |
| 835 | None, |
| 836 | "pack-quantized", |
| 837 | CompressedTensorsWNA16, |
| 838 | id="w2_pack", |
| 839 | ), |
| 840 | pytest.param( |
| 841 | 3, |
| 842 | QuantizationStrategy.GROUP.value, |
| 843 | None, |
| 844 | None, |
| 845 | "pack-quantized", |
| 846 | CompressedTensorsWNA16, |
| 847 | id="w3_pack", |
| 848 | ), |
| 849 | pytest.param( |
| 850 | 5, |
| 851 | QuantizationStrategy.GROUP.value, |
| 852 | None, |
| 853 | None, |
| 854 | "pack-quantized", |
| 855 | CompressedTensorsWNA16, |
| 856 | id="w5_pack", |
| 857 | ), |
| 858 | pytest.param( |
| 859 | 6, |
| 860 | QuantizationStrategy.GROUP.value, |
| 861 | None, |
| 862 | None, |
| 863 | "pack-quantized", |
| 864 | CompressedTensorsWNA16, |
| 865 | id="w6_pack", |
| 866 | ), |
| 867 | pytest.param( |
| 868 | 7, |
| 869 | QuantizationStrategy.GROUP.value, |
| 870 | None, |
| 871 | None, |
| 872 | "pack-quantized", |
| 873 | CompressedTensorsWNA16, |
| 874 | id="w7_pack", |
| 875 | ), |
| 876 | ], |
| 877 | ) |
| 878 | def test_scheme_selection( |
| 879 | weight_bits, weight_strategy, input_act, output_act, format, expected_scheme |
| 880 | ): |
| 881 | """Test that _get_scheme_from_parts selects the correct scheme. |
| 882 | |
| 883 | This parametrized test verifies scheme selection for various combinations |
| 884 | of weight bits, quantization strategies, input/output activations, and |
| 885 | compression formats. |
| 886 | |
| 887 | Key regression test: W8A8 int-quantized models with channel-wise weights |
| 888 | should use W8A8Int8 (true int8 gemm), not WNA8O8Int (fake-quant). |
| 889 | WNA8O8Int should only match when BOTH input and output activations are |
| 890 | present. |
| 891 | """ |
| 892 | weight_quant = QuantizationArgs( |
| 893 | num_bits=weight_bits, |
| 894 | type=QuantizationType.INT, |
| 895 | strategy=weight_strategy, |
| 896 | symmetric=True, |
| 897 | dynamic=False, |
| 898 | group_size=128 if weight_strategy == QuantizationStrategy.GROUP.value else None, |
| 899 | ) |
| 900 | |
| 901 | config = CompressedTensorsConfig( |
| 902 | target_scheme_map={}, |
| 903 | ignore=[], |
| 904 | quant_format=format, |
| 905 | ) |
| 906 | |
| 907 | scheme = config._get_scheme_from_parts( |
| 908 | weight_quant=weight_quant, |
| 909 | input_quant=input_act, |
| 910 | output_quant=output_act, |
| 911 | format=format, |
| 912 | ) |
| 913 | |
| 914 | assert isinstance(scheme, expected_scheme), ( |
| 915 | f"Expected {expected_scheme.__name__} for " |
| 916 | f"W{weight_bits} {weight_strategy} + " |
| 917 | f"input_act={input_act} + output_act={output_act} + " |
| 918 | f"format={format}, got {type(scheme).__name__}" |
| 919 | ) |
| 920 | |
| 921 | |
| 922 | @pytest.mark.skipif( |
| 923 | not current_platform.is_cuda() or not current_platform.has_device_capability(75), |
| 924 | reason="MXFP8 requires Turing (sm_75+) or newer.", |
| 925 | ) |
| 926 | def test_compressed_tensors_mxfp8_moe_setup(vllm_runner): |
| 927 | """Verify MXFP8 scheme, dtypes, and generation for a MoE model.""" |
| 928 | model_path = "AliEdalati97/Qwen3-30B-A3B-MXFP8" |
| 929 | with vllm_runner( |
| 930 | model_path, |
| 931 | enforce_eager=True, |
| 932 | load_format="dummy", |
| 933 | hf_overrides={"num_hidden_layers": 4}, |
| 934 | ) as llm: |
| 935 | |
| 936 | def check_model(model): |
| 937 | from vllm.model_executor.layers.fused_moe import MoERunner |
| 938 | from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_mxfp8 import ( # noqa: E501 |
| 939 | CompressedTensorsW8A8Mxfp8MoEMethod, |
| 940 | ) |
| 941 | |
| 942 | layer = model.model.layers[0] |
| 943 | |
| 944 | qkv = layer.self_attn.qkv_proj |
| 945 | assert isinstance(qkv.quant_method, CompressedTensorsLinearMethod) |
| 946 | assert isinstance(qkv.scheme, CompressedTensorsW8A8Mxfp8) |
| 947 | |
| 948 | experts = layer.mlp.experts |
| 949 | assert isinstance(experts, MoERunner) |
| 950 | assert isinstance( |
| 951 | experts._quant_method, CompressedTensorsW8A8Mxfp8MoEMethod |
| 952 | ) |
| 953 | |
| 954 | llm.apply_model(check_model) |
| 955 | output = llm.generate_greedy("Hello my name is", max_tokens=4) |
| 956 | assert output |
| 957 | |
| 958 | |
| 959 | @pytest.mark.parametrize("num_bits", range(2, 9)) |
| 960 | def test_humming_supports_compressed_tensors_wna16_quant_key(num_bits): |
| 961 | from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( |
| 962 | HummingExpertsBase, |
| 963 | ) |
| 964 | from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa: E501 |
| 965 | WNA16_SUPPORTED_TYPES_MAP, |
| 966 | ) |
| 967 | |
| 968 | weight_key = QuantKey( |
| 969 | dtype=WNA16_SUPPORTED_TYPES_MAP[num_bits], |
| 970 | scale=ScaleDesc( |
| 971 | dtype=torch.float16, |
| 972 | static=True, |
| 973 | group_shape=GroupShape(row=1, col=128), |
| 974 | ), |
| 975 | symmetric=True, |
| 976 | ) |
| 977 | |
| 978 | assert HummingExpertsBase._supports_quant_scheme(weight_key, None) |
| 979 | |
| 980 | |
| 981 | def test_quant_key_str_supports_scalar_type_dtypes(): |
| 982 | quant_key = QuantKey( |
| 983 | dtype=scalar_types.uint2b2, |
| 984 | scale=ScaleDesc( |
| 985 | dtype=torch.float16, |
| 986 | static=True, |
| 987 | group_shape=GroupShape(row=1, col=128), |
| 988 | ), |
| 989 | symmetric=True, |
| 990 | ) |
| 991 | |
| 992 | assert str(quant_key) == ( |
| 993 | "QuantKey(uint2b2,scale(f16,static,GroupShape(row=1, col=128)),symmetric)" |
| 994 | ) |
| 995 | |
| 996 | |
| 997 | @pytest.mark.skipif( |
| 998 | not current_platform.is_cuda() or not current_platform.has_device_capability(80), |
| 999 | reason="MXFP4 requires ampere or newer", |
| 1000 | ) |
| 1001 | def test_compressed_tensors_mxfp4(vllm_runner): |
| 1002 | model_path = "nm-testing/TinyLlama-1.1B-Chat-v1.0-MXFP4" |
| 1003 | with vllm_runner(model_path, enforce_eager=True) as llm: |
| 1004 | |
| 1005 | def check_model(model): |
| 1006 | layer = model.model.layers[0] |
| 1007 | |
| 1008 | qkv_proj = layer.self_attn.qkv_proj |
| 1009 | o_proj = layer.self_attn.o_proj |
| 1010 | gate_up_proj = layer.mlp.gate_up_proj |
| 1011 | down_proj = layer.mlp.down_proj |
| 1012 | |
| 1013 | for proj in (qkv_proj, o_proj, gate_up_proj, down_proj): |
| 1014 | assert isinstance(proj.quant_method, CompressedTensorsLinearMethod) |
| 1015 | assert isinstance(proj.scheme, CompressedTensorsW4A4Mxfp4) |
| 1016 | |
| 1017 | # Verify group size |
| 1018 | assert proj.scheme.group_size == 32 |
| 1019 | |
| 1020 | llm.apply_model(check_model) |
| 1021 | output = llm.generate_greedy("Hello my name is", max_tokens=4) |
| 1022 | assert output |