diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py index 06e5923a30f..ff02e24b6ab 100644 --- a/modelopt/torch/export/convert_hf_config.py +++ b/modelopt/torch/export/convert_hf_config.py @@ -104,6 +104,19 @@ def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) }, "weights": {"dynamic": False, "num_bits": 8, "type": "float", "group_size": gs}, } + elif quant_algo == "FP8_PB": + # Block-wise FP8 (W8A8). Weights in gsxgs blocks; activations quantized + # dynamically per-token at runtime (no input_activations entry). + gs = group_size or 128 + return { + "weights": { + "dynamic": False, + "num_bits": 8, + "type": "float", + "strategy": "block", + "block_structure": [gs, gs], + }, + } else: warnings.warn( f"Unsupported quantization algorithm '{quant_algo}' in " @@ -166,6 +179,29 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An original_quantization_details = input_config.get("quantization", {}) quant_algo_value = original_quantization_details.get("quant_algo") + # FP8_PB (block-wise FP8, W8A8): emit the native ``quant_method: fp8`` config + # vLLM/SGLang expect (weight_scale_inv + dynamic activations), matching the + # official Qwen3.5 FP8 checkpoint. + if quant_algo_value == "FP8_PB": + kv_cache_quant_algo = original_quantization_details.get("kv_cache_quant_algo") + assert not kv_cache_quant_algo, ( + "FP8_PB export does not support kv_cache quantization yet " + f"(got kv_cache_quant_algo={kv_cache_quant_algo!r})." + ) + group_size = original_quantization_details.get("group_size") or 128 + exclude_modules = original_quantization_details.get("exclude_modules") or [] + fp8_config: dict[str, Any] = { + "quant_method": "fp8", + "fmt": "e4m3", + "activation_scheme": "dynamic", + "weight_block_size": [group_size, group_size], + "modules_to_not_convert": exclude_modules, + } + producer_info = input_config.get("producer") + if producer_info: + fp8_config["producer"] = producer_info + return fp8_config + # This structure is derived based on the example for "FP8" and "NVFP4" # TODO: Handle other quantization algorithms if quant_algo_value == "FP8": diff --git a/modelopt/torch/export/model_config.py b/modelopt/torch/export/model_config.py index 5f92cc2e5dc..bb6200d9dc1 100755 --- a/modelopt/torch/export/model_config.py +++ b/modelopt/torch/export/model_config.py @@ -42,6 +42,7 @@ QUANTIZATION_NVFP4_AWQ = "nvfp4_awq" QUANTIZATION_FP8_PB_REAL = "fp8_pb_real" QUANTIZATION_FP8_PB_WO = "fp8_pb_wo" +QUANTIZATION_FP8_PB_W8A8 = "fp8_pb_w8a8" QUANTIZATION_FP8_PC_PT = "fp8_pc_pt" KV_CACHE_FP8 = "FP8" diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 2af5f6eab0b..3c457e288d6 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -54,6 +54,7 @@ KV_CACHE_NVFP4_AFFINE, QUANTIZATION_FP8, QUANTIZATION_FP8_PB_REAL, + QUANTIZATION_FP8_PB_W8A8, QUANTIZATION_FP8_PB_WO, QUANTIZATION_FP8_PC_PT, QUANTIZATION_INT4_AWQ, @@ -537,10 +538,12 @@ def _get_quantization_from_layer(layer, quantizer_attr_names: QuantizerAttrNames and block_sizes.get("scale_bits") == (8, 0) ): return QUANTIZATION_MXFP8 - if weight_quantizer.fake_quant: - return QUANTIZATION_FP8_PB_WO - else: + # Block FP8: input quantizer enabled -> W8A8, else weight-only. + if not weight_quantizer.fake_quant: return QUANTIZATION_FP8_PB_REAL + if input_quantizer is not None and input_quantizer.is_enabled: + return QUANTIZATION_FP8_PB_W8A8 + return QUANTIZATION_FP8_PB_WO if weight_quantizer.axis == 0: return QUANTIZATION_FP8_PC_PT return QUANTIZATION_FP8 @@ -758,6 +761,12 @@ def process_layer_quant_config(layer_config_dict): "quant_algo": "MXFP8", "group_size": block_size_value, } + elif v == "fp8_pb_w8a8": + # Block-wise FP8, W8A8 at serve time. + layer_config = { + "quant_algo": "FP8_PB", + "group_size": block_size_value, + } else: layer_config = {"quant_algo": v} @@ -865,7 +874,7 @@ def to_quantized_weight( if quantization == QUANTIZATION_MXFP8: return MXFP8QTensor.quantize_with_scale(weight, weights_scaling_factor) - if quantization == QUANTIZATION_FP8_PB_WO: + if quantization in (QUANTIZATION_FP8_PB_WO, QUANTIZATION_FP8_PB_W8A8): return FP8QTensor.quantize( weight, weights_scaling_factor.squeeze(), block_sizes={-1: block_size, -2: block_size} )[0]._quantized_data diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 8bc92ed5eb9..813226b1761 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -78,6 +78,7 @@ from .model_config import ( QUANTIZATION_FP8, QUANTIZATION_FP8_PB_REAL, + QUANTIZATION_FP8_PB_W8A8, QUANTIZATION_FP8_PC_PT, QUANTIZATION_MXFP8, QUANTIZATION_NONE, @@ -741,9 +742,22 @@ def _export_quantized_weight( setattr(sub_module, weight_name, nn.Parameter(quantized_weight, requires_grad=False)) - # Register the corrected weight_scale as a buffer + # Register the corrected weight scale as a buffer. if weight_scale is not None: - sub_module.register_buffer(quantizer_attrs.weight_scale, weight_scale) + if quantization_format == QUANTIZATION_FP8_PB_W8A8: + # Store per-block scale as 2-D weight_scale_inv (amax/448, not + # inverted); squeeze the keepdim block-amax [out, 1, in, 1]. + if ( + weight_scale.dim() == 4 + and weight_scale.shape[1] == 1 + and weight_scale.shape[3] == 1 + ): + weight_scale = weight_scale.squeeze(3).squeeze(1) + sub_module.register_buffer(quantizer_attrs.weight_scale_inv, weight_scale) + if quantizer_attrs.weight_scale in sub_module._buffers: + del sub_module._buffers[quantizer_attrs.weight_scale] + else: + sub_module.register_buffer(quantizer_attrs.weight_scale, weight_scale) # Tied-weight dedup: if a previously-processed module shared the same # source weight memory, alias the packed weight + scale buffers so the @@ -758,6 +772,7 @@ def _export_quantized_weight( setattr(sub_module, weight_name, getattr(_prior, weight_name)) for _attr in ( quantizer_attrs.weight_scale, + quantizer_attrs.weight_scale_inv, quantizer_attrs.weight_scale_2, quantizer_attrs.input_scale, ): diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 070a4478838..5b353d010e7 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -41,6 +41,7 @@ KV_CACHE_NVFP4, QUANTIZATION_FP8, QUANTIZATION_FP8_PB_REAL, + QUANTIZATION_FP8_PB_W8A8, QUANTIZATION_FP8_PB_WO, QUANTIZATION_NONE, QUANTIZATION_NVFP4, @@ -296,6 +297,11 @@ def save_pretrained( quantization = "NVFP4" elif quantization_format == QUANTIZATION_W4A16_NVFP4: quantization = "W4A16_NVFP4" + elif quantization_format == QUANTIZATION_FP8_PB_W8A8: + raise NotImplementedError( + "Block-wise FP8 W8A8 (FP8_PB_W8A8) export is not supported on the " + "Megatron path; export via the HF path (unified_export_hf)." + ) # We use the last PP rank and the 1st EP rank to write the config because # medusa_heads and eagle_module only exist in the last stage. @@ -860,6 +866,11 @@ def _get_quantized_state( """ name_to_value = {} qformat: str = self._get_quantization_format(module) + if qformat == QUANTIZATION_FP8_PB_W8A8: + raise NotImplementedError( + "Block-wise FP8 W8A8 (FP8_PB_W8A8) export is not supported on the " + "Megatron path; export via the HF path (unified_export_hf)." + ) if qformat is None and "norm" not in prefix: self._record_excluded_module(prefix) block_size = get_weight_block_size(module) diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index b0049b5a08d..34d0bc82703 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -272,6 +272,7 @@ def weight_attr_names(module: nn.Module) -> "Generator[str, None, None]": "input_quantizer", "output_quantizer", "weight_scale", + "weight_scale_inv", "weight_scale_2", "input_scale", "output_scale", @@ -287,6 +288,7 @@ def quantizer_attr_names(weight_name: str = "weight") -> QuantizerAttrNames: input_quantizer=f"{prefix}input_quantizer", output_quantizer=f"{prefix}output_quantizer", weight_scale=f"{prefix}weight_scale", + weight_scale_inv=f"{prefix}weight_scale_inv", weight_scale_2=f"{prefix}weight_scale_2", input_scale=f"{prefix}input_scale", output_scale=f"{prefix}output_scale", diff --git a/modelopt_recipes/configs/ptq/presets/model/fp8_2d_blockwise_w8a8_dynamic.yaml b/modelopt_recipes/configs/ptq/presets/model/fp8_2d_blockwise_w8a8_dynamic.yaml new file mode 100644 index 00000000000..7ce0df50d9b --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/fp8_2d_blockwise_w8a8_dynamic.yaml @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# QuantizeConfig preset for FP8 E4M3 2D blockwise weights + dynamic per-token FP8 +# activations (W8A8). The dynamic input quantizer makes PTQ calibrate as W8A8; +# it stores no input_scale, so exported weights match the weight-only preset. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + fp8: configs/numerics/fp8 + +algorithm: max +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*weight_quantizer' + cfg: + $import: fp8 + block_sizes: + -1: 128 + -2: 128 + - quantizer_name: '*input_quantizer' + cfg: + $import: fp8 + block_sizes: + -1: 128 + type: dynamic + - $import: default_disabled_quantizers diff --git a/tests/gpu/torch/export/test_export.py b/tests/gpu/torch/export/test_export.py index cac0a9a9aef..c48aa0824ac 100644 --- a/tests/gpu/torch/export/test_export.py +++ b/tests/gpu/torch/export/test_export.py @@ -141,6 +141,21 @@ def test_get_quantization_format(config, expected): "exclude_modules": ["layer8"], }, ), + ( + { + "layer1.quantization": "fp8_pb_w8a8", # 128x128 block-wise FP8 (W8A8) + "layer1.awq_block_size": 128, + "layer2.quantization": "fp8_pb_w8a8", + "layer2.awq_block_size": 128, + "layer8.quantization": None, + }, + { + "quant_algo": "FP8_PB", + "kv_cache_quant_algo": None, + "group_size": 128, + "exclude_modules": ["layer8"], + }, + ), ], ) def test_process_layer_quant_config(layer_config_dict, expected_processed_dict): @@ -520,3 +535,66 @@ def is_excluded(module_name: str) -> bool: assert not is_excluded("model.layers.0.mlp.experts.0.down_proj"), ( f"Routed experts should not be excluded, got patterns: {exclude_modules}" ) + + +def test_fp8_pb_w8a8_export_uses_weight_scale_inv(tmp_path): + """W8A8 block-FP8 (FP8_PB) export stores per-block scales as weight_scale_inv + (DeepSeek/Qwen convention) and drops the plain weight_scale; activations are + dynamic so no input_scale is stored.""" + from safetensors import safe_open + + model = get_tiny_qwen3_moe().to("cuda") + model.config.architectures = ["Qwen3MoeForCausalLM"] + + cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": {"num_bits": (4, 3), "block_sizes": {-1: 128, -2: 128}, "axis": None}, + "enable": True, + }, + { + "quantizer_name": "*input_quantizer", + "cfg": { + "num_bits": (4, 3), + "block_sizes": {-1: 128, "type": "dynamic"}, + "axis": None, + }, + "enable": True, + }, + {"quantizer_name": "*lm_head*", "enable": False}, + ], + "algorithm": "max", + } + dummy_inputs = {k: v.to("cuda") for k, v in model.dummy_inputs.items()} + mtq.quantize(model, cfg, lambda m: m(**dummy_inputs)) + + export_dir = tmp_path / "fp8_pb_w8a8" + export_hf_checkpoint(model, export_dir=export_dir) + + keys = set() + for st in export_dir.glob("*.safetensors"): + with safe_open(st, framework="pt") as f: + keys.update(f.keys()) + + assert any(k.endswith(".weight_scale_inv") for k in keys), ( + "block FP8 must store weight_scale_inv" + ) + assert not any(k.endswith(".weight_scale") for k in keys), "plain weight_scale must be dropped" + assert not any(k.endswith(".input_scale") for k in keys), ( + "dynamic activations store no input_scale" + ) + + # Per-block scales must be 2-D [out_blocks, in_blocks] (DeepSeek/Qwen + # convention) -- not the 4-D [out_blocks, 1, in_blocks, 1] block-amax shape. + # This is what ModelOpt's _QuantFP8Linear reload path and vLLM/SGLang's stock + # block-FP8 loader expect. + for st in export_dir.glob("*.safetensors"): + with safe_open(st, framework="pt") as f: + for k in list(f.keys()): + if k.endswith(".weight_scale_inv"): + scale_inv = f.get_tensor(k) + assert scale_inv.ndim == 2, ( + f"{k} must be 2-D [out_blocks, in_blocks], got {tuple(scale_inv.shape)}" + ) diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index 1199f4c7cf0..448152a973b 100644 --- a/tests/unit/torch/export/test_get_quantization.py +++ b/tests/unit/torch/export/test_get_quantization.py @@ -26,16 +26,39 @@ from modelopt.torch.export.layer_utils import get_quantization_format from modelopt.torch.export.model_config import ( QUANTIZATION_FP8, + QUANTIZATION_FP8_PB_W8A8, QUANTIZATION_NVFP4, QUANTIZATION_W4A8_AWQ, ) from modelopt.torch.export.quant_utils import get_quant_config from modelopt.torch.quantization.nn import NVFP4StaticQuantizer +# Block-wise FP8 W8A8 on the ".1" linear: weight block + enabled dynamic input quantizer. +_fp8_pb_w8a8_config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*.1.weight_quantizer", + "cfg": {"num_bits": (4, 3), "block_sizes": {-1: 128, -2: 128}, "axis": None}, + "enable": True, + }, + { + "quantizer_name": "*.1.input_quantizer", + "cfg": {"num_bits": (4, 3), "block_sizes": {-1: 128, "type": "dynamic"}, "axis": None}, + "enable": True, + }, + ], + "algorithm": "max", +} + @pytest.mark.parametrize( ("config", "expected"), - [(partial_fp8_config, QUANTIZATION_FP8), (partial_w4a8_config, QUANTIZATION_W4A8_AWQ)], + [ + (partial_fp8_config, QUANTIZATION_FP8), + (partial_w4a8_config, QUANTIZATION_W4A8_AWQ), + (_fp8_pb_w8a8_config, QUANTIZATION_FP8_PB_W8A8), + ], ) def test_get_quantization_format(config, expected): model = ToyModel()