Skip to content

feat(qwen3_5_moe): support compressed-tensors NVFP4 experts with block-fp8 dense side - #296

Open
JIAQI13 wants to merge 1 commit into
FlashML-org:mainfrom
JIAQI13:feat/ct-nvfp4-expert-banks
Open

feat(qwen3_5_moe): support compressed-tensors NVFP4 experts with block-fp8 dense side#296
JIAQI13 wants to merge 1 commit into
FlashML-org:mainfrom
JIAQI13:feat/ct-nvfp4-expert-banks

Conversation

@JIAQI13

@JIAQI13 JIAQI13 commented Aug 30, 2026

Copy link
Copy Markdown

Fixes #263
Also addresses #236 (uncensored/abliterated Qwen3.6-35B-A3B NVFP4 support). Related: #165, #252, #274.

Problem

Community abliterated/uncensored Qwen3.6-35B-A3B NVFP4 checkpoints exported with llm-compressor (quant_method: compressed-tensors, format: mixed-precision) cannot be converted:

  • routed experts: NVFP4 stored as weight_packed (uint8) + weight_scale (fp8-e4m3 block) + scalar weight_global_scale
  • dense side (attn/GDN + shared expert): 128×128 block FP8 (.weight_scale, W8A8)

FreeToken only recognized the nvidia-modelopt NVFP4 expert layout (weight / weight_scale / weight_scale_2), so:

The official nvidia/Qwen3.6-35B-A3B-NVFP4 (modelopt layout) and block-FP8 checkpoints are unaffected.

Changes

  • models/nvfp4_banks.py: add load_nvfp4_ct_expert_source_banks / ..._parallel — a compressed-tensors NVFP4 expert source builder for the offload banks. Parses weight_packed / weight_scale / weight_global_scale (the stored global is the quant-side scale; the dequant global is its reciprocal, broadcast per output row — same convention vLLM uses), and builds the single-file safetensors weight map from the safetensors header when no model.safetensors.index.json is present.
  • models/qwen3_5_moe/config.py: add _compressed_tensors_mixed_fp8block() detection (8-bit block group + 4-bit group_size=16 tensor_group group in config_groups); when matched, route experts to NVFP4 banks and the dense side (attn/GDN/shared expert) to native 128×128 block FP8 (attn_quant/dense_quant = "fp8_block").
  • models/qwen3_5_moe/weight.py: register the CT expert key pattern/source spec; route the dense pass of these checkpoints through the existing block-FP8 loader (_iter_weights_fp8), renaming the block scale suffix .weight_scale.weight_scale_inv and skipping dynamic-activation scales.
  • models/qwen3_5_moe/moe.py, gdn.py: enable the block-FP8 modules for the shared expert and GDN when dense_quant/attn_quant == "fp8_block" (previously gated on the experts being block-FP8).
  • tests/models/test_qwen3_5_moe_ct_mixed_detect.py: detection tests — mixed CT flagged, pure-CT-NVFP4 / per-channel-FP8 / modelopt mixed / no-quant not flagged.

Pure modelopt NVFP4, block-FP8, dense NVFP4, and pure compressed-tensors NVFP4 paths are unchanged.

Tested on (real hardware)

  • GPU: RTX 3080 Laptop, 16 GB VRAM (sm_86), driver 610.57.04; CPU: Intel i7-10870H, 32 GB RAM + swap; CachyOS
  • FreeToken 0.1.2, torch 2.11.0+cu130
  • Checkpoint: kyaky/Qwen3.6-35B-A3B-Uncensored-NVFP4 (compressed-tensors mixed-precision, single 23 GB model.safetensors)
  • A/B with the same command
    ft checkpoint --model <HF dir> --out <out> --dtype bfloat16 --moe-backend offload --shard-gib 8:
  • ft serve --model-path <FTW> --moe-backend offload --moe-cache-auto: loads via the FTW fast path (3.25 GiB dense + 16.9 GiB expert banks), API server is ready, /v1/chat/completions returns correct generations.
  • Detection unit tests pass standalone (5/5, no weights required).

Scope note

This covers the compressed-tensors expert layout (also used by sakamakismile/Huihui-...-abliterated-NVFP4 in #263). The other failing checkpoint in that report, joshebbs/qwen3.6-35b-abliterated-nvfp4-modelopt, uses an unusual doubled key prefix (language_model.model.layers.N.mlp.experts.*) that neither the existing modelopt regex nor this change matches — that looks like a separate export bug and is not handled here.

…k-fp8 dense side

llm-compressor mixed-precision exports (e.g. kyaky/Qwen3.6-35B-A3B-Uncensored-NVFP4 and the sakamakismile abliterated checkpoint in FlashML-org#263) store routed experts as compressed-tensors NVFP4 (weight_packed/weight_scale/weight_global_scale) while the dense side is 128x128 block FP8. Add a CT NVFP4 expert bank loader for offload, detect the format from config_groups, and route the dense pass through the existing block-fp8 loader (gating shared-expert/GDN block-fp8 modules on dense_quant/attn_quant).

Fixes FlashML-org#263, addresses FlashML-org#236.
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.

[Bug] ft checkpoint fails on Qwen3.6-35B-A3B abliterated NVFP4 checkpoints: Missing MoE expert source layers

1 participant