Skip to content

[Bugfix] Consolidate Gemma2/3 GGUF fixes for correctness on Blackwell #37220

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kitaekatt wants to merge 5 commits into
base: main
Choose a base branch
from
Open

[Bugfix] Consolidate Gemma2/3 GGUF fixes for correctness on Blackwell #37220

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
fff7e08
fix(gguf): Consolidate Gemma2/3 GGUF fixes for correctness on Blackwell
kitaekatt Mar 16, 2026
c3ac087
Merge branch 'main' into fix/gemma-gguf-consolidated
kitaekatt Mar 25, 2026
c7a1a30
refactor: Extract _get_norm_cls helper to reduce duplication
kitaekatt Mar 25, 2026
5befed8
Merge remote-tracking branch 'upstream/main' into fix/gemma-gguf-cons…
kitaekatt Apr 18, 2026
86fa040
refactor(gguf): Extract maybe_patch_gguf_tokenizer helper
kitaekatt Apr 18, 2026
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
31 changes: 19 additions & 12 deletions vllm/model_executor/models/gemma2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
from vllm.logger import init_logger
from vllm.model_executor.layers.activation import GeluAndMul
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear,
QKVParallelLinear,
Expand Down Expand Up @@ -214,16 +214,15 @@ def __init__(
quant_config=quant_config,
prefix=f"{prefix}.mlp",
)
self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = GemmaRMSNorm(
config.hidden_size, eps=config.rms_norm_eps
)
self.pre_feedforward_layernorm = GemmaRMSNorm(
config.hidden_size, eps=config.rms_norm_eps
)
self.post_feedforward_layernorm = GemmaRMSNorm(
config.hidden_size, eps=config.rms_norm_eps
)
# GGUF stores RMSNorm weights with +1 baked in (llama.cpp convention).
# GemmaRMSNorm adds 1 in its forward pass, so use plain RMSNorm for GGUF.
quant_name = quant_config.get_name() if quant_config else None
rms_norm_cls = RMSNorm if quant_name == "gguf" else GemmaRMSNorm

@gemini-code-assist gemini-code-assist Bot Mar 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The logic to determine the normalization class based on the quantization config is duplicated in Gemma2DecoderLayer and Gemma2Model within this file, and also multiple times in vllm/model_executor/models/gemma3.py. This code duplication increases maintenance overhead and the risk of introducing inconsistencies in the future if the logic needs to be updated.

To improve this, you could extract the logic into a module-level helper function. For example:

from typing import Type
from torch import nn
from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
from vllm.model_executor.layers.quantization import QuantizationConfig

def _get_norm_cls(quant_config: QuantizationConfig | None) -> Type[nn.Module]:
    quant_name = quant_config.get_name() if quant_config else None
    return RMSNorm if quant_name == "gguf" else GemmaRMSNorm

This helper can then be called from all the locations where this logic is needed, making the code more DRY and easier to maintain.

Suggested change
quant_name = quant_config.get_name() if quant_config else None
rms_norm_cls = RMSNorm if quant_name == "gguf" else GemmaRMSNorm
rms_norm_cls = _get_norm_cls(quant_config)

@Isotr0py Isotr0py Mar 17, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you address this?

@kitaekatt kitaekatt Mar 25, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in c7a1a30 — extracted _get_norm_cls() helper in both gemma2.py and gemma3.py. All 5 occurrences now use it.

rms_norm_kwargs = dict(hidden_size=config.hidden_size, eps=config.rms_norm_eps)
self.input_layernorm = rms_norm_cls(**rms_norm_kwargs)
self.post_attention_layernorm = rms_norm_cls(**rms_norm_kwargs)
self.pre_feedforward_layernorm = rms_norm_cls(**rms_norm_kwargs)
self.post_feedforward_layernorm = rms_norm_cls(**rms_norm_kwargs)

def forward(
self,
Expand Down Expand Up @@ -263,6 +262,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
self.embed_tokens = VocabParallelEmbedding(
config.vocab_size,
config.hidden_size,
quant_config=quant_config,
prefix=f"{prefix}.embed_tokens",
)
self.start_layer, self.end_layer, self.layers = make_layers(
config.num_hidden_layers,
Expand All @@ -271,7 +272,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
),
prefix=f"{prefix}.layers",
)
self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
quant_name = quant_config.get_name() if quant_config else None
rms_norm_cls = RMSNorm if quant_name == "gguf" else GemmaRMSNorm
self.norm = rms_norm_cls(config.hidden_size, eps=config.rms_norm_eps)

# Normalize the embedding by sqrt(hidden_size)
# The normalizer's data type should be downcasted to the model's
Expand Down Expand Up @@ -361,6 +364,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
continue
if is_pp_missing_parameter(name, self):
continue
# Skip parameters not in the model (e.g., GGUF quantization
# metadata like qweight_type for embeddings)
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
Expand Down
40 changes: 17 additions & 23 deletions vllm/model_executor/models/gemma3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
Attention,
EncoderOnlyAttention,
)
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear,
QKVParallelLinear,
Expand Down Expand Up @@ -156,8 +156,12 @@ def __init__(
prefix=f"{prefix}.o_proj",
)

self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
# GGUF stores RMSNorm weights with +1 baked in (llama.cpp convention).
# GemmaRMSNorm adds 1 in its forward pass, so use plain RMSNorm for GGUF.
quant_name = quant_config.get_name() if quant_config else None
rms_norm_cls = RMSNorm if quant_name == "gguf" else GemmaRMSNorm
self.q_norm = rms_norm_cls(self.head_dim, eps=config.rms_norm_eps)
self.k_norm = rms_norm_cls(self.head_dim, eps=config.rms_norm_eps)

layer_idx = extract_layer_index(prefix)
layer_type = config.layer_types[layer_idx]
Expand Down Expand Up @@ -261,16 +265,13 @@ def __init__(
quant_config=quant_config,
prefix=f"{prefix}.mlp",
)
self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = GemmaRMSNorm(
config.hidden_size, eps=config.rms_norm_eps
)
self.pre_feedforward_layernorm = GemmaRMSNorm(
config.hidden_size, eps=config.rms_norm_eps
)
self.post_feedforward_layernorm = GemmaRMSNorm(
config.hidden_size, eps=config.rms_norm_eps
)
quant_name = quant_config.get_name() if quant_config else None
rms_norm_cls = RMSNorm if quant_name == "gguf" else GemmaRMSNorm
rms_norm_kwargs = dict(hidden_size=config.hidden_size, eps=config.rms_norm_eps)
self.input_layernorm = rms_norm_cls(**rms_norm_kwargs)
self.post_attention_layernorm = rms_norm_cls(**rms_norm_kwargs)
self.pre_feedforward_layernorm = rms_norm_cls(**rms_norm_kwargs)
self.post_feedforward_layernorm = rms_norm_cls(**rms_norm_kwargs)

def forward(
self,
Expand Down Expand Up @@ -322,7 +323,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
),
prefix=f"{prefix}.layers",
)
self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
quant_name = quant_config.get_name() if quant_config else None
rms_norm_cls = RMSNorm if quant_name == "gguf" else GemmaRMSNorm
self.norm = rms_norm_cls(config.hidden_size, eps=config.rms_norm_eps)

# Normalize the embedding by sqrt(hidden_size)
# The normalizer's data type should be downcasted to the model's
Expand Down Expand Up @@ -383,15 +386,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
for name, loaded_weight in weights:
# Revert +1 during llama.cpp conversion
# see: https://github.com/ggml-org/llama.cpp/blob/be7c3034108473beda214fd1d7c98fd6a7a3bdf5/convert_hf_to_gguf.py#L3397-L3400
if (
self.quant_config
and self.quant_config.get_name() == "gguf"
and name.endswith("norm.weight")
):
loaded_weight -= 1

if self.quant_config is not None and (
scale_name := self.quant_config.get_cache_scale(name)
):
Expand Down
26 changes: 26 additions & 0 deletions vllm/tokenizers/hf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,16 @@

from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast

from vllm.logger import init_logger
from vllm.transformers_utils.config import get_sentence_transformer_tokenizer_config
from vllm.transformers_utils.gguf_utils import extract_eos_token_id_from_gguf

from .protocol import TokenizerLike

HfTokenizer: TypeAlias = PreTrainedTokenizer | PreTrainedTokenizerFast

logger = init_logger(__name__)


def get_cached_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer:
"""
Expand Down Expand Up @@ -81,6 +85,9 @@ def from_pretrained(
download_dir: str | None = None,
**kwargs,
) -> HfTokenizer:
# Save gguf_file before AutoTokenizer.from_pretrained() pops it from kwargs
gguf_file = kwargs.get("gguf_file")

try:
tokenizer = AutoTokenizer.from_pretrained(
path_or_repo_id,
Expand Down Expand Up @@ -122,4 +129,23 @@ def from_pretrained(
}
tokenizer.add_special_tokens(special_tokens_map)

# Patch EOS token ID from GGUF metadata if available
# GGUF files may have a different EOS token ID than HF tokenizer config
# (e.g., Gemma uses <end_of_turn> ID 106 as EOS, but HF reports <eos> ID 1)
# Note: gguf_file was saved above before
# AutoTokenizer.from_pretrained() popped it
if gguf_file:
gguf_path = Path(path_or_repo_id) / gguf_file
gguf_eos_id = extract_eos_token_id_from_gguf(str(gguf_path))
if gguf_eos_id is not None:
hf_eos_id = tokenizer.eos_token_id
if hf_eos_id != gguf_eos_id:
logger.info(
"Patching tokenizer eos_token_id from %d to %d "
"(using GGUF metadata)",
hf_eos_id,
gguf_eos_id,
)
tokenizer.eos_token_id = gguf_eos_id

@Isotr0py Isotr0py Mar 27, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What if we provided a repo id instead of local path here?

BTW, I prefer to add an extra maybe_patch_gguf_tokenizer function at gguf_utils.py instead, because it's not very elegant to explicitly handle gguf file here:

def maybe_patch_gguf_tokenizer(
    tokenizer,
    path_or_repo_id: str,
    **kwargs,
):
    ...

maybe_patch_gguf_tokenizer(...)


return get_cached_tokenizer(tokenizer)
40 changes: 40 additions & 0 deletions vllm/transformers_utils/gguf_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -220,6 +220,46 @@ def extract_vision_config_from_gguf(mmproj_path: str) -> "SiglipVisionConfig | N
return config


def extract_eos_token_id_from_gguf(model: str) -> int | None:
"""Extract EOS token ID from GGUF metadata.

GGUF files store the EOS token ID in tokenizer.ggml.eos_token_id field.
This may differ from HuggingFace's tokenizer config (e.g., Gemma models
use <end_of_turn> token ID 106 as EOS in GGUF, but HF tokenizer reports
<eos> token ID 1).

Args:
model: Path to GGUF model file

Returns:
EOS token ID from GGUF metadata, or None if not found
"""
# Note: We don't check for .gguf extension here because HuggingFace Hub
# stores GGUF files as blob hashes without extensions. The caller is
# responsible for ensuring this is a valid GGUF file (via check_gguf_file).
try:
model_path = Path(model)
if not model_path.is_file():
return None

reader = gguf.GGUFReader(str(model_path))

eos_field = reader.get_field(Keys.Tokenizer.EOS_ID)
if eos_field is not None:
eos_token_id = int(eos_field.parts[-1][0])
logger.debug(
"Extracted eos_token_id=%d from GGUF metadata",
eos_token_id,
)
return eos_token_id

return None

except Exception as e:
logger.debug("Error extracting EOS token ID from GGUF: %s", e)
return None


def maybe_patch_hf_config_from_gguf(
model: str,
hf_config: PretrainedConfig,
Expand Down
Loading