[PP] Add pipeline parallelism (#5724)
This commit is contained in:
Ying Sheng
@@ -16,6 +16,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import bisect
|
||||
import inspect
|
||||
import os
|
||||
from contextlib import contextmanager
|
||||
from typing import TYPE_CHECKING, Callable
|
||||
@@ -33,12 +34,14 @@ from sglang.srt.model_executor.forward_batch_info import (
|
||||
CaptureHiddenMode,
|
||||
ForwardBatch,
|
||||
ForwardMode,
|
||||
PPProxyTensors,
|
||||
)
|
||||
from sglang.srt.patch_torch import monkey_patch_torch_compile
|
||||
from sglang.srt.utils import (
|
||||
get_available_gpu_memory,
|
||||
get_device_memory_capacity,
|
||||
is_hip,
|
||||
rank0_log,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -188,10 +191,11 @@ class CudaGraphRunner:
|
||||
self.speculative_algorithm = model_runner.server_args.speculative_algorithm
|
||||
self.tp_size = model_runner.server_args.tp_size
|
||||
self.dp_size = model_runner.server_args.dp_size
|
||||
self.pp_size = model_runner.server_args.pp_size
|
||||
|
||||
# Batch sizes to capture
|
||||
self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(model_runner)
|
||||
|
||||
rank0_log(f"Capture cuda graph bs {self.capture_bs}")
|
||||
self.capture_forward_mode = ForwardMode.DECODE
|
||||
self.capture_hidden_mode = CaptureHiddenMode.NULL
|
||||
self.num_tokens_per_bs = 1
|
||||
@@ -234,6 +238,19 @@ class CudaGraphRunner:
|
||||
self.positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
|
||||
self.mrope_positions = torch.zeros((3, self.max_bs), dtype=torch.int64)
|
||||
|
||||
# pipeline parallelism
|
||||
if self.pp_size > 1:
|
||||
self.pp_proxy_tensors = {
|
||||
"hidden_states": torch.zeros(
|
||||
(self.max_bs, self.model_runner.model_config.hidden_size),
|
||||
dtype=torch.bfloat16,
|
||||
),
|
||||
"residual": torch.zeros(
|
||||
(self.max_bs, self.model_runner.model_config.hidden_size),
|
||||
dtype=torch.bfloat16,
|
||||
),
|
||||
}
|
||||
|
||||
# Speculative_inference
|
||||
if (
|
||||
model_runner.spec_algorithm.is_eagle3()
|
||||
@@ -384,6 +401,12 @@ class CudaGraphRunner:
|
||||
encoder_lens = None
|
||||
mrope_positions = self.mrope_positions[:, :bs]
|
||||
|
||||
# pipeline parallelism
|
||||
if self.pp_size > 1:
|
||||
pp_proxy_tensors = PPProxyTensors(
|
||||
{k: v[:num_tokens] for k, v in self.pp_proxy_tensors.items()}
|
||||
)
|
||||
|
||||
if self.enable_dp_attention or self.enable_sp_layernorm:
|
||||
self.global_num_tokens_gpu.copy_(
|
||||
torch.tensor(
|
||||
@@ -456,8 +479,20 @@ class CudaGraphRunner:
|
||||
# Clean intermediate result cache for DP attention
|
||||
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
|
||||
|
||||
logits_output = forward(input_ids, forward_batch.positions, forward_batch)
|
||||
return logits_output.next_token_logits, logits_output.hidden_states
|
||||
kwargs = {}
|
||||
if (
|
||||
self.pp_size > 1
|
||||
and "pp_proxy_tensors" in inspect.signature(forward).parameters
|
||||
):
|
||||
kwargs["pp_proxy_tensors"] = pp_proxy_tensors
|
||||
|
||||
logits_output_or_pp_proxy_tensors = forward(
|
||||
input_ids,
|
||||
forward_batch.positions,
|
||||
forward_batch,
|
||||
**kwargs,
|
||||
)
|
||||
return logits_output_or_pp_proxy_tensors
|
||||
|
||||
for _ in range(2):
|
||||
torch.cuda.synchronize()
|
||||
@@ -490,7 +525,11 @@ class CudaGraphRunner:
|
||||
self.capture_hidden_mode = hidden_mode_from_spec_info
|
||||
self.capture()
|
||||
|
||||
def replay_prepare(self, forward_batch: ForwardBatch):
|
||||
def replay_prepare(
|
||||
self,
|
||||
forward_batch: ForwardBatch,
|
||||
pp_proxy_tensors: Optional[PPProxyTensors] = None,
|
||||
):
|
||||
self.recapture_if_needed(forward_batch)
|
||||
|
||||
raw_bs = forward_batch.batch_size
|
||||
@@ -519,6 +558,11 @@ class CudaGraphRunner:
|
||||
self.seq_lens_cpu.fill_(1)
|
||||
self.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
|
||||
|
||||
if pp_proxy_tensors:
|
||||
for key in self.pp_proxy_tensors.keys():
|
||||
dim = pp_proxy_tensors[key].shape[0]
|
||||
self.pp_proxy_tensors[key][:dim].copy_(pp_proxy_tensors[key])
|
||||
|
||||
if self.is_encoder_decoder:
|
||||
self.encoder_lens[:raw_bs].copy_(forward_batch.encoder_lens)
|
||||
if forward_batch.mrope_positions is not None:
|
||||
@@ -547,10 +591,13 @@ class CudaGraphRunner:
|
||||
self.bs = bs
|
||||
|
||||
def replay(
|
||||
self, forward_batch: ForwardBatch, skip_attn_backend_init: bool = False
|
||||
) -> LogitsProcessorOutput:
|
||||
self,
|
||||
forward_batch: ForwardBatch,
|
||||
skip_attn_backend_init: bool = False,
|
||||
pp_proxy_tensors: Optional[PPProxyTensors] = None,
|
||||
) -> Union[LogitsProcessorOutput, PPProxyTensors]:
|
||||
if not skip_attn_backend_init:
|
||||
self.replay_prepare(forward_batch)
|
||||
self.replay_prepare(forward_batch, pp_proxy_tensors)
|
||||
else:
|
||||
# In speculative decoding, these two fields are still needed.
|
||||
self.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids)
|
||||
@@ -558,17 +605,19 @@ class CudaGraphRunner:
|
||||
|
||||
# Replay
|
||||
self.graphs[self.bs].replay()
|
||||
next_token_logits, hidden_states = self.output_buffers[self.bs]
|
||||
|
||||
logits_output = LogitsProcessorOutput(
|
||||
next_token_logits=next_token_logits[: self.raw_num_token],
|
||||
hidden_states=(
|
||||
hidden_states[: self.raw_num_token]
|
||||
if hidden_states is not None
|
||||
else None
|
||||
),
|
||||
)
|
||||
return logits_output
|
||||
output = self.output_buffers[self.bs]
|
||||
if isinstance(output, LogitsProcessorOutput):
|
||||
return LogitsProcessorOutput(
|
||||
next_token_logits=output.next_token_logits[: self.raw_num_token],
|
||||
hidden_states=(
|
||||
output.hidden_states[: self.raw_num_token]
|
||||
if output.hidden_states is not None
|
||||
else None
|
||||
),
|
||||
)
|
||||
else:
|
||||
assert isinstance(output, PPProxyTensors)
|
||||
return PPProxyTensors({k: v[: self.bs] for k, v in output.tensors.items()})
|
||||
|
||||
def get_spec_info(self, num_tokens: int):
|
||||
spec_info = None
|
||||
|
||||
Reference in New Issue
Block a user