53 lines
2.2 KiB
Diff
53 lines
2.2 KiB
Diff
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,
|