From 165931cf9121ff84c1819756a3b855fe227dd815 Mon Sep 17 00:00:00 2001 From: Stas Alekseev <100800+salekseev@users.noreply.github.com> Date: Mon, 7 Sep 2026 20:35:01 -0400 Subject: [PATCH] fix(qwen3_5_moe): accept a per-channel fp8 weight_scale in the dense 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. --- python/freetoken/models/qwen3_5_moe/weight.py | 27 ++++++++-- tests/models/test_qwen3_5_moe_weight.py | 54 +++++++++++++++++++ 2 files changed, 78 insertions(+), 3 deletions(-) create mode 100644 tests/models/test_qwen3_5_moe_weight.py diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index d07f18cd7..8324a956d 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -311,9 +311,30 @@ def iter_weights( } -def _per_row_scale(scalar: torch.Tensor, rows: int) -> torch.Tensor: - """Per-tensor scalar -> per-output-row fp32 vector ``[rows]`` (exact broadcast).""" - return scalar.reshape(1).to(torch.float32).expand(rows) +def _per_row_scale(scale: torch.Tensor, rows: int) -> torch.Tensor: + """FP8 weight scale -> per-output-row fp32 vector ``[rows]``, which is what + ``Fp8PerTensorLinear.weight_scale`` wants whichever on-disk shape arrives: + + * **per-tensor** (``strategy: "tensor"``) -- one scalar for the whole weight, broadcast + across the rows. Exact. + * **per-channel** (``strategy: "channel"``) -- ``[rows, 1]``, already one scalar per row, + so the reshape alone is enough. This case is why the per-tensor ``reshape(1)`` cannot be + unconditional: on a per-channel checkpoint it raises ``RuntimeError: shape '[1]' is + invalid for input of size ``. + + Any other element count is raised on rather than broadcast: a silently mis-shaped scale + would be applied to the wrong output rows and produce plausible garbage. + """ + flat = scale.reshape(-1).to(torch.float32) + if flat.numel() == 1: + return flat.expand(rows) + if flat.numel() != rows: + raise ValueError( + f"fp8 weight_scale has {flat.numel()} elements for a weight with {rows} output " + f"rows (shape {tuple(scale.shape)}); expected either 1 (per-tensor) or {rows} " + "(per-channel)" + ) + return flat def _pt_fp8_fuse(base: str, weight: torch.Tensor, scalar: torch.Tensor, diff --git a/tests/models/test_qwen3_5_moe_weight.py b/tests/models/test_qwen3_5_moe_weight.py new file mode 100644 index 000000000..ffd4f6120 --- /dev/null +++ b/tests/models/test_qwen3_5_moe_weight.py @@ -0,0 +1,54 @@ +"""qwen3_5_moe dense weight-loader helpers: the fp8 ``weight_scale`` shape contract. + +``_per_row_scale`` is what turns whatever ``weight_scale`` a compressed-tensors or modelopt +checkpoint puts on disk into the per-output-row fp32 vector ``Fp8PerTensorLinear.weight_scale`` +declares. Both on-disk granularities are exercised, plus the refusal that keeps a mis-shaped +scale from being applied to the wrong output rows. +""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.models.qwen3_5_moe.weight import _per_row_scale + +_ROWS = 7 + + +@pytest.mark.parametrize("shape", [(), (1,), (1, 1)]) +def test_per_tensor_scale_broadcasts_to_every_row(shape): + """``strategy: "tensor"`` -- one scalar for the whole weight.""" + out = _per_row_scale(torch.full(shape, 0.25), _ROWS) + assert out.shape == (_ROWS,) + assert out.dtype == torch.float32 + assert out.tolist() == [0.25] * _ROWS + + +@pytest.mark.parametrize("shape", [(_ROWS, 1), (1, _ROWS), (_ROWS,)]) +def test_per_channel_scale_keeps_row_order(shape): + """``strategy: "channel"`` -- already one scalar per row; order must survive the reshape. + + Order is asserted with distinct values rather than a set/sum: permuting the scales would + keep every aggregate identical while silently scaling each output row by another row's + factor. + """ + values = [1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0] + out = _per_row_scale(torch.tensor(values).reshape(shape), _ROWS) + assert out.shape == (_ROWS,) + assert out.dtype == torch.float32 + assert out.tolist() == values + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64]) +def test_scale_is_promoted_to_fp32_from_any_storage_dtype(dtype): + out = _per_row_scale(torch.ones(_ROWS, 1, dtype=dtype), _ROWS) + assert out.dtype == torch.float32 + + +@pytest.mark.parametrize("shape", [(3, 1), (_ROWS + 1, 1), (2, 3)]) +def test_a_mismatched_scale_raises_instead_of_broadcasting(shape): + """Neither 1 nor ``rows`` elements: raise. Broadcasting row 0 over every output row would + load without error and serve fluent, wrong tokens.""" + with pytest.raises(ValueError, match="expected either 1"): + _per_row_scale(torch.ones(shape), _ROWS)