A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 21:00 UTC
Fix `head_dim` not existing in all model configs (Transformers backend) (#14141)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Mar 3, 2025 at 17:48 UTC
91373a0d15f3d47a45ab5145467940a1664900fc
1 file changed
+15
-12
vllm/model_executor/models/transformers.py
+15
-12
@@ -25,7 +25,6 @@ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
25
from vllm.attention import Attention
26
from vllm.config import VllmConfig
27
from vllm.distributed import get_tensor_model_parallel_world_size
28
-from vllm.distributed.utils import divide
28
from vllm.logger import init_logger
29
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
30
ReplicatedLinear,
@@ -128,10 +127,12 @@ class TransformersModel(nn.Module, SupportsQuant, SupportsLoRA):
127
128
config = vllm_config.model_config.hf_config
129
cache_config = vllm_config.cache_config
130
+ model_config = vllm_config.model_config
131
+ parallel_config = vllm_config.parallel_config
132
133
self.config = config
133
- self.vocab_size = config.vocab_size
134
- self.unpadded_vocab_size = config.vocab_size
134
+ self.vocab_size = model_config.get_vocab_size()
135
+ self.unpadded_vocab_size = model_config.get_vocab_size()
136
137
self.model: PreTrainedModel = AutoModel.from_config(
138
self.config,
@@ -145,15 +146,17 @@ class TransformersModel(nn.Module, SupportsQuant, SupportsLoRA):
146
self.apply_base_model_tp_plan(self.model)
147
148
# Attention modifications (assumes 1 attention op per hidden layer)
148
- tp_size = get_tensor_model_parallel_world_size()
149
+ num_heads = model_config.get_num_attention_heads(parallel_config)
150
+ head_size = model_config.get_head_size()
151
+ num_kv_heads = model_config.get_num_kv_heads(parallel_config)
152
self.attention_instances = [
153
Attention(
151
- num_heads=divide(config.num_attention_heads, tp_size),
152
- head_size=config.head_dim,
154
+ num_heads=num_heads,
155
+ head_size=head_size,
156
# NOTE: We use Llama scale as default, if it's set by
157
# Transformers, it's updated in vllm_flash_attention_forward
155
- scale=config.head_dim**-0.5,
156
- num_kv_heads=divide(config.num_key_value_heads, tp_size),
158
+ scale=head_size**-0.5,
159
+ num_kv_heads=num_kv_heads,
160
cache_config=cache_config,
161
quant_config=self.quant_config,
162
prefix=f"{i}.attn") for i in range(config.num_hidden_layers)
@@ -163,7 +166,7 @@ class TransformersModel(nn.Module, SupportsQuant, SupportsLoRA):
166
self.replace_vocab_embed_class(self.model)
167
168
# ForCausalLM modifications
166
- self.lm_head = ParallelLMHead(config.vocab_size,
169
+ self.lm_head = ParallelLMHead(self.vocab_size,
170
config.hidden_size,
171
quant_config=self.quant_config,
172
prefix=maybe_prefix(prefix, "lm_head"))
@@ -172,7 +175,7 @@ class TransformersModel(nn.Module, SupportsQuant, SupportsLoRA):
175
176
logit_scale = getattr(config, "logit_scale", 1.0)
177
self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
175
- config.vocab_size, logit_scale)
178
+ self.vocab_size, logit_scale)
179
self.sampler = get_sampler()
180
181
def apply_base_model_tp_plan(self, module: nn.Module, prefix: str = ""):
@@ -203,12 +206,12 @@ class TransformersModel(nn.Module, SupportsQuant, SupportsLoRA):
206
new_module = VocabParallelEmbedding(
207
self.vocab_size,
208
self.config.hidden_size,
206
- org_num_embeddings=self.config.vocab_size,
209
+ org_num_embeddings=self.vocab_size,
210
quant_config=None,
211
)
212
log_replacement("input embedding", self.model.get_input_embeddings(),
213
new_module)
211
- self.model.set_input_embeddings(new_module)
214
+ module.set_input_embeddings(new_module)
215
216
def forward(
217
self,