wehub-resource-sync
286 lines
11 KiB
Python
286 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import tempfile
|
|
|
|
import huggingface_hub.constants
|
|
import pytest
|
|
from huggingface_hub.utils import LocalEntryNotFoundError
|
|
|
|
from vllm.model_executor.model_loader.weight_utils import (
|
|
download_weights_from_hf,
|
|
maybe_remap_kv_scale_name,
|
|
)
|
|
|
|
|
|
def test_download_weights_from_hf():
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
# assert LocalEntryNotFoundError error is thrown
|
|
# if offline is set and model is not cached
|
|
huggingface_hub.constants.HF_HUB_OFFLINE = True
|
|
with pytest.raises(LocalEntryNotFoundError):
|
|
download_weights_from_hf(
|
|
"facebook/opt-125m",
|
|
allow_patterns=["*.safetensors", "*.bin"],
|
|
cache_dir=tmpdir,
|
|
)
|
|
|
|
# download the model
|
|
huggingface_hub.constants.HF_HUB_OFFLINE = False
|
|
download_weights_from_hf(
|
|
"facebook/opt-125m",
|
|
allow_patterns=["*.safetensors", "*.bin"],
|
|
cache_dir=tmpdir,
|
|
)
|
|
|
|
# now it should work offline
|
|
huggingface_hub.constants.HF_HUB_OFFLINE = True
|
|
assert (
|
|
download_weights_from_hf(
|
|
"facebook/opt-125m",
|
|
allow_patterns=["*.safetensors", "*.bin"],
|
|
cache_dir=tmpdir,
|
|
)
|
|
is not None
|
|
)
|
|
|
|
|
|
class TestMaybeRemapKvScaleName:
|
|
"""Tests for maybe_remap_kv_scale_name covering all checkpoint formats."""
|
|
|
|
PARAMS_DICT = {
|
|
"model.layers.0.self_attn.attn.k_scale": None,
|
|
"model.layers.0.self_attn.attn.v_scale": None,
|
|
"model.layers.0.self_attn.attn.q_scale": None,
|
|
"model.layers.0.self_attn.qkv_proj.weight": None,
|
|
}
|
|
|
|
def test_qkv_proj_k_scale(self):
|
|
"""Qwen3-MoE / llm-compressor format: qkv_proj.k_scale -> attn.k_scale
|
|
Regression test for https://github.com/vllm-project/vllm/issues/25047"""
|
|
result = maybe_remap_kv_scale_name(
|
|
"model.layers.0.self_attn.qkv_proj.k_scale", self.PARAMS_DICT
|
|
)
|
|
assert result == "model.layers.0.self_attn.attn.k_scale"
|
|
|
|
def test_qkv_proj_v_scale(self):
|
|
"""Qwen3-MoE / llm-compressor format: qkv_proj.v_scale -> attn.v_scale
|
|
Regression test for https://github.com/vllm-project/vllm/issues/25047"""
|
|
result = maybe_remap_kv_scale_name(
|
|
"model.layers.0.self_attn.qkv_proj.v_scale", self.PARAMS_DICT
|
|
)
|
|
assert result == "model.layers.0.self_attn.attn.v_scale"
|
|
|
|
def test_modelopt_k_proj_k_scale(self):
|
|
"""ModelOpt format: k_proj.k_scale -> attn.k_scale"""
|
|
result = maybe_remap_kv_scale_name(
|
|
"model.layers.0.self_attn.k_proj.k_scale", self.PARAMS_DICT
|
|
)
|
|
assert result == "model.layers.0.self_attn.attn.k_scale"
|
|
|
|
def test_modelopt_v_proj_v_scale(self):
|
|
"""ModelOpt format: v_proj.v_scale -> attn.v_scale"""
|
|
result = maybe_remap_kv_scale_name(
|
|
"model.layers.0.self_attn.v_proj.v_scale", self.PARAMS_DICT
|
|
)
|
|
assert result == "model.layers.0.self_attn.attn.v_scale"
|
|
|
|
def test_deprecated_kv_scale(self):
|
|
"""Old format: kv_scale -> attn.k_scale (deprecated)"""
|
|
result = maybe_remap_kv_scale_name(
|
|
"model.layers.0.self_attn.kv_scale", self.PARAMS_DICT
|
|
)
|
|
assert result == "model.layers.0.self_attn.attn.k_scale"
|
|
|
|
def test_default_bare_k_scale(self):
|
|
"""Default format: .k_scale -> .attn.k_scale"""
|
|
result = maybe_remap_kv_scale_name(
|
|
"model.layers.0.self_attn.k_scale", self.PARAMS_DICT
|
|
)
|
|
assert result == "model.layers.0.self_attn.attn.k_scale"
|
|
|
|
def test_non_scale_name_unchanged(self):
|
|
"""Non-scale names should be returned unchanged."""
|
|
name = "model.layers.0.self_attn.qkv_proj.weight"
|
|
result = maybe_remap_kv_scale_name(name, self.PARAMS_DICT)
|
|
assert result == name
|
|
|
|
def test_nvfp4_modelopt_k_proj_k_scale(self):
|
|
"""ModelOpt NVFP4 format (e.g. nvidia/Qwen3-30B-A3B-NVFP4):
|
|
k_proj.k_scale -> attn.k_scale.
|
|
Validates that NVFP4 checkpoints are not broken by this change."""
|
|
result = maybe_remap_kv_scale_name(
|
|
"model.layers.0.self_attn.k_proj.k_scale", self.PARAMS_DICT
|
|
)
|
|
assert result == "model.layers.0.self_attn.attn.k_scale"
|
|
|
|
def test_nvfp4_modelopt_v_proj_v_scale(self):
|
|
"""ModelOpt NVFP4 format (e.g. nvidia/Qwen3-30B-A3B-NVFP4):
|
|
v_proj.v_scale -> attn.v_scale.
|
|
Validates that NVFP4 checkpoints are not broken by this change."""
|
|
result = maybe_remap_kv_scale_name(
|
|
"model.layers.0.self_attn.v_proj.v_scale", self.PARAMS_DICT
|
|
)
|
|
assert result == "model.layers.0.self_attn.attn.v_scale"
|
|
|
|
def test_qwen3_vl_moe_qkv_proj_k_scale(self):
|
|
"""Qwen3-VL-MoE uses the same fused qkv_proj naming as Qwen3-MoE.
|
|
Regression test for qwen3_vl_moe.py fix (same bug as #25047)."""
|
|
result = maybe_remap_kv_scale_name(
|
|
"model.layers.0.self_attn.qkv_proj.k_scale", self.PARAMS_DICT
|
|
)
|
|
assert result == "model.layers.0.self_attn.attn.k_scale"
|
|
|
|
def test_qwen3_vl_moe_qkv_proj_v_scale(self):
|
|
"""Qwen3-VL-MoE uses the same fused qkv_proj naming as Qwen3-MoE.
|
|
Regression test for qwen3_vl_moe.py fix (same bug as #25047)."""
|
|
result = maybe_remap_kv_scale_name(
|
|
"model.layers.0.self_attn.qkv_proj.v_scale", self.PARAMS_DICT
|
|
)
|
|
assert result == "model.layers.0.self_attn.attn.v_scale"
|
|
|
|
def test_nvfp4_weight_scale_not_remapped(self):
|
|
"""NVFP4 weight_scale should not be touched by remap (not a kv scale)."""
|
|
name = "model.layers.0.self_attn.k_proj.weight_scale"
|
|
result = maybe_remap_kv_scale_name(name, self.PARAMS_DICT)
|
|
assert result == name
|
|
|
|
def test_nvfp4_input_scale_not_remapped(self):
|
|
"""NVFP4 input_scale should not be touched by remap (not a kv scale)."""
|
|
name = "model.layers.0.self_attn.k_proj.input_scale"
|
|
result = maybe_remap_kv_scale_name(name, self.PARAMS_DICT)
|
|
assert result == name
|
|
|
|
def test_missing_target_returns_none(self):
|
|
"""If remapped name not in params_dict, return None."""
|
|
empty_params: dict[str, None] = {}
|
|
result = maybe_remap_kv_scale_name(
|
|
"model.layers.0.self_attn.qkv_proj.k_scale", empty_params
|
|
)
|
|
assert result is None
|
|
|
|
|
|
class TestKvCacheScaleMapper:
|
|
"""The `WeightsMapper` returned by `get_cache_scale_mapper` replaces the
|
|
per-model `maybe_remap_kv_scale_name` calls. It must remap the same set of
|
|
checkpoint formats (the non-`params_dict`-dependent ones) and be idempotent
|
|
so it composes safely with a model's own qkv/gate_up `hf_to_vllm_mapper`."""
|
|
|
|
def _mapper(self):
|
|
# `get_cache_scale_mapper` does not use `self`; call it on the base
|
|
# class to get the default (non-config-specific) mapper.
|
|
from vllm.model_executor.layers.quantization.base_config import (
|
|
QuantizationConfig,
|
|
)
|
|
|
|
return QuantizationConfig.get_cache_scale_mapper()
|
|
|
|
def _map(self, name: str) -> str | None:
|
|
return self._mapper()._map_name(name)
|
|
|
|
@pytest.mark.parametrize(
|
|
"name,expected",
|
|
[
|
|
# Qwen3-MoE / llm-compressor fused qkv_proj
|
|
(
|
|
"model.layers.0.self_attn.qkv_proj.k_scale",
|
|
"model.layers.0.self_attn.attn.k_scale",
|
|
),
|
|
(
|
|
"model.layers.0.self_attn.qkv_proj.v_scale",
|
|
"model.layers.0.self_attn.attn.v_scale",
|
|
),
|
|
# ModelOpt / NVFP4 k_proj/v_proj
|
|
(
|
|
"model.layers.0.self_attn.k_proj.k_scale",
|
|
"model.layers.0.self_attn.attn.k_scale",
|
|
),
|
|
(
|
|
"model.layers.0.self_attn.v_proj.v_scale",
|
|
"model.layers.0.self_attn.attn.v_scale",
|
|
),
|
|
# deprecated fused kv_scale and bare scales
|
|
(
|
|
"model.layers.0.self_attn.kv_scale",
|
|
"model.layers.0.self_attn.attn.k_scale",
|
|
),
|
|
(
|
|
"model.layers.0.self_attn.k_scale",
|
|
"model.layers.0.self_attn.attn.k_scale",
|
|
),
|
|
# NemotronH mixer
|
|
(
|
|
"model.layers.0.mixer.k_proj.k_scale",
|
|
"model.layers.0.mixer.attn.k_scale",
|
|
),
|
|
# already in vLLM form -> unchanged (idempotent)
|
|
(
|
|
"model.layers.0.self_attn.attn.k_scale",
|
|
"model.layers.0.self_attn.attn.k_scale",
|
|
),
|
|
# non-kv scales must not be touched
|
|
(
|
|
"model.layers.0.self_attn.k_proj.weight_scale",
|
|
"model.layers.0.self_attn.k_proj.weight_scale",
|
|
),
|
|
(
|
|
"model.layers.0.self_attn.k_proj.input_scale",
|
|
"model.layers.0.self_attn.k_proj.input_scale",
|
|
),
|
|
# regular weights untouched
|
|
(
|
|
"model.layers.0.self_attn.q_proj.weight",
|
|
"model.layers.0.self_attn.q_proj.weight",
|
|
),
|
|
],
|
|
)
|
|
def test_remap(self, name, expected):
|
|
assert self._map(name) == expected
|
|
|
|
@pytest.mark.parametrize(
|
|
"name",
|
|
[
|
|
"model.layers.0.self_attn.k_scale",
|
|
"model.layers.0.self_attn.k_proj.k_scale",
|
|
"model.layers.0.self_attn.qkv_proj.v_scale",
|
|
"model.layers.0.mixer.k_proj.k_scale",
|
|
],
|
|
)
|
|
def test_idempotent(self, name):
|
|
once = self._map(name)
|
|
assert once is not None
|
|
assert self._map(once) == once
|
|
|
|
def test_composes_with_qkv_mapper(self):
|
|
"""Applied together with a model's qkv/gate_up mapper, the regex scale
|
|
rules run before the substr rename, so scales are normalized to `.attn.`
|
|
and regular projections are still fused correctly."""
|
|
from vllm.model_executor.models.utils import WeightsMapper
|
|
|
|
model_mapper = WeightsMapper(
|
|
orig_to_new_substr={
|
|
".q_proj": ".qkv_proj.q",
|
|
".k_proj": ".qkv_proj.k",
|
|
".v_proj": ".qkv_proj.v",
|
|
}
|
|
)
|
|
# AutoWeightsLoader does `mapper |= cache_scale_mapper`
|
|
combined = model_mapper | self._mapper()
|
|
|
|
assert (
|
|
combined._map_name("model.layers.0.self_attn.q_proj.weight")
|
|
== "model.layers.0.self_attn.qkv_proj.q.weight"
|
|
)
|
|
assert (
|
|
combined._map_name("model.layers.0.self_attn.k_proj.k_scale")
|
|
== "model.layers.0.self_attn.attn.k_scale"
|
|
)
|
|
assert (
|
|
combined._map_name("model.layers.0.self_attn.k_scale")
|
|
== "model.layers.0.self_attn.attn.k_scale"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_download_weights_from_hf()
|