diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index 6a08110774..bf1554ad72 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -2,13 +2,18 @@ # 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._custom_ops import scaled_fp4_quant from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm.model_executor.layers.quantization.online.fp8 import ( @@ -16,11 +21,22 @@ from vllm.model_executor.layers.quantization.online.fp8 import ( 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.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 @@ -186,6 +202,169 @@ def test_online_nvfp4_per_token_moe(vllm_runner, monkeypatch) -> None: 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="cuda", 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) + ): + 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="cuda", dtype=torch.bfloat16) + weight[:, -1, -1] = torch.tensor([32.0, 64.0], device="cuda") + + 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) diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index 1ec6836059..3e98cfef5c 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -34,13 +34,17 @@ from vllm.model_executor.layers.quantization.online.moe_base import ( ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, + amax_for_moe_weight_quant, + amax_for_tp_weight_quant, create_fp8_quant_key, + get_fp8_min_max, kFp8Dynamic128Sym, kFp8DynamicTensorSym, kFp8DynamicTokenSym, kFp8Static128BlockSym, kFp8StaticChannelSym, kFp8StaticTensorSym, + weight_amax, ) from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( cutlass_fp8_supported, @@ -59,6 +63,54 @@ from vllm.utils.math_utils import round_up # --------------------------------------------------------------------------- +def _fp8_max(device: torch.device) -> torch.Tensor: + # A 0-d tensor divisor keeps the true division the quant kernels do; + # dividing by a Python float lowers to multiply-by-reciprocal instead. + return torch.tensor(get_fp8_min_max()[1], device=device, dtype=torch.float32) + + +def _fp8_scale(amax: torch.Tensor) -> torch.Tensor: + amax = amax.to(torch.float32) + return amax / _fp8_max(amax.device) + + +def _fp8_channel_scale(amax: torch.Tensor) -> torch.Tensor: + fp8_max = _fp8_max(amax.device) + return _fp8_scale(amax).clamp_min(1.0 / (fp8_max * 512)) + + +# Chunk so the fp32 intermediate from the divide stays under ~64 MB, regardless +# of how wide the weight is. +_QUANT_CHUNK_ELEMS = 16 * 1024 * 1024 + + +def _fp8_quant_per_channel(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + # True division, matching the dynamic per-token kernel; the static kernel + # multiplies by the reciprocal instead and can round a value that sits on a + # representable midpoint the other way. + fp8_min, fp8_max = get_fp8_min_max() + out = torch.empty_like(weight, dtype=current_platform.fp8_dtype()) + chunk_rows = max(1, _QUANT_CHUNK_ELEMS // weight.shape[-1]) + for start in range(0, weight.shape[0], chunk_rows): + rows = slice(start, start + chunk_rows) + out[rows] = (weight[rows] / scale[rows]).clamp_(fp8_min, fp8_max).to(out.dtype) + return out + + +def _is_tp_sharded(layer: Module, *, reduces_output_dim: bool = True) -> bool: + """Whether the weight is sharded along a dim its ``amax`` reduces over. + + Row parallel shards the input dim, column parallel the output dim. + Per-output-channel scales only reduce the input dim, so they are already + unsharded-equivalent on column-parallel layers. + """ + if layer.tp_size == 1: + return False + is_row_parallel = layer.input_size != layer.input_size_per_partition + is_column_parallel = layer.output_size != layer.output_size_per_partition + return is_row_parallel or (reduces_output_dim and is_column_parallel) + + class _Fp8OnlineLinearBase(LinearMethodBase): """Shared base for online FP8 linear methods. Loads fp16/bf16 checkpoint weights onto meta device and materializes them just-in-time.""" @@ -156,7 +208,10 @@ class Fp8PerTensorOnlineLinearMethod(_Fp8OnlineLinearBase): return layer.input_scale = None - qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None) + amax = weight_amax(layer.weight).reshape(1) + amax = amax_for_tp_weight_quant(amax, _is_tp_sharded(layer)) + weight_scale = _fp8_scale(amax) + qweight, _ = ops.scaled_fp8_quant(layer.weight, scale=weight_scale) # Update layer with new values. replace_parameter(layer, "weight", qweight.t().data) @@ -335,9 +390,12 @@ class Fp8PtpcOnlineLinearMethod(_Fp8OnlineLinearBase): return layer.input_scale = None - qweight, weight_scale = ops.scaled_fp8_quant( - layer.weight, scale=None, use_per_token_if_dynamic=True + amax = weight_amax(layer.weight, dim=-1, keepdim=True) + amax = amax_for_tp_weight_quant( + amax, _is_tp_sharded(layer, reduces_output_dim=False) ) + weight_scale = _fp8_channel_scale(amax) + qweight = _fp8_quant_per_channel(layer.weight, weight_scale) replace_parameter(layer, "weight", qweight.t()) replace_parameter(layer, "weight_scale", weight_scale) @@ -514,19 +572,23 @@ class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase): fp8_dtype = current_platform.fp8_dtype() w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype) w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype) - w13_scale = torch.ones( - layer.num_experts, device=w13.device, dtype=torch.float32 - ) - w2_scale = torch.ones(layer.num_experts, device=w2.device, dtype=torch.float32) layer.w13_input_scale = None layer.w2_input_scale = None + moe_tp_size = self.moe.tp_size + w13_amax = weight_amax(layer.w13_weight.flatten(1), dim=-1) + w13_amax = amax_for_moe_weight_quant(w13_amax, moe_tp_size) + w13_scale = _fp8_scale(w13_amax) + w2_amax = weight_amax(layer.w2_weight.flatten(1), dim=-1) + w2_amax = amax_for_moe_weight_quant(w2_amax, moe_tp_size) + w2_scale = _fp8_scale(w2_amax) + for expert in range(layer.local_num_experts): - w13[expert, :, :], w13_scale[expert] = ops.scaled_fp8_quant( - layer.w13_weight[expert, :, :] + w13[expert, :, :], _ = ops.scaled_fp8_quant( + layer.w13_weight[expert, :, :], scale=w13_scale[expert] ) - w2[expert, :, :], w2_scale[expert] = ops.scaled_fp8_quant( - layer.w2_weight[expert, :, :] + w2[expert, :, :], _ = ops.scaled_fp8_quant( + layer.w2_weight[expert, :, :], scale=w2_scale[expert] ) # Shuffle weights to runtime format and setup kernel. @@ -723,26 +785,24 @@ class Fp8PtpcOnlineMoEMethod(_Fp8OnlineMoEBase): # construction, so it cannot drift from the weight's expert count # under EP / padded MoE. n_w13 = layer.w13_weight.shape[1] - n_w2 = layer.w2_weight.shape[1] w13_scale = torch.ones( w13.shape[0], n_w13, 1, device=w13.device, dtype=torch.float32 ) - w2_scale = torch.ones( - w2.shape[0], n_w2, 1, device=w2.device, dtype=torch.float32 - ) layer.w13_input_scale = None layer.w2_input_scale = None + w2_amax = weight_amax(layer.w2_weight, dim=-1, keepdim=True) + w2_amax = amax_for_moe_weight_quant(w2_amax, self.moe.tp_size) + w2_scale = _fp8_channel_scale(w2_amax) + for expert in range(layer.local_num_experts): w13[expert], w13_scale[expert] = ops.scaled_fp8_quant( layer.w13_weight[expert], scale=None, use_per_token_if_dynamic=True, ) - w2[expert], w2_scale[expert] = ops.scaled_fp8_quant( - layer.w2_weight[expert], - scale=None, - use_per_token_if_dynamic=True, + w2[expert] = _fp8_quant_per_channel( + layer.w2_weight[expert], w2_scale[expert] ) self._setup_kernel( diff --git a/vllm/model_executor/layers/quantization/online/int8.py b/vllm/model_executor/layers/quantization/online/int8.py index 870e4d2a0f..661da7f997 100644 --- a/vllm/model_executor/layers/quantization/online/int8.py +++ b/vllm/model_executor/layers/quantization/online/int8.py @@ -22,8 +22,10 @@ from vllm.model_executor.layers.quantization.online.moe_base import ( OnlineMoEMethodBase, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( + amax_for_moe_weight_quant, kInt8DynamicTokenSym, kInt8StaticChannelSym, + weight_amax, ) from vllm.model_executor.utils import replace_parameter @@ -72,6 +74,9 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase): dtype=torch.float32, ) + w2_amax = weight_amax(layer.w2_weight, dim=-1) + w2_amax = amax_for_moe_weight_quant(w2_amax, self.moe.tp_size) + for expert in range(layer.local_num_experts): # w13: per-row quantization over hidden_size dim w = layer.w13_weight[expert, :, :] @@ -82,7 +87,7 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase): # w2: per-row quantization over intermediate_size dim w = layer.w2_weight[expert, :, :] - scales = w.abs().amax(dim=1) / vmax + scales = w2_amax[expert] / vmax q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax) w2[expert, :, :] = q.to(torch.int8) w2_scale[expert, :] = scales diff --git a/vllm/model_executor/layers/quantization/online/nvfp4.py b/vllm/model_executor/layers/quantization/online/nvfp4.py index eaa0069942..0a8a60a73a 100644 --- a/vllm/model_executor/layers/quantization/online/nvfp4.py +++ b/vllm/model_executor/layers/quantization/online/nvfp4.py @@ -20,8 +20,10 @@ from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import FLOAT4_E2M1_MAX, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( + amax_for_moe_weight_quant, kNvfp4Dynamic, kNvfp4Static, + weight_amax, ) from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform @@ -31,6 +33,7 @@ FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max def _quantize_moe_weight_to_nvfp4( weight: torch.Tensor, + moe_tp_size: int = 1, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Quantize stacked MoE expert weights ``(E, N, K)`` to NVFP4. @@ -43,7 +46,8 @@ def _quantize_moe_weight_to_nvfp4( k = weight.shape[-1] assert k % 16 == 0, f"last dim must be a multiple of 16, got {k}" - amax = weight.abs().amax(dim=(1, 2)).to(torch.float32).clamp_min(1e-8) + amax = weight_amax(weight.flatten(1), dim=-1).to(torch.float32) + amax = amax_for_moe_weight_quant(amax, moe_tp_size).clamp_min(1e-8) global_scale = (FLOAT4_E2M1_MAX * FLOAT8_E4M3_MAX) / amax weight_scale_2 = (1.0 / global_scale).to(torch.float32) @@ -106,8 +110,13 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase): layer._already_called_process_weights_after_loading = True def _quantize_weights(self, layer: Module) -> None: - w13, w13_scale, w13_scale_2 = _quantize_moe_weight_to_nvfp4(layer.w13_weight) - w2, w2_scale, w2_scale_2 = _quantize_moe_weight_to_nvfp4(layer.w2_weight) + moe_tp_size = self.moe.tp_size + w13, w13_scale, w13_scale_2 = _quantize_moe_weight_to_nvfp4( + layer.w13_weight, moe_tp_size + ) + w2, w2_scale, w2_scale_2 = _quantize_moe_weight_to_nvfp4( + layer.w2_weight, moe_tp_size + ) replace_parameter(layer, "w13_weight", w13) replace_parameter(layer, "w13_weight_scale", w13_scale) diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 6678aa15e2..a63db87fe6 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -11,6 +11,7 @@ import numpy import torch from torch import fx +from vllm.distributed.parallel_state import get_ep_group, get_tp_group from vllm.platforms import current_platform from vllm.scalar_type import ScalarType, scalar_types @@ -24,6 +25,43 @@ INT4_DTYPE = scalar_types.uint4b8 INT8_DTYPE = scalar_types.uint8b128 +def weight_amax( + weight: torch.Tensor, *, dim: int | None = None, keepdim: bool = False +) -> torch.Tensor: + """``max(|weight|)``, without materializing a full-size ``abs()``.""" + lo, hi = weight.aminmax(dim=dim, keepdim=keepdim) + return torch.maximum(lo.abs(), hi.abs()) + + +def amax_for_tp_weight_quant(amax: torch.Tensor, is_sharded: bool) -> torch.Tensor: + """Reduce a weight ``amax`` over the TP group when the weight is sharded + along a dim the ``amax`` reduces over, so each shard derives the same scale + it would as part of the whole weight. + """ + if is_sharded: + torch.distributed.all_reduce( + amax, + op=torch.distributed.ReduceOp.MAX, + group=get_tp_group().device_group, + ) + return amax + + +def amax_for_moe_weight_quant(amax: torch.Tensor, moe_tp_size: int) -> torch.Tensor: + """Reduce a per-expert weight ``amax`` over the ranks that tensor-shard the + MoE weights. That sharding is flattened over DP x PCP x TP, exactly the EP + group's span. Under EP ``moe_tp_size`` is 1 and each rank owns whole + experts, so no reduction is needed. + """ + if moe_tp_size > 1: + torch.distributed.all_reduce( + amax, + op=torch.distributed.ReduceOp.MAX, + group=get_ep_group().device_group, + ) + return amax + + def get_fp8_min_max() -> tuple[float, float]: """Get the min and max values for FP8 quantization.""" # Using the default value (240.0) from pytorch will cause accuracy