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

[Fix] Fix logprob and normalized_logprob (#1428)

Browse Source
This commit is contained in:
Lianmin Zheng
2024-09-15 06:36:06 -07:00
committed by GitHub
parent 282681b8a1
commit 9ba1f09760
22 changed files with 314 additions and 215 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/io_struct.py
+2 -1
View File
@@ -20,7 +20,7 @@ processes (TokenizerManager, DetokenizerManager, Controller).
import copy
import uuid
from dataclasses import dataclass, field
from dataclasses import dataclass
from typing import Dict, List, Optional, Union
from sglang.srt.managers.schedule_batch import BaseFinishReason
@@ -43,6 +43,7 @@ class GenerateReqInput:
# Whether to return logprobs.
return_logprob: Optional[Union[List[bool], bool]] = None
# If return logprobs, the start location in the prompt for returning logprobs.
# By default, this value is "-1", which means it will only return logprobs for output tokens.
logprob_start_len: Optional[Union[List[int], int]] = None
# If return logprobs, the number of top logprobs to return at each position.
top_logprobs_num: Optional[Union[List[int], int]] = None
python/sglang/srt/managers/schedule_batch.py
+61 -33
View File
@@ -19,7 +19,7 @@ limitations under the License.
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING, List, Optional, Union
from typing import List, Optional, Tuple, Union
import torch
@@ -53,7 +53,7 @@ class BaseFinishReason:
self.is_error = is_error
def to_json(self):
raise NotImplementedError("Subclasses must implement this method")
raise NotImplementedError()
class FINISH_MATCHED_TOKEN(BaseFinishReason):
@@ -105,7 +105,13 @@ class FINISH_ABORT(BaseFinishReason):
class Req:
"""Store all inforamtion of a request."""
def __init__(self, rid, origin_input_text, origin_input_ids, lora_path=None):
def __init__(
self,
rid: str,
origin_input_text: str,
origin_input_ids: Tuple[int],
lora_path: Optional[str] = None,
):
# Input and output info
self.rid = rid
self.origin_input_text = origin_input_text
@@ -118,6 +124,10 @@ class Req:
# Memory info
self.req_pool_idx = None
# Check finish
self.tokenizer = None
self.finished_reason = None
# For incremental decoding
# ----- | --------- read_ids -------|
# ----- | surr_ids |
@@ -136,7 +146,7 @@ class Req:
# this does not include the jump forward tokens.
self.completion_tokens_wo_jump_forward = 0
# For vision input
# For vision inputs
self.pixel_values = None
self.image_sizes = None
self.image_offsets = None
@@ -144,31 +154,35 @@ class Req:
self.modalities = None
# Prefix info
self.extend_input_len = 0
self.prefix_indices = []
self.extend_input_len = 0
self.last_node = None
# Sampling parameters
self.sampling_params = None
self.stream = False
# Check finish
self.tokenizer = None
self.finished_reason = None
# Logprobs
# Logprobs (arguments)
self.return_logprob = False
self.embedding = None
self.logprob_start_len = 0
self.top_logprobs_num = 0
# Logprobs (return value)
self.normalized_prompt_logprob = None
self.input_token_logprobs = None
self.input_top_logprobs = None
self.output_token_logprobs = []
self.output_top_logprobs = []
# Logprobs (internal values)
# The tokens is prefilled but need to be considered as decode tokens
# and should be updated for the decode logprobs
self.last_update_decode_tokens = 0
# The relative logprob_start_len in an extend batch
self.extend_logprob_start_len = 0
# Embedding
self.embedding = None
# Constrained decoding
self.regex_fsm: RegexGuide = None
@@ -363,9 +377,13 @@ class ScheduleBatch:
return_logprob: bool = False
top_logprobs_nums: List[int] = None
# Stream
has_stream: bool = False
@classmethod
def init_new(cls, reqs, req_to_token_pool, token_to_kv_pool, tree_cache):
return_logprob = any(req.return_logprob for req in reqs)
has_stream = any(req.stream for req in reqs)
return cls(
reqs=reqs,
@@ -373,18 +391,15 @@ class ScheduleBatch:
token_to_kv_pool=token_to_kv_pool,
tree_cache=tree_cache,
return_logprob=return_logprob,
has_stream=has_stream,
)
def batch_size(self):
return len(self.reqs) if self.reqs else 0
return len(self.reqs)
def is_empty(self):
return len(self.reqs) == 0
def has_stream(self) -> bool:
# Return whether batch has at least 1 streaming request
return any(r.stream for r in self.reqs)
def alloc_req_slots(self, num_reqs):
req_pool_indices = self.req_to_token_pool.alloc(num_reqs)
if req_pool_indices is None:
@@ -427,8 +442,8 @@ class ScheduleBatch:
for i, req in enumerate(reqs):
req.req_pool_idx = req_pool_indices_cpu[i]
pre_len, seq_len = len(req.prefix_indices), len(req.fill_ids)
ext_len = seq_len - pre_len
seq_lens.append(seq_len)
assert seq_len - pre_len == req.extend_input_len
if pre_len > 0:
self.req_to_token_pool.req_to_token[req.req_pool_idx][
@@ -436,9 +451,19 @@ class ScheduleBatch:
] = req.prefix_indices
self.req_to_token_pool.req_to_token[req.req_pool_idx][pre_len:seq_len] = (
out_cache_loc[pt : pt + ext_len]
out_cache_loc[pt : pt + req.extend_input_len]
)
pt += ext_len
# Compute the relative logprob_start_len in an extend batch
if req.logprob_start_len >= pre_len:
extend_logprob_start_len = min(
req.logprob_start_len - pre_len, req.extend_input_len - 1
)
else:
extend_logprob_start_len = req.extend_input_len - 1
req.extend_logprob_start_len = extend_logprob_start_len
pt += req.extend_input_len
# Set fields
with torch.device("cuda"):
@@ -451,21 +476,13 @@ class ScheduleBatch:
self.out_cache_loc = out_cache_loc
self.top_logprobs_nums = [r.top_logprobs_num for r in reqs]
self.prefix_lens_cpu = [len(r.prefix_indices) for r in reqs]
self.extend_lens_cpu = [r.extend_input_len for r in reqs]
self.extend_logprob_start_lens_cpu = [r.extend_logprob_start_len for r in reqs]
self.sampling_info = SamplingBatchInfo.from_schedule_batch(self, vocab_size)
def mix_with_running(self, running_batch: "ScheduleBatch"):
self.forward_mode = ForwardMode.MIXED
self.running_bs = running_batch.batch_size()
# NOTE: prefix_indices is what has been cached, but we don't cache each decode step
prefix_lens_cpu = [len(r.prefix_indices) for r in self.reqs]
prefix_lens_cpu.extend(
[
len(r.origin_input_ids) + len(r.output_ids) - 1
for r in running_batch.reqs
]
)
running_bs = running_batch.batch_size()
for req in running_batch.reqs:
req.fill_ids = req.origin_input_ids + req.output_ids
@@ -473,12 +490,22 @@ class ScheduleBatch:
input_ids = torch.cat([self.input_ids, running_batch.input_ids])
out_cache_loc = torch.cat([self.out_cache_loc, running_batch.out_cache_loc])
extend_num_tokens = self.extend_num_tokens + running_batch.batch_size()
extend_num_tokens = self.extend_num_tokens + running_bs
self.merge(running_batch)
self.input_ids = input_ids
self.out_cache_loc = out_cache_loc
self.extend_num_tokens = extend_num_tokens
self.prefix_lens_cpu = prefix_lens_cpu
# NOTE: prefix_indices is what has been cached, but we don't cache each decode step
self.prefix_lens_cpu.extend(
[
len(r.origin_input_ids) + len(r.output_ids) - 1
for r in running_batch.reqs
]
)
self.extend_lens_cpu.extend([1] * running_bs)
self.extend_logprob_start_lens_cpu.extend([0] * running_bs)
def check_decode_mem(self):
bs = self.batch_size()
@@ -685,6 +712,7 @@ class ScheduleBatch:
self.out_cache_loc = None
self.top_logprobs_nums = [self.top_logprobs_nums[i] for i in unfinished_indices]
self.return_logprob = any(req.return_logprob for req in self.reqs)
self.has_stream = any(req.stream for req in self.reqs)
self.sampling_info.filter(unfinished_indices, new_indices)
@@ -695,7 +723,6 @@ class ScheduleBatch:
self.sampling_info.merge(other.sampling_info)
self.reqs.extend(other.reqs)
self.req_pool_indices = torch.concat(
[self.req_pool_indices, other.req_pool_indices]
)
@@ -706,3 +733,4 @@ class ScheduleBatch:
self.out_cache_loc = None
self.top_logprobs_nums.extend(other.top_logprobs_nums)
self.return_logprob = any(req.return_logprob for req in self.reqs)
self.has_stream = any(req.stream for req in self.reqs)
python/sglang/srt/managers/tokenizer_manager.py
-6
View File
@@ -197,8 +197,6 @@ class TokenizerManager:
if not_use_index
else obj.logprob_start_len[index]
)
if return_logprob and logprob_start_len == -1:
logprob_start_len = len(input_ids) - 1
top_logprobs_num = (
obj.top_logprobs_num
if not_use_index
@@ -251,8 +249,6 @@ class TokenizerManager:
# Send to the controller
if self.is_generation:
if return_logprob and logprob_start_len == -1:
logprob_start_len = len(input_ids) - 1
tokenized_obj = TokenizedGenerateReqInput(
rid,
input_text,
@@ -349,8 +345,6 @@ class TokenizerManager:
sampling_params = self._get_sampling_params(obj.sampling_params[index])
if self.is_generation:
if obj.return_logprob[index] and obj.logprob_start_len[index] == -1:
obj.logprob_start_len[index] = len(input_ids) - 1
pixel_values, image_hashes, image_sizes = (
await self._get_pixel_values(obj.image_data[index])
)
python/sglang/srt/managers/tp_worker.py
+44 -24
View File
@@ -278,7 +278,7 @@ class ModelTpServer:
self.running_batch = None
break
if self.out_pyobjs and self.running_batch.has_stream():
if self.out_pyobjs and self.running_batch.has_stream:
break
else:
self.check_memory()
@@ -360,9 +360,13 @@ class ModelTpServer:
# Only when pixel values is not None we have modalities
req.modalities = recv_req.modalites
req.return_logprob = recv_req.return_logprob
req.logprob_start_len = recv_req.logprob_start_len
req.top_logprobs_num = recv_req.top_logprobs_num
req.stream = recv_req.stream
req.logprob_start_len = recv_req.logprob_start_len
if req.logprob_start_len == -1:
# By default, only return the logprobs for output tokens
req.logprob_start_len = len(recv_req.input_ids) - 1
# Init regex FSM
if (
@@ -384,7 +388,7 @@ class ModelTpServer:
# Truncate prompts that are too long
if len(req.origin_input_ids) >= self.max_req_input_len:
logger.warn(
logger.warning(
"Request length is longer than the KV cache pool size or "
"the max context length. Truncated!!!"
)
@@ -583,7 +587,7 @@ class ModelTpServer:
next_token_ids = [self.tokenizer.eos_token_id] * len(batch.reqs)
# Check finish conditions
pt = 0
logprob_pt = 0
for i, req in enumerate(batch.reqs):
if req is not self.current_inflight_req:
# Inflight reqs' prefill is not finished
@@ -607,10 +611,9 @@ class ModelTpServer:
self.req_to_token_pool.free(req.req_pool_idx)
if req.return_logprob:
self.add_logprob_return_values(
i, req, pt, next_token_ids, logits_output
logprob_pt += self.add_logprob_return_values(
i, req, logprob_pt, next_token_ids, logits_output
)
pt += req.extend_input_len
else:
assert batch.extend_num_tokens != 0
logits_output = self.model_runner.forward(batch)
@@ -638,48 +641,63 @@ class ModelTpServer:
def add_logprob_return_values(
self,
i,
i: int,
req: Req,
pt: int,
next_token_ids: List[int],
output: LogitsProcessorOutput,
):
"""Attach logprobs to the return values."""
req.output_token_logprobs.append(
(output.next_token_logprobs[i], next_token_ids[i])
)
# If logprob_start_len > 0, then first logprob_start_len prompt tokens will be ignored.
num_input_logprobs = req.extend_input_len - req.extend_logprob_start_len
if req.normalized_prompt_logprob is None:
req.normalized_prompt_logprob = output.normalized_prompt_logprobs[i]
if req.input_token_logprobs is None:
# If logprob_start_len > 0, then first logprob_start_len prompt tokens will be ignored.
req.input_token_logprobs = list(
zip(
output.input_token_logprobs[pt : pt + req.extend_input_len - 1],
req.fill_ids[-req.extend_input_len + 1 :],
)
)
if req.logprob_start_len == 0:
input_token_logprobs = output.input_token_logprobs[
pt : pt + num_input_logprobs - 1 - req.last_update_decode_tokens
]
input_token_ids = req.fill_ids[
len(req.fill_ids)
- num_input_logprobs
+ 1 : len(req.fill_ids)
- req.last_update_decode_tokens
]
req.input_token_logprobs = list(zip(input_token_logprobs, input_token_ids))
if (
req.logprob_start_len == 0
): # The first token does not have logprob, pad it.
req.input_token_logprobs = [
(None, req.fill_ids[0])
] + req.input_token_logprobs
if req.last_update_decode_tokens != 0:
# Some decode tokens are re-computed in an extend batch
req.output_token_logprobs.extend(
list(
zip(
output.input_token_logprobs[
pt
+ req.extend_input_len
+ num_input_logprobs
- 1
- req.last_update_decode_tokens : pt
+ req.extend_input_len
+ num_input_logprobs
- 1
],
req.fill_ids[-req.last_update_decode_tokens + 1 :],
req.fill_ids[
len(req.fill_ids)
- req.last_update_decode_tokens : len(req.fill_ids)
],
)
)
)
req.output_token_logprobs.append(
(output.next_token_logprobs[i], next_token_ids[i])
)
if req.top_logprobs_num > 0:
if req.input_top_logprobs is None:
req.input_top_logprobs = output.input_top_logprobs[i]
@@ -688,10 +706,12 @@ class ModelTpServer:
if req.last_update_decode_tokens != 0:
req.output_top_logprobs.extend(
output.input_top_logprobs[i][-req.last_update_decode_tokens + 1 :]
output.input_top_logprobs[i][-req.last_update_decode_tokens :]
)
req.output_top_logprobs.append(output.output_top_logprobs[i])
return num_input_logprobs
def forward_decode_batch(self, batch: ScheduleBatch):
# Check if decode out of memory
if not batch.check_decode_mem():