wehub-resource-sync
80 lines
2.6 KiB
Python
80 lines
2.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""E2E tests for online FP8 per-channel quantization.
|
|
|
|
Loads a BF16 model with ``--quantization fp8_per_channel`` (online
|
|
quantization) and compares log-probabilities against the same model served in
|
|
BF16 without quantization. This exercises the full pipeline: config parsing,
|
|
``Fp8PtpcOnlineLinearMethod``, ``Fp8PtpcOnlineMoEMethod``, weight
|
|
loading, online quantization / shuffling, and inference.
|
|
|
|
``example_prompts`` is a pytest fixture (from conftest.py) that loads 8
|
|
diverse prompts from ``tests/prompts/example.txt``.
|
|
"""
|
|
|
|
import pytest
|
|
|
|
from tests.quantization.utils import is_quant_method_supported
|
|
|
|
from ..utils import check_logprobs_close
|
|
|
|
# Small MoE model that fits on a single GPU and exercises both linear + MoE.
|
|
MOE_MODEL = "allenai/OLMoE-1B-7B-0125-Instruct"
|
|
# Small dense model (no MoE) to validate the linear-only path.
|
|
DENSE_MODEL = "Qwen/Qwen3-0.6B"
|
|
|
|
MAX_MODEL_LEN = 1024
|
|
MAX_TOKENS = 4
|
|
NUM_LOG_PROBS = 8
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not is_quant_method_supported("fp8"),
|
|
reason="fp8 is not supported on this GPU type.",
|
|
)
|
|
@pytest.mark.quant_model
|
|
@pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"])
|
|
def test_fp8_per_channel_logprobs(
|
|
vllm_runner,
|
|
example_prompts,
|
|
model: str,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Compare BF16 baseline logprobs against online per-channel-quantized
|
|
model.
|
|
|
|
Runs the same model twice -- once in BF16 (baseline) and once with online
|
|
FP8 per-channel quantization -- then checks that the top log-probabilities
|
|
are close. Only 4 tokens are generated to keep the test fast while still
|
|
catching numerical divergence beyond expected per-channel error.
|
|
"""
|
|
with monkeypatch.context() as m:
|
|
m.setenv("TOKENIZERS_PARALLELISM", "true")
|
|
|
|
with vllm_runner(
|
|
model,
|
|
max_model_len=MAX_MODEL_LEN,
|
|
enforce_eager=True,
|
|
) as vllm_model:
|
|
baseline_outputs = vllm_model.generate_greedy_logprobs(
|
|
example_prompts, MAX_TOKENS, NUM_LOG_PROBS
|
|
)
|
|
|
|
with vllm_runner(
|
|
model,
|
|
max_model_len=MAX_MODEL_LEN,
|
|
enforce_eager=True,
|
|
quantization="fp8_per_channel",
|
|
) as vllm_model:
|
|
test_outputs = vllm_model.generate_greedy_logprobs(
|
|
example_prompts, MAX_TOKENS, NUM_LOG_PROBS
|
|
)
|
|
|
|
check_logprobs_close(
|
|
outputs_0_lst=baseline_outputs,
|
|
outputs_1_lst=test_outputs,
|
|
name_0="bf16",
|
|
name_1="fp8_per_channel",
|
|
)
|