[PP] Add pipeline parallelism (#5724)

python/sglang/srt/managers/tp_worker.py

@@ -15,11 +15,12 @@
 
 import logging
 import threading
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Union
 
 import torch
 
 from sglang.srt.configs.model_config import ModelConfig
+from sglang.srt.distributed import get_pp_group, get_tp_group, get_world_group
 from sglang.srt.hf_transformers_utils import get_processor, get_tokenizer
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.managers.io_struct import (
@@ -31,7 +32,7 @@ from sglang.srt.managers.io_struct import (
 )
 from sglang.srt.managers.schedule_batch import ModelWorkerBatch, global_server_args_dict
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool, TokenToKVPoolAllocator
-from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import MultiprocessingSerializer, broadcast_pyobj, set_random_seed
@@ -47,6 +48,7 @@ class TpModelWorker:
         server_args: ServerArgs,
         gpu_id: int,
         tp_rank: int,
+        pp_rank: int,
         dp_rank: Optional[int],
         nccl_port: int,
         is_draft_worker: bool = False,
@@ -54,7 +56,9 @@ class TpModelWorker:
         token_to_kv_pool_allocator: Optional[TokenToKVPoolAllocator] = None,
     ):
         # Parse args
+        self.tp_size = server_args.tp_size
         self.tp_rank = tp_rank
+        self.pp_rank = pp_rank
 
         # Init model and tokenizer
         self.model_config = ModelConfig(
@@ -73,12 +77,15 @@ class TpModelWorker:
             quantization=server_args.quantization,
             is_draft_model=is_draft_worker,
         )
+
         self.model_runner = ModelRunner(
             model_config=self.model_config,
             mem_fraction_static=server_args.mem_fraction_static,
             gpu_id=gpu_id,
             tp_rank=tp_rank,
             tp_size=server_args.tp_size,
+            pp_rank=pp_rank,
+            pp_size=server_args.pp_size,
             nccl_port=nccl_port,
             server_args=server_args,
             is_draft_worker=is_draft_worker,
@@ -105,6 +112,10 @@ class TpModelWorker:
                 )
         self.device = self.model_runner.device
 
+        # Init nccl groups
+        self.pp_group = get_pp_group()
+        self.world_group = get_world_group()
+
         # Profile number of tokens
         self.max_total_num_tokens = self.model_runner.max_total_num_tokens
         self.max_prefill_tokens = server_args.max_prefill_tokens
@@ -130,8 +141,9 @@ class TpModelWorker:
         # Sync random seed across TP workers
         self.random_seed = broadcast_pyobj(
             [server_args.random_seed],
-            self.tp_rank,
-            self.model_runner.tp_group.cpu_group,
+            self.tp_size * self.pp_rank + tp_rank,
+            self.world_group.cpu_group,
+            src=self.world_group.ranks[0],
         )[0]
         set_random_seed(self.random_seed)
 
@@ -156,11 +168,14 @@ class TpModelWorker:
     def get_pad_input_ids_func(self):
         return getattr(self.model_runner.model, "pad_input_ids", None)
 
-    def get_tp_cpu_group(self):
-        return self.model_runner.tp_group.cpu_group
+    def get_tp_group(self):
+        return self.model_runner.tp_group
+
+    def get_attention_tp_group(self):
+        return self.model_runner.attention_tp_group
 
     def get_attention_tp_cpu_group(self):
-        return self.model_runner.attention_tp_group.cpu_group
+        return getattr(self.model_runner.attention_tp_group, "cpu_group", None)
 
     def get_memory_pool(self):
         return (
@@ -172,19 +187,38 @@ class TpModelWorker:
         self,
         model_worker_batch: ModelWorkerBatch,
         skip_sample: bool = False,
-    ) -> Tuple[LogitsProcessorOutput, Optional[torch.Tensor]]:
+    ) -> Tuple[Union[LogitsProcessorOutput, torch.Tensor], Optional[torch.Tensor]]:
         forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
-        logits_output = self.model_runner.forward(forward_batch)
 
-        if model_worker_batch.launch_done is not None:
-            model_worker_batch.launch_done.set()
+        pp_proxy_tensors = None
+        if not self.pp_group.is_first_rank:
+            pp_proxy_tensors = PPProxyTensors(
+                self.pp_group.recv_tensor_dict(
+                    all_gather_group=self.get_attention_tp_group()
+                )
+            )
 
-        if skip_sample:
-            next_token_ids = None
+        if self.pp_group.is_last_rank:
+            logits_output = self.model_runner.forward(
+                forward_batch, pp_proxy_tensors=pp_proxy_tensors
+            )
+            if model_worker_batch.launch_done is not None:
+                model_worker_batch.launch_done.set()
+
+            if skip_sample:
+                next_token_ids = None
+            else:
+                next_token_ids = self.model_runner.sample(
+                    logits_output, model_worker_batch
+                )
+
+            return logits_output, next_token_ids
         else:
-            next_token_ids = self.model_runner.sample(logits_output, model_worker_batch)
-
-        return logits_output, next_token_ids
+            pp_proxy_tensors = self.model_runner.forward(
+                forward_batch,
+                pp_proxy_tensors=pp_proxy_tensors,
+            )
+            return pp_proxy_tensors.tensors, None
 
     def forward_batch_embedding(self, model_worker_batch: ModelWorkerBatch):
         forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)