From ac406d4301377d6c7212b5936e5ef1b7a0b2b915 Mon Sep 17 00:00:00 2001 From: AlphaBaby Date: Mon, 17 Nov 2025 19:54:23 +0800 Subject: [PATCH] fix uneven PP layer indices (#13282) Co-authored-by: Xuchun Shang --- python/sglang/srt/distributed/utils.py | 21 ++++++++++++++------- python/sglang/srt/models/deepseek_v2.py | 21 ++++++++++++++------- 2 files changed, 28 insertions(+), 14 deletions(-) diff --git a/python/sglang/srt/distributed/utils.py b/python/sglang/srt/distributed/utils.py index bfe54b9d4..519fedca0 100644 --- a/python/sglang/srt/distributed/utils.py +++ b/python/sglang/srt/distributed/utils.py @@ -65,7 +65,7 @@ def get_pp_indices( ) -> Tuple[int, int]: """Try to evenly distribute layers across partitions. If the number of layers is not divisible by the number of partitions, - the last partition will have the remaining layers. + the first N partitions will have one extra layer, where N = remainder. """ # partition_list_str can be set to None in sglang partition_list_str = os.getenv("SGLANG_PP_LAYER_PARTITION", None) @@ -83,12 +83,19 @@ def get_pp_indices( start_layer = sum(partitions[:pp_rank]) end_layer = start_layer + partitions[pp_rank] else: - layers_per_partition = num_hidden_layers // pp_size - start_layer = pp_rank * layers_per_partition - end_layer = start_layer + layers_per_partition - - if pp_rank == pp_size - 1: - end_layer = num_hidden_layers + base_layers = num_hidden_layers // pp_size + remainder = num_hidden_layers % pp_size + # Distribute the extra layers to the first 'remainder' partitions + if pp_rank < remainder: + # This partition gets one extra layer + start_layer = pp_rank * (base_layers + 1) + end_layer = start_layer + (base_layers + 1) + else: + # This partition gets only base layers + start_layer = ( + remainder * (base_layers + 1) + (pp_rank - remainder) * base_layers + ) + end_layer = start_layer + base_layers return (start_layer, end_layer) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 3193a1a07..54f6caa02 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -3254,13 +3254,20 @@ class DeepseekV2ForCausalLM(nn.Module): self.model = DeepseekV2Model( config, quant_config, prefix=add_prefix("model", prefix) ) - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=add_prefix("lm_head", prefix), - use_attn_tp_group=get_global_server_args().enable_dp_lm_head, - ) + if self.pp_group.is_last_rank: + if self.pp_group.world_size == 1 and config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + use_attn_tp_group=get_global_server_args().enable_dp_lm_head, + ) + else: + # ranks other than the last rank will have a placeholder layer + self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor(config) self._routed_experts_weights_of_layer = LazyValue(