fix(models): support llm-compressor NVFP4 MoE export variants - #413
Open
Romeriz wants to merge 1 commit into
Open
fix(models): support llm-compressor NVFP4 MoE export variants#413Romeriz wants to merge 1 commit into
Romeriz wants to merge 1 commit into
Conversation
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
4 tasks
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.
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
maincould not convertor serve: per-expert vs stacked per-layer routed experts, a GDN left bf16
or fp8 by the export's
ignorelist or arecipe: "all,gdn:fp8,-router",format-only configs with no
config_groups, native NVFP4 lm_head / embedtokens, 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_nvfp4falls back to the exacttop-level
format: nvfp4-pack-quantizedwhenconfig_groupsis absent;ModelConfig.gdn_quant(GDNout_projquant, independent ofattn_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-projectionlayer-global (reciprocal) into the
*_globalbanks. Same 6-bank dict as theper-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_nvfp4read the export'signorelist andrecipestring;_expert_quantrecognizes CT-MoE ->nvfp4; CTlm_head_quantis
nvfp4unless the export skippedlm_head;gdn_quantwired.models/qwen3_5_moe/gdn.py,model.py-- GDNout_projbuilt fromgdn_quant(defaults to
attn_quant, preserving the modelopt/fp8 paths).models/qwen3_5_moe/weight.py--_dequant_fp8_weightbroadcasts per-output-rowscales; per-expert CT + stacked CT expert specs and a spec selector that probes the
weight map (index or header);
_iter_weights_compressed_tensorsnow serves MoEcheckpoints: skips per-expert and stacked experts (they ride the offload banks),
dequantizes an fp8/bf16 GDN (fused
in_proj), dequantizes NVFP4embed_tokens,keeps
lm_headnative when the model wants it, and keeps NVFP4 attn / shared-expertnative. Parallel bank read raises
NotImplementedErrorfor stacked experts (theserial read already saturates and the auto-pick already selects serial).
tests/models/test_qwen3_5_moe_ct_nvfp4.py-- 16 weight-free tests: configdetection (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
nvidia/Qwen3.6-35B-A3B-NVFP4), block-fp8, and densecompressed-tensors NVFP4 (Qwen3.6-27B) take unchanged code paths.
bf16 GDN
out_proj, native NVFP4 shared-expert/attn).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_FUSEshared_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)
f8d6a56onmain(af71ba4); torch 2.11.0+cu130AEON-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: nvfp4doth4580/Kwaipilot-KAT-Coder-V2.5-Dev-NVFP4-MIXED(format-only, stacked experts,gdn:fp8recipe, NVFP4 lm_head/embed_tokens)Ttimms/KAT-Coder-V2.5-Dev-REAP-50-NVFP4A16(per-expert experts, GDN bf16, 128 experts)ft checkpoint --model <checkpoint_dir> --out <ftw_dir>thenft serve --model-path <ftw_dir>mainall 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).