fix uneven PP layer indices (#13282)

Co-authored-by: Xuchun Shang <xuchun.shang@linux.alibaba.com>

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
+        base_layers = num_hidden_layers // pp_size
-        start_layer = pp_rank * layers_per_partition
+        remainder = num_hidden_layers % pp_size
-        end_layer = start_layer + layers_per_partition
+        # Distribute the extra layers to the first 'remainder' partitions
-
+        if pp_rank < remainder:
-        if pp_rank == pp_size - 1:
+            # This partition gets one extra layer
-            end_layer = num_hidden_layers
+            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)
 

python/sglang/srt/models/deepseek_v2.py

@@ -3254,6 +3254,10 @@ class DeepseekV2ForCausalLM(nn.Module):
         self.model = DeepseekV2Model(
             config, quant_config, prefix=add_prefix("model", prefix)
         )
+        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,
@@ -3261,6 +3265,9 @@ class DeepseekV2ForCausalLM(nn.Module):
                     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(