diff --git a/docs/features/quantization/online.md b/docs/features/quantization/online.md index a81ea219f7..e70be35ed9 100644 --- a/docs/features/quantization/online.md +++ b/docs/features/quantization/online.md @@ -111,7 +111,10 @@ Currently supported combinations are: | Checkpoint layer | Online target | Requirements | | ---------------- | ------------- | ------------ | -| Serialized ModelOpt MXFP8 linear | any | Target-specific hardware requirements | +| Serialized ModelOpt MXFP8 linear | any online linear method | Target-specific hardware requirements | +| Serialized compressed-tensors NVFP4 MoE | any online MoE method | Target-specific hardware requirements | +| Serialized ModelOpt NVFP4 MoE | any online MoE method | Target-specific hardware requirements | +| Serialized Quark NVFP4 MoE | any online MoE method | Target-specific hardware requirements | For example, serialized ModelOpt MXFP8 linears can be converted to per-channel-weight/per-token-activation FP8 at load time: @@ -122,10 +125,18 @@ vllm serve MiniMaxAI/MiniMax-M3-MXFP8 \ --quantization-config.linear fp8_per_channel ``` +Similarly, for example, NVFP4 MoE checkpoints can be re-quantized to any online MoE method, +for example compressed-tensors NVFP4 experts to online MXFP4: + +```bash +vllm serve RedHatAI/Qwen3-30B-A3B-NVFP4 \ + --quantization-config '{"targets":{"*mlp.experts*":"mxfp4"}}' +``` + !!! info - This is only available for linear layers at the moment. - + Re-quantization is available for both linear and MoE layers. + Checkpoint-exclusion policy is not enforced by online quantization/re-quantization. Only exclusion through `--quantization-config.targets`/`--quantization-config.ignore` is enforced. ### Online quantization on unquantized layers from partially-quantized checkpoints diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index ce13aa9a65..f18814dcd8 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -42,11 +42,13 @@ from vllm.model_executor.kernels.linear.scaled_mm import ( MarlinFP8ScaledMMLinearKernel, ) from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import FusedMoEFactory +from vllm.model_executor.layers.fused_moe import FusedMoEFactory, RoutedExperts from vllm.model_executor.layers.fused_moe.oracle import nvfp4 as nvfp4_oracle from vllm.model_executor.layers.linear import ( ColumnParallelLinear, LinearBase, + MergedColumnParallelLinear, + RowParallelLinear, UnquantizedLinearMethod, ) from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig @@ -57,10 +59,15 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 CompressedTensorsMoEMethod, ) +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4 import ( # noqa: E501 + CompressedTensorsW4A4Nvfp4MoEMethod, +) from vllm.model_executor.layers.quantization.modelopt import ( ModelOptFp8Config, ModelOptLinearMethod, ModelOptMxFp8Config, + ModelOptNvFp4Config, + ModelOptNvFp4FusedMoE, ) from vllm.model_executor.layers.quantization.online.base import ( OnlineQuantizationConfig, @@ -84,6 +91,7 @@ from vllm.model_executor.layers.quantization.online.mxfp4 import ( ) from vllm.model_executor.layers.quantization.online.mxfp8 import ( Mxfp8OnlineLinearMethod, + Mxfp8OnlineMoEMethod, ) from vllm.model_executor.layers.quantization.online.nvfp4 import ( Nvfp4OnlineMoEMethod, @@ -93,16 +101,23 @@ from vllm.model_executor.layers.quantization.quark.quark import ( QuarkConfig, QuarkLinearMethod, ) +from vllm.model_executor.layers.quantization.quark.quark_moe import ( + QuarkNvfp4MoEMethod, +) from vllm.model_executor.layers.quantization.utils import quant_utils from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( MXFP8_SCALE_DTYPE, MXFP8_VALUE_DTYPE, ) +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( + ref_nvfp4_quant, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( amax_for_moe_weight_quant, amax_for_tp_weight_quant, kMxfp8Dynamic, kMxfp8Static, + kNvfp4Dynamic, kNvfp4DynamicToken, kNvfp4Static, weight_amax, @@ -113,6 +128,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( from vllm.model_executor.model_loader import weight_utils from vllm.model_executor.model_loader.base_loader import log_online_quantization from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader +from vllm.model_executor.model_loader.reload.layerwise import ( + finalize_layerwise_reload, + initialize_layerwise_reload, + record_metadata_for_reloading, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.granitemoe import ( GraniteMoeModel, ) @@ -136,6 +157,8 @@ GRANITE_MODEL_NAME = "ibm-granite/granite-3.0-1b-a400m-base" PARTIALLY_PREQUANTIZED_MODEL_NAME = ( "nm-testing/tinysmokeqwen3moe-W4A16-first-only-CTstable" ) +NVFP4_MOE_MODEL_NAME = "RedHatAI/Qwen3-30B-A3B-NVFP4" +MODELOPT_NVFP4_MOE_MODEL_NAME = "nvidia/Qwen3-30B-A3B-NVFP4" @pytest.mark.parametrize( @@ -507,35 +530,749 @@ def test_online_ignore_keeps_checkpoint_quantization_linear( assert isinstance(layer.quant_method, QuarkLinearMethod) -def test_online_quantization_rejects_prequantized_moe( - default_vllm_config, dist_init, monkeypatch +def test_online_quantization_configures_prequantized_moe( + default_vllm_config, dist_init ) -> None: - """Reject pre-quantized MoE before constructing an online target backend.""" + """Configure the checkpoint method as the online MoE's weight source.""" default_vllm_config.model_config = ModelConfig() prefix = "model.layers.0.mlp.experts" quant_config = _fully_quantized_quark_config() quant_config.online_quantization_config = OnlineQuantizationConfig( quant_config_args(linear="mxfp4", moe="mxfp4") ) - monkeypatch.setattr( - quant_config.online_quantization_config, - "get_quant_method", - lambda *args: pytest.fail("unsupported online backend must not be constructed"), + + layer = FusedMoEFactory( + num_experts=4, + top_k=2, + hidden_size=32, + intermediate_size=64, + params_dtype=torch.bfloat16, + quant_config=quant_config, + prefix=prefix, ) - with pytest.raises( - NotImplementedError, - match="Requantizing checkpoint-quantized MoE layers is not supported", - ): - FusedMoEFactory( - num_experts=4, - top_k=2, - hidden_size=32, - intermediate_size=64, - params_dtype=torch.bfloat16, - quant_config=quant_config, - prefix=prefix, + assert isinstance(layer._quant_method, Mxfp4OnlineMoEMethod) + assert layer._quant_method.requantization_source is not None + + +def _pack_e2m1(values: torch.Tensor) -> torch.Tensor: + codes = torch.zeros_like(values, dtype=torch.uint8) + for code, magnitude in enumerate((0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)): + codes[values.abs() == magnitude] = code + codes |= (values < 0).to(torch.uint8) << 3 + return codes[..., 0::2] | (codes[..., 1::2] << 4) + + +def _serialize_nvfp4_experts( + weights: torch.Tensor, global_scales: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + packed_weights = [] + block_scales = [] + reconstructed_weights = [] + for weight, global_scale in zip(weights, global_scales, strict=True): + fp4_values, block_scale = ref_nvfp4_quant(weight, global_scale, block_size=16) + packed_weights.append(_pack_e2m1(fp4_values)) + block_scales.append(block_scale.to(torch.float8_e4m3fn)) + reconstructed_weights.append( + ( + fp4_values.reshape(*fp4_values.shape[:-1], -1, 16) + * (block_scale / global_scale).unsqueeze(-1) + ) + .reshape_as(weight) + .to(weight.dtype) ) + return ( + torch.stack(packed_weights), + torch.stack(block_scales), + torch.stack(reconstructed_weights), + ) + + +def _make_nvfp4_moe_source(source_format: str, moe, monkeypatch): + if source_format == "compressed_tensors": + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.compressed_tensors." + "compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4." + "select_nvfp4_moe_backend", + lambda **kwargs: (object(), None), + ) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.compressed_tensors." + "compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4." + "is_global_sf_supported_for_nvfp4_backend", + lambda backend: False, + ) + return CompressedTensorsW4A4Nvfp4MoEMethod(moe) + + if source_format == "quark": + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.quark.quark_moe." + "select_nvfp4_moe_backend", + lambda **kwargs: (object(), None), + ) + return QuarkNvfp4MoEMethod( + moe=moe, + quant_config=QuarkConfig({}), + activation_quant_key=kNvfp4Dynamic, + ) + + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.modelopt.select_nvfp4_moe_backend", + lambda **kwargs: (object(), None), + ) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.modelopt." + "is_global_sf_supported_for_nvfp4_backend", + lambda backend: False, + ) + config = ModelOptNvFp4Config( + is_checkpoint_nvfp4_serialized=True, + kv_cache_quant_algo=None, + exclude_modules=[], + ) + return ModelOptNvFp4FusedMoE(config, moe) + + +def _nvfp4_moe_checkpoint_tensors( + source_format: str, + w1: torch.Tensor, + w3: torch.Tensor, + w2: torch.Tensor, + w1_global_scale: torch.Tensor, + w3_global_scale: torch.Tensor, + w2_global_scale: torch.Tensor, +) -> tuple[dict[str, torch.Tensor], torch.Tensor, torch.Tensor]: + w1_packed, w1_scale, w1_reference = _serialize_nvfp4_experts(w1, w1_global_scale) + w3_packed, w3_scale, w3_reference = _serialize_nvfp4_experts(w3, w3_global_scale) + w2_packed, w2_scale, w2_reference = _serialize_nvfp4_experts(w2, w2_global_scale) + tensors = { + "w13_weight_scale": torch.cat((w1_scale, w3_scale), dim=1), + "w2_weight_scale": w2_scale, + } + if source_format == "compressed_tensors": + tensors.update( + { + "w13_weight_packed": torch.cat((w1_packed, w3_packed), dim=1), + "w2_weight_packed": w2_packed, + "w13_weight_global_scale": torch.stack( + (w1_global_scale, w3_global_scale), dim=1 + ), + "w2_weight_global_scale": w2_global_scale, + "w13_input_global_scale": torch.ones_like( + torch.stack((w1_global_scale, w3_global_scale), dim=1) + ), + "w2_input_global_scale": torch.ones_like(w2_global_scale), + } + ) + else: + tensors.update( + { + "w13_weight": torch.cat((w1_packed, w3_packed), dim=1), + "w2_weight": w2_packed, + "w13_weight_scale_2": 1.0 + / torch.stack((w1_global_scale, w3_global_scale), dim=1), + "w2_weight_scale_2": 1.0 / w2_global_scale, + ( + "w13_input_scale_2" + if source_format == "quark" + else "w13_input_scale" + ): torch.ones_like( + torch.stack((w1_global_scale, w3_global_scale), dim=1) + ), + ( + "w2_input_scale_2" if source_format == "quark" else "w2_input_scale" + ): torch.ones_like(w2_global_scale), + } + ) + return tensors, torch.cat((w1_reference, w3_reference), dim=1), w2_reference + + +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize("tp_rank", [0, 1]) +@pytest.mark.parametrize("projection", ["column", "row", "merged"]) +@pytest.mark.parametrize("source_format", ["modelopt_mxfp8"]) +def test_weight_loading_requantization( + source_format: str, + projection: str, + tp_rank: int, + default_vllm_config, + monkeypatch, +) -> None: + """Checkpoint sources retain parameter-specific loaders under online overrides.""" + default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16) + for module in ( + "vllm.model_executor.layers.linear", + "vllm.model_executor.parameter", + ): + monkeypatch.setattr(f"{module}.get_tensor_model_parallel_rank", lambda: tp_rank) + monkeypatch.setattr(f"{module}.get_tensor_model_parallel_world_size", lambda: 2) + if source_format == "modelopt_mxfp8": + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.modelopt.init_mxfp8_linear_kernel", + lambda **kwargs: Mock(), + ) + config = _fully_quantized_modelopt_mxfp8_config() + expected_source_cls = ModelOptLinearMethod + else: + pytest.fail(f"Unsupported checkpoint source: {source_format}") + config.online_quantization_config = OnlineQuantizationConfig( + quant_config_args(linear="fp8_per_channel") + ) + kwargs = dict( + input_size=128, + bias=False, + params_dtype=torch.bfloat16, + quant_config=config, + prefix="model.layers.0.self_attn.o_proj", + ) + if projection == "merged": + layer = MergedColumnParallelLinear(output_sizes=[64, 64], **kwargs) + elif projection == "row": + layer = RowParallelLinear(output_size=64, **kwargs) + else: + layer = ColumnParallelLinear(output_size=64, **kwargs) + target = layer.quant_method + # Online processing must remain assigned to the requested target. + assert isinstance(target, Fp8PtpcOnlineLinearMethod) + source = target.requantization_source + # Checkpoint parameters must belong to the selected source method. + assert isinstance(source, expected_source_cls) + for name in ("weight", "weight_scale"): + # Source registration must choose V2 rather than the online target's loader. + assert getattr(layer, name).weight_loader == layer.weight_loader_v2 + + torch.manual_seed(11) + rows = 128 if projection == "merged" else 64 + if source_format == "modelopt_mxfp8": + checkpoint = { + "weight": torch.randn(rows, 128, dtype=torch.bfloat16).to( + MXFP8_VALUE_DTYPE + ), + "weight_scale": torch.randint(124, 131, (rows, 4), dtype=MXFP8_SCALE_DTYPE), + } + else: + pytest.fail(f"Unsupported checkpoint source: {source_format}") + expected = {} + for name, tensor in checkpoint.items(): + param = getattr(layer, name) + if projection == "merged": + for shard, part in enumerate(tensor.chunk(2, dim=0)): + param.weight_loader(param, part, shard) + expected[name] = torch.cat( + [part.chunk(2, dim=0)[tp_rank] for part in tensor.chunk(2, dim=0)] + ) + else: + param.weight_loader(param, tensor) + expected[name] = tensor.chunk(2, dim=1 if projection == "row" else 0)[ + tp_rank + ] + # TP loading must preserve each source weight/scale shard exactly. + torch.testing.assert_close(param.cpu(), expected[name], rtol=0, atol=0) + actual_weight = source.dequantize_weight(layer) + reference = source.dequantize_weight(SimpleNamespace(**expected)) + # Loaded source shards must dequantize identically to their checkpoint slices. + torch.testing.assert_close(actual_weight.cpu(), reference.cpu(), rtol=0, atol=0) + # Source loading must leave the online processing method attached. + assert layer.quant_method is target + + +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize( + "source_format", + [ + pytest.param("compressed_tensors", id="compressed_tensors_nvfp4"), + pytest.param("modelopt", id="modelopt_nvfp4"), + pytest.param("quark", id="quark_nvfp4"), + ], +) +def test_weight_loading_requantization_moe( + source_format: str, default_vllm_config, monkeypatch +) -> None: + """Online re-quantization uses the original quant method for loading.""" + from tests.kernels.moe.utils import make_dummy_moe_config + + # Loading local experts needs parameter TP metadata, not a process group. + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0 + ) + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1 + ) + default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.online.fp8.select_fp8_moe_backend", + lambda **kwargs: (None, None), + ) + config = make_dummy_moe_config(num_experts=2, hidden_dim=32, intermediate_size=64) + source = _make_nvfp4_moe_source(source_format, config, monkeypatch) + quant_config = Mock() + quant_config.packed_modules_mapping = {} + quant_config.get_name.return_value = source_format + quant_config.get_quant_method.return_value = source + quant_config.online_quantization_config = OnlineQuantizationConfig( + quant_config_args(moe="fp8_per_tensor") + ) + prefix = "model.layers.0.mlp.experts" + layer = RoutedExperts( + prefix, + torch.bfloat16, + config, + quant_config, + expert_map_manager=SimpleNamespace( + local_num_experts=2, + num_fused_shared_experts=0, + placement_strategy="linear", + expert_map=None, + expert_mask=None, + routing_tables=None, + map_global_to_local=lambda index: index, + ), + ) + target = layer.quant_method + # Construction must select the override, not execute the checkpoint method. + assert isinstance(target, Fp8PerTensorOnlineMoEMethod) + # Source allocation and loading must use the original checkpoint method. + assert target.requantization_source is source + # Requantization requires a checkpoint source for dequantization. + assert source is not None + torch.manual_seed(11) + w1 = torch.randn(2, 64, 32, dtype=torch.bfloat16) + w3 = torch.randn_like(w1) + w2 = torch.randn(2, 32, 64, dtype=torch.bfloat16) + tensors, expected_w13, expected_w2 = _nvfp4_moe_checkpoint_tensors( + source_format, + w1, + w3, + w2, + torch.tensor([96.0, 160.0]), + torch.tensor([40.0, 72.0]), + torch.tensor([128.0, 56.0]), + ) + # Distinct input scales expose gate/up overwrites independently of dequantization. + # Quark's existing input-scale loader broadcasts each shard across both slots. + input_name = { + "compressed_tensors": "w13_input_global_scale", + "modelopt": "w13_input_scale", + "quark": "w13_input_scale_2", + }[source_format] + tensors[input_name] = torch.tensor( + [[0.25, 0.25], [0.75, 0.75]] + if source_format == "quark" + else [[0.25, 0.5], [0.75, 1.0]] + ) + checkpoint = [] + for name, tensor in tensors.items(): + stem, suffix = name.split("_", 1) + projections = ("gate_proj", "up_proj") if stem == "w13" else ("down_proj",) + for expert in range(2): + parts = ( + tensor[expert].chunk(2, dim=0) if stem == "w13" else (tensor[expert],) + ) + for projection, part in zip(projections, parts, strict=True): + checkpoint.append((f"{expert}.{projection}.{suffix}", part)) + loaded = set(layer.load_weights(checkpoint)) + + # Every source parameter must be reached through checkpoint-name mapping. + assert loaded == set(tensors) + + for name, expected in tensors.items(): + # Public loading must preserve all packed weights and scale conventions. + torch.testing.assert_close(getattr(layer, name).cpu(), expected, rtol=0, atol=0) + actual_w13, actual_w2 = source.dequantize_weight(layer) + + # Gate/up ordering and distinct global scales must survive source loading. + torch.testing.assert_close(actual_w13.cpu(), expected_w13, rtol=0, atol=0) + # Down-projection block and global scales must survive source loading. + torch.testing.assert_close(actual_w2.cpu(), expected_w2, rtol=0, atol=0) + + # Source-aware loading must not replace the target used for processing/inference. + assert layer.quant_method is target + + +@pytest.mark.parametrize("is_act_and_mul", [False, True]) +@pytest.mark.parametrize( + "target_hidden,target_intermediate,has_bias,error", + [ + pytest.param(16, 32, False, None, id="unpadded"), + pytest.param(24, 48, False, None, id="padded"), + pytest.param(24, 48, True, NotImplementedError, id="bias-with-padding"), + pytest.param(8, 32, False, ValueError, id="shrinking-hidden"), + pytest.param(16, 16, False, ValueError, id="shrinking-intermediate"), + ], +) +def test_online_requantized_moe_padding( + is_act_and_mul: bool, + target_hidden: int, + target_intermediate: int, + has_bias: bool, + error: type[Exception] | None, +) -> None: + """Preserve logical expert weights and shard boundaries when padding.""" + num_shards = 2 if is_act_and_mul else 1 + w13 = torch.arange( + 2 * num_shards * 32 * 16, device="cpu", dtype=torch.float32 + ).reshape(2, num_shards * 32, 16) + w2 = torch.arange(2 * 16 * 32, device="cpu", dtype=torch.float32).reshape(2, 16, 32) + method = Fp8PerTensorOnlineMoEMethod.__new__(Fp8PerTensorOnlineMoEMethod) + method.moe = SimpleNamespace( + w13_num_shards=num_shards, is_act_and_mul=is_act_and_mul, has_bias=has_bias + ) + method.requantization_source = Mock() + method.requantization_source.dequantize_weight.return_value = (w13, w2) + layer = SimpleNamespace( + local_num_experts=2, + moe_config=SimpleNamespace( + hidden_dim_unpadded=16, + intermediate_size_per_partition_unpadded=32, + hidden_dim=target_hidden, + intermediate_size_per_partition=target_intermediate, + ), + w13_bias=torch.zeros(2, num_shards * 32, device="cpu") if has_bias else None, + w2_bias=torch.zeros(2, 16, device="cpu") if has_bias else None, + ) + if error is not None: + message = "biases is unsupported" if has_bias else "cannot shrink weights" + # Unsupported padding must fail explicitly rather than truncate weights. + with pytest.raises(error, match=message): + method.get_weights_for_quantization(layer) + return + + padded_w13, padded_w2 = method.get_weights_for_quantization(layer) + if (target_hidden, target_intermediate) == (16, 32): + # Unpadded w13 must not allocate a replacement tensor. + assert padded_w13 is w13 + # Unpadded w2 must not allocate a replacement tensor. + assert padded_w2 is w2 + + expected_w13 = w13.new_zeros(2, num_shards * target_intermediate, target_hidden) + + for shard in range(num_shards): + start = shard * target_intermediate + expected_w13[:, start : start + 32, :16] = w13[:, shard * 32 : (shard + 1) * 32] + + expected_w2 = w2.new_zeros(2, target_hidden, target_intermediate) + expected_w2[:, :16, :32] = w2 + + # Preserve each gate/up shard and zero all added rows and columns. + torch.testing.assert_close(padded_w13, expected_w13, rtol=0, atol=0) + + # Preserve the logical down projection and zero both padding dimensions. + torch.testing.assert_close(padded_w2, expected_w2, rtol=0, atol=0) + + +@pytest.mark.skipif( + not (current_platform.is_cuda() or current_platform.is_rocm()) + or not is_quant_method_supported("fp8"), + reason="Requires FP8 support on CUDA or ROCm.", +) +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize("source_format", ["compressed_tensors", "modelopt", "quark"]) +@pytest.mark.parametrize( + "online_target", + [ + "fp8_per_tensor", + "fp8_per_block", + "fp8_per_channel", + "mxfp4", + "mxfp8", + "int8", + ], +) +def test_online_requantized_moe_reload( + source_format: str, + online_target: str, + default_vllm_config, + monkeypatch, +) -> None: + """Reload each online target in place and match a fresh checkpoint conversion. + + Source loading is covered separately; this test isolates target reload processing. + """ + from vllm.model_executor.layers.quantization.online import fp8 as online_fp8 + + # Test conversion and reload behavior, not torch.compile. + monkeypatch.setattr( + online_fp8, + "per_block_cast_to_fp8", + online_fp8.per_block_cast_to_fp8.__wrapped__, + ) + + default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16) + + # This is faster than `dist_init` arg. + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0 + ) + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1 + ) + + moe = SimpleNamespace( + w13_num_shards=2, + has_bias=False, + is_act_and_mul=True, + tp_size=1, + use_mxfp4_w4a4_dsv4=False, + ) + target_methods = { + "fp8_per_tensor": (Fp8PerTensorOnlineMoEMethod, "fp8"), + "fp8_per_block": (Fp8PerBlockOnlineMoEMethod, "fp8"), + "fp8_per_channel": (Fp8PtpcOnlineMoEMethod, "fp8"), + "mxfp4": (Mxfp4OnlineMoEMethod, "mxfp4"), + "mxfp8": (Mxfp8OnlineMoEMethod, "mxfp8"), + "int8": (Int8OnlineMoEMethod, "int8"), + } + method_cls, module_name = target_methods[online_target] + monkeypatch.setattr( + f"vllm.model_executor.layers.quantization.online.{module_name}." + f"select_{module_name}_moe_backend", + lambda **kwargs: (object(), object()), + ) + source = _make_nvfp4_moe_source(source_format, moe, monkeypatch) + + if module_name in ("fp8", "mxfp8"): + kernel_module = "vllm.model_executor.layers.fused_moe.oracle.fp8" + converter_name = "convert_to_fp8_moe_kernel_format" + factory_name = "make_fp8_moe_kernel" + weight_names = ("w13", "w2", "w13_scale", "w2_scale") + elif module_name == "mxfp4": + kernel_module = "vllm.model_executor.layers.quantization.online.mxfp4" + converter_name = "convert_weight_to_mxfp4_moe_kernel_format" + factory_name = "make_mxfp4_moe_kernel" + weight_names = ( + "w13_weight", + "w2_weight", + "w13_weight_scale", + "w2_weight_scale", + "w13_bias", + "w2_bias", + ) + else: + kernel_module = "vllm.model_executor.layers.quantization.online.int8" + converter_name = "convert_to_int8_moe_kernel_format" + factory_name = "make_int8_moe_kernel" + weight_names = ("w13", "w2") + + converter = Mock( + side_effect=lambda **kwargs: tuple(kwargs[name] for name in weight_names) + ) + kernel = Mock() + kernel_factory = Mock(return_value=kernel) + monkeypatch.setattr(f"{kernel_module}.{converter_name}", converter) + monkeypatch.setattr(f"{kernel_module}.{factory_name}", kernel_factory) + device = current_platform.device_type + size = 128 + + def create_layer() -> tuple[torch.nn.Module, Any]: + method = method_cls(moe=moe) + monkeypatch.setattr( + method, "get_fused_moe_quant_config", lambda layer: object() + ) + method.set_requantization_source(source) + layer = torch.nn.Module() + layer.local_num_experts = 2 + layer.num_experts = 2 + layer._expert_routing_tables = lambda: None + layer.moe_config = SimpleNamespace( + hidden_dim_unpadded=size, + intermediate_size_per_partition_unpadded=size, + hidden_dim=size, + intermediate_size_per_partition=size, + is_act_and_mul=True, + moe_parallel_config=SimpleNamespace(tp_size=1), + ) + layer._get_hidden_dim = RoutedExperts._get_hidden_dim + layer._narrow_expert_data_for_padding = ( + RoutedExperts._narrow_expert_data_for_padding + ) + layer.quant_method = method + layer.quant_config = SimpleNamespace(online_quantization_config=None) + method.create_weights( + layer, + num_experts=2, + hidden_size=size, + intermediate_size_per_partition=size, + params_dtype=torch.bfloat16, + weight_loader=default_weight_loader, + ) + record_metadata_for_reloading(layer) + return layer, method + + with torch.device(device): + layer, method = create_layer() + source_parameters = dict(method.requantization_source_parameters) + source_names = set(source_parameters) + packed_w13_name = ( + "w13_weight_packed" + if source_format == "compressed_tensors" + else "w13_weight" + ) + # Source allocation must retain the packed NVFP4 checkpoint layout. + assert source_parameters[packed_w13_name].shape == ( + 2, + 2 * size, + size // 2, + ) + + def serialized_weights(value: float) -> dict[str, torch.Tensor]: + shape = (2, size, size) + pattern = ( + torch.linspace( + -1.0, 1.0, size * size, device=device, dtype=torch.bfloat16 + ) + .reshape(1, size, size) + .expand(shape) + ) + w1 = pattern * value + value**2 + w3 = pattern.flip(-1) * (value * 2) - value**2 + w2 = pattern.flip(-2) * (value * 3) + value**2 + w1_global_scale = torch.tensor([96.0, 160.0], device=device) + w3_global_scale = torch.tensor([40.0, 72.0], device=device) + w2_global_scale = torch.tensor([128.0, 56.0], device=device) + tensors, _, _ = _nvfp4_moe_checkpoint_tensors( + source_format, + w1, + w3, + w2, + w1_global_scale, + w3_global_scale, + w2_global_scale, + ) + return tensors + + def load_serialized_weights( + layer: torch.nn.Module, weights: dict[str, torch.Tensor], *, reload: bool + ) -> None: + for name, loaded_weight in weights.items(): + parameter = getattr(layer, name) + if reload: + parameter.weight_loader(parameter, loaded_weight) + elif name in (packed_w13_name, "w2_weight", "w2_weight_packed"): + for expert in range(2): + if name == packed_w13_name: + gate, up = loaded_weight[expert].chunk(2, dim=0) + for shard_id, shard_weight in (("w1", gate), ("w3", up)): + RoutedExperts._load_w13( + layer, + parameter.data[expert], + 0, + shard_id, + shard_weight, + 0, + ) + else: + RoutedExperts._load_w2( + layer, + parameter.data[expert], + 1, + loaded_weight[expert], + 0, + ) + else: + parameter.data.copy_(loaded_weight) + + # Initial loading replaces NVFP4 staging weights with online target weights. + load_serialized_weights(layer, serialized_weights(0.25), reload=False) + method.process_weights_after_loading(layer) + finalize_layerwise_reload(layer, default_vllm_config.model_config) + # Initial processing must invoke backend-format conversion exactly once. + converter.assert_called_once() + # Initial processing must construct the target kernel exactly once. + kernel_factory.assert_called_once() + # The method must attach the kernel returned by the factory. + assert method.moe_kernel is kernel + + # Initial conversion must release references to the NVFP4 parameters. + assert not method.requantization_source_parameters + + for name, parameter in source_parameters.items(): + # No original NVFP4 parameter object may remain attached to the layer. + assert getattr(layer, name, None) is not parameter + + expected_dtype = { + "mxfp4": torch.uint8, + "int8": torch.int8, + "mxfp8": torch.float8_e4m3fn, + }.get(online_target, current_platform.fp8_dtype()) + + # Gate/up weights must use the requested target's quantized dtype. + assert layer.w13_weight.dtype == expected_dtype + # Down-projection weights must use the same target dtype. + assert layer.w2_weight.dtype == expected_dtype + + packed_factor = 2 if online_target == "mxfp4" else 1 + # Target packing must preserve both gate/up shards and expert count. + assert layer.w13_weight.shape == ( + 2, + 2 * size, + size // packed_factor, + ) + # Target packing must preserve the down-projection dimensions. + assert layer.w2_weight.shape == (2, size, size // packed_factor) + # Initial processing must materialize gate/up weights on the device. + assert not layer.w13_weight.is_meta + # Initial processing must materialize down-projection weights on the device. + assert not layer.w2_weight.is_meta + + target_parameters = dict(layer.named_parameters()) + target_storage = { + name: parameter.data_ptr() for name, parameter in target_parameters.items() + } + first_weights = { + name: parameter.clone() for name, parameter in target_parameters.items() + } + second_checkpoint = serialized_weights(0.75) + fresh_layer, fresh_method = create_layer() + load_serialized_weights(fresh_layer, second_checkpoint, reload=False) + fresh_method.process_weights_after_loading(fresh_layer) + finalize_layerwise_reload(fresh_layer, default_vllm_config.model_config) + expected_parameters = dict(fresh_layer.named_parameters()) + # Fresh conversion must produce the same target parameter schema. + assert expected_parameters.keys() == target_parameters.keys() + # The reference conversion must also release its NVFP4 staging references. + assert not fresh_method.requantization_source_parameters + for name in ("w13_weight", "w2_weight"): + # Checkpoint B must change both matrices so stale reloads cannot pass. + assert not torch.equal(expected_parameters[name], first_weights[name]) + converter.reset_mock() + kernel_factory.reset_mock() + + # Reload restores the original NVFP4 checkpoint-facing layout, not the + # already-converted target runtime layout. + initialize_layerwise_reload(layer) + + # Reload must recreate every NVFP4 source parameter on the meta device. + assert all(getattr(layer, name).is_meta for name in source_names) + + load_serialized_weights(layer, second_checkpoint, reload=True) + finalize_layerwise_reload(layer, default_vllm_config.model_config) + # Reload must reconvert checkpoint B rather than reuse checkpoint A's weights. + converter.assert_called_once() + # Reload must rebuild the kernel using the newly processed target tensors. + kernel_factory.assert_called_once() + # The method must attach the kernel returned by the reload factory call. + assert method.moe_kernel is kernel + + # Reload must not retain references to temporary NVFP4 source parameters. + assert not method.requantization_source_parameters + + source_only_names = source_names - target_parameters.keys() + + # NVFP4-only staging parameters must be absent or reset to the target's + # non-parameter sentinel after reload processing. + assert all(getattr(layer, name, None) is None for name in source_only_names) + + # Reload must restore exactly the target parameters, without extra staging weights. + assert dict(layer.named_parameters()).keys() == target_parameters.keys() + + for name, parameter in target_parameters.items(): + # Reload must update target weights and scales without replacing Parameters. + assert getattr(layer, name) is parameter + + assert parameter.data_ptr() == target_storage[name] + assert not parameter.is_meta + + torch.testing.assert_close(parameter, expected_parameters[name], rtol=0, atol=0) def test_activation_only_override_applies_to_checkpoint_method( @@ -939,6 +1676,24 @@ def test_nvfp4_one_sided_rejects_unsupported_input_dtype(input_dtype: torch.dtyp {}, id="requantization_mxfp8_ptcp_fp8", ), + pytest.param( + NVFP4_MOE_MODEL_NAME, + None, + {"targets": {"*mlp.experts*": "mxfp4"}}, + CompressedTensorsLinearMethod, + Mxfp4OnlineMoEMethod, + {}, + id="requantization_nvfp4_mxfp4_moe", + ), + pytest.param( + MODELOPT_NVFP4_MOE_MODEL_NAME, + None, + {"targets": {"*mlp.experts*": "fp8_per_tensor"}}, + ModelOptLinearMethod, + Fp8PerTensorOnlineMoEMethod, + {}, + id="requantization_modelopt_nvfp4_fp8_moe", + ), ], ) @pytest.mark.parametrize( @@ -964,7 +1719,9 @@ def test_online_quantization( """ # 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()): + if (quant_scheme == "mxfp4" or expected_moe_cls is Mxfp4OnlineMoEMethod) and not ( + on_gfx950() or on_gfx942() + ): pytest.skip("mxfp4 online quantization is only tested on AMD gfx942, gfx950.") if current_platform.is_rocm(): @@ -991,6 +1748,35 @@ def test_online_quantization( moe_backend=runner_kwargs.get("moe_backend", "auto"), ), } + if model_name in (NVFP4_MOE_MODEL_NAME, MODELOPT_NVFP4_MOE_MODEL_NAME): + original_process = expected_moe_cls.process_weights_after_loading + expected_source_cls = ( + CompressedTensorsW4A4Nvfp4MoEMethod + if model_name == NVFP4_MOE_MODEL_NAME + else ModelOptNvFp4FusedMoE + ) + + def assert_source_weights_released(method, layer) -> None: + source_parameters = dict(method.requantization_source_parameters) + assert isinstance(method.requantization_source, expected_source_cls) + original_process(method, layer) + assert not method.requantization_source_parameters + for name, parameter in source_parameters.items(): + assert getattr(layer, name, None) is not parameter + + if expected_moe_cls is Mxfp4OnlineMoEMethod: + assert layer.w13_weight_scale.dtype == torch.uint8 + assert layer.w2_weight_scale.dtype == torch.uint8 + else: + assert layer.w13_weight.dtype == current_platform.fp8_dtype() + assert layer.w2_weight.dtype == current_platform.fp8_dtype() + + monkeypatch.setattr( + expected_moe_cls, + "process_weights_after_loading", + assert_source_weights_released, + ) + if model_name == "mgoin/Qwen3-0.6B-MXFP8": original_process = Fp8PtpcOnlineLinearMethod.process_weights_after_loading @@ -1068,6 +1854,9 @@ def test_online_quantization( elif model_name == "mgoin/Qwen3-0.6B-MXFP8": o_proj = model.model.layers[0].self_attn.o_proj moe = None + elif model_name in (NVFP4_MOE_MODEL_NAME, MODELOPT_NVFP4_MOE_MODEL_NAME): + o_proj = model.model.layers[0].self_attn.o_proj + moe = model.model.layers[0].mlp.experts else: o_proj = model.model.layers[0].self_attn.o_proj moe = model.model.layers[0].block_sparse_moe.experts @@ -1087,8 +1876,23 @@ def test_online_quantization( assert o_proj.weight.dtype == torch.bfloat16 elif model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME: assert o_proj.weight.dtype == MXFP8_VALUE_DTYPE - elif quant_scheme == "mxfp4": - assert o_proj.weight.dtype == torch.uint8 + elif ( + model_name + in ( + NVFP4_MOE_MODEL_NAME, + MODELOPT_NVFP4_MOE_MODEL_NAME, + ) + or quant_scheme == "mxfp4" + ): + if ( + model_name in (NVFP4_MOE_MODEL_NAME, MODELOPT_NVFP4_MOE_MODEL_NAME) + and current_platform.is_cuda() + and not current_platform.has_device_capability(100) + ): + # NVFP4 linear falls back to Marlin, which repacks weights to int32. + assert o_proj.weight.dtype == torch.int32 + else: + assert o_proj.weight.dtype == torch.uint8 elif current_platform.is_cuda() or current_platform.is_xpu(): if current_platform.supports_fp8() and not force_marlin: assert o_proj.weight.dtype == torch.float8_e4m3fn @@ -1548,6 +2352,104 @@ def _quantize_moe(weight, scheme, moe_tp_size): return qweight, scale +@pytest.mark.skipif( + not is_quant_method_supported("fp8"), + reason="FP8 is not supported on this GPU type.", +) +@pytest.mark.parametrize("source_format", ["compressed_tensors", "modelopt", "quark"]) +def test_nvfp4_moe_dequantization_and_fp8_requantization_numerics( + source_format: str, + monkeypatch, +) -> None: + """NVFP4 checkpoint weights survive dequantization and FP8 requantization.""" + torch.manual_seed(11) + num_experts = 2 + hidden_size = intermediate_size = 32 + device = current_platform.device_type + + w1 = ( + torch.randn( + num_experts, + intermediate_size, + hidden_size, + device=device, + dtype=torch.bfloat16, + ) + * 0.35 + + 0.2 + ) + w3 = torch.randn_like(w1) * 1.7 - 0.15 + w2 = ( + torch.randn( + num_experts, + hidden_size, + intermediate_size, + device=device, + dtype=torch.bfloat16, + ) + * 0.8 + + 0.05 + ) + + # The reference quantizer takes encode scales; the checkpoint helper converts + # them to each source format's stored convention. Distinct values for every + # expert and w13 shard catch reciprocal, ordering, and chunking bugs. + w1_global_scale = torch.tensor([96.0, 160.0], device=device) + w3_global_scale = torch.tensor([40.0, 72.0], device=device) + w2_global_scale = torch.tensor([128.0, 56.0], device=device) + + checkpoint_tensors, w13_reference, w2_reference = _nvfp4_moe_checkpoint_tensors( + source_format, + w1, + w3, + w2, + w1_global_scale, + w3_global_scale, + w2_global_scale, + ) + moe = SimpleNamespace(is_act_and_mul=True, has_bias=False) + source = _make_nvfp4_moe_source(source_format, moe, monkeypatch) + source_layer = SimpleNamespace( + **checkpoint_tensors, + params_dtype=torch.bfloat16, + ) + w13_dequantized, w2_dequantized = source.dequantize_weight(source_layer) + + w13 = torch.cat((w1, w3), dim=1) + torch.testing.assert_close(w13_dequantized, w13_reference, rtol=0, atol=0) + torch.testing.assert_close(w2_dequantized, w2_reference, rtol=0, atol=0) + + w13_fp8, w13_fp8_scale = _quantize_moe(w13_dequantized, "per_tensor", 1) + w2_fp8, w2_fp8_scale = _quantize_moe(w2_dequantized, "per_tensor", 1) + w13_requantized = w13_fp8.float() * w13_fp8_scale[:, None, None] + w2_requantized = w2_fp8.float() * w2_fp8_scale[:, None, None] + + inputs = ( + torch.randn(num_experts, 8, hidden_size, device=device, dtype=torch.bfloat16) + * 0.4 + ) + + def expert_output( + w13_weight: torch.Tensor, w2_weight: torch.Tensor + ) -> torch.Tensor: + gate_up = torch.einsum("eth,eih->eti", inputs.float(), w13_weight.float()) + gate, up = gate_up.chunk(2, dim=-1) + hidden = torch.nn.functional.silu(gate) * up + return torch.einsum("eti,ehi->eth", hidden, w2_weight.float()) + + reference_output = expert_output(w13, w2) + requantized_output = expert_output(w13_requantized, w2_requantized) + relative_error = torch.linalg.vector_norm( + requantized_output - reference_output + ) / torch.linalg.vector_norm(reference_output) + + # This includes NVFP4 error in both expert matrices plus a second FP8 + # quantization. A 20% relative L2 bound allows for that compounded error + # while still rejecting scale inversion or w1/w3 layout errors by orders + # of magnitude. + assert relative_error < 0.2 + + @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.""" diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 71f4206008..6fb49314fc 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -700,12 +700,29 @@ class RoutedExperts(PluggableLayer): param.data[:, :dim1, :dim2].copy_(loaded_weight) return True if return_success else None - quant_method_name = self.quant_method.__class__.__name__ + from vllm.model_executor.layers.quantization.online.moe_base import ( + OnlineMoEMethodBase, + ) + + base_quant_method = self.quant_method + + # In case online MOE re-quantization is used, loading is handled + # by the original quant method. + if ( + isinstance(base_quant_method, OnlineMoEMethodBase) + and base_quant_method.requantization_source is not None + ): + assert isinstance( + base_quant_method.requantization_source, FusedMoEMethodBase + ) + base_quant_method = base_quant_method.requantization_source + + quant_method_name = base_quant_method.__class__.__name__ global_expert_id = expert_id expert_id = self._map_global_expert_id_to_local_expert_id(global_expert_id) use_global_sf = ( - getattr(self.quant_method, "use_global_sf", False) + getattr(base_quant_method, "use_global_sf", False) and "input_scale" in weight_name ) @@ -842,7 +859,7 @@ class RoutedExperts(PluggableLayer): if "ModelOpt" in quant_method_name: # Determine per-tensor weight scale patterns based on variant # Use the dedicated method instead of brittle string matching - uses_weight_scale_2 = self.quant_method.uses_weight_scale_2_pattern() + uses_weight_scale_2 = base_quant_method.uses_weight_scale_2_pattern() quant_method = getattr(param, "quant_method", None) # Call _load_per_tensor_weight_scale() to load per-tensor (scalar) @@ -921,7 +938,7 @@ class RoutedExperts(PluggableLayer): FusedMoeWeightScaleSupported.GROUP.value, FusedMoeWeightScaleSupported.BLOCK.value, ]: - scale_refine = getattr(self.quant_method, "weight_scale_refine", None) + scale_refine = getattr(base_quant_method, "weight_scale_refine", None) if ( quant_method == FusedMoeWeightScaleSupported.BLOCK.value and scale_refine is not None diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 65cda325a2..5b77caa256 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -516,6 +516,20 @@ class ColumnParallelLinear(LinearBase): self._maybe_allow_fp8_block_shape_mismatch() self.gather_output = gather_output + from vllm.model_executor.layers.quantization.online.fp8 import OnlineLinearBase + + quant_method = self.quant_method + if ( + isinstance(quant_method, OnlineLinearBase) + and quant_method.requantization_source is not None + ): + quant_method = quant_method.requantization_source + weight_loader = ( + self.weight_loader_v2 + if quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED + else self.weight_loader + ) + self.quant_method.create_weights( layer=self, input_size_per_partition=self.input_size_per_partition, @@ -523,11 +537,7 @@ class ColumnParallelLinear(LinearBase): input_size=self.input_size, output_size=self.output_size, params_dtype=self.params_dtype, - weight_loader=( - self.weight_loader_v2 - if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED - else self.weight_loader - ), + weight_loader=weight_loader, ) if bias: @@ -1697,6 +1707,20 @@ class RowParallelLinear(LinearBase): self.input_is_parallel = input_is_parallel self.reduce_results = reduce_results + from vllm.model_executor.layers.quantization.online.fp8 import OnlineLinearBase + + quant_method = self.quant_method + if ( + isinstance(quant_method, OnlineLinearBase) + and quant_method.requantization_source is not None + ): + quant_method = quant_method.requantization_source + weight_loader = ( + self.weight_loader_v2 + if quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED + else self.weight_loader + ) + self.quant_method.create_weights( layer=self, input_size_per_partition=self.input_size_per_partition, @@ -1704,11 +1728,7 @@ class RowParallelLinear(LinearBase): input_size=self.input_size, output_size=self.output_size, params_dtype=self.params_dtype, - weight_loader=( - self.weight_loader_v2 - if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED - else self.weight_loader - ), + weight_loader=weight_loader, ) if not reduce_results and (bias and not skip_bias_add): raise ValueError( diff --git a/vllm/model_executor/layers/quantization/base_config.py b/vllm/model_executor/layers/quantization/base_config.py index 3ec2e728e1..5bc1e42aa5 100644 --- a/vllm/model_executor/layers/quantization/base_config.py +++ b/vllm/model_executor/layers/quantization/base_config.py @@ -84,11 +84,14 @@ class QuantizeMethodBase(ABC): """ return - def dequantize_weight(self, layer: nn.Module) -> torch.Tensor: + def dequantize_weight( + self, layer: nn.Module + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """Materialize a serialized quantized weight for requantization.""" raise NotImplementedError( f"The quantization method {type(self)} does not implement " - "dequantize_weight. Please open an issue." + "dequantize_weight, so re-quantization from this checkpoint is " + "not supported. Please open an issue." ) @@ -302,6 +305,9 @@ def resolve_quant_method( UnquantizedLinearMethod, ) from vllm.model_executor.layers.quantization.online.fp8 import OnlineLinearBase + from vllm.model_executor.layers.quantization.online.moe_base import ( + OnlineMoEMethodBase, + ) base_quant_method = quant_config.get_quant_method(layer, prefix) if quant_config.online_quantization_config is None: @@ -328,11 +334,6 @@ def resolve_quant_method( # The checkpoint quant method is applied as there is no online override. return base_quant_method - if isinstance(layer, RoutedExperts): - raise NotImplementedError( - "Requantizing checkpoint-quantized MoE layers is not supported." - ) - online_quant_method = quant_config.online_quantization_config.get_quant_method( layer, prefix ) @@ -341,7 +342,7 @@ def resolve_quant_method( base_quant_method, (UnquantizedLinearMethod, UnquantizedFusedMoEMethod) ) - assert isinstance(online_quant_method, OnlineLinearBase) + assert isinstance(online_quant_method, (OnlineLinearBase, OnlineMoEMethodBase)) online_quant_method.set_requantization_source(base_quant_method) # The online method dequantizes the checkpoint method before requantizing. diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py index dff0c62669..e7fc3f1df3 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py @@ -26,6 +26,9 @@ from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501 CompressedTensorsMoEMethod, ) +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( + dequantize_nvfp4_moe_weights, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( kNvfp4Dynamic, kNvfp4Static, @@ -185,6 +188,22 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): ) set_weight_attrs(w2_input_scale, extra_weight_attrs) + def dequantize_weight( + self, layer: torch.nn.Module + ) -> tuple[torch.Tensor, torch.Tensor]: + """Materialize serialized NVFP4 expert weights for requantization.""" + return dequantize_nvfp4_moe_weights( + layer.w13_weight_packed, + layer.w13_weight_scale, + layer.w13_weight_global_scale, + layer.w2_weight_packed, + layer.w2_weight_scale, + layer.w2_weight_global_scale, + layer.params_dtype, + group_size=self.group_size, + invert_global_scales=True, + ) + def process_weights_after_loading(self, layer: RoutedExperts) -> None: """Convert NVFP4 MoE weights into kernel format and setup the kernel.""" # NOTE(rob): wN_weight_packed -> wN_weight is because ModularKernelMethod diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 8354f885d9..ef16e8c5dc 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -81,6 +81,9 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( MXFP8_VALUE_DTYPE, dequant_mxfp8_to_bf16, ) +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( + dequantize_nvfp4_moe_weights, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( FP4_DTYPE, QuantKey, @@ -889,6 +892,22 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): """FP4 variants use 'weight_scale_2' pattern for per-tensor weight scales.""" return True + def dequantize_weight( + self, layer: torch.nn.Module + ) -> tuple[torch.Tensor, torch.Tensor]: + """Materialize serialized ModelOpt NVFP4 expert weights.""" + return dequantize_nvfp4_moe_weights( + layer.w13_weight, + layer.w13_weight_scale, + layer.w13_weight_scale_2, + layer.w2_weight, + layer.w2_weight_scale, + layer.w2_weight_scale_2, + layer.params_dtype, + group_size=self.quant_config.group_size, + invert_global_scales=False, + ) + def create_weights( self, layer: RoutedExperts, diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index 8d8a7eca60..905155f5d8 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -34,10 +34,12 @@ from vllm.model_executor.layers.fusion.quant_activation import ( from vllm.model_executor.layers.linear import ( LinearMethodBase, ) -from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.model_executor.layers.quantization.online.moe_base import ( OnlineMoEMethodBase, ) +from vllm.model_executor.layers.quantization.online.requantization import ( + OnlineRequantizationMixin, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, QuantKey, @@ -118,7 +120,7 @@ def _is_tp_sharded(layer: Module, *, reduces_output_dim: bool = True) -> bool: return is_row_parallel or (reduces_output_dim and is_column_parallel) -class OnlineLinearBase(LinearMethodBase): +class OnlineLinearBase(OnlineRequantizationMixin, LinearMethodBase): """Shared base for online FP8 linear methods. Loads fp16/bf16 checkpoint weights onto meta device and materializes them just-in-time.""" @@ -126,15 +128,17 @@ class OnlineLinearBase(LinearMethodBase): activation_quant_key: QuantKey | None def __init__(self): + super().__init__() self.out_dtype = torch.get_default_dtype() self.input_dtype = get_current_vllm_config().model_config.dtype - self.requantization_source: QuantizeMethodBase | None = None - self.requantization_source_parameters: dict[str, torch.nn.Parameter] = {} - def set_requantization_source(self, source_method: QuantizeMethodBase) -> None: - """Configure serialized-weight conversion before online quantization.""" - self.requantization_source = source_method - self.uses_meta_device = False + def get_weight_for_quantization(self, layer: torch.nn.Module) -> torch.Tensor: + if self.requantization_source is None: + return layer.weight + weight = self.requantization_source.dequantize_weight(layer) + if not isinstance(weight, torch.Tensor): + raise TypeError("Linear requantization requires a single weight tensor.") + return weight def create_weights( self, @@ -146,22 +150,15 @@ class OnlineLinearBase(LinearMethodBase): params_dtype: torch.dtype, **extra_weight_attrs, ): - if self.requantization_source is not None: - existing_parameter_names = set(layer._parameters) - self.requantization_source.create_weights( - layer, - input_size_per_partition, - output_partition_sizes, - input_size, - output_size, - params_dtype, - **extra_weight_attrs, - ) - self.requantization_source_parameters = { - name: parameter - for name, parameter in layer._parameters.items() - if name not in existing_parameter_names and parameter is not None - } + if self.create_requantization_source_weights( + layer, + input_size_per_partition, + output_partition_sizes, + input_size, + output_size, + params_dtype, + **extra_weight_attrs, + ): return output_size_per_partition = sum(output_partition_sizes) @@ -187,19 +184,6 @@ class OnlineLinearBase(LinearMethodBase): initialize_online_processing(layer) - def release_requantization_source_weights(self, layer: torch.nn.Module) -> None: - """Release checkpoint parameters after successful requantization.""" - for name, source_parameter in self.requantization_source_parameters.items(): - if layer._parameters.get(name) is source_parameter: - delattr(layer, name) - self.requantization_source_parameters.clear() - - def get_weight_for_quantization(self, layer: Module) -> torch.Tensor: - """Return checkpoint weights materialized for online quantization.""" - if self.requantization_source is None: - return layer.weight - return self.requantization_source.dequantize_weight(layer) - class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase): """Online tensorwise FP8 linear quantization. @@ -268,6 +252,7 @@ class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase): if self.use_marlin and hasattr(self.fp8_linear, "marlin_input_dtype"): self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype self.fp8_linear.process_weights_after_loading(layer) + self.release_requantization_source_weights(layer) # Prevent duplicate processing (e.g., during weight reload) layer._already_called_process_weights_after_loading = True @@ -368,6 +353,7 @@ class Fp8PerBlockOnlineLinearMethod(OnlineLinearBase): del layer.weight_scale self.fp8_linear.process_weights_after_loading(layer) + self.release_requantization_source_weights(layer) # Prevent duplicate processing (e.g., during weight reload) layer._already_called_process_weights_after_loading = True @@ -624,27 +610,29 @@ class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase): if getattr(layer, "_already_called_process_weights_after_loading", False): return + w13_weight, w2_weight = self.get_weights_for_quantization(layer) + # If checkpoint is fp16, quantize in place. 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 = torch.empty_like(w13_weight, dtype=fp8_dtype) + w2 = torch.empty_like(w2_weight, dtype=fp8_dtype) 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 = weight_amax(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 = weight_amax(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, :, :], _ = ops.scaled_fp8_quant( - layer.w13_weight[expert, :, :], scale=w13_scale[expert] + w13_weight[expert, :, :], scale=w13_scale[expert] ) w2[expert, :, :], _ = ops.scaled_fp8_quant( - layer.w2_weight[expert, :, :], scale=w2_scale[expert] + w2_weight[expert, :, :], scale=w2_scale[expert] ) # Shuffle weights to runtime format and setup kernel. @@ -657,6 +645,7 @@ class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase): w13_input_scale=layer.w13_input_scale, w2_input_scale=layer.w2_input_scale, ) + self.release_requantization_source_weights(layer) # Prevent duplicate processing (e.g., during weight reload) layer._already_called_process_weights_after_loading = True @@ -700,11 +689,12 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase): if getattr(layer, "_already_called_process_weights_after_loading", False): return - self._zero_padding(layer) + w13_weight, w2_weight = self.get_weights_for_quantization(layer) + self._zero_padding(layer, w13_weight, w2_weight) 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 = torch.empty_like(w13_weight, dtype=fp8_dtype) + w2 = torch.empty_like(w2_weight, dtype=fp8_dtype) block_size = self.weight_block_size assert block_size is not None @@ -713,8 +703,8 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase): # Create block-shaped scales (computed here rather than in # create_weights because online quant doesn't need them until now). num_experts = layer.local_num_experts - _, w13_out, w13_in = layer.w13_weight.shape - _, w2_out, w2_in = layer.w2_weight.shape + _, w13_out, w13_in = w13_weight.shape + _, w2_out, w2_in = w2_weight.shape w13_scale = torch.ones( num_experts, @@ -733,18 +723,22 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase): for expert in range(num_experts): w13[expert], w13_scale[expert] = per_block_cast_to_fp8( - layer.w13_weight[expert], + w13_weight[expert], block_size=block_size, use_ue8m0=False, ) w2[expert], w2_scale[expert] = per_block_cast_to_fp8( - layer.w2_weight[expert], + w2_weight[expert], block_size=block_size, use_ue8m0=False, ) layer.weight_block_size = block_size + # Activation quantization is dynamic (including in re-quantization) + layer.w13_input_scale = None + layer.w2_input_scale = None + # Shuffle weights to runtime format and setup kernel. self._setup_kernel( layer, @@ -752,9 +746,10 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase): w2, w13_scale, w2_scale, - layer.w13_input_scale, - layer.w2_input_scale, + w13_input_scale=None, + w2_input_scale=None, ) + self.release_requantization_source_weights(layer) # Prevent duplicate processing (e.g., during weight reload) layer._already_called_process_weights_after_loading = True @@ -806,32 +801,32 @@ class Fp8PtpcOnlineMoEMethod(_Fp8OnlineMoEBase): if getattr(layer, "_already_called_process_weights_after_loading", False): return + w13_weight, w2_weight = self.get_weights_for_quantization(layer) + 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 = torch.empty_like(w13_weight, dtype=fp8_dtype) + w2 = torch.empty_like(w2_weight, dtype=fp8_dtype) # Scale's leading dim is taken from the fp8 weight tensor by # construction, so it cannot drift from the weight's expert count # under EP / padded MoE. - n_w13 = layer.w13_weight.shape[1] + n_w13 = w13_weight.shape[1] w13_scale = torch.ones( w13.shape[0], n_w13, 1, device=w13.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 = weight_amax(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], + w13_weight[expert], scale=None, use_per_token_if_dynamic=True, ) - w2[expert] = _fp8_quant_per_channel( - layer.w2_weight[expert], w2_scale[expert] - ) + w2[expert] = _fp8_quant_per_channel(w2_weight[expert], w2_scale[expert]) self._setup_kernel( layer, @@ -842,5 +837,6 @@ class Fp8PtpcOnlineMoEMethod(_Fp8OnlineMoEBase): w13_input_scale=None, w2_input_scale=None, ) + self.release_requantization_source_weights(layer) layer._already_called_process_weights_after_loading = True diff --git a/vllm/model_executor/layers/quantization/online/int8.py b/vllm/model_executor/layers/quantization/online/int8.py index 268345ca09..5d358488f7 100644 --- a/vllm/model_executor/layers/quantization/online/int8.py +++ b/vllm/model_executor/layers/quantization/online/int8.py @@ -52,42 +52,51 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase): if getattr(layer, "_already_called_process_weights_after_loading", False): return - self._quantize_weights(layer) + w13_weight, w2_weight = self.get_weights_for_quantization(layer) + self._quantize_weights(layer, w13_weight, w2_weight) self._setup_kernel(layer) + self.release_requantization_source_weights(layer) layer._already_called_process_weights_after_loading = True - def _quantize_weights(self, layer: Module) -> None: + def _quantize_weights( + self, + layer: Module, + w13_weight: torch.Tensor | None = None, + w2_weight: torch.Tensor | None = None, + ) -> None: + w13_weight = layer.w13_weight if w13_weight is None else w13_weight + w2_weight = layer.w2_weight if w2_weight is None else w2_weight vmax = torch.iinfo(torch.int8).max - w13 = torch.empty_like(layer.w13_weight, dtype=torch.int8) - w2 = torch.empty_like(layer.w2_weight, dtype=torch.int8) + w13 = torch.empty_like(w13_weight, dtype=torch.int8) + w2 = torch.empty_like(w2_weight, dtype=torch.int8) w13_scale = torch.zeros( layer.num_experts, - layer.w13_weight.shape[1], + w13_weight.shape[1], device=w13.device, dtype=torch.float32, ) w2_scale = torch.zeros( layer.num_experts, - layer.w2_weight.shape[1], + w2_weight.shape[1], device=w2.device, dtype=torch.float32, ) - w2_amax = weight_amax(layer.w2_weight, dim=-1) + w2_amax = weight_amax(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, :, :] + w = w13_weight[expert, :, :] scales = w.abs().amax(dim=1) / vmax q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax) w13[expert, :, :] = q.to(torch.int8) w13_scale[expert, :] = scales # w2: per-row quantization over intermediate_size dim - w = layer.w2_weight[expert, :, :] + w = w2_weight[expert, :, :] scales = w2_amax[expert] / vmax q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax) w2[expert, :, :] = q.to(torch.int8) diff --git a/vllm/model_executor/layers/quantization/online/moe_base.py b/vllm/model_executor/layers/quantization/online/moe_base.py index e296eab373..eb175ebd33 100644 --- a/vllm/model_executor/layers/quantization/online/moe_base.py +++ b/vllm/model_executor/layers/quantization/online/moe_base.py @@ -10,26 +10,26 @@ from vllm.model_executor.layers.fused_moe import ( RoutedExperts, SharedExperts, ) +from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput -from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase +from vllm.model_executor.layers.quantization.online.requantization import ( + OnlineRequantizationMixin, +) from vllm.model_executor.model_loader.reload.layerwise import ( initialize_online_processing, ) from vllm.model_executor.utils import set_weight_attrs -class OnlineMoEMethodBase(FusedMoEMethodBase): +class OnlineMoEMethodBase(OnlineRequantizationMixin, FusedMoEMethodBase): """Base for MoE methods that load full-precision weights on meta device and quantize them after loading via the QeRL layerwise processing system. """ uses_meta_device: bool = True - def set_requantization_source(self, source_method: QuantizeMethodBase) -> None: - """Reject requantization from a checkpoint-quantized MoE method.""" - raise NotImplementedError( - "Requantizing checkpoint-quantized MoE layers is not supported." - ) + def __init__(self, moe: FusedMoEConfig): + super().__init__(moe) def create_weights( self, @@ -40,10 +40,22 @@ class OnlineMoEMethodBase(FusedMoEMethodBase): params_dtype: torch.dtype, **extra_weight_attrs, ): + # Target conversion needs this metadata even when the source creates weights. layer.num_experts = num_experts layer.orig_dtype = params_dtype layer.weight_block_size = None + if self.requantization_source is not None: + self.create_requantization_source_weights( + layer, + num_experts, + layer.moe_config.hidden_dim_unpadded, + layer.moe_config.intermediate_size_per_partition_unpadded, + params_dtype, + **extra_weight_attrs, + ) + return + # Fused gate_up_proj (column parallel) — full precision on meta device w13_weight = torch.nn.Parameter( torch.empty( @@ -103,22 +115,95 @@ class OnlineMoEMethodBase(FusedMoEMethodBase): initialize_online_processing(layer) - def _zero_padding(self, layer: torch.nn.Module) -> None: + def get_weights_for_quantization( + self, layer: torch.nn.Module + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.requantization_source is None: + return layer.w13_weight, layer.w2_weight + + weights = self.requantization_source.dequantize_weight(layer) + if not isinstance(weights, tuple): + raise TypeError("MoE requantization requires w13 and w2 weight tensors.") + + w13, w2 = weights + + config = layer.moe_config + hidden_size = config.hidden_dim_unpadded + intermediate_size = config.intermediate_size_per_partition_unpadded + target_hidden_size = config.hidden_dim + target_intermediate_size = config.intermediate_size_per_partition + + num_experts = layer.local_num_experts + num_shards = self.moe.w13_num_shards + + if w13.shape != (num_experts, num_shards * intermediate_size, hidden_size): + raise ValueError("Requantization requires unpadded source w13 weights.") + if w2.shape != (num_experts, hidden_size, intermediate_size): + raise ValueError("Requantization requires unpadded source w2 weights.") + + if ( + target_hidden_size < hidden_size + or target_intermediate_size < intermediate_size + ): + raise ValueError("Requantization target dimensions cannot shrink weights.") + + # No padding required. + if (hidden_size, intermediate_size) == ( + target_hidden_size, + target_intermediate_size, + ): + return w13, w2 + + if any( + getattr(layer, name, None) is not None for name in ("w13_bias", "w2_bias") + ): + raise NotImplementedError( + "Padded requantization with biases is unsupported." + ) + + def pad_moe_weights( + w13: torch.Tensor, w2: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + padded_w13 = w13.new_zeros( + (num_experts, num_shards * target_intermediate_size, target_hidden_size) + ) + padded_w2 = w2.new_zeros( + (num_experts, target_hidden_size, target_intermediate_size) + ) + for shard in range(num_shards): + source_start = shard * intermediate_size + target_start = shard * target_intermediate_size + padded_w13[ + :, target_start : target_start + intermediate_size, :hidden_size + ] = w13[:, source_start : source_start + intermediate_size, :] + padded_w2[:, :hidden_size, :intermediate_size] = w2 + return padded_w13, padded_w2 + + return pad_moe_weights(w13, w2) + + def _zero_padding( + self, + layer: torch.nn.Module, + w13_weight: torch.Tensor | None = None, + w2_weight: torch.Tensor | None = None, + ) -> None: + w13_weight = layer.w13_weight if w13_weight is None else w13_weight + w2_weight = layer.w2_weight if w2_weight is None else w2_weight hidden_size = layer.moe_config.hidden_dim_unpadded intermediate_size = layer.moe_config.intermediate_size_per_partition_unpadded - w13_shard = layer.w13_weight.shape[1] // self.moe.w13_num_shards + w13_shard = w13_weight.shape[1] // self.moe.w13_num_shards if w13_shard > intermediate_size: for shard in range(self.moe.w13_num_shards): start = shard * w13_shard + intermediate_size - layer.w13_weight[:, start : (shard + 1) * w13_shard, :] = 0 - if layer.w13_weight.shape[2] > hidden_size: - layer.w13_weight[:, :, hidden_size:] = 0 + w13_weight[:, start : (shard + 1) * w13_shard, :] = 0 + if w13_weight.shape[2] > hidden_size: + w13_weight[:, :, hidden_size:] = 0 - if layer.w2_weight.shape[1] > hidden_size: - layer.w2_weight[:, hidden_size:, :] = 0 - if layer.w2_weight.shape[2] > intermediate_size: - layer.w2_weight[:, :, intermediate_size:] = 0 + if w2_weight.shape[1] > hidden_size: + w2_weight[:, hidden_size:, :] = 0 + if w2_weight.shape[2] > intermediate_size: + w2_weight[:, :, intermediate_size:] = 0 if getattr(layer, "w13_bias", None) is not None: w13_bias_shard = layer.w13_bias.shape[1] // self.moe.w13_num_shards diff --git a/vllm/model_executor/layers/quantization/online/mxfp4.py b/vllm/model_executor/layers/quantization/online/mxfp4.py index b949c9b7ff..70a64340f6 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp4.py +++ b/vllm/model_executor/layers/quantization/online/mxfp4.py @@ -127,6 +127,7 @@ class Mxfp4OnlineLinearMethod(OnlineLinearBase): replace_parameter(layer, "weight_scale", weight_scale.data) self.kernel.process_weights_after_loading(layer) + self.release_requantization_source_weights(layer) layer._already_called_process_weights_after_loading = True @@ -295,7 +296,8 @@ class Mxfp4OnlineMoEMethod(OnlineMoEMethodBase): if getattr(layer, "_already_called_process_weights_after_loading", False): return - self._zero_padding(layer) + w13_weight, w2_weight = self.get_weights_for_quantization(layer) + self._zero_padding(layer, w13_weight, w2_weight) if self.mxfp4_backend == Mxfp4MoeBackend.NONE: layer._already_called_process_weights_after_loading = True @@ -304,8 +306,8 @@ class Mxfp4OnlineMoEMethod(OnlineMoEMethodBase): layer.w13_input_scale = None layer.w2_input_scale = None - w13, w13_scale = _quantize_mxfp4_moe_weight(layer.w13_weight) - w2, w2_scale = _quantize_mxfp4_moe_weight(layer.w2_weight) + w13, w13_scale = _quantize_mxfp4_moe_weight(w13_weight) + w2, w2_scale = _quantize_mxfp4_moe_weight(w2_weight) self._setup_kernel( layer, @@ -316,5 +318,6 @@ class Mxfp4OnlineMoEMethod(OnlineMoEMethodBase): getattr(layer, "w13_bias", None), getattr(layer, "w2_bias", None), ) + self.release_requantization_source_weights(layer) layer._already_called_process_weights_after_loading = True diff --git a/vllm/model_executor/layers/quantization/online/mxfp8.py b/vllm/model_executor/layers/quantization/online/mxfp8.py index ecc196e9f1..f16f2b9fbf 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp8.py +++ b/vllm/model_executor/layers/quantization/online/mxfp8.py @@ -34,7 +34,6 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( mxfp8_e4m3_quantize, ) from vllm.model_executor.utils import replace_parameter -from vllm.platforms import current_platform class Mxfp8OnlineLinearMethod(OnlineLinearBase): @@ -86,6 +85,7 @@ class Mxfp8OnlineLinearMethod(OnlineLinearBase): replace_parameter(layer, "weight_scale", weight_scale.data) self.kernel.process_weights_after_loading(layer) + self.release_requantization_source_weights(layer) layer._already_called_process_weights_after_loading = True @@ -237,14 +237,12 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase): if getattr(layer, "_already_called_process_weights_after_loading", False): return - 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_weight, w2_weight = self.get_weights_for_quantization(layer) layer.w13_input_scale = None layer.w2_input_scale = None - w13, w13_scale = self._quantize_mxfp8_moe_weight(layer.w13_weight) - w2, w2_scale = self._quantize_mxfp8_moe_weight(layer.w2_weight) + w13, w13_scale = self._quantize_mxfp8_moe_weight(w13_weight) + w2, w2_scale = self._quantize_mxfp8_moe_weight(w2_weight) self._setup_kernel( layer, @@ -255,5 +253,6 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase): layer.w13_input_scale, layer.w2_input_scale, ) + self.release_requantization_source_weights(layer) layer._already_called_process_weights_after_loading = True diff --git a/vllm/model_executor/layers/quantization/online/nvfp4.py b/vllm/model_executor/layers/quantization/online/nvfp4.py index 6c780c2c46..1d79a53fd3 100644 --- a/vllm/model_executor/layers/quantization/online/nvfp4.py +++ b/vllm/model_executor/layers/quantization/online/nvfp4.py @@ -116,19 +116,26 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase): if getattr(layer, "_already_called_process_weights_after_loading", False): return - self._quantize_weights(layer) + w13_weight, w2_weight = self.get_weights_for_quantization(layer) + self._quantize_weights(layer, w13_weight, w2_weight) self._setup_kernel(layer) + self.release_requantization_source_weights(layer) layer._already_called_process_weights_after_loading = True - def _quantize_weights(self, layer: Module) -> None: + def _quantize_weights( + self, + layer: Module, + w13_weight: torch.Tensor | None = None, + w2_weight: torch.Tensor | None = None, + ) -> None: + w13_weight = layer.w13_weight if w13_weight is None else w13_weight + w2_weight = layer.w2_weight if w2_weight is None else 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 + w13_weight, moe_tp_size ) + w2, w2_scale, w2_scale_2 = _quantize_moe_weight_to_nvfp4(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/online/requantization.py b/vllm/model_executor/layers/quantization/online/requantization.py new file mode 100644 index 0000000000..80ff9f6789 --- /dev/null +++ b/vllm/model_executor/layers/quantization/online/requantization.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase + + +class OnlineRequantizationMixin: + """Shared lifecycle for converting serialized weights before quantization.""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.requantization_source: QuantizeMethodBase | None = None + self.requantization_source_parameters: dict[str, torch.nn.Parameter] = {} + + def set_requantization_source(self, source_method: QuantizeMethodBase) -> None: + """Configure serialized-weight conversion before online quantization.""" + self.requantization_source = source_method + self.uses_meta_device = False + + def create_requantization_source_weights( + self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs + ) -> bool: + """Delegate weight creation to the source and track its parameters.""" + if self.requantization_source is None: + return False + + existing_parameter_names = set(layer._parameters) + self.requantization_source.create_weights( + layer, *weight_args, **extra_weight_attrs + ) + self.requantization_source_parameters = { + name: parameter + for name, parameter in layer._parameters.items() + if name not in existing_parameter_names and parameter is not None + } + return True + + def release_requantization_source_weights(self, layer: torch.nn.Module) -> None: + """Release serialized parameters after successful requantization.""" + for name, source_parameter in self.requantization_source_parameters.items(): + if layer._parameters.get(name) is source_parameter: + delattr(layer, name) + self.requantization_source_parameters.clear() diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 262f9bdb0a..849b0c9638 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -69,6 +69,9 @@ from vllm.model_executor.layers.quantization.quark.utils import ( from vllm.model_executor.layers.quantization.utils.fp8_utils import ( validate_fp8_block_shape_moe, ) +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( + dequantize_nvfp4_moe_weights, +) from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( _ACTIVATION_QUANT_KEY_MAP, _WEIGHT_QUANT_KEY_MAP, @@ -2110,6 +2113,22 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod): layer.register_parameter("w2_input_scale_2", w2_input_scale_2) set_weight_attrs(w2_input_scale_2, extra_weight_attrs) + def dequantize_weight( + self, layer: torch.nn.Module + ) -> tuple[torch.Tensor, torch.Tensor]: + """Dequantize Quark NVFP4 expert weights for requantization.""" + return dequantize_nvfp4_moe_weights( + layer.w13_weight, + layer.w13_weight_scale, + layer.w13_weight_scale_2, + layer.w2_weight, + layer.w2_weight_scale, + layer.w2_weight_scale_2, + layer.params_dtype, + group_size=self.group_size, + invert_global_scales=False, + ) + def process_weights_after_loading(self, layer: RoutedExperts) -> None: """Convert NVFP4 MoE weights into kernel format and setup the kernel.""" # Match existing NVFP4 MoE paths: fused w13 uses the w1 global scale. diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py index 568ae3fd1b..cad1412066 100644 --- a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py +++ b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py @@ -11,6 +11,7 @@ from vllm.utils.torch_utils import direct_register_custom_op __all__ = [ "break_fp4_bytes", + "dequantize_nvfp4_moe_weights", "dequantize_to_dtype", "nvfp4_gathered_bias", "ref_nvfp4_quant", @@ -629,6 +630,58 @@ def dequantize_to_dtype( return out.to(dtype) +def dequantize_nvfp4_moe_weights( + w13_weight_packed: torch.Tensor, + w13_weight_scale: torch.Tensor, + w13_weight_global_scale: torch.Tensor, + w2_weight_packed: torch.Tensor, + w2_weight_scale: torch.Tensor, + w2_weight_global_scale: torch.Tensor, + dtype: torch.dtype, + group_size: int = 16, + *, + invert_global_scales: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + """Materialize serialized NVFP4 MoE weights for requantization. + + Compressed-tensors stores global encode scales, while ModelOpt stores + ``weight_scale_2`` as the corresponding decode multiplier. Set + ``invert_global_scales`` according to the source checkpoint convention. + """ + if invert_global_scales: + w13_weight_global_scale = 1.0 / w13_weight_global_scale + w2_weight_global_scale = 1.0 / w2_weight_global_scale + + w13_num_shards = w13_weight_global_scale.shape[1] + w13_weights = torch.chunk(w13_weight_packed, w13_num_shards, dim=1) + w13_scales = torch.chunk(w13_weight_scale, w13_num_shards, dim=1) + w13 = torch.cat( + [ + dequantize_to_dtype( + weight, + scale, + w13_weight_global_scale[:, shard_idx].contiguous(), + dtype, + block_size=group_size, + swizzle=False, + ) + for shard_idx, (weight, scale) in enumerate( + zip(w13_weights, w13_scales, strict=True) + ) + ], + dim=1, + ) + w2 = dequantize_to_dtype( + w2_weight_packed, + w2_weight_scale, + w2_weight_global_scale, + dtype, + block_size=group_size, + swizzle=False, + ) + return w13, w2 + + def get_reciprocal(x): if isinstance(x, torch.Tensor): # torch.where yields operation not permitted when stream is capturing.