diff --git a/src/mobius/_configs/_quantization.py b/src/mobius/_configs/_quantization.py index 40556e0b..69432ebf 100644 --- a/src/mobius/_configs/_quantization.py +++ b/src/mobius/_configs/_quantization.py @@ -53,6 +53,27 @@ def from_transformers(cls, hf_config) -> QuantizationConfig | None: if not isinstance(qc, dict): return None method = qc.get("quant_method", "none") + # NVIDIA ModelOpt NVFP4/FP8 checkpoints (e.g. quantized Qwen3.6) encode + # weights as packed E2M1 (fp4) / float8 with block + global scales — a + # layout the INT4 ``QuantizedLinear``/``MatMulNBits`` path below would + # silently mis-dequantize. Checked before the ``quant_method == "none"`` + # early-return because ModelOpt may name its scheme only via + # ``quant_algo``/``quant_cfg`` (no ``quant_method``). Fail loudly: the + # reconstruction math is available in :mod:`mobius.integrations.modelopt`, + # but the full weight-load integration + native routed-expert NVFP4 QMoE + # emission (CUDA/Blackwell, ``onnxruntime_USE_FP4_QMOE=ON``) are not wired. + from mobius.integrations.modelopt import is_modelopt_quant_config + + if is_modelopt_quant_config(qc): + raise NotImplementedError( + "NVIDIA ModelOpt NVFP4/FP8 checkpoints are not yet fully " + "supported for export. The weight-reconstruction core lives in " + "mobius.integrations.modelopt (dequantize_nvfp4 / dequantize_fp8); " + "remaining work is checkpoint weight-loading and native " + "routed-expert NVFP4 QMoE emission (CUDA/Blackwell, " + "onnxruntime_USE_FP4_QMOE=ON). Export the unquantized (bf16) " + "checkpoint instead, or quantize the bf16 export via Olive." + ) if method == "none": return None # FP8 per-tensor quantization (float8_e4m3fn + scalar scale) diff --git a/src/mobius/_configs_test.py b/src/mobius/_configs_test.py index da97a82e..ba1af0ab 100644 --- a/src/mobius/_configs_test.py +++ b/src/mobius/_configs_test.py @@ -837,6 +837,24 @@ def test_from_transformers_fp8_returns_none(self): )() assert QuantizationConfig.from_transformers(hf) is None + def test_from_transformers_modelopt_nvfp4_raises(self): + """ModelOpt NVFP4/FP8 checkpoints fail loudly rather than mis-quantizing. + + The INT4 path would silently mis-dequantize packed E2M1/float8 weights, + so ``from_transformers`` raises until the ModelOpt loader is wired. + """ + import pytest + + for qc in ( + {"quant_method": "modelopt"}, + {"quant_method": "modelopt", "quant_algo": "NVFP4"}, + {"quant_algo": "W4A16_NVFP4"}, + {"quant_algo": "FP8"}, + ): + hf = type("HFConfig", (), {"quantization_config": qc})() + with pytest.raises(NotImplementedError, match="ModelOpt"): + QuantizationConfig.from_transformers(hf) + def test_from_transformers_to_dict_object(self): """HF QuantizationConfig objects have a to_dict() method.""" inner = type( diff --git a/src/mobius/integrations/modelopt/__init__.py b/src/mobius/integrations/modelopt/__init__.py new file mode 100644 index 00000000..a7a2c099 --- /dev/null +++ b/src/mobius/integrations/modelopt/__init__.py @@ -0,0 +1,37 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""NVIDIA TensorRT Model Optimizer (ModelOpt) checkpoint support. + +ModelOpt exports mixed-precision NVFP4 + FP8 checkpoints (e.g. Qwen3.6). This +package provides the weight-reconstruction math used to consume those +checkpoints: + +- NVFP4 (``W4A16_NVFP4``): block-16 E2M1 4-bit weights with FP8-E4M3 block + scales and a per-tensor FP32 global scale (``weight_scale_2``). +- FP8 (``E4M3``): per-tensor scaled float8 weights. + +The dequantization functions reconstruct BF16 weights so the standard mobius +build path (plain ``Linear`` / ``bf16`` graph) can consume ModelOpt checkpoints +without a native FP8/NVFP4 kernel. Native routed-expert NVFP4 QMoE emission +(the CUDA-only, Blackwell ``QMoE`` ``quant_type="nvfp4"`` op) is intentionally +out of scope here — see the module docstring in :mod:`._dequant`. +""" + +from __future__ import annotations + +from mobius.integrations.modelopt._dequant import ( + FP4_E2M1_LUT, + dequantize_fp8, + dequantize_nvfp4, + is_modelopt_quant_config, + unpack_nvfp4_codes, +) + +__all__ = [ + "FP4_E2M1_LUT", + "dequantize_fp8", + "dequantize_nvfp4", + "is_modelopt_quant_config", + "unpack_nvfp4_codes", +] diff --git a/src/mobius/integrations/modelopt/_dequant.py b/src/mobius/integrations/modelopt/_dequant.py new file mode 100644 index 00000000..ce1255e0 --- /dev/null +++ b/src/mobius/integrations/modelopt/_dequant.py @@ -0,0 +1,131 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""NVFP4 / FP8 weight reconstruction for NVIDIA ModelOpt checkpoints. + +NVIDIA TensorRT Model Optimizer (ModelOpt) exports mixed-precision checkpoints +where different modules use different numeric formats. For Qwen3.6 the layout is: + +- **Routed MoE experts, shared expert, lm_head**: ``W4A16_NVFP4`` — block-16 + E2M1 (fp4) 4-bit weights, FP8-E4M3 per-block scales, and a per-tensor FP32 + global scale stored as ``weight_scale_2``. +- **Attention / linear-attention projections**: ``FP8`` (E4M3) — float8 weights + with a per-tensor ``weight_scale``. + +ONNX Runtime has no FP8 attention GEMM, so — following ORT GenAI's ModelOpt +loader — the dense FP8 projections and the NVFP4 shared-expert / lm_head are +**dequantized back to BF16**, which reconstructs them exactly and lets the +standard (BF16) build path consume the checkpoint. Only the *routed* MoE experts +are consumed natively by the CUDA QMoE ``quant_type="nvfp4"`` op; that native +emission is Blackwell/``onnxruntime_USE_FP4_QMOE=ON``-only and is intentionally +NOT implemented here (it cannot be built or verified without that runtime). + +The functions below are the numeric core of the loader and are format-faithful +to ModelOpt's on-disk encoding (verified in ``_dequant_test.py``). + +E2M1 (fp4) magnitude table, indexed by the low 3 bits of each 4-bit code +(bit 3 is the sign): ``{0, 0.5, 1, 1.5, 2, 3, 4, 6}``. +""" + +from __future__ import annotations + +import ml_dtypes +import numpy as np + +# E2M1 (fp4) decoded magnitudes, indexed by ``code & 0x7`` (bit 3 is the sign). +FP4_E2M1_LUT = np.array([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=np.float32) + +# NVFP4 pins the weight block size to 16 K-elements per FP8-E4M3 block scale. +NVFP4_BLOCK_SIZE = 16 + + +def unpack_nvfp4_codes(packed_nk2: np.ndarray) -> np.ndarray: + """Unpack a ModelOpt NVFP4 weight tensor to per-element E2M1 codes. + + ``packed_nk2`` is uint8 ``[N, K/2]`` where each byte holds two adjacent + K-axis E2M1 codes for the same output row ``N`` (low nibble = even ``K``, + high nibble = odd ``K``) — the layout ModelOpt writes. Returns uint8 codes + ``[N, K]`` in ``0..15``. + """ + packed = np.ascontiguousarray(packed_nk2).astype(np.uint8) + low = packed & 0x0F + high = packed >> 4 + n = packed.shape[0] + # Interleave (low, high) along a new trailing axis, then flatten to [N, K]. + codes = np.stack((low, high), axis=-1).reshape(n, -1) + return np.ascontiguousarray(codes) + + +def dequantize_nvfp4( + weight_u8: np.ndarray, + block_scale_e4m3: np.ndarray, + global_scale: float | np.floating, +) -> np.ndarray: + """Reconstruct a BF16 weight from ModelOpt NVFP4 tensors. + + Args: + weight_u8: uint8 ``[N, K/2]`` packed E2M1 codes (low nibble = even + ``K``, high nibble = odd ``K``). + block_scale_e4m3: FP8-E4M3 ``[N, K/16]`` per-block scales + (``ml_dtypes.float8_e4m3fn`` or a uint8 raw-code view). + global_scale: per-tensor FP32 scalar. + + Returns: + ``ml_dtypes.bfloat16`` array ``[N, K]`` where + ``w = e2m1(code) * e4m3(block_scale[n, k // 16]) * global_scale``. + """ + codes = unpack_nvfp4_codes(weight_u8).astype(np.int64) # [N, K] + mag = FP4_E2M1_LUT[codes & 0x7] # [N, K] + val = np.where((codes & 0x8) > 0, -mag, mag).astype(np.float32) # [N, K] + + block_scale = _to_float32(block_scale_e4m3) # [N, K/16] + k = codes.shape[1] + n_blocks = block_scale.shape[1] + if k % n_blocks != 0: + raise ValueError(f"NVFP4 K={k} is not divisible by the block count {n_blocks}.") + block_scale = np.repeat(block_scale, k // n_blocks, axis=1) # [N, K] + + dequant = val * block_scale * np.float32(global_scale) + return dequant.astype(ml_dtypes.bfloat16) + + +def dequantize_fp8( + weight_f8: np.ndarray, + weight_scale: float | np.floating, +) -> np.ndarray: + """Reconstruct a BF16 weight from an FP8 (E4M3) weight + per-tensor scale. + + ``w = e4m3(weight) * weight_scale``. + """ + weight = _to_float32(weight_f8) + return (weight * np.float32(weight_scale)).astype(ml_dtypes.bfloat16) + + +def is_modelopt_quant_config(quantization_config: dict | None) -> bool: + """Return ``True`` if a HF ``quantization_config`` is a ModelOpt export. + + ModelOpt writes ``quant_method="modelopt"`` (newer HF serialisations) or a + ``quant_algo`` / ``quant_cfg`` naming an ``NVFP4`` / ``FP8`` scheme. This + keeps detection tolerant of both spellings. + """ + if not isinstance(quantization_config, dict): + return False + method = str(quantization_config.get("quant_method", "")).lower() + if method == "modelopt": + return True + algo = str( + quantization_config.get("quant_algo") or quantization_config.get("quant_cfg") or "" + ).upper() + return "NVFP4" in algo or "FP8" in algo + + +def _to_float32(arr: np.ndarray) -> np.ndarray: + """View/convert an FP8-E4M3 (or raw uint8 code) tensor as float32. + + Raw ``uint8`` inputs are reinterpreted as ``float8_e4m3fn`` byte codes + before widening, matching how ModelOpt stores block scales. + """ + arr = np.asarray(arr) + if arr.dtype == np.uint8: + arr = arr.view(ml_dtypes.float8_e4m3fn) + return arr.astype(np.float32) diff --git a/src/mobius/integrations/modelopt/_dequant_test.py b/src/mobius/integrations/modelopt/_dequant_test.py new file mode 100644 index 00000000..0a26b79f --- /dev/null +++ b/src/mobius/integrations/modelopt/_dequant_test.py @@ -0,0 +1,101 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for ModelOpt NVFP4 / FP8 weight reconstruction.""" + +from __future__ import annotations + +import ml_dtypes +import numpy as np + +from mobius.integrations.modelopt import ( + FP4_E2M1_LUT, + dequantize_fp8, + dequantize_nvfp4, + is_modelopt_quant_config, + unpack_nvfp4_codes, +) + + +def _e2m1_code(sign: int, mag_index: int) -> int: + """Build a 4-bit E2M1 code from a sign bit and a magnitude index (0..7).""" + return (sign << 3) | (mag_index & 0x7) + + +def test_fp4_lut_values(): + assert FP4_E2M1_LUT.tolist() == [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] + + +def test_unpack_nvfp4_codes_splits_nibbles(): + # byte = low | (high << 4); low nibble is the even-K code. + packed = np.array([[0x21, 0xC6]], dtype=np.uint8) # (1, 0x21=33, 0xC6=198) + codes = unpack_nvfp4_codes(packed) + # 0x21 -> low=1, high=2 ; 0xC6 -> low=6, high=12 + assert codes.tolist() == [[1, 2, 6, 12]] + assert codes.dtype == np.uint8 + + +def test_dequantize_nvfp4_known_values(): + # K=4 codes: k0=+1.0, k1=-0.5, k2=+4.0, k3=-2.0 + k0 = _e2m1_code(0, 2) # mag 1.0 + k1 = _e2m1_code(1, 1) # -0.5 + k2 = _e2m1_code(0, 6) # 4.0 + k3 = _e2m1_code(1, 4) # -2.0 + byte0 = k0 | (k1 << 4) + byte1 = k2 | (k3 << 4) + weight_u8 = np.array([[byte0, byte1]], dtype=np.uint8) + + block_scale = np.array([[2.0]], dtype=ml_dtypes.float8_e4m3fn) # one block + global_scale = 0.5 + + out = dequantize_nvfp4(weight_u8, block_scale, global_scale) + assert out.dtype == ml_dtypes.bfloat16 + # val * 2.0 (block) * 0.5 (global) == val + expected = np.array([[1.0, -0.5, 4.0, -2.0]], dtype=np.float32) + np.testing.assert_array_equal(out.astype(np.float32), expected) + + +def test_dequantize_nvfp4_block_scale_repeat(): + # Two 16-element blocks with distinct block scales exercise repeat_interleave. + n, k = 1, 32 + # All codes = +1.0 (index 2) -> byte 0x22 packs two such codes. + weight_u8 = np.full((n, k // 2), 0x22, dtype=np.uint8) + block_scale = np.array([[1.0, 4.0]], dtype=ml_dtypes.float8_e4m3fn) + out = dequantize_nvfp4(weight_u8, block_scale, 1.0).astype(np.float32) + assert out.shape == (1, 32) + # First 16 elements scaled by 1.0, next 16 by 4.0. + np.testing.assert_array_equal(out[0, :16], np.ones(16, dtype=np.float32)) + np.testing.assert_array_equal(out[0, 16:], np.full(16, 4.0, dtype=np.float32)) + + +def test_dequantize_nvfp4_raw_uint8_block_scale(): + # Block scales may arrive as raw uint8 e4m3 code bytes; result must match a + # typed float8 view. + k0 = _e2m1_code(0, 2) + weight_u8 = np.array([[k0 | (k0 << 4)]], dtype=np.uint8) # K=2, both +1.0 + typed = np.array([[3.0]], dtype=ml_dtypes.float8_e4m3fn) + raw = typed.view(np.uint8) + a = dequantize_nvfp4(weight_u8, typed, 1.0).astype(np.float32) + b = dequantize_nvfp4(weight_u8, raw, 1.0).astype(np.float32) + np.testing.assert_array_equal(a, b) + np.testing.assert_array_equal(a, np.array([[3.0, 3.0]], dtype=np.float32)) + + +def test_dequantize_fp8_per_tensor_scale(): + weight = np.array([[1.0, -2.0, 0.5, 3.0]], dtype=ml_dtypes.float8_e4m3fn) + out = dequantize_fp8(weight, 0.25) + assert out.dtype == ml_dtypes.bfloat16 + expected = np.array([[0.25, -0.5, 0.125, 0.75]], dtype=np.float32) + np.testing.assert_array_equal(out.astype(np.float32), expected) + + +def test_is_modelopt_quant_config(): + assert is_modelopt_quant_config({"quant_method": "modelopt"}) + assert is_modelopt_quant_config({"quant_algo": "NVFP4"}) + assert is_modelopt_quant_config({"quant_cfg": "W4A16_NVFP4"}) + assert is_modelopt_quant_config({"quant_algo": "FP8"}) + # Non-ModelOpt schemes and empty configs are rejected. + assert not is_modelopt_quant_config({"quant_method": "gptq"}) + assert not is_modelopt_quant_config({"quant_method": "awq"}) + assert not is_modelopt_quant_config({}) + assert not is_modelopt_quant_config(None)