Skip to content

fix(models): support llm-compressor NVFP4 MoE export variants - #413

Open
Romeriz wants to merge 1 commit into
FlashML-org:mainfrom
Romeriz:fix/models-ct-nvfp4-moe-variants
Open

fix(models): support llm-compressor NVFP4 MoE export variants#413
Romeriz wants to merge 1 commit into
FlashML-org:mainfrom
Romeriz:fix/models-ct-nvfp4-moe-variants

Conversation

@Romeriz

@Romeriz Romeriz commented Sep 8, 2026

Copy link
Copy Markdown

Fixes #410
Related: #390 (complementary -- FP8-attn/shared CT exports; see the #390 comment)

Summary

llm-compressor Qwen3.5-MoE NVFP4 exports come in layouts main could not convert
or serve: per-expert vs stacked per-layer routed experts, a GDN left bf16
or fp8
by the export's ignore list or a recipe: "all,gdn:fp8,-router",
format-only configs with no config_groups, native NVFP4 lm_head / embed
tokens
, and single-file checkpoints without a shard index. Detect each variant
from the on-disk layout / recipe, build the matching offload banks (per-expert and
stacked loaders), and emit the state dict the model expects.

Changes

  • models/config.py -- detect_compressed_tensors_nvfp4 falls back to the exact
    top-level format: nvfp4-pack-quantized when config_groups is absent;
    ModelConfig.gdn_quant (GDN out_proj quant, independent of attn_quant).
  • models/nvfp4_banks.py -- _weight_map() (index OR single-file header);
    Nvfp4ExpertSourceSpec.stacked; load_nvfp4_stacked_expert_sources() reshapes each
    [E*rows, cols] bank tensor to [E, rows, cols] and broadcasts the per-projection
    layer-global (reciprocal) into the *_global banks. Same 6-bank dict as the
    per-expert loader, so the marlin/b12x repack and the offload cache are unchanged.
  • models/qwen3_5_moe/config.py -- _ct_expert_groups_nvfp4 (+format fallback);
    _ct_recipe / _ct_ignored / _ct_gdn_nvfp4 read the export's ignore list and
    recipe string; _expert_quant recognizes CT-MoE -> nvfp4; CT lm_head_quant
    is nvfp4 unless the export skipped lm_head; gdn_quant wired.
  • models/qwen3_5_moe/gdn.py, model.py -- GDN out_proj built from gdn_quant
    (defaults to attn_quant, preserving the modelopt/fp8 paths).
  • models/qwen3_5_moe/weight.py -- _dequant_fp8_weight broadcasts per-output-row
    scales; per-expert CT + stacked CT expert specs and a spec selector that probes the
    weight map (index or header); _iter_weights_compressed_tensors now serves MoE
    checkpoints: skips per-expert and stacked experts (they ride the offload banks),
    dequantizes an fp8/bf16 GDN (fused in_proj), dequantizes NVFP4 embed_tokens,
    keeps lm_head native when the model wants it, and keeps NVFP4 attn / shared-expert
    native. Parallel bank read raises NotImplementedError for stacked experts (the
    serial read already saturates and the auto-pick already selects serial).
  • tests/models/test_qwen3_5_moe_ct_nvfp4.py -- 16 weight-free tests: config
    detection (config_groups / format-only + recipe), recipe GDN override, per-row fp8
    dequant, per-expert + stacked spec and keying detection, per-expert and stacked
    single-file bank-loader placement with distinct gate_up/down globals, and full
    dense-pass fixtures for the AEON and Kwaipilot layouts.

Backwards compatibility

  • Pure modelopt NVFP4 (nvidia/Qwen3.6-35B-A3B-NVFP4), block-fp8, and dense
    compressed-tensors NVFP4 (Qwen3.6-27B) take unchanged code paths.
  • AEON-7 re-read after the change is identical to before (663-tensor state dict,
    bf16 GDN out_proj, native NVFP4 shared-expert/attn).
  • Ornith-style CT exports (per-tensor-FP8 attn/shared-expert group + NVFP4 experts)
    are intentionally out of scope -- that is PR fix(models): better support for mixed-precision compressed-tensors NVFP4  #390; the two are complementary and
    overlap in shared plumbing (_CT_NVFP4_FUSE shared_expert fusion, _expert_quant,
    per-expert llm-compressor bank naming), so they should be reconciled (per-variant
    MoE-CT dense routing) rather than merged independently.

Tested on (real hardware)

  • GPU: NVIDIA GeForce RTX 5060 Ti 16Gb, driver KMD Version: 610.57.04 CUDA UMD Version: 13.3 ; OS: DEbian 13 (trixie)
  • FreeToken commit f8d6a56 on main (af71ba4); torch 2.11.0+cu130
  • Checkpoints (all convert + serve):
    • AEON-7/Qwen3.6-35B-A3B-heretic-NVFP4 (CT NVFP4, per-expert experts, GDN bf16,
      single 23 GB no-index file) -- FTW: 663 weight + 240 experts_bank, 20.96 GiB,
      quant_format: nvfp4
    • doth4580/Kwaipilot-KAT-Coder-V2.5-Dev-NVFP4-MIXED (format-only, stacked experts,
      gdn:fp8 recipe, NVFP4 lm_head/embed_tokens)
    • Ttimms/KAT-Coder-V2.5-Dev-REAP-50-NVFP4A16 (per-expert experts, GDN bf16, 128 experts)
  • Commands: ft checkpoint --model <checkpoint_dir> --out <ftw_dir> then
    ft serve --model-path <ftw_dir>
  • On main all three fail to convert or serve; on this branch they convert and serve,
    and the dense pass matches the model state dict exactly (AEON: 663 tensors,
    Kwaipilot: 665 tensors; 0 missing / 0 extra each).

llm-compressor Qwen3.5-MoE NVFP4 exports come in layouts main could not
convert or serve: per-expert vs stacked per-layer routed experts, a GDN
left bf16 / fp8 by the ignore list or a ``gdn:fp8`` recipe, format-only
configs with no ``config_groups``, native NVFP4 lm_head/embed_tokens, and
single-file checkpoints without a shard index. Detect each from the
on-disk layout / recipe, build the matching offload banks (per-expert and
stacked loaders), and emit the state dict the model expects.

Assisted-by: opencode
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.

[Model request] KAT-Coder-V2.5-Dev NVFP4 community exports (compressed-tensors MoE): dense-branch misroute + GDN granularity

1 participant