Add NVFP4/FP8 (ModelOpt) weight reconstruction foundation for Qwen3.6 - #446
Open
justinchuby wants to merge 1 commit into
Open
Add NVFP4/FP8 (ModelOpt) weight reconstruction foundation for Qwen3.6#446justinchuby wants to merge 1 commit into
justinchuby wants to merge 1 commit 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. |
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.