Skip to content

fix(qwen3_5_moe): accept a per-channel fp8 weight_scale in the dense loader - #415

Open
salekseev wants to merge 1 commit into
FlashML-org:mainfrom
salekseev:upstream/ct-fp8-per-channel-scale
Open

fix(qwen3_5_moe): accept a per-channel fp8 weight_scale in the dense loader#415
salekseev wants to merge 1 commit into
FlashML-org:mainfrom
salekseev:upstream/ct-fp8-per-channel-scale

Conversation

@salekseev

Copy link
Copy Markdown

Hit this while trying to load a compressed-tensors FP8 checkpoint that uses per-channel weight scales. _per_row_scale assumes the scale is always a single scalar:

return scalar.reshape(1).to(torch.float32).expand(rows)

That's right for strategy: "tensor", but llm-compressor also emits strategy: "channel" — one scalar per output row, stored as [rows, 1]. On that shape the reshape blows up:

RuntimeError: shape '[1]' is invalid for input of size 248320

Since Fp8PerTensorLinear.weight_scale is already declared per output row, there's nothing to change downstream — a genuine per-tensor weight just stores the same scalar in every row. So this reshapes to [-1] and branches on the element count: 1 broadcasts exactly as before, rows passes straight through.

The part I'd push back on if I were reviewing: the third branch raises instead of doing something lenient. That's deliberate. If a mis-shaped scale broadcast row 0's factor across all 248320 output rows, the model would load happily and generate fluent nonsense. I'd much rather it fail at load with a message naming both counts.

Reachability, since it affects how you want to review this

You can't trigger this on main today. The only checkpoints that reach _iter_weights_attn_fp8 are ModelOpt ones, and those use a per-tensor scale. So on main this is hardening plus a prerequisite — it becomes load-bearing once a per-channel checkpoint can actually route there, which needs the compressed-tensors detection from #390 plus a one-token widening to accept "channel".

I mention it so you can weigh it as "small safe change that unblocks a checkpoint class" rather than "fixes a bug users are hitting", because the latter would be overselling it.

Tests

Adds tests/models/test_qwen3_5_moe_weight.py, modelled on tests/models/test_glm5_next_config.py. Covers both on-disk granularities, fp32 promotion from every storage dtype, and the mismatch raising. The per-channel case asserts row order with distinct values rather than a sum or a set — permuting the scales would keep every aggregate identical while silently scaling each row by the wrong factor.

There were no test_qwen3_5_moe_* files before this, so nothing is replaced.

Testing done

Verified against unsloth/Qwen3.6-35B-A3B-NVFP4-Fast and nvidia/Qwen3.6-35B-A3B-NVFP4 on an RTX 4080 SUPER (16 GiB, sm_89, TP=1). Full background and measurements in #252.

…loader

_per_row_scale unconditionally did scale.reshape(1), which is only valid for a
per-tensor scale. compressed-tensors also emits strategy: "channel", one scalar
per output row stored as [rows, 1], and on such a checkpoint the reshape raises

    RuntimeError: shape '[1]' is invalid for input of size 248320

Reshape to [-1] and branch on the element count instead: 1 broadcasts as before,
rows passes through, anything else raises rather than being broadcast -- a
mis-shaped scale that loaded would apply one row's factor to every output row and
serve fluent, wrong tokens.

Fp8PerTensorLinear.weight_scale is already declared per-output-row, so nothing
downstream changes and no kernel work follows.
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.

1 participant