Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 32 additions & 0 deletions tests/quantization/test_modelopt.py
Original file line number Diff line number Diff line change
Expand Up @@ -165,6 +165,38 @@ def test_modelopt_mixed_precision_does_not_quantize_unlisted_fused_sibling():
assert config._resolve_quant_algo("model.layers.0.linear_attn.in_proj_ba") is None


def test_modelopt_mixed_precision_composes_gemma4_mappers():
from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
from vllm.model_executor.models.gemma4_mm import (
Gemma4ForConditionalGeneration,
)

config = _mixed_precision_config(
{
"model.language_model.layers.0.experts": {
"quant_algo": "NVFP4",
"group_size": 16,
},
"model.language_model.layers.1.moe.experts.gate_up_proj": {
"quant_algo": "NVFP4",
"group_size": 16,
},
}
)

config.apply_vllm_mapper(
Gemma4ForConditionalGeneration.hf_to_vllm_mapper.get_unstacked_mapper()
)
config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_unstacked_mapper())

expected_prefix = "language_model.model.layers.0.moe.experts"
assert set(config.quantized_layers) == {
expected_prefix,
"language_model.model.layers.1.moe.gate_up_proj",
}
assert config._resolve_quant_algo(expected_prefix) == "NVFP4"


def test_modelopt_mixed_precision_infers_fused_gate_up_projection():
from vllm.model_executor.layers.linear import LinearBase

Expand Down
8 changes: 7 additions & 1 deletion vllm/model_executor/models/gemma4.py
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,12 @@

logger = init_logger(__name__)

_GEMMA4_EXPERT_PARENT_MAPPER = WeightsMapper(
orig_to_new_regex={
re.compile(r"(?<!\.moe)\.experts$"): ".moe.experts",
}
)


def _remap_gemma4_expert_weight_name(name: str) -> str:
return re.sub(r"(?<!\.moe)\.experts\.(\d+)\.", r".moe.experts.\1.", name)
Expand Down Expand Up @@ -1508,7 +1514,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
class Gemma4ForCausalLM(
nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3
):
hf_to_vllm_mapper = WeightsMapper(
hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
orig_to_new_prefix={
# Gemma4ForConditionalGeneration already loads the text stack
# from `model.language_model.*`. We reuse that same checkpoint
Expand Down
9 changes: 6 additions & 3 deletions vllm/model_executor/models/gemma4_mm.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,10 @@
from vllm.logger import init_logger
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import ReplicatedLinear
from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
from vllm.model_executor.models.gemma4 import (
_GEMMA4_EXPERT_PARENT_MAPPER,
Gemma4ForCausalLM,
)
from vllm.model_executor.models.module_mapping import MultiModelKeys
from vllm.model_executor.models.transformers.utils import recursive_replace_linear
from vllm.multimodal import MULTIMODAL_REGISTRY
Expand Down Expand Up @@ -998,7 +1001,7 @@ class Gemma4ForConditionalGeneration(
}

# Maps checkpoint prefixes to vLLM module paths.
hf_to_vllm_mapper = WeightsMapper(
hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
orig_to_new_prefix={
# vision tower
"model.vision_tower": "vision_tower",
Expand All @@ -1010,7 +1013,7 @@ class Gemma4ForConditionalGeneration(
"model.language_model.": "language_model.model.",
"lm_head.": "language_model.lm_head.",
"model": "language_model.model",
}
},
)

def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Expand Down
Loading