diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index d1498f1..745d1d3 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -26,7 +26,7 @@ import torch from torch import nn from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, VllmConfig +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config from vllm.config.utils import getattr_iter from vllm.distributed import ( get_pp_group, @@ -491,9 +491,12 @@ class Gemma4Attention(nn.Module): f"{kv_shared_layer_index}.self_attn.attn" ) + # Gemma 4's rope types (default, proportional) do not depend on max_position, so + # size the cos/sin caches for max_model_len: 262144 positions take ~0.4 GiB. + max_model_len = get_current_vllm_config().model_config.max_model_len self.rotary_emb = get_rope( self.head_dim, - max_position=max_position_embeddings, + max_position=min(max_position_embeddings, max_model_len), rope_parameters=rope_parameters, is_neox_style=True, ) @@ -1444,14 +1447,17 @@ class Gemma4ForCausalLM( prefix=maybe_prefix(prefix, "model"), ) - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) if config.tie_word_embeddings: - self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) + # Share embed_tokens: a ParallelLMHead tied after construction would hold the + # vocab matrix (1.4 GiB) twice while loading, which overflows 8 GB GPUs. + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) self.logits_processor = LogitsProcessor( config.vocab_size,