470 lines
16 KiB
Python
470 lines
16 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Tests online quantization."""
|
|
|
|
from types import SimpleNamespace
|
|
from typing import cast
|
|
|
|
import pytest
|
|
import torch
|
|
from torch.distributed import ProcessGroup
|
|
|
|
from tests.quantization.utils import (
|
|
_test_online_quant_peak_mem_impl,
|
|
is_quant_method_supported,
|
|
)
|
|
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.model_executor.layers.linear import UnquantizedLinearMethod
|
|
from vllm.model_executor.layers.quantization.online.fp8 import (
|
|
Fp8PerBlockOnlineLinearMethod,
|
|
Fp8PerBlockOnlineMoEMethod,
|
|
Fp8PerTensorOnlineLinearMethod,
|
|
Fp8PerTensorOnlineMoEMethod,
|
|
_fp8_channel_scale,
|
|
_fp8_quant_per_channel,
|
|
_fp8_scale,
|
|
_is_tp_sharded,
|
|
)
|
|
from vllm.model_executor.layers.quantization.online.int8 import Int8OnlineMoEMethod
|
|
from vllm.model_executor.layers.quantization.online.mxfp4 import (
|
|
Mxfp4OnlineLinearMethod,
|
|
Mxfp4OnlineMoEMethod,
|
|
)
|
|
from vllm.model_executor.layers.quantization.online.nvfp4 import (
|
|
Nvfp4OnlineMoEMethod,
|
|
_quantize_moe_weight_to_nvfp4,
|
|
)
|
|
from vllm.model_executor.layers.quantization.utils import quant_utils
|
|
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
|
amax_for_moe_weight_quant,
|
|
amax_for_tp_weight_quant,
|
|
weight_amax,
|
|
)
|
|
from vllm.platforms import current_platform
|
|
from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
|
|
|
|
if current_platform.is_rocm():
|
|
from vllm.platforms.rocm import on_gfx942, on_gfx950
|
|
else:
|
|
|
|
def on_gfx950() -> bool:
|
|
return False
|
|
|
|
def on_gfx942() -> bool:
|
|
return False
|
|
|
|
|
|
DEVICE = current_platform.device_type
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not is_quant_method_supported("fp8"),
|
|
reason="FP8 is not supported on this GPU type.",
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"quant_scheme,online_quant_args,expected_linear_cls,expected_moe_cls",
|
|
[
|
|
# simple case - quantization='fp8_per_tensor'
|
|
(
|
|
"fp8_per_tensor",
|
|
None,
|
|
Fp8PerTensorOnlineLinearMethod,
|
|
Fp8PerTensorOnlineMoEMethod,
|
|
),
|
|
# simple case - quantization='fp8_per_block'
|
|
(
|
|
"fp8_per_block",
|
|
None,
|
|
Fp8PerBlockOnlineLinearMethod,
|
|
Fp8PerBlockOnlineMoEMethod,
|
|
),
|
|
# quantization='online' with per-layer-kind overrides
|
|
(
|
|
"online",
|
|
{
|
|
"linear": "fp8_per_block",
|
|
"moe": "fp8_per_tensor",
|
|
},
|
|
Fp8PerBlockOnlineLinearMethod,
|
|
Fp8PerTensorOnlineMoEMethod,
|
|
),
|
|
# ignore with direct layer name
|
|
(
|
|
"fp8_per_tensor",
|
|
# qkv_proj is fused from q_proj/k_proj/v_proj, so currently the
|
|
# ignore regex must match the unfused shard names
|
|
# TODO(future PR): also make 're:.*qkv_proj.*' work
|
|
{"ignore": ["model.layers.1.self_attn.o_proj", "re:.*[qkv]_proj"]},
|
|
Fp8PerTensorOnlineLinearMethod,
|
|
Fp8PerTensorOnlineMoEMethod,
|
|
),
|
|
(
|
|
"mxfp4",
|
|
None,
|
|
Mxfp4OnlineLinearMethod,
|
|
Mxfp4OnlineMoEMethod,
|
|
),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"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,
|
|
) -> None:
|
|
"""
|
|
Tests that online quantization frontend configuration works -
|
|
selecting quant schemes, overriding quant schemes by type, ignoring
|
|
layers.
|
|
|
|
Does not test performance, peak memory usage, etc.
|
|
"""
|
|
|
|
# TODO: Relax this condition once there is a native MXFP4_MXFP4
|
|
# linear/moe backend supported on cuda.
|
|
if quant_scheme == "mxfp4" and not (on_gfx950() or on_gfx942()):
|
|
pytest.skip("mxfp4 online quantization is only tested on AMD gfx942, gfx950.")
|
|
|
|
if current_platform.is_rocm():
|
|
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0")
|
|
rocm_aiter_ops.refresh_env_variables()
|
|
|
|
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(
|
|
quantization=quant_scheme,
|
|
enforce_eager=True,
|
|
)
|
|
if online_quant_args is not None:
|
|
runner_kwargs["quantization_config"] = online_quant_args
|
|
|
|
with vllm_runner(
|
|
model_name,
|
|
**runner_kwargs,
|
|
) 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)
|
|
|
|
llm.apply_model(check_model)
|
|
|
|
outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4)
|
|
print(outputs[0][1])
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not (
|
|
current_platform.is_cuda()
|
|
and current_platform.is_device_capability_family(100)
|
|
and has_flashinfer_trtllm_fused_moe()
|
|
),
|
|
reason="nvfp4_per_token needs a Blackwell (SM100) GPU + FlashInfer TRTLLM MoE.",
|
|
)
|
|
def test_online_nvfp4_per_token_moe(vllm_runner, monkeypatch) -> None:
|
|
"""Online NVFP4 quantizes the MoE and leaves dense layers unquantized."""
|
|
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
|
|
|
with vllm_runner(
|
|
"ibm-granite/granite-3.0-1b-a400m-base",
|
|
quantization="nvfp4_per_token",
|
|
enforce_eager=True,
|
|
) as llm:
|
|
|
|
def check_model(model):
|
|
layer = model.model.layers[0]
|
|
assert isinstance(
|
|
layer.block_sparse_moe.experts._quant_method, Nvfp4OnlineMoEMethod
|
|
)
|
|
assert isinstance(
|
|
layer.self_attn.o_proj.quant_method, UnquantizedLinearMethod
|
|
)
|
|
|
|
llm.apply_model(check_model)
|
|
outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4)
|
|
print(outputs[0][1])
|
|
|
|
|
|
def _patch_max_reduce(monkeypatch, full_amax) -> None:
|
|
"""Stand in for the TP/EP MAX all-reduce, returning the unsharded amax."""
|
|
expected = cast(ProcessGroup, object())
|
|
stub = SimpleNamespace(device_group=expected)
|
|
monkeypatch.setattr(quant_utils, "get_tp_group", lambda: stub)
|
|
monkeypatch.setattr(quant_utils, "get_ep_group", lambda: stub)
|
|
|
|
def fake_all_reduce(tensor, op, group):
|
|
assert op == torch.distributed.ReduceOp.MAX
|
|
assert group is expected
|
|
tensor.copy_(full_amax)
|
|
|
|
monkeypatch.setattr(torch.distributed, "all_reduce", fake_all_reduce)
|
|
|
|
|
|
def test_is_tp_sharded_false_when_scale_is_already_global() -> None:
|
|
"""Replicated and column-parallel-with-channel-scales need no collective."""
|
|
replicated = SimpleNamespace(
|
|
tp_size=4,
|
|
input_size=64,
|
|
output_size=32,
|
|
input_size_per_partition=64,
|
|
output_size_per_partition=32,
|
|
)
|
|
assert not _is_tp_sharded(replicated)
|
|
|
|
column = SimpleNamespace(
|
|
tp_size=4,
|
|
input_size=64,
|
|
output_size=32,
|
|
input_size_per_partition=64,
|
|
output_size_per_partition=8,
|
|
)
|
|
assert not _is_tp_sharded(column, reduces_output_dim=False)
|
|
assert _is_tp_sharded(column)
|
|
|
|
|
|
def _quantize_linear(weight, scheme, is_sharded):
|
|
if scheme == "per_tensor":
|
|
amax = weight_amax(weight).reshape(1)
|
|
scale = _fp8_scale(amax_for_tp_weight_quant(amax, is_sharded))
|
|
return ops.scaled_fp8_quant(weight, scale=scale)[0], scale
|
|
amax = weight_amax(weight, dim=-1, keepdim=True)
|
|
scale = _fp8_channel_scale(amax_for_tp_weight_quant(amax, is_sharded))
|
|
return _fp8_quant_per_channel(weight, scale), scale
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not is_quant_method_supported("fp8"),
|
|
reason="FP8 is not supported on this GPU type.",
|
|
)
|
|
@pytest.mark.parametrize("scheme", ["per_tensor", "per_channel"])
|
|
@pytest.mark.parametrize("shard_dim", [0, 1])
|
|
def test_online_linear_tp_weight_quant_matches_unsharded(
|
|
monkeypatch, scheme: str, shard_dim: int
|
|
) -> None:
|
|
"""TP shards pack the same FP8 values and scales as the unsharded weight."""
|
|
torch.manual_seed(0)
|
|
weight = torch.randn(32, 64, device=DEVICE, dtype=torch.bfloat16)
|
|
weight[-1, -1] = 64.0
|
|
|
|
full_weight, full_scale = _quantize_linear(weight, scheme, False)
|
|
|
|
# Per-channel scales reduce only the input dim, so a column (dim 0) shard
|
|
# already matches without a collective.
|
|
is_sharded = scheme == "per_tensor" or shard_dim == 1
|
|
if is_sharded:
|
|
full_amax = (
|
|
weight_amax(weight).reshape(1)
|
|
if scheme == "per_tensor"
|
|
else weight_amax(weight, dim=-1, keepdim=True)
|
|
)
|
|
_patch_max_reduce(monkeypatch, full_amax)
|
|
|
|
shard_size = weight.shape[shard_dim] // 2
|
|
shard = weight.narrow(shard_dim, 0, shard_size).contiguous()
|
|
tp_weight, tp_scale = _quantize_linear(shard, scheme, is_sharded)
|
|
|
|
assert torch.equal(tp_weight, full_weight.narrow(shard_dim, 0, shard_size))
|
|
if scheme == "per_channel" and shard_dim == 0:
|
|
assert torch.equal(tp_scale, full_scale.narrow(0, 0, shard_size))
|
|
else:
|
|
assert torch.equal(tp_scale, full_scale)
|
|
|
|
|
|
def _quantize_moe(weight, scheme, moe_tp_size):
|
|
if scheme == "nvfp4":
|
|
return _quantize_moe_weight_to_nvfp4(weight, moe_tp_size)
|
|
if scheme == "per_tensor":
|
|
amax = weight_amax(weight.flatten(1), dim=-1)
|
|
scale = _fp8_scale(amax_for_moe_weight_quant(amax, moe_tp_size))
|
|
quant = lambda w, s: ops.scaled_fp8_quant(w, scale=s)[0] # noqa: E731
|
|
else:
|
|
amax = weight_amax(weight, dim=-1, keepdim=True)
|
|
scale = _fp8_channel_scale(amax_for_moe_weight_quant(amax, moe_tp_size))
|
|
quant = _fp8_quant_per_channel
|
|
qweight = torch.stack([quant(w, s) for w, s in zip(weight, scale)])
|
|
return qweight, scale
|
|
|
|
|
|
@pytest.mark.parametrize("scheme", ["per_tensor", "per_channel", "nvfp4"])
|
|
def test_online_moe_tp_weight_quant_matches_ep(monkeypatch, scheme: str) -> None:
|
|
"""TP shards of w2 pack the same values and scales as full experts."""
|
|
if scheme == "nvfp4":
|
|
if (
|
|
not (
|
|
current_platform.is_cuda()
|
|
and current_platform.is_device_capability_family(100)
|
|
)
|
|
or current_platform.is_xpu()
|
|
):
|
|
pytest.skip("NVFP4 weight quantization needs a Blackwell (SM100) GPU.")
|
|
elif not is_quant_method_supported("fp8"):
|
|
pytest.skip("FP8 is not supported on this GPU type.")
|
|
|
|
torch.manual_seed(0)
|
|
weight = torch.randn(2, 32, 32, device=DEVICE, dtype=torch.bfloat16)
|
|
weight[:, -1, -1] = torch.tensor([32.0, 64.0], device=DEVICE)
|
|
|
|
ep_out = _quantize_moe(weight, scheme, 1)
|
|
|
|
full_amax = (
|
|
weight_amax(weight, dim=-1, keepdim=True)
|
|
if scheme == "per_channel"
|
|
else weight_amax(weight.flatten(1), dim=-1).to(torch.float32)
|
|
)
|
|
_patch_max_reduce(monkeypatch, full_amax)
|
|
|
|
# w2 is sharded along its last (intermediate) dim.
|
|
shard_size = weight.shape[2] // 2
|
|
tp_out = _quantize_moe(weight[:, :, :shard_size], scheme, 2)
|
|
|
|
packing = 2 if scheme == "nvfp4" else 1
|
|
assert torch.equal(tp_out[0], ep_out[0][:, :, : shard_size // packing])
|
|
if scheme == "nvfp4":
|
|
assert torch.equal(tp_out[1], ep_out[1][:, :, : shard_size // 16])
|
|
assert torch.equal(tp_out[-1], ep_out[-1])
|
|
|
|
|
|
def test_online_int8_moe_w2_scale_matches_unsharded(monkeypatch) -> None:
|
|
"""Int8 MoE w2 reduces over the sharded intermediate dim."""
|
|
torch.manual_seed(0)
|
|
w13 = torch.randn(2, 16, 8, dtype=torch.bfloat16)
|
|
w2 = torch.randn(2, 8, 16, dtype=torch.bfloat16)
|
|
w2[:, -1, -1] = 64.0
|
|
|
|
def quantize(w2_in, moe_tp_size):
|
|
layer = torch.nn.Module()
|
|
layer.w13_weight = torch.nn.Parameter(w13, requires_grad=False)
|
|
layer.w2_weight = torch.nn.Parameter(w2_in, requires_grad=False)
|
|
layer.num_experts = layer.local_num_experts = w13.shape[0]
|
|
method = SimpleNamespace(moe=SimpleNamespace(tp_size=moe_tp_size))
|
|
Int8OnlineMoEMethod._quantize_weights(method, layer)
|
|
return layer.w2_weight, layer.w2_scale
|
|
|
|
full_weight, full_scale = quantize(w2, 1)
|
|
|
|
_patch_max_reduce(monkeypatch, weight_amax(w2, dim=-1))
|
|
|
|
shard_size = w2.shape[2] // 2
|
|
tp_weight, tp_scale = quantize(w2[:, :, :shard_size].contiguous(), 2)
|
|
|
|
assert torch.equal(tp_weight, full_weight[:, :, :shard_size])
|
|
assert torch.equal(tp_scale, full_scale)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not (
|
|
current_platform.is_cuda() and current_platform.is_device_capability_family(100)
|
|
),
|
|
reason="NVFP4 weight quantization needs a Blackwell (SM100) GPU.",
|
|
)
|
|
def test_online_nvfp4_quantizes_original_expert_weights() -> None:
|
|
torch.manual_seed(0)
|
|
weight = torch.randn(2, 32, 32, device="cuda", dtype=torch.bfloat16)
|
|
|
|
quantized, block_scale, global_decode_scale = _quantize_moe_weight_to_nvfp4(weight)
|
|
global_encode_scale = 1.0 / global_decode_scale
|
|
expected = [
|
|
scaled_fp4_quant(
|
|
expert_weight,
|
|
expert_scale,
|
|
is_sf_swizzled_layout=False,
|
|
)
|
|
for expert_weight, expert_scale in zip(
|
|
weight,
|
|
global_encode_scale,
|
|
strict=True,
|
|
)
|
|
]
|
|
|
|
assert torch.equal(
|
|
quantized,
|
|
torch.stack([expert_weight for expert_weight, _ in expected]),
|
|
)
|
|
assert torch.equal(
|
|
block_scale,
|
|
torch.stack([expert_scale for _, expert_scale in expected]),
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not is_quant_method_supported("fp8"),
|
|
reason="FP8 is not supported on this GPU type.",
|
|
)
|
|
def test_online_quant_peak_mem(
|
|
vllm_runner,
|
|
caplog_mp_spawn,
|
|
monkeypatch,
|
|
) -> None:
|
|
_test_online_quant_peak_mem_impl(
|
|
"fp8_per_tensor", vllm_runner, caplog_mp_spawn, monkeypatch
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not is_quant_method_supported("fp8"),
|
|
reason="FP8 is not supported on this GPU type.",
|
|
)
|
|
def test_online_quant_load_format_dummy(
|
|
vllm_runner,
|
|
monkeypatch,
|
|
caplog,
|
|
) -> None:
|
|
with vllm_runner(
|
|
"ibm-granite/granite-3.0-1b-a400m-base",
|
|
quantization="fp8_per_tensor",
|
|
enforce_eager=True,
|
|
load_format="dummy",
|
|
) as llm:
|
|
outputs = llm.generate_greedy(["The future of AI is"], max_tokens=4)
|
|
print(outputs[0][1])
|