prain: tune searches the auto CPU threshold and writes --cpu-threads; vLLM patch: fit Gemma 4 in 8 GB (tied lm_head without a copy, rope caches sized by max_model_len)

2026-10-27 23:52:53 +01:00
parent d69fd34492
commit 01dce2ffc6
4 changed files with 75 additions and 3 deletions
@@ -0,0 +1,52 @@
+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,