[Feature] Migrate qwen3_5_moe onto layers/quantization: the new config layer resolves compressed-tensors checkpoints correctly, the legacy detectors don't (offering a PR)

[salekseev]

Applies to

qwen3_5_moe (engine), following the quantization refactor in #418 / #427. Verified against main @ fb7f732.

What problem does this solve

layers/quantization/ already describes the compressed-tensors checkpoints people keep filing issues about -- correctly -- but qwen3_5_moe doesn't consume it, so those checkpoints still don't load.

Probing QuantConfig.from_hf(...).scheme_for_name(...) on today's main:

unsloth/Qwen3.6-35B-A3B-NVFP4-Fast
  self_attn.q_proj            -> fp8_tensor, group=(1, -1)     # per-channel
  lm_head                     -> fp8_tensor, group=(1, -1)
  mlp.shared_expert.gate_proj -> nvfp4,      group=(1, 16)

primitive-ai/Ornith-1.5-35B-A3B-mixed-NVFP4-FP8
  self_attn.q_proj            -> fp8_tensor, group=(-1, -1)    # per-tensor
  lm_head                     -> None                          # correctly excluded via `ignore`

nvidia/Qwen3.6-35B-A3B-NVFP4   (modelopt)
  self_attn.q_proj / experts / shared_expert / lm_head -> fp8_tensor / nvfp4 / nvfp4 / nvfp4

All three resolve exactly right. But qwen3_5_moe/config.py still computes its own verdicts with hand-rolled detectors -- _expert_quant (L42), _lm_head_quant (L71), _attn_quant (L107), consumed at L169/172/180 -- and _iter_weights_attn_fp8 in weight.py still routes off those strings. Those detectors return attn_quant='none' for both compressed-tensors checkpoints above, so the reader never reaches the fp8 dense path even though the config layer describes it perfectly.

model.py is already migrated for the head -- it passes quant_config=config.quant, prefix="lm_head" to a single ParallelLMHead -- so the model half and the config half currently disagree about the same checkpoint.

This is the shared root cause under a cluster of open reports: #252 (per-channel FP8, targets: ["Linear"]), #381 (model.language_model prefix), #344, #410, and the routing discussion on #390 / #413.

Proposed solution

Port qwen3_5_moe to derive its verdicts from the QuantConfig, the way qwen4_exp/config.py already does:

expert_scheme = QuantConfig.from_hf(hf_config).scheme_for_name("...layers.0.mlp.experts.0.gate_proj")
expert_quant  = "none" if expert_scheme is None else str(expert_scheme.kind)

Keep the verdict strings -- engine.py reads expert_quant for MoE-strategy and bench-format decisions (L1162, L1316, L1552) and attn_quant/dense_quant at L1397, and checkpoint/convert.py:44 hashes expert_quant -- and just stop hand-deriving them. Each module already takes its own scheme when it's built.

The reader is the harder half: qwen3_5_moe/weight.py only imports QuantKind and never receives a QuantConfig, so #427 didn't reach it. _iter_weights_attn_fp8's dispatch would need to become scheme-driven rather than keyed on attn_quant == "fp8_pertensor", and _per_row_scale becomes redundant once scales come from Fp8TensorLinearMethod.create_weights, which already allocates [out_features] fp32.

Scoped deliberately to fp8/NVFP4, with the legacy detectors retained as the fallback for AWQ and compressed-tensors int4. That is not tidiness -- a wholesale port would break three checkpoints today, because layers/quantization has no dialect for them:

checkpoint	quant_method	QuantConfig.from_hf
cyankiwi/Qwen3.6-35B-A3B-AWQ-4bit	awq	NotImplementedError: quantization method 'awq' is not supported
Avesed/Qwen3.6-35B-A3B-INT4-W4A16	compressed-tensors	NotImplementedError: weight scheme ...
cyankiwi/Qwen-AgentWorld-35B-A3B-AWQ-INT4	compressed-tensors	NotImplementedError

Adding AWQ and CT-int4 dialects to layers/quantization/configs/ is the natural follow-up, and it's the same ground as #396 -- I'd rather do that as a separate change than smuggle it in here.

Alternatives considered

• Patch the legacy detectors instead. That's what I was doing: fix(qwen3_5_moe): accept a per-channel fp8 weight_scale in the dense loader #415 (per-channel weight_scale) and feat(qwen3_5_moe): keep a natively-fp8 lm_head native instead of failing to load #416, plus a strategy: "channel" widening offered on fix(models): better support for mixed-precision compressed-tensors NVFP4  #390. I closed feat(qwen3_5_moe): keep a natively-fp8 lm_head native instead of failing to load #416 once refactor(quant): config, scheme and method layers for quantization #418/refactor(quant): hand the checkpoint QuantConfig to the weight readers #427 landed -- the new config layer covers it more completely than my patch, including class-name targets in ignore, which is the actual blocker for [Feature Request] Support compressed-tensors per-channel FP8 checkpoints for Qwen3.5 MoE #252's checkpoint. Continuing down that road means maintaining two detection systems that disagree.
• Wait for the migration to happen incidentally. Reasonable, but fix(models): better support for mixed-precision compressed-tensors NVFP4  #390/fix(models): support llm-compressor NVFP4 MoE export variants #413 are actively designing per-variant dense routing right now, and if the variant is already a QuantScheme per module that routing may not need to exist. Worth settling before that work lands.
• Migrate everything including int4 in one PR. Bigger, unreviewable, and it would couple an int4 dialect to an fp8 fix.

Offering to do it

Happy to take this if nobody has claimed it. I can validate against every family involved -- unsloth/...-NVFP4-Fast (CT per-channel fp8 + NVFP4), primitive-ai/Ornith-1.5-...-mixed-NVFP4-FP8 (CT per-tensor fp8), nvidia/...-NVFP4 (modelopt, the regression case), plus the three int4 checkpoints above to prove the fallback still works -- on a 4080 SUPER with 62 GB host RAM and load / needle-recall / BFCL harnesses already wired up for this engine.

Before writing code I would want your call on the following.

1. Is retaining the legacy detectors as an AWQ/int4 fallback acceptable as an interim state, or would you rather the int4 dialects land in layers/quantization first and the migration be one clean cut?
2. Should weight.py receive the QuantConfig the same way refactor(quant): hand the checkpoint QuantConfig to the weight readers #427 did it for other readers, or is there a preferred shape for MoE readers that split dense and expert passes?

[jason-fxz]

Thanks, the diagnosis matches ours. The migration is up as #438 config.py drops the string detectors, weight.py replaces the four readers with one pass that asks scheme_for for every Linear module, the same call the model builds its buffers from.

On your questions:

1. No fallback. The legacy detectors go. They never protected the AWQ / int4 checkpoints: on main all three already fail in checkpoint_quant_config ('awq' is not supported, compressed-tensors weight scheme ... is not supported) before parse_config runs. Int4 support belongs to [Feature] int4 routed experts: AWQ -> Q4_1 and compressed-tensors pack-quantized -> Q4_0 #396 and has to be complete when it comes: an integer W4A16 kind with zero points, an AWQ dialect plus the type: int branch in compressed-tensors, expert bank layout and kernels, and a repack step (AutoAWQ and compressed-tensors pack int4 differently). All three of those checkpoints quantize only the experts; the dense side is bf16.
2. Readers get the config from get_quant_config() (refactor(quant): hand the checkpoint QuantConfig to the weight readers #427). Each module resolves through the family NameMap and a per-dialect table naming the checkpoint tensor behind each role; the expert reader uses the same table.

Two details: _per_row_scale stays, the layer allocates [out] but does not broadcast, so the reader still builds the per-row vector (your #415 logic, credited in the PR). And compressed-tensors ignore entries name the module alone, not its children; our matcher used ancestor semantics, which is what broke unsloth's linear_attn entry and Ornith's experts.N entries. Fixed.

AIME e2e passes on an H100 for unsloth/Qwen3.6-35B-A3B-NVFP4-Fast, RedHatAI/Qwen3.6-35B-A3B-NVFP4, primitive-ai/Ornith-1.5-35B-A3B-mixed-NVFP4-FP8, kyaky/Qwen3.6-35B-A3B-Uncensored-NVFP4, unsloth/Qwen3.8-27B-NVFP4, sakamakismile/Qwen3.6-27B-NVFP4 and the nvidia ModelOpt checkpoints; the full table is in the PR.