Skip to content

fix(nvfp4): close global-scale load gaps in QKV, VLM and DeepSeek - #2449

Merged
inureyes merged 3 commits into
mainfrom
fix/issue-2444-nvfp4-global-scale-gaps
Oct 19, 2026
Merged

inureyes merged 3 commits into
mainfrom
fix/issue-2444-nvfp4-global-scale-gaps

Conversation

@inureyes

@inureyes inureyes commented Oct 18, 2026 •

Copy link
Copy Markdown
Member

Closes #2444

nvidia/Qwen3-30B-A3B-NVFP4 (ModelOpt) failed to load on main with 1-bit weights are affine-only ... (prefix: model.layers.0.self_attn.q_proj). The fused QKV loader hardcoded affine, so the bias-free nvfp4 planes solved to a 1-bit affine layout, and the concatenated plane had no slot for each plane's weight_scale_2.

Changes

  • FusedQKVLinear::from_weights_separate infers the mode from q_proj with infer_quantization_mode, as UnifiedLinear::from_weights does.
  • Qwen3 and Qwen3-MoE group_size() fall back to modelopt_nvfp4_group_size (16), as PhiMoE does. This adds a #[serde(default)] quantization_config field to both ModelArgs and one line to each of three struct-literal test fixtures.
  • finish_vlm_weights_common and load_gemma4_unified (raw keys, before the sanitizer) run the native global-scale flag pass.
    • The filtered VLM loaders end in finish_vlm_weights_common too. There a subset variant skips a flagged plane with neither its weight nor its scale in the map. A loaded flagged plane without its scale still fails.
    • split_fused_moe_key handles .native_global_scale and .global_scale. The [num_experts] vector is copied unsplit to gate_proj and up_proj and renamed under down_proj, so an unflagged scale fails at its switch_glu.gate_proj prefix instead of being orphaned.
  • DeepSeek v1 SwitchLinear::from_stacked_parts refuses a plane that has only one of .scales / .biases, naming the prefix. Previously it built Regular and aborted in gather_mm.

Validation (GB10, CUDA, every GPU run under gpu-lock)

  • MLXCEL_SDPA_DETERMINISTIC=1 mlxcel generate -m ~/models_archive/hf-src/Qwen3-30B-A3B-NVFP4 -p "The capital of France is" -n 80 --temp 0 --show-reasoning loads (main fails at load) and generates coherent text: "... Since the question is about the capital of France, the answer is Paris. ...". No mlx-lm teacher-forced trace: mlx-lm 0.32.2 refuses this checkpoint (55968 unexpected parameters, unstacked ModelOpt keys).
  • Greedy output is token-identical to main (MLXCEL_SDPA_DETERMINISTIC=1, same prompts) for ~/models_archive/mlx/laguna-xs-2.1-nvfp4 (with --show-reasoning), olmoe-1b-7b-0924-nvfp4-gs and gemma-4-31b-it-nvfp4 (two prompts). This was re-checked on the final commit.
  • 10 new tests:
    • layers: nvfp4 split QKV with per-plane scales, its forward checked against x @ dequantize(w).T * s; mxfp4 QKV fused at group 32.
    • qwen3 / qwen3_moe: ModelOpt-only group_size() == 16.
    • Gemma 4 Unified: flagged and unflagged fused experts.
    • VLM: a shipped .native_global_scale is refused, a flagged scale is resolved, and a filtered load skips planes outside the filter.
    • DeepSeek: the half-quantized plane.
    • Each test fails with its fix reverted and passes with it restored.
  • These filters pass: layers::, runtime_lora, one_bit (mlxcel-core) and qwen3 vlm deepseek:: switch_layers sanitize (mlxcel). Clippy --lib --tests -D warnings is clean for both crates. fmt and the five contract checks pass.

Not addressed here (pre-existing): the Kimi K3 VLM loader bypasses finish_vlm_weights_common, so a flagged export is refused cleanly at SwitchLinear. A dense .global_scale is not shape-validated, which is the same as o_proj and MLP on main.

@inureyes inureyes added type:bug Bug fixes, error corrections, or issue resolutions priority:medium Medium priority area:models Model architectures, weights, loading, metadata status:review Under review status:done Completed and removed status:review Under review labels Oct 18, 2026
inureyes added a commit that referenced this pull request Oct 19, 2026
nvidia/Qwen3-30B-A3B-NVFP4 failed to load: FusedQKVLinear::from_weights_separate hardcoded affine, so its bias-free nvfp4 q/k/v solved to a 1-bit affine layout, and the fused plane had no slot for each plane's weight_scale_2.

- from_weights_separate infers the mode from q_proj as UnifiedLinear::from_weights does. When any of q/k/v carries a .global_scale sidecar the layer loads as QkvProjection::Split through UnifiedLinear::from_weights_with_mode, so each plane applies its own scalar after the matmul and claims its own width-pinned runtime-LoRA terms. No struct grows.
- Qwen3 and Qwen3-MoE group_size() fall back to the ModelOpt NVFP4 block size (16) when only quantization_config is present.
- finish_vlm_weights_common and the Gemma 4 Unified loader run apply_native_global_scale_flags on the raw keys. split_fused_moe_key carries .native_global_scale and .global_scale through the fused expert split (shared by gate and up, renamed under down), so an unflagged scale fails at its switch_glu prefix instead of being orphaned.
- DeepSeek v1 SwitchLinear refuses a stacked expert plane with only one of .scales/.biases instead of loading it as Regular and aborting in gather_mm.

Closes #2444
Review follow-ups for #2444.

- The fused q/k/v arm now reconciles a block-float plane's group size from its shapes and checks the backend, as the single-plane loader does. Mode inference made this arm reachable for bias-free exports, and it used to store the declared group size (for example 64 for mxfp4 packed at 32), which aborts in quantized_matmul.
- The global-scale split builds each plane at its reconciled layout, like the mixed-width split, so no plane logs a spurious divergence warning.
- The filtered VLM loaders share finish_vlm_weights_common, so the flag pass there uses a subset variant: a flagged plane with neither its weight nor its scale in the map was filtered out and is skipped. A loaded flagged plane without its scale still fails.
- The nvfp4 QKV test follows the backend's NVFP4 support. Stale Jamba and modelopt_nvfp4_group_size comments are updated.

New tests: mxfp4 q/k/v fuse at group 32 and match the standalone planes, and a filtered VLM load skips a flagged plane outside the filter. Both fail with the fix reverted.

Refs #2444
@inureyes
inureyes force-pushed the fix/issue-2444-nvfp4-global-scale-gaps branch from 405e30d to 56ba8b7 Compare October 19, 2026 01:45
@inureyes
inureyes merged commit e772cf2 into main Oct 19, 2026
27 checks passed
@inureyes
inureyes deleted the fix/issue-2444-nvfp4-global-scale-gaps branch October 19, 2026 02:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:models Model architectures, weights, loading, metadata priority:medium Medium priority status:done Completed type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix: Close NVFP4 global-scale load gaps left after #2441

1 participant