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 11, 00:00 UTC
[Bugfix] Fix Jais2ForCausalLM (#31198)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Dec 23, 2025 at 15:44 UTC
6b16fff01b671840420fef304a62a77279d56533
1 file changed
+4
-25
vllm/model_executor/models/jais2.py
+4
-25
@@ -48,7 +48,6 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
48
from vllm.model_executor.layers.quantization import QuantizationConfig
49
from vllm.model_executor.layers.rotary_embedding import get_rope
50
from vllm.model_executor.layers.vocab_parallel_embedding import (
51
- DEFAULT_VOCAB_PADDING_SIZE,
51
ParallelLMHead,
52
VocabParallelEmbedding,
53
)
@@ -167,7 +166,6 @@ class Jais2Attention(nn.Module):
166
167
self.rotary_emb = get_rope(
168
self.head_dim,
170
- rotary_dim=self.head_dim,
169
max_position=max_position_embeddings,
170
rope_parameters=getattr(config, "rope_parameters", None),
171
is_neox_style=is_neox_style,
@@ -304,17 +302,12 @@ class Jais2Model(nn.Module):
302
303
config = vllm_config.model_config.hf_config
304
quant_config = vllm_config.quant_config
307
- lora_config = vllm_config.lora_config
305
306
self.config = config
307
self.quant_config = quant_config
308
self.padding_idx = config.pad_token_id
312
- lora_vocab = (
313
- (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1))
314
- if lora_config
315
- else 0
316
- )
317
- self.vocab_size = config.vocab_size + lora_vocab
309
+
310
+ self.vocab_size = config.vocab_size
311
self.org_vocab_size = config.vocab_size
312
if get_pp_group().is_first_rank or (
313
config.tie_word_embeddings and get_pp_group().is_last_rank
@@ -456,29 +449,15 @@ class Jais2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
449
super().__init__()
450
config = vllm_config.model_config.hf_config
451
quant_config = vllm_config.quant_config
459
- lora_config = vllm_config.lora_config
452
self.config = config
461
- self.lora_config = lora_config
462
-
453
self.model = self._init_model(
454
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
455
)
456
457
if get_pp_group().is_last_rank:
468
- self.unpadded_vocab_size = config.vocab_size
469
- if lora_config:
470
- self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
458
self.lm_head = ParallelLMHead(
472
- self.unpadded_vocab_size,
459
+ config.vocab_size,
460
config.hidden_size,
474
- org_num_embeddings=config.vocab_size,
475
- padding_size=(
476
- DEFAULT_VOCAB_PADDING_SIZE
477
- # We need bigger padding if using lora for kernel
478
- # compatibility
479
- if not lora_config
480
- else lora_config.lora_vocab_padding_size
481
- ),
461
quant_config=quant_config,
462
prefix=maybe_prefix(prefix, "lm_head"),
463
)
@@ -487,7 +466,7 @@ class Jais2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
466
467
logit_scale = getattr(config, "logit_scale", 1.0)
468
self.logits_processor = LogitsProcessor(
490
- self.unpadded_vocab_size, config.vocab_size, logit_scale
469
+ config.vocab_size, scale=logit_scale
470
)
471
else:
472
self.lm_head = PPMissingLayer()