Add NVFP4/FP8 (ModelOpt) weight reconstruction foundation for Qwen3.6 - #446
Open
justinchuby wants to merge 2 commits into
Open
Add NVFP4/FP8 (ModelOpt) weight reconstruction foundation for Qwen3.6#446justinchuby wants to merge 2 commits into
justinchuby wants to merge 2 commits into
Conversation
NVIDIA TensorRT Model Optimizer (ModelOpt) exports mixed-precision NVFP4 + FP8 checkpoints (e.g. quantized Qwen3.6): routed MoE experts / shared expert / lm_head as W4A16_NVFP4 (block-16 E2M1 weights, FP8-E4M3 block scales, per-tensor FP32 global scale weight_scale_2), and attention/linear-attn projections as FP8 (E4M3, per-tensor scale). Add the numeric core of the ModelOpt loader under mobius/integrations/modelopt: - dequantize_nvfp4(): E2M1 codes x FP8-E4M3 block scale x FP32 global -> BF16 - dequantize_fp8(): FP8 (E4M3) x per-tensor scale -> BF16 - unpack_nvfp4_codes(), is_modelopt_quant_config(), FP4_E2M1_LUT Following ORT GenAI's loader, dense FP8 and NVFP4 shared-expert/lm_head weights are reconstructed to BF16 exactly (ORT has no FP8 attention GEMM), so the standard BF16 build path can consume the checkpoint. Fully unit-tested against hand-computed references (numeric-only; no runtime dependency). Guard QuantizationConfig.from_transformers: ModelOpt schemes now raise a clear NotImplementedError instead of being silently routed through the INT4 MatMulNBits path (which would mis-dequantize packed E2M1/float8 weights). The check runs before the quant_method=="none" early-return because ModelOpt may name its scheme only via quant_algo/quant_cfg. Out of scope (documented, follow-up): full checkpoint weight-load integration and native routed-expert NVFP4 QMoE emission (CUDA/Blackwell, onnxruntime_USE_FP4_QMOE=ON) — neither buildable nor verifiable without that runtime and a real ModelOpt checkpoint. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds a new mobius.integrations.modelopt package that implements the numeric core for reconstructing NVIDIA ModelOpt NVFP4 and FP8 checkpoint weights back to BF16, and adds an explicit guard in QuantizationConfig.from_transformers() to fail loudly on ModelOpt quantization configs to avoid silently mis-dequantizing them via the INT4 path.
Changes:
- Add NVFP4/FP8 dequantization utilities (
unpack_nvfp4_codes,dequantize_nvfp4,dequantize_fp8) plus ModelOpt quant-config detection. - Add unit tests covering dequant math and config detection.
- Add a
NotImplementedErrorguard in HF quantization-config parsing to prevent misrouting ModelOpt schemes through the existing INT4 quantization path.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| src/mobius/integrations/modelopt/_dequant.py | Adds NVFP4/FP8 unpack/dequant math and ModelOpt quant-config detection helpers. |
| src/mobius/integrations/modelopt/_dequant_test.py | Adds unit tests for the new dequantization functions and config detection. |
| src/mobius/integrations/modelopt/init.py | Exposes the ModelOpt integration API via package exports. |
| src/mobius/_configs/_quantization.py | Adds a loud failure path for ModelOpt quantization configs to avoid silent mis-dequantization. |
| src/mobius/_configs_test.py | Adds a regression test ensuring ModelOpt configs raise NotImplementedError. |
Suppressed comments (1)
src/mobius/integrations/modelopt/_dequant.py:85
dequantize_nvfp4only checksK % n_blocks == 0and then repeats scales. Ifblock_scale_e4m3is accidentally loaded with the wrong second dimension (e.g. K/32 or K/8), this will still run and silently produce incorrect dequantized weights. Since NVFP4 is specified as fixed 16-element blocks, enforceK % NVFP4_BLOCK_SIZE == 0andblock_scale.shape[1] == K // NVFP4_BLOCK_SIZE. Also, expanding scales withnp.repeatcreates a large temporary; multiplying in[N, blocks, 16]form avoids that extra allocation.
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}.")
Comment on lines
+50
to
+56
| 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) |
Comment on lines
+38
to
+55
| 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) |
Comment on lines
+71
to
+81
| 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_nvfp4_block_scale_repeat(): | ||
| # Two 16-element blocks with distinct block scales exercise repeat_interleave. |
Address PR review feedback on the NVFP4 dequant path: - unpack_nvfp4_codes now rejects non-2D input with a clear ValueError so a loader shape/dtype mistake surfaces instead of being silently reinterpreted via the uint8 cast. - dequantize_nvfp4 enforces the fixed NVFP4 16-element block size; a derived size other than 16 means mismatched weight/scale shapes and now raises. - Rewrite the known-value and raw-uint8 block-scale tests to use full K=16 blocks (spec-realistic), asserting trailing zero-code reconstruction. - Fix stale "repeat_interleave" comment to "np.repeat". - Add regression tests for the two new validation errors. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.
Suppressed comments (8)
src/mobius/integrations/modelopt/_dequant.py:88
- After making
ml_dtypesoptional,dequantize_nvfp4()should explicitly fail fast with a clear message whenml_dtypesis unavailable; otherwise this will error later with anAttributeErrorwhen trying to accessml_dtypes.bfloat16.
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]
src/mobius/integrations/modelopt/_dequant.py:148
_to_float32()reinterprets rawuint8bytes as FP8, but ifml_dtypesisn't installed this will currently fail with a crypticNameError/AttributeError. Raise a clearImportErrorwhen this conversion is requested withoutml_dtypes.
arr = np.asarray(arr)
if arr.dtype == np.uint8:
arr = arr.view(ml_dtypes.float8_e4m3fn)
return arr.astype(np.float32)
src/mobius/integrations/modelopt/_dequant.py:118
dequantize_fp8()returnsml_dtypes.bfloat16but doesn't check whetherml_dtypesis available (unlike the suggested guard fordequantize_nvfp4). With an optionalml_dtypesimport, this will otherwise fail with anAttributeError.
weight = _to_float32(weight_f8)
return (weight * np.float32(weight_scale)).astype(ml_dtypes.bfloat16)
src/mobius/integrations/modelopt/_dequant.py:60
unpack_nvfp4_codes()currently casts inputs touint8, which can silently reinterpret incorrect dtypes (e.g. int8/float) and contradicts the docstring intent to surface upstream dtype mistakes. Validatedtype==uint8instead of coercing.
if packed_nk2.ndim != 2:
raise ValueError(
f"NVFP4 packed codes must be 2D [N, K/2], got shape {packed_nk2.shape}."
)
packed = np.ascontiguousarray(packed_nk2).astype(np.uint8)
low = packed & 0x0F
src/mobius/integrations/modelopt/_dequant.py:106
dequantize_nvfp4()assumesblock_scale_e4m3is 2D and has matchingN, but it never validates either. A mismatched shape can silently broadcast and reconstruct wrong weights. Also,np.repeatmaterializes a full[N,K]scale tensor, doubling peak memory for large weights; reshape/broadcast avoids that extra allocation.
block_scale = _to_float32(block_scale_e4m3) # [N, K/16]
k = codes.shape[1]
n_blocks = block_scale.shape[1]
if n_blocks == 0 or k % n_blocks != 0:
raise ValueError(f"NVFP4 K={k} is not divisible by the block count {n_blocks}.")
block_size = k // n_blocks
if block_size != NVFP4_BLOCK_SIZE:
# NVFP4 pins the block size to 16; a different derived size means the
# weight/scale shapes are mismatched (silently-wrong reconstruction).
raise ValueError(
f"NVFP4 block size must be {NVFP4_BLOCK_SIZE}, got {block_size} "
f"(K={k}, block scales={n_blocks})."
)
block_scale = np.repeat(block_scale, block_size, axis=1) # [N, K]
dequant = val * block_scale * np.float32(global_scale)
return dequant.astype(ml_dtypes.bfloat16)
src/mobius/_configs/_quantization.py:55
quant_methodis used for branching but isn't normalized. If HuggingFace ever serializes values like "FP8"/"None" with different casing, this logic can mis-route into the INT4 path. Normalizing to lowercase makes this parser more robust.
method = qc.get("quant_method", "none")
src/mobius/integrations/modelopt/_dequant.py:33
mobius.integrations.modeloptimportsml_dtypesat module import time. Sinceml_dtypesis not a declared core dependency, this can raiseImportErrorwhenQuantizationConfig.from_transformers()tries to importis_modelopt_quant_config, preventing the intendedNotImplementedErrorguard from firing. Make theml_dtypesimport optional at module import time and fail with a clear error only when dequantization is actually invoked.
This issue also appears in the following locations of the same file:
- line 86
- line 117
- line 145
import ml_dtypes
import numpy as np
src/mobius/integrations/modelopt/_dequant_test.py:9
- This test module imports
ml_dtypesunconditionally. Sinceml_dtypesis not a declared core dependency, test collection can fail in environments where it isn't installed. Preferpytest.importorskip("ml_dtypes")so the suite can still run without this optional dependency.
import ml_dtypes
import numpy as np
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Foundation for consuming NVIDIA ModelOpt NVFP4/FP8 mixed-precision checkpoints (e.g. quantized Qwen3.6), porting the verifiable numeric core of ORT GenAI #2350 / #2351.
ModelOpt Qwen3.6 layout:
W4A16_NVFP4: block-16 E2M1 (fp4) weights, FP8-E4M3 block scales, per-tensor FP32 global scale (weight_scale_2).FP8(E4M3, per-tensorweight_scale).How
New
mobius/integrations/modelopt/package with the loader's numeric core:dequantize_nvfp4()—e2m1(code) × e4m3(block_scale[k//16]) × global_scale → bf16dequantize_fp8()—e4m3(weight) × weight_scale → bf16unpack_nvfp4_codes(),is_modelopt_quant_config(),FP4_E2M1_LUTFollowing ORT GenAI's loader, dense FP8 and NVFP4 shared-expert/lm_head weights are reconstructed to BF16 exactly (ORT has no FP8 attention GEMM), so the standard BF16 build path can consume the checkpoint. Fully unit-tested against hand-computed references — numeric-only, no runtime/GPU dependency.
QuantizationConfig.from_transformersnow fails loudly (NotImplementedError) on ModelOpt schemes instead of silently routing packed E2M1/float8 weights through the INT4MatMulNBitspath (which would mis-dequantize them). The check runs before thequant_method=="none"early-return since ModelOpt may name its scheme only viaquant_algo/quant_cfg.Scope / verification
_configs_test.py+ module).onnxruntime_USE_FP4_QMOE=ON). That path is not buildable or verifiable in this environment (no Blackwell GPU, no FP4-QMOE ORT build, no real ModelOpt checkpoint), so it is intentionally gated rather than shipped unverified — the safe increment given the repo's "no silently-wrong weights" bar.Please review — do not merge yet.