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, 22:15 UTC
DeepSeek V2/V3/R1 only place `lm_head` on last pp rank (#13833)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Feb 26, 2025 at 01:24 UTC
24679788ed389f63106a9733ae0c87559bd11604
1 file changed
+6
-3
vllm/model_executor/models/deepseek_v2.py
+6
-3
@@ -636,9 +636,12 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP):
636
self.quant_config = quant_config
637
self.model = DeepseekV2Model(vllm_config=vllm_config,
638
prefix=maybe_prefix(prefix, "model"))
639
- self.lm_head = ParallelLMHead(config.vocab_size,
640
- config.hidden_size,
641
- quant_config=quant_config)
639
+ if get_pp_group().is_last_rank:
640
+ self.lm_head = ParallelLMHead(config.vocab_size,
641
+ config.hidden_size,
642
+ quant_config=quant_config)
643
+ else:
644
+ self.lm_head = PPMissingLayer()
645
self.logits_processor = LogitsProcessor(config.vocab_size)
646
self.sampler = get_sampler()
647
self.make_empty_intermediate_tensors = (