diff --git a/python/freetoken/kernel/triton/fp8_pertensor_linear.py b/python/freetoken/kernel/triton/fp8_pertensor_linear.py index d1d4de5b4..2be284407 100644 --- a/python/freetoken/kernel/triton/fp8_pertensor_linear.py +++ b/python/freetoken/kernel/triton/fp8_pertensor_linear.py @@ -423,8 +423,45 @@ def __init__(self, in_features: int, output_sizes: list[int], has_bias: bool = F super().__init__(in_features, sum(output_sizes), has_bias) +class Fp8LMHead(Fp8PerTensorLinear): + """FP8 (W8A16) LM head: ``Fp8PerTensorLinear`` with ``ParallelLMHead``'s prefill + behaviour at TP=1 -- slice to the last token per sequence, then the W8A16 GEMV/GEMM + instead of a bf16 ``F.linear`` over the dequantized weight. Weight buffers, the + ``input_scale`` dance and the uniform-scale/segment precompute are all inherited; the + prefill slice is the only new code. + + Exists because a checkpoint can store ``lm_head`` as fp8 natively and no head class could + consume that: ``Nvfp4LMHead`` is FP4, and glm_moe_dsa/glm5_next's fp8 heads subclass the + bf16 ``ParallelLMHead`` and quantize at load. So a natively-fp8 head had to be + dequantized at load. For unsloth/Qwen3.6-35B-A3B-NVFP4-Fast that matrix is + ``[248320, 2048]``: 0.474 GiB kept native versus 0.947 GiB as bf16. On a 16 GiB card the + difference is not only decode traffic, it is the headroom an 8k prefill needs -- the + dequantized head left 118 MiB free and the GDN chunked-prefill workspace then OOM'd. + + ``weight_scale`` is per output row ``[vocab]``, which covers both compressed-tensors + strategies: "channel" ships one scalar per row, "tensor" ships one for the whole matrix + and the loader broadcasts it. TP=1 and untied embeddings, same as ``Nvfp4LMHead``.""" + + def __init__(self, num_embeddings: int, embedding_dim: int): + self.num_embeddings = num_embeddings + self.embedding_dim = embedding_dim + super().__init__(in_features=embedding_dim, out_features=num_embeddings) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + from freetoken.core import get_global_ctx + + batch = get_global_ctx().batch + if batch.is_prefill: + # Only the last position of each sequence produces logits, so slicing here keeps + # the [vocab, hidden] GEMM at M=batch rather than M=prompt_tokens. + indices = batch.attn_metadata.get_last_indices(batch.size) + x = x[indices].contiguous() + return super().forward(x) + + __all__ = [ "FP8", + "Fp8LMHead", "Fp8PerTensorLinear", "Fp8PerTensorColMerged", "fp8_pertensor_linear", diff --git a/python/freetoken/models/qwen3_5_moe/config.py b/python/freetoken/models/qwen3_5_moe/config.py index 2ac4b607f..813973777 100644 --- a/python/freetoken/models/qwen3_5_moe/config.py +++ b/python/freetoken/models/qwen3_5_moe/config.py @@ -71,7 +71,9 @@ def _expert_quant(hf_config: Any) -> str: def _lm_head_quant(hf_config: Any) -> str: """Whether the checkpoint stores ``lm_head`` as NVFP4. modelopt MIXED_PRECISION lists it in the per-layer ``quantized_layers`` map (``W4A16_NVFP4``); pure-NVFP4 checkpoints have no - per-layer map and leave lm_head bf16. Returns ``"nvfp4"`` or ``"none"``.""" + per-layer map and leave lm_head bf16. llm-compressor mixed-precision instead puts + ``lm_head`` in an fp8 ``config_groups`` entry. Returns ``"nvfp4"``, ``"fp8"`` + or ``"none"``.""" get = _quant_accessor(hf_config) if get is None: return "none" @@ -82,6 +84,23 @@ def _lm_head_quant(hf_config: Any) -> str: if name == "lm_head" or name.endswith(".lm_head"): if "fp4" in str((spec or {}).get("quant_algo", "")).lower(): return "nvfp4" + # compressed-tensors ``ignore`` wins over a group's ``targets``. A head listed there is + # bf16 on disk however broadly the group matches, so claiming it as fp8 would build an + # Fp8LMHead for a bf16 weight and fail the load on the dtype check. + if any("lm_head" in str(x) for x in (get("ignore") or [])): + return "none" + # llm-compressor mixed-precision puts lm_head in the fp8 group (unsloth NVFP4-Fast), with + # no per-layer `quantized_layers` map at all. Keeping it fp8 rather than dequantizing is + # worth 0.473 GiB on a [248320, 2048] head -- decode traffic, and the headroom an 8k + # prefill needs. Gated on the same weights geometry _attn_quant uses. + for g in (get("config_groups") or {}).values(): + if not g or not any("lm_head" in t for t in (g.get("targets") or [])): + continue + w = g.get("weights") or {} + if str(w.get("type", "")).lower() != "float" or int(w.get("num_bits", 0) or 0) != 8: + continue + if w.get("group_size") is None and w.get("strategy") in ("tensor", "channel"): + return "fp8" return "none" diff --git a/python/freetoken/models/qwen3_5_moe/model.py b/python/freetoken/models/qwen3_5_moe/model.py index eba7fd24f..3f3bfc67d 100644 --- a/python/freetoken/models/qwen3_5_moe/model.py +++ b/python/freetoken/models/qwen3_5_moe/model.py @@ -100,6 +100,15 @@ def __init__(self, config: ModelConfig): self.lm_head = Nvfp4LMHead( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size ) + elif getattr(config, "lm_head_quant", "none") == "fp8": + # checkpoint stores the (untied) lm_head as fp8 (llm-compressor mixed-precision): + # keep it native (W8A16). Halves this matrix versus dequantizing to bf16. + from freetoken.kernel.triton.fp8_pertensor_linear import Fp8LMHead + + assert not config.tie_word_embeddings, "fp8 lm_head assumes untied embeddings" + self.lm_head = Fp8LMHead( + num_embeddings=config.vocab_size, embedding_dim=config.hidden_size + ) else: self.lm_head = ParallelLMHead( num_embeddings=config.vocab_size, diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index d07f18cd7..bbaabf142 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -207,6 +207,7 @@ def iter_weights( include_non_moe=include_non_moe, include_moe_experts=include_moe_experts, dense_nvfp4=config.dense_quant == "nvfp4", lmhead_nvfp4=config.lm_head_quant == "nvfp4", + lmhead_fp8=config.lm_head_quant == "fp8", ) return tp_info = get_tp_info() @@ -311,9 +312,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, @@ -426,7 +448,7 @@ def _dense_nvfp4_emit( def _iter_weights_attn_fp8( model_path: str, device: torch.device, *, include_non_moe: bool, include_moe_experts: bool, - dense_nvfp4: bool = False, lmhead_nvfp4: bool = False, + dense_nvfp4: bool = False, lmhead_nvfp4: bool = False, lmhead_fp8: bool = False, ) -> Iterator[tuple[str, torch.Tensor]]: """Dense pass for the modelopt MIXED_PRECISION Qwen3.5 checkpoint. @@ -474,7 +496,14 @@ def _iter_weights_attn_fp8( raw_base = raw_name[: -len(".weight")] has_s2 = raw_base + ".weight_scale_2" in keyset has_s = raw_base + ".weight_scale" in keyset - if has_s and not has_s2: # per-tensor FP8 dense projection + # An fp8 lm_head is only safe to keep native if the model built an + # Fp8LMHead for it (lm_head_quant == "fp8"). The bf16 ParallelLMHead has + # no weight_scale buffer, so emitting fp8 into it fails the load with + # "Unexpected keys in state_dict: ['lm_head.weight_scale']". Keep it + # native only when the model asked for an fp8 head; otherwise fall + # through and dequantize. + is_lmhead = base == "lm_head" or base.endswith(".lm_head") + if has_s and not has_s2 and (lmhead_fp8 or not is_lmhead): w = f.get_tensor(raw_name) # fp8-e4m3, kept verbatim sc = f.get_tensor(raw_base + ".weight_scale") # modelopt's calibrated activation scale: kept (not dropped with the diff --git a/tests/models/test_qwen3_5_moe_fp8_lm_head.py b/tests/models/test_qwen3_5_moe_fp8_lm_head.py new file mode 100644 index 000000000..282aaa8f9 --- /dev/null +++ b/tests/models/test_qwen3_5_moe_fp8_lm_head.py @@ -0,0 +1,99 @@ +"""qwen3_5_moe fp8 ``lm_head``: detection, layer choice, and the loader gate. + +A compressed-tensors mixed-precision checkpoint can put ``lm_head`` in its fp8 group +(unsloth/Qwen3.6-35B-A3B-NVFP4-Fast). Three things have to agree for that to load: the config +detector reports ``"fp8"``, the model builds an ``Fp8LMHead``, and the dense loader keeps the +weight fp8 instead of dequantizing it. The loader gate is the subtle one -- emitting fp8 into +the bf16 ``ParallelLMHead`` fails with an unexpected ``lm_head.weight_scale`` key. + +Runs off trimmed copies of the real ``quantization_config`` blocks. +""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kernel.triton.fp8_pertensor_linear import Fp8LMHead, Fp8PerTensorLinear +from freetoken.models.qwen3_5_moe.config import _lm_head_quant + + +def _quant_config(*, lm_head_in_targets: bool, strategy: str = "channel") -> dict: + """Trimmed from unsloth/Qwen3.6-35B-A3B-NVFP4-Fast: fp8 group + NVFP4 expert group.""" + targets = [".self_attn.q_proj", ".linear_attn.in_proj_qkv"] + if lm_head_in_targets: + targets = targets + ["lm_head"] + return { + "quant_method": "compressed-tensors", + "format": "mixed-precision", + "config_groups": { + "group_0": { + "targets": targets, + "weights": {"num_bits": 8, "type": "float", + "strategy": strategy, "group_size": None}, + }, + "group_1": { + "targets": ["re:.*mlp.experts.*"], + "weights": {"num_bits": 4, "type": "float", + "strategy": "tensor_group", "group_size": 16}, + }, + }, + } + + +class _Cfg: + def __init__(self, quantization_config): + self.quantization_config = quantization_config + + +@pytest.mark.parametrize("strategy", ["tensor", "channel"]) +def test_lm_head_in_the_fp8_group_is_detected(strategy): + cfg = _Cfg(_quant_config(lm_head_in_targets=True, strategy=strategy)) + assert _lm_head_quant(cfg) == "fp8" + + +def test_lm_head_outside_the_fp8_group_is_not_claimed(): + """Only the group that targets ``lm_head`` may decide the head's dtype; a checkpoint whose + fp8 group covers the attention projections alone must leave the head bf16.""" + cfg = _Cfg(_quant_config(lm_head_in_targets=False)) + assert _lm_head_quant(cfg) == "none" + + +def test_an_ignored_lm_head_is_not_claimed_even_when_a_group_targets_it(): + """compressed-tensors ``ignore`` wins over ``targets``. This is the common llm-compressor + shape: a group matches broadly, then ``ignore`` carves modules back out. Such a head is + bf16 on disk, so claiming it as fp8 would fail the load on the dtype check.""" + cfg = _Cfg(_quant_config(lm_head_in_targets=True)) + cfg.quantization_config["ignore"] = ["re:.*lm_head"] + assert _lm_head_quant(cfg) == "none" + + +def test_an_unrelated_ignore_entry_does_not_suppress_the_fp8_head(): + cfg = _Cfg(_quant_config(lm_head_in_targets=True)) + cfg.quantization_config["ignore"] = ["re:.*mlp.experts.*", "model.embed_tokens"] + assert _lm_head_quant(cfg) == "fp8" + + +def test_a_4bit_lm_head_group_is_not_read_as_fp8(): + cfg = _Cfg(_quant_config(lm_head_in_targets=True)) + cfg.quantization_config["config_groups"]["group_0"]["weights"]["num_bits"] = 4 + assert _lm_head_quant(cfg) == "none" + + +def test_a_grouped_fp8_lm_head_is_not_read_as_per_tensor(): + """Fp8LMHead consumes a per-output-row scale; a block/group scale is a different layout + and must not be routed here.""" + cfg = _Cfg(_quant_config(lm_head_in_targets=True, strategy="group")) + cfg.quantization_config["config_groups"]["group_0"]["weights"]["group_size"] = 128 + assert _lm_head_quant(cfg) == "none" + + +def test_fp8_lm_head_buffers_match_the_per_output_row_contract(): + head = Fp8LMHead(num_embeddings=64, embedding_dim=8) + assert isinstance(head, Fp8PerTensorLinear) + assert head.weight.shape == (64, 8) + assert head.weight.dtype == torch.float8_e4m3fn + # Per output row, fp32 -- covers "channel" directly and "tensor" once broadcast. + assert head.weight_scale.shape == (64,) + assert head.weight_scale.dtype == torch.float32 + assert head.bias is None 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)