Skip to content

Commit bc27359

Browse files
committed
feat(qwen3_5_moe): read llm-compressor packed NVFP4 on the mixed-fp8 dense pass
A compressed-tensors mixed-precision checkpoint stores its dense NVFP4 projections as .weight_packed / .weight_scale / .weight_global_scale, not as modelopt's .weight / .weight_scale / .weight_scale_2. The mixed-fp8 dense pass only looked for the modelopt names, so on such a checkpoint the packed shared expert was never emitted and the load failed with KeyError: 'model.layers.0.mlp.shared_expert.gate_up_proj.weight'. Add a .weight_packed branch to that pass. Layers the model keeps native NVFP4 are emitted through the existing _dense_nvfp4_emit with llm-compressor's parts reader; layers it does not keep native are dequantized to bf16 and routed through the same bf16 fusion the .weight path already uses. The two parts readers are NOT interchangeable and must stay distinct: _nvfp4_parts (modelopt) uses the global scale directly, while nvfp4_parts_ct (llm-compressor) takes its reciprocal. Swapping them produces garbled output rather than an error, so each call site names the reader it needs. Verified: unsloth/Qwen3.6-35B-A3B-NVFP4-Fast resolves and loads with attn_quant='fp8_pertensor' expert_quant='nvfp4' dense_quant='nvfp4' lm_head_quant='fp8'; nvidia/Qwen3.6-35B-A3B-NVFP4's four values are unchanged.
1 parent b1596f4 commit bc27359

1 file changed

Lines changed: 46 additions & 5 deletions

File tree

‎python/freetoken/models/qwen3_5_moe/weight.py‎

Lines changed: 46 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -423,8 +423,17 @@ def _nvfp4_parts(f, raw_base: str):
423423
_NOT_DENSE_NVFP4 = object()
424424

425425

426+
def _emit_ct_packed_as_bf16(base: str, w, sc, g, bf16_buf: dict):
427+
"""Dequantize an llm-compressor packed NVFP4 projection to bf16 and route it through the
428+
same bf16 fusion the ``.weight`` path uses, for layers the model does NOT keep native."""
429+
bf16 = _dequant_nvfp4_weight(w, sc, g[:1])
430+
emit = _ct_bf16_fuse(base, bf16, bf16_buf, _PT_BF16_FUSE)
431+
return emit if emit is not None else [(base + ".weight", bf16)]
432+
433+
426434
def _dense_nvfp4_emit(
427-
f, base: str, raw_base: str, *, shared_nvfp4: bool, lmhead_nvfp4: bool, shared_buf: dict
435+
f, base: str, raw_base: str, *, shared_nvfp4: bool, lmhead_nvfp4: bool, shared_buf: dict,
436+
parts=_nvfp4_parts,
428437
):
429438
"""For a dense ``.weight`` whose checkpoint has a ``weight_scale_2`` (NVFP4), return the list
430439
of ``(key, tensor)`` to yield as native FP4 -- ``(.weight uint8, .weight_scale fp8 block,
@@ -437,19 +446,27 @@ def _dense_nvfp4_emit(
437446
438447
Returns ``[]`` while a gate/up merge is still buffered, or ``_NOT_DENSE_NVFP4`` if the model
439448
does not keep this layer native (the caller dequantizes to bf16 exactly as before). Shared by
440-
the mixed-FP8 dense pass and the default (pure-NVFP4) dense pass."""
449+
the mixed-FP8 dense pass and the default (pure-NVFP4) dense pass.
450+
451+
``parts`` reads the on-disk NVFP4 triple; it defaults to modelopt's ``_nvfp4_parts``
452+
(``.weight`` + ``.weight_scale`` + ``.weight_scale_2``), and callers pass
453+
``_nvfp4_parts_ct`` for llm-compressor (``.weight_packed`` + ``.weight_scale`` +
454+
``.weight_global_scale``). The two are NOT interchangeable: llm-compressor stores the
455+
QUANT-side global scale where modelopt stores the DEQUANT-side divisor, so
456+
``_nvfp4_parts_ct`` takes its reciprocal. Reading one convention with the other's reader
457+
garbles output rather than raising -- see upstream #390."""
441458
is_lmhead = base == "lm_head" or base.endswith(".lm_head")
442459
if lmhead_nvfp4 and is_lmhead:
443-
w, s, g = _nvfp4_parts(f, raw_base)
460+
w, s, g = parts(f, raw_base)
444461
return [(base + ".weight", w), (base + ".weight_scale", s), (base + ".weight_global", g)]
445462
if not shared_nvfp4:
446463
return _NOT_DENSE_NVFP4
447464
for gate_b, up_b, down_b, infix in _NVFP4_MLP_LAYOUTS:
448465
if base.endswith(down_b):
449-
w, s, g = _nvfp4_parts(f, raw_base)
466+
w, s, g = parts(f, raw_base)
450467
return [(base + ".weight", w), (base + ".weight_scale", s), (base + ".weight_global", g)]
451468
if base.endswith(gate_b) or base.endswith(up_b):
452-
w, s, g = _nvfp4_parts(f, raw_base)
469+
w, s, g = parts(f, raw_base)
453470
prefix = base.rsplit(infix, 1)[0] + infix
454471
slots = shared_buf.setdefault(prefix, {})
455472
slots["gate" if base.endswith(gate_b) else "up"] = (w, s, g)
@@ -512,6 +529,30 @@ def _iter_weights_attn_fp8(
512529
if _PACKED_EXPERT_PATTERN.match(name) is not None:
513530
continue # no packed experts in this checkpoint; guard anyway
514531

532+
if name.endswith(".weight_packed"):
533+
# llm-compressor stores dense NVFP4 as ``weight_packed`` with no plain
534+
# ``.weight``, so this projection would otherwise never reach
535+
# _dense_nvfp4_emit and the model would KeyError on the fused
536+
# ``shared_expert.gate_up_proj.weight`` it expects. Routed experts are
537+
# already excluded above (offload cache), so what arrives here is the
538+
# dense side: shared_expert gate/up/down, and lm_head on checkpoints that
539+
# pack it. Uses the CT parts reader for the reciprocal global scale.
540+
base = name[: -len(".weight_packed")]
541+
raw_base = raw_name[: -len(".weight_packed")]
542+
emit = _dense_nvfp4_emit(
543+
f, base, raw_base, shared_nvfp4=dense_nvfp4,
544+
lmhead_nvfp4=lmhead_nvfp4, shared_buf=nvfp4_shared_buf,
545+
parts=_nvfp4_parts_ct,
546+
)
547+
if emit is not _NOT_DENSE_NVFP4:
548+
yield from emit
549+
continue
550+
# Not a layer the model keeps native: dequantize to bf16 like the
551+
# ``.weight`` path below does for its own non-native cases.
552+
w, sc, g = _nvfp4_parts_ct(f, raw_base)
553+
yield from _emit_ct_packed_as_bf16(base, w, sc, g, bf16_buf)
554+
continue
555+
515556
if name.endswith(".weight"):
516557
base = name[: -len(".weight")]
517558
raw_base = raw_name[: -len(".weight")]

0 commit comments

Comments
 (0)