[Qwen3.5] Build MTP head unquantized when checkpoint ships it dense (PR #58813)
Fixes #58807. A compressed-tensors pack-quantized checkpoint that ships its MTP head dense (plain mtp.*.weight names in a separate bf16 shard) crashes at weight load: Qwen3_5MultiTokenPredictor builds its Linears from the target quant config, so the loader expects packed names (weight_packed) the checkpoint never had, and the weight-map walk dies with "no module or parameter named fc.weight". Detect the dense MTP head from the checkpoint itself: read the safetensors index via the same repo-file channel used for recipe.yaml (local cache first, Hub download as fallback). If any mtp.* tensor carries a quantization artifact suffix the head is properly quantized and nothing changes; if every mtp.* entry is a plain .weight the whole predictor builds unquantized (fc and layers), mirroring the existing modelopt_fp4 and GPTQ dynamic carve-outs. Single-file checkpoints have no index and keep current behavior; NVFP4 checkpoints are unaffected (their mtp.layers.* carry artifacts). Co-Authored-By: Hinata <hinata@idihore.id> Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
@@ -0,0 +1,78 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Dense-MTP-head detection for pack-quantized checkpoints (#58807).
|
||||
|
||||
The whole fix hinges on reading the checkpoint's actual tensor names, so the
|
||||
tests pin the naming contract: dense means plain `mtp.*.weight` keys with no
|
||||
quantization artifact suffix anywhere under `mtp.`.
|
||||
"""
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.model_executor.models.qwen3_5_mtp import (
|
||||
_checkpoint_ships_dense_mtp,
|
||||
_mtp_weight_names,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
|
||||
|
||||
def _model_config_with_index(monkeypatch, names):
|
||||
monkeypatch.setattr(
|
||||
"vllm.model_executor.models.qwen3_5_mtp._mtp_weight_names",
|
||||
lambda _config: names,
|
||||
)
|
||||
return SimpleNamespace()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"names,expected",
|
||||
[
|
||||
# The reporter's checkpoint: every mtp.* tensor dense.
|
||||
(["mtp.fc.weight", "mtp.layers.0.mlp.down_proj.weight"], True),
|
||||
# Biases are dense too and must not count as artifacts.
|
||||
(["mtp.fc.weight", "mtp.fc.bias"], True),
|
||||
# A properly pack-quantized MTP head keeps the target quantization.
|
||||
(["mtp.fc.weight_packed", "mtp.fc.weight_scale"], False),
|
||||
# Channel-wise quantized MTP head: dense name plus scale companion.
|
||||
(["mtp.fc.weight", "mtp.fc.weight_scale"], False),
|
||||
# No MTP head in the checkpoint at all.
|
||||
(["model.layers.0.self_attn.q_proj.weight"], False),
|
||||
# Layout unknown (no index): keep current behavior.
|
||||
(None, False),
|
||||
([], False),
|
||||
],
|
||||
)
|
||||
def test_dense_mtp_detection(monkeypatch, names, expected):
|
||||
config = _model_config_with_index(monkeypatch, names)
|
||||
assert _checkpoint_ships_dense_mtp(config) is expected
|
||||
|
||||
|
||||
def test_weight_names_read_from_local_index(tmp_path, monkeypatch):
|
||||
index = {
|
||||
"weight_map": {
|
||||
"model.layers.0.self_attn.q_proj.weight_packed": "model-00001.safetensors",
|
||||
"mtp.fc.weight": "model-mtp-bf16.safetensors",
|
||||
"mtp.layers.0.mlp.down_proj.weight": "model-mtp-bf16.safetensors",
|
||||
}
|
||||
}
|
||||
(tmp_path / "model.safetensors.index.json").write_text(json.dumps(index))
|
||||
config = SimpleNamespace(model=str(tmp_path), revision=None)
|
||||
|
||||
names = _mtp_weight_names(config)
|
||||
|
||||
assert names == ["mtp.fc.weight", "mtp.layers.0.mlp.down_proj.weight"]
|
||||
|
||||
|
||||
def test_weight_names_none_without_index(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"vllm.model_executor.models.qwen3_5_mtp.huggingface_hub.hf_hub_download",
|
||||
lambda *args, **kwargs: (_ for _ in ()).throw(FileNotFoundError()),
|
||||
)
|
||||
config = SimpleNamespace(model=str(tmp_path), revision=None)
|
||||
|
||||
assert _mtp_weight_names(config) is None
|
||||
@@ -2,13 +2,15 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Inference-only Qwen3_5 MTP model."""
|
||||
|
||||
import json
|
||||
from collections.abc import Iterable
|
||||
|
||||
import huggingface_hub
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm.compilation.decorators import support_torch_compile
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.config import ModelConfig, VllmConfig
|
||||
from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fused_moe.utils import (
|
||||
@@ -34,6 +36,7 @@ from vllm.model_executor.models.utils import sequence_parallel_chunk
|
||||
from vllm.sequence import IntermediateTensors
|
||||
from vllm.transformers_utils.configs.qwen3_5 import Qwen3_5TextConfig
|
||||
from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeTextConfig
|
||||
from vllm.transformers_utils.repo_utils import try_get_local_file
|
||||
|
||||
from .interfaces import (
|
||||
MultiModalEmbeddings,
|
||||
@@ -52,6 +55,64 @@ from .utils import (
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Parameter suffixes a quantized MTP head would carry in the checkpoint.
|
||||
_MTP_QUANT_ARTIFACT_SUFFIXES = (
|
||||
"weight_packed",
|
||||
"weight_scale",
|
||||
"weight_zero_point",
|
||||
"input_scale",
|
||||
"input_zero_point",
|
||||
)
|
||||
|
||||
|
||||
def _mtp_weight_names(model_config: ModelConfig) -> list[str] | None:
|
||||
"""Return the checkpoint's mtp.* tensor names, or None when unknown.
|
||||
|
||||
Reads the safetensors index from the local cache or the Hub, the same way
|
||||
config-time helpers read recipe.yaml. A single-file checkpoint has no
|
||||
index, so its layout stays unknown (callers keep the default behavior).
|
||||
"""
|
||||
file_path = try_get_local_file(
|
||||
model=model_config.model,
|
||||
file_name="model.safetensors.index.json",
|
||||
revision=model_config.revision,
|
||||
)
|
||||
if file_path is None:
|
||||
try:
|
||||
file_path = huggingface_hub.hf_hub_download(
|
||||
model_config.model,
|
||||
"model.safetensors.index.json",
|
||||
revision=model_config.revision,
|
||||
)
|
||||
except Exception:
|
||||
return None
|
||||
try:
|
||||
with open(file_path) as f:
|
||||
weight_map = json.load(f).get("weight_map", {})
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
return [name for name in weight_map if name.startswith("mtp.")]
|
||||
|
||||
|
||||
def _checkpoint_ships_dense_mtp(model_config: ModelConfig) -> bool:
|
||||
"""Whether the checkpoint ships the MTP head as dense weights.
|
||||
|
||||
Quantizers sometimes leave the MTP head unquantized without declaring it
|
||||
in the quant config's ignore list; the tensors then arrive as plain
|
||||
`mtp.*.weight` with no packed or scale companions (#58807). Detection is
|
||||
name-based, so a properly quantized MTP head (any artifact suffix on any
|
||||
mtp.* tensor) keeps the target's quantization.
|
||||
"""
|
||||
names = _mtp_weight_names(model_config)
|
||||
if not names:
|
||||
return False
|
||||
mtp_names = [name for name in names if name.startswith("mtp.")]
|
||||
if not mtp_names:
|
||||
return False
|
||||
return not any(
|
||||
name.rsplit(".", 1)[-1] in _MTP_QUANT_ARTIFACT_SUFFIXES for name in mtp_names
|
||||
)
|
||||
|
||||
|
||||
@support_torch_compile(
|
||||
dynamic_arg_dims={
|
||||
@@ -91,9 +152,19 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
|
||||
# missing from hf_quant_config.json exclude_modules. Force unquantized.
|
||||
# Ref: https://github.com/vllm-project/vllm/pull/38650
|
||||
# Ref: https://github.com/NVIDIA/Model-Optimizer/pull/1124
|
||||
# compressed-tensors pack-quantized checkpoints may likewise ship the
|
||||
# whole MTP head dense without declaring it in the ignore list (#58807).
|
||||
mtp_dense = (
|
||||
quant_config is not None
|
||||
and quant_config.get_name() == "compressed-tensors"
|
||||
and _checkpoint_ships_dense_mtp(model_config)
|
||||
)
|
||||
fc_quant = (
|
||||
None
|
||||
if (quant_config and quant_config.get_name() == "modelopt_fp4")
|
||||
if (
|
||||
quant_config
|
||||
and (quant_config.get_name() == "modelopt_fp4" or mtp_dense)
|
||||
)
|
||||
else quant_config
|
||||
)
|
||||
self.fc = ColumnParallelLinear(
|
||||
@@ -109,12 +180,17 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
|
||||
# GPTQ: quantized checkpoints may exclude MTP from quantization via
|
||||
# quantization_config.dynamic with "-:pattern" entries. When detected,
|
||||
# disable quantization for MTP layers so they use unquantized params.
|
||||
# The dense-MTP detection above covers the same need for
|
||||
# compressed-tensors checkpoints that declare nothing at all.
|
||||
original_quant = vllm_config.quant_config
|
||||
if quant_config and quant_config.get_name() not in ("modelopt_fp4",):
|
||||
hf_qc = getattr(model_config.hf_config, "quantization_config", None)
|
||||
if isinstance(hf_qc, dict):
|
||||
dynamic = hf_qc.get("dynamic", {})
|
||||
if any(k.startswith("-:") and "mtp" in k for k in dynamic):
|
||||
if mtp_dense or (
|
||||
isinstance(hf_qc, dict)
|
||||
and any(
|
||||
k.startswith("-:") and "mtp" in k for k in hf_qc.get("dynamic", {})
|
||||
)
|
||||
):
|
||||
vllm_config.quant_config = None
|
||||
self.layers = torch.nn.ModuleList(
|
||||
Qwen3_5DecoderLayer(
|
||||
|
||||
Reference in New Issue
Block a user