chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Split the overlapped version of TpModelWorkerClient into a separate file (#1726)

Browse Source
This commit is contained in:
Lianmin Zheng
2024-10-20 00:29:29 -07:00
committed by GitHub
parent 593b19f29d
commit b48edff67f
7 changed files with 217 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files

109
python/sglang/srt/managers/tp_worker.py
View File

@@ -17,13 +17,8 @@ limitations under the License.
import json
import logging
import threading
import time
from queue import Queue
from typing import Optional
import torch
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.hf_transformers_utils import get_processor, get_tokenizer
from sglang.srt.managers.io_struct import UpdateWeightReqInput
@@ -108,9 +103,6 @@ class TpModelWorker:
)[0]
set_random_seed(self.random_seed)
if server_args.enable_overlap_schedule:
self.init_overlap_status()
def get_worker_info(self):
return (
self.max_total_num_tokens,
@@ -137,81 +129,6 @@ class TpModelWorker:
self.model_runner.token_to_kv_pool,
)
def init_overlap_status(self):
self.future_logits_output_dict = dict()
self.future_logits_output_ct = 0
self.future_token_ids_ct = 0
self.future_token_ids_map = torch.empty(
(self.max_running_requests * 5,), dtype=torch.int32, device=self.device
)
self.future_token_ids_limit = self.max_running_requests * 3
self.future_token_ids_output = dict()
self.future_event_map = dict()
self.forward_queue = Queue()
self.forward_stream = torch.cuda.Stream()
self.forward_thread = threading.Thread(
target=self.forward_thread_func,
)
self.forward_thread.start()
def forward_thread_func(self):
with torch.cuda.stream(self.forward_stream):
self.forward_thread_func_()
@torch.inference_mode()
def forward_thread_func_(self):
while True:
tic1 = time.time()
model_worker_batch, future_logits_output, future_next_token_ids = (
self.forward_queue.get()
)
# Resolve future tokens in the input
tic2 = time.time()
resolved_input_ids = model_worker_batch.input_ids
future_mask = resolved_input_ids < 0
resolved_input_ids[future_mask] = self.future_token_ids_map[
-resolved_input_ids[future_mask]
]
# Run forward
logits_output, next_token_ids = self.forward_batch_generation(
model_worker_batch
)
# Set future values
if model_worker_batch.return_logprob:
self.future_logits_output_dict[future_logits_output] = logits_output
# logger.info(f"set output {future_next_token_ids=}, {next_token_ids=}")
self.future_token_ids_map[-future_next_token_ids] = next_token_ids.to(
torch.int32
)
# logger.info("Set event")
self.future_token_ids_output[model_worker_batch.bid] = (
next_token_ids.tolist()
)
self.future_event_map[model_worker_batch.bid].set()
if False:
tic3 = time.time()
self.acc_time_with_waiting += tic3 - tic1
self.acc_time_without_waiting += tic3 - tic2
if self.forward_queue.qsize() == 0:
logger.info(
f"{self.acc_time_with_waiting=:.3f}, {self.acc_time_without_waiting=:.3f}, {self.forward_queue.qsize()=}"
)
def resolve_future_token_ids(self, bid: int):
self.future_event_map[bid].wait()
ret = self.future_token_ids_output[bid]
del self.future_event_map[bid]
return ret
def resolve_future_logits_output(self, future_obj):
return self.future_logits_output_dict.pop(future_obj)
def forward_batch_generation(self, model_worker_batch: ModelWorkerBatch):
forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
logits_output = self.model_runner.forward(forward_batch)
@@ -224,32 +141,6 @@ class TpModelWorker:
embeddings = logits_output.embeddings
return embeddings
def forward_batch_generation_non_blocking(
self, model_worker_batch: ModelWorkerBatch
):
# Allocate output future objects
future_logits_output = self.future_logits_output_ct
self.future_logits_output_ct += 1
bs = len(model_worker_batch.seq_lens)
with torch.cuda.stream(self.forward_stream):
future_next_token_ids = -torch.arange(
self.future_token_ids_ct + 1,
self.future_token_ids_ct + 1 + bs,
dtype=torch.int32,
device=self.device,
)
self.future_token_ids_ct = (
self.future_token_ids_ct + bs
) % self.future_token_ids_limit
ret = future_logits_output, future_next_token_ids
self.future_event_map[model_worker_batch.bid] = threading.Event()
self.forward_queue.put(
(model_worker_batch.copy(), future_logits_output, future_next_token_ids)
)
return ret
def update_weights(self, recv_req: UpdateWeightReqInput):
success, message = self.model_runner.update_weights(
recv_req.model_path, recv_req.load_format