Skip to content

fix(models): support nvidia/Qwen3.8-Flash-Next-NVFP4 in qwen4_exp - #404

Open
Eng-Ahmd wants to merge 1 commit into
FlashML-org:mainfrom
Eng-Ahmd:fix/qwen4-exp-nvidia-nvfp4
Open

fix(models): support nvidia/Qwen3.8-Flash-Next-NVFP4 in qwen4_exp#404
Eng-Ahmd wants to merge 1 commit into
FlashML-org:mainfrom
Eng-Ahmd:fix/qwen4-exp-nvidia-nvfp4

Conversation

@Eng-Ahmd

@Eng-Ahmd Eng-Ahmd commented Sep 7, 2026

Copy link
Copy Markdown

Problem: nvidia/Qwen3.8-Flash-Next-NVFP4 fails to load. First with missing MoE expert sources on every layer, and after working around the quant flag, with KeyError on model.layers.0.linear_attn.out_proj.weight_scale. RadixArk/Qwen3.8-Flash-Next-NVFP4 loads fine.

Root cause (python/freetoken/models/qwen4_exp/config.py):

  • Expert detection only checks the top-level quant_algo for fp4. The nvidia checkpoint labels it MIXED_PRECISION and carries the NVFP4 experts in quantized_layers/config_groups, so experts resolved to none and fell into the BF16 offload path.
  • attn_quant probes layers.0.self_attn.q_proj, but layer 0 is linear_attention and has no self_attn module. The ignore match always misses, so attn resolved to nvfp4 and GDN out_proj was built as NVFP4 while the checkpoint stores it BF16.

Change:

  • Add _is_nvfp4_experts(): top-level fp4, or mixed plus an FP4 experts entry in quantized_layers or a 4-bit float experts group in config_groups.
  • Compute layer split before quant parsing. Probe self_attn on full_ids[0] and linear_attn on linear_ids[0]; attn stays none unless either side is actually quantized. Empty groups are treated as ignored.
  • docs/models.md: list the nvidia checkpoint. weight.py docstring updated (loader regex already matched both layouts).

Verification:

  • tests/models/qwen4_exp/test_config.py: 8 passed.
  • tests/models/qwen4_exp/test_weight_ckpt.py against a local nvidia/Qwen3.8-Flash-Next-NVFP4 copy: 41 passed (dense names equal the model state dict, sampled tensors match).
  • parse_config on the nvidia checkpoint now yields expert=nvfp4 attn=none dense=none lm_head=none; RadixArk synthetic config unchanged.

Scope: qwen4_exp only. No Laguna changes.

Detect NVFP4 experts behind MIXED_PRECISION via quantized_layers/config_groups. Probe self_attn on a full_attention layer and linear_attn on a linear layer instead of hardcoded layers.0, which has no self_attn and misrouted attn_quant to nvfp4.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant