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

[Minor] Many cleanup (#1357)

Browse Source
This commit is contained in:
Lianmin Zheng
2024-09-09 04:14:11 -07:00
committed by GitHub
parent c9b75917d5
commit e4d68afcf0
24 changed files with 416 additions and 296 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/constrained/fsm_cache.py
+29 -38
View File
@@ -16,6 +16,7 @@ limitations under the License.
"""Cache for the compressed finite state machine."""
from outlines.fsm.json_schema import build_regex_from_schema
from transformers import AutoTokenizer
from sglang.srt.constrained import RegexGuide, TransformerTokenizer
from sglang.srt.constrained.base_tool_cache import BaseToolCache
@@ -28,12 +29,9 @@ class FSMCache(BaseToolCache):
tokenizer_args_dict,
enable=True,
skip_tokenizer_init=False,
json_schema_mode=False,
):
super().__init__(enable=enable)
self.json_schema_mode = json_schema_mode
if (
skip_tokenizer_init
or tokenizer_path.endswith(".json")
@@ -42,44 +40,37 @@ class FSMCache(BaseToolCache):
# Do not support TiktokenTokenizer or SentencePieceTokenizer
return
from importlib.metadata import version
tokenizer_args_dict.setdefault("padding_side", "left")
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, **tokenizer_args_dict)
try:
self.outlines_tokenizer = TransformerTokenizer(tokenizer)
except AttributeError:
# FIXME: tmp fix for chatglm2 & chatglm3 (pad_token_id=0)
origin_pad_token_id = tokenizer.pad_token_id
if version("outlines") >= "0.0.35":
from transformers import AutoTokenizer
def fset(self, value):
self._value = value
tokenizer_args_dict.setdefault("padding_side", "left")
tokenizer = AutoTokenizer.from_pretrained(
tokenizer_path, **tokenizer_args_dict
type(tokenizer).pad_token_id = property(
fget=type(tokenizer).pad_token_id.fget, fset=fset
)
try:
self.outlines_tokenizer = TransformerTokenizer(tokenizer)
except AttributeError:
# FIXME: tmp fix for chatglm2 & chatglm3 (pad_token_id=0)
origin_pad_token_id = tokenizer.pad_token_id
def fset(self, value):
self._value = value
type(tokenizer).pad_token_id = property(
fget=type(tokenizer).pad_token_id.fget, fset=fset
)
self.outlines_tokenizer = TransformerTokenizer(tokenizer)
self.outlines_tokenizer.tokenizer.pad_token_id = origin_pad_token_id
self.outlines_tokenizer.pad_token_id = origin_pad_token_id
self.outlines_tokenizer.pad_token = (
self.outlines_tokenizer.tokenizer.pad_token
)
self.outlines_tokenizer.vocabulary = (
self.outlines_tokenizer.tokenizer.get_vocab()
)
else:
self.outlines_tokenizer = TransformerTokenizer(
tokenizer_path, **tokenizer_args_dict
self.outlines_tokenizer = TransformerTokenizer(tokenizer)
self.outlines_tokenizer.tokenizer.pad_token_id = origin_pad_token_id
self.outlines_tokenizer.pad_token_id = origin_pad_token_id
self.outlines_tokenizer.pad_token = (
self.outlines_tokenizer.tokenizer.pad_token
)
self.outlines_tokenizer.vocabulary = (
self.outlines_tokenizer.tokenizer.get_vocab()
)
def init_value(self, value):
if self.json_schema_mode:
regex = build_regex_from_schema(value, whitespace_pattern=r"[\n\t ]*")
return RegexGuide(regex, self.outlines_tokenizer), regex
def init_value(self, key):
key_type, key_string = key
if key_type == "json":
regex = build_regex_from_schema(key_string, whitespace_pattern=r"[\n\t ]*")
elif key_type == "regex":
regex = key_string
else:
return RegexGuide(value, self.outlines_tokenizer)
raise ValueError(f"Invalid key_type: {key_type}")
return RegexGuide(regex, self.outlines_tokenizer), regex
python/sglang/srt/managers/controller_multi.py
+1 -5
View File
@@ -71,12 +71,10 @@ class ControllerMulti:
self,
server_args: ServerArgs,
port_args: PortArgs,
model_override_args,
):
# Parse args
self.server_args = server_args
self.port_args = port_args
self.model_override_args = model_override_args
self.load_balance_method = LoadBalanceMethod.from_str(
server_args.load_balance_method
)
@@ -114,7 +112,6 @@ class ControllerMulti:
self.server_args,
self.port_args,
pipe_controller_writer,
self.model_override_args,
True,
gpu_ids,
dp_worker_id,
@@ -189,14 +186,13 @@ def start_controller_process(
server_args: ServerArgs,
port_args: PortArgs,
pipe_writer,
model_override_args: dict,
):
"""Start a controller process."""
configure_logger(server_args)
try:
controller = ControllerMulti(server_args, port_args, model_override_args)
controller = ControllerMulti(server_args, port_args)
except Exception:
pipe_writer.send(get_exception_traceback())
raise
python/sglang/srt/managers/controller_single.py
-5
View File
@@ -40,7 +40,6 @@ class ControllerSingle:
self,
server_args: ServerArgs,
port_args: PortArgs,
model_override_args: dict,
gpu_ids: List[int],
is_data_parallel_worker: bool,
dp_worker_id: int,
@@ -76,7 +75,6 @@ class ControllerSingle:
tp_rank_range,
server_args,
port_args.nccl_ports[dp_worker_id],
model_override_args,
)
# Launch tp rank 0
@@ -85,7 +83,6 @@ class ControllerSingle:
0,
server_args,
port_args.nccl_ports[dp_worker_id],
model_override_args,
)
self.tp_cpu_group = self.tp_server.model_runner.tp_group.cpu_group
@@ -126,7 +123,6 @@ def start_controller_process(
server_args: ServerArgs,
port_args: PortArgs,
pipe_writer: multiprocessing.connection.Connection,
model_override_args: dict,
is_data_parallel_worker: bool = False,
gpu_ids: List[int] = None,
dp_worker_id: int = None,
@@ -149,7 +145,6 @@ def start_controller_process(
controller = ControllerSingle(
server_args,
port_args,
model_override_args,
gpu_ids,
is_data_parallel_worker,
dp_worker_id,
python/sglang/srt/managers/tokenizer_manager.py
+2 -2
View File
@@ -18,6 +18,7 @@ limitations under the License.
import asyncio
import concurrent.futures
import dataclasses
import json
import logging
import multiprocessing as mp
import os
@@ -77,7 +78,6 @@ class TokenizerManager:
self,
server_args: ServerArgs,
port_args: PortArgs,
model_override_args: dict = None,
):
self.server_args = server_args
@@ -95,7 +95,7 @@ class TokenizerManager:
self.hf_config = get_config(
self.model_path,
trust_remote_code=server_args.trust_remote_code,
model_override_args=model_override_args,
model_override_args=json.loads(server_args.json_model_override_args),
)
self.is_generation = is_generation_model(
self.hf_config.architectures, self.server_args.is_embedding
python/sglang/srt/managers/tp_worker.py
+79 -76
View File
@@ -15,13 +15,14 @@ limitations under the License.
"""A tensor parallel worker."""
import json
import logging
import multiprocessing
import os
import pickle
import time
import warnings
from typing import Any, List, Optional, Union
from typing import Any, List, Optional
import torch
import torch.distributed
@@ -66,6 +67,7 @@ from sglang.utils import get_exception_traceback
logger = logging.getLogger(__name__)
# Crash on warning if we are running CI tests
crash_on_warning = os.getenv("SGLANG_IS_IN_CI", "false") == "true"
@@ -76,11 +78,10 @@ class ModelTpServer:
tp_rank: int,
server_args: ServerArgs,
nccl_port: int,
model_override_args: dict,
):
suppress_other_loggers()
# Copy arguments
# Parse arguments
self.gpu_id = gpu_id
self.tp_rank = tp_rank
self.tp_size = server_args.tp_size
@@ -93,9 +94,8 @@ class ModelTpServer:
server_args.model_path,
server_args.trust_remote_code,
context_length=server_args.context_length,
model_override_args=model_override_args,
model_override_args=json.loads(server_args.json_model_override_args),
)
self.model_runner = ModelRunner(
model_config=self.model_config,
mem_fraction_static=server_args.mem_fraction_static,
@@ -136,7 +136,7 @@ class ModelTpServer:
self.max_total_num_tokens - 1,
)
# Sync random seed
# Sync random seed across TP workers
server_args.random_seed = broadcast_recv_input(
[server_args.random_seed],
self.tp_rank,
@@ -144,7 +144,7 @@ class ModelTpServer:
)[0]
set_random_seed(server_args.random_seed)
# Print info
# Print debug info
logger.info(
f"max_total_num_tokens={self.max_total_num_tokens}, "
f"max_prefill_tokens={self.max_prefill_tokens}, "
@@ -181,7 +181,7 @@ class ModelTpServer:
self.num_generated_tokens = 0
self.last_stats_tic = time.time()
# Chunked prefill
# Init chunked prefill
self.chunked_prefill_size = server_args.chunked_prefill_size
self.current_inflight_req = None
self.is_mixed_chunk = (
@@ -197,16 +197,6 @@ class ModelTpServer:
"trust_remote_code": server_args.trust_remote_code,
},
skip_tokenizer_init=server_args.skip_tokenizer_init,
json_schema_mode=False,
)
self.json_fsm_cache = FSMCache(
server_args.tokenizer_path,
{
"tokenizer_mode": server_args.tokenizer_mode,
"trust_remote_code": server_args.trust_remote_code,
},
skip_tokenizer_init=server_args.skip_tokenizer_init,
json_schema_mode=True,
)
self.jump_forward_cache = JumpForwardCache()
@@ -227,11 +217,12 @@ class ModelTpServer:
try:
# Recv requests
for recv_req in recv_reqs:
if isinstance(
recv_req, (TokenizedGenerateReqInput, TokenizedEmbeddingReqInput)
):
if isinstance(recv_req, TokenizedGenerateReqInput):
self.handle_generate_request(recv_req)
self.do_not_get_new_batch = False
elif isinstance(recv_req, TokenizedEmbeddingReqInput):
self.handle_embedding_request(recv_req)
self.do_not_get_new_batch = False
elif isinstance(recv_req, FlushCacheReq):
self.flush_cache()
elif isinstance(recv_req, AbortReq):
@@ -331,57 +322,56 @@ class ModelTpServer:
def handle_generate_request(
self,
recv_req: Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput],
recv_req: TokenizedGenerateReqInput,
):
req = Req(recv_req.rid, recv_req.input_text, recv_req.input_ids)
req.tokenizer = self.tokenizer
req.sampling_params = recv_req.sampling_params
if self.model_runner.is_generation:
req.pixel_values = recv_req.pixel_values
if req.pixel_values is not None:
# Use image hash as fake token_ids, which is then used
# for prefix matching
image_hash = hash(tuple(recv_req.image_hashes))
req.pad_value = [
(image_hash) % self.model_config.vocab_size,
(image_hash >> 16) % self.model_config.vocab_size,
(image_hash >> 32) % self.model_config.vocab_size,
(image_hash >> 64) % self.model_config.vocab_size,
]
req.image_sizes = recv_req.image_sizes
(
req.origin_input_ids,
req.image_offsets,
) = self.model_runner.model.pad_input_ids(
req.origin_input_ids_unpadded,
req.pad_value,
req.pixel_values,
req.image_sizes,
)
# 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.pixel_values = recv_req.pixel_values
if req.pixel_values is not None:
# Use image hash as fake token_ids, which is then used
# for prefix matching
image_hash = hash(tuple(recv_req.image_hashes))
req.pad_value = [
(image_hash) % self.model_config.vocab_size,
(image_hash >> 16) % self.model_config.vocab_size,
(image_hash >> 32) % self.model_config.vocab_size,
(image_hash >> 64) % self.model_config.vocab_size,
]
req.image_sizes = recv_req.image_sizes
(
req.origin_input_ids,
req.image_offsets,
) = self.model_runner.model.pad_input_ids(
req.origin_input_ids_unpadded,
req.pad_value,
req.pixel_values,
req.image_sizes,
)
# 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
# Init regex fsm fron json
# Init regex FSM
if (
req.sampling_params.json_schema is not None
or req.sampling_params.regex is not None
):
if req.sampling_params.json_schema is not None:
req.regex_fsm, computed_regex_string = self.json_fsm_cache.query(
req.sampling_params.json_schema
req.regex_fsm, computed_regex_string = self.regex_fsm_cache.query(
("json", req.sampling_params.json_schema)
)
if not self.disable_regex_jump_forward:
req.jump_forward_map = self.jump_forward_cache.query(
computed_regex_string
)
# Init regex fsm
elif req.sampling_params.regex is not None:
req.regex_fsm = self.regex_fsm_cache.query(req.sampling_params.regex)
if not self.disable_regex_jump_forward:
req.jump_forward_map = self.jump_forward_cache.query(
req.sampling_params.regex
)
req.regex_fsm, computed_regex_string = self.regex_fsm_cache.query(
("regex", req.sampling_params.regex)
)
if not self.disable_regex_jump_forward:
req.jump_forward_map = self.jump_forward_cache.query(
computed_regex_string
)
# Truncate prompts that are too long
if len(req.origin_input_ids) >= self.max_req_input_len:
@@ -390,16 +380,32 @@ class ModelTpServer:
"the max context length. Truncated!!!"
)
req.origin_input_ids = req.origin_input_ids[: self.max_req_input_len]
req.sampling_params.max_new_tokens = min(
(
req.sampling_params.max_new_tokens
if req.sampling_params.max_new_tokens is not None
else 1 << 30
),
self.max_req_input_len - 1 - len(req.origin_input_ids),
)
if self.model_runner.is_generation:
req.sampling_params.max_new_tokens = min(
(
req.sampling_params.max_new_tokens
if req.sampling_params.max_new_tokens is not None
else 1 << 30
),
self.max_req_input_len - 1 - len(req.origin_input_ids),
self.waiting_queue.append(req)
def handle_embedding_request(
self,
recv_req: TokenizedEmbeddingReqInput,
):
req = Req(recv_req.rid, recv_req.input_text, recv_req.input_ids)
req.tokenizer = self.tokenizer
req.sampling_params = recv_req.sampling_params
# Truncate prompts that are too long
if len(req.origin_input_ids) >= self.max_req_input_len:
logger.warn(
"Request length is longer than the KV cache pool size or "
"the max context length. Truncated!!!"
)
req.origin_input_ids = req.origin_input_ids[: self.max_req_input_len]
self.waiting_queue.append(req)
@@ -892,7 +898,6 @@ def run_tp_server(
tp_rank: int,
server_args: ServerArgs,
nccl_port: int,
model_override_args: dict,
):
"""Run a tensor parallel model server."""
configure_logger(server_args, prefix=f" TP{tp_rank}")
@@ -903,7 +908,6 @@ def run_tp_server(
tp_rank,
server_args,
nccl_port,
model_override_args,
)
tp_cpu_group = model_server.model_runner.tp_group.cpu_group
@@ -920,14 +924,13 @@ def launch_tp_servers(
tp_rank_range: List[int],
server_args: ServerArgs,
nccl_port: int,
model_override_args: dict,
):
"""Launch multiple tensor parallel servers."""
procs = []
for i in tp_rank_range:
proc = multiprocessing.Process(
target=run_tp_server,
args=(gpu_ids[i], i, server_args, nccl_port, model_override_args),
args=(gpu_ids[i], i, server_args, nccl_port),
)
proc.start()
procs.append(proc)
python/sglang/srt/model_executor/model_runner.py
+1
View File
@@ -18,6 +18,7 @@ limitations under the License.
import gc
import importlib
import importlib.resources
import json
import logging
import pkgutil
from functools import lru_cache
python/sglang/srt/server.py
+3 -6
View File
@@ -272,7 +272,6 @@ async def retrieve_file_content(file_id: str):
def launch_server(
server_args: ServerArgs,
model_override_args: Optional[dict] = None,
pipe_finish_writer: Optional[mp.connection.Connection] = None,
):
"""Launch an HTTP server."""
@@ -317,7 +316,6 @@ def launch_server(
tp_rank_range,
server_args,
ports[3],
model_override_args,
)
try:
@@ -328,7 +326,7 @@ def launch_server(
return
# Launch processes
tokenizer_manager = TokenizerManager(server_args, port_args, model_override_args)
tokenizer_manager = TokenizerManager(server_args, port_args)
if server_args.chat_template:
load_chat_template_for_openai_api(tokenizer_manager, server_args.chat_template)
pipe_controller_reader, pipe_controller_writer = mp.Pipe(duplex=False)
@@ -341,7 +339,7 @@ def launch_server(
proc_controller = mp.Process(
target=start_controller_process,
args=(server_args, port_args, pipe_controller_writer, model_override_args),
args=(server_args, port_args, pipe_controller_writer),
)
proc_controller.start()
@@ -501,7 +499,6 @@ class Runtime:
def __init__(
self,
log_level: str = "error",
model_override_args: Optional[dict] = None,
*args,
**kwargs,
):
@@ -525,7 +522,7 @@ class Runtime:
proc = mp.Process(
target=launch_server,
args=(self.server_args, model_override_args, pipe_writer),
args=(self.server_args, pipe_writer),
)
proc.start()
pipe_writer.close()
python/sglang/srt/server_args.py
+18 -22
View File
@@ -76,6 +76,14 @@ class ServerArgs:
dp_size: int = 1
load_balance_method: str = "round_robin"
# Distributed args
nccl_init_addr: Optional[str] = None
nnodes: int = 1
node_rank: Optional[int] = None
# Model override args in JSON
json_model_override_args: str = "{}"
# Optimization/debug options
disable_flashinfer: bool = False
disable_flashinfer_sampling: bool = False
@@ -91,14 +99,6 @@ class ServerArgs:
enable_mla: bool = False
triton_attention_reduce_in_fp32: bool = False
# Distributed args
nccl_init_addr: Optional[str] = None
nnodes: int = 1
node_rank: Optional[int] = None
# Model override args in JSON
json_model_override_args: Optional[dict] = None
def __post_init__(self):
if self.tokenizer_path is None:
self.tokenizer_path = self.model_path
@@ -385,6 +385,14 @@ class ServerArgs:
)
parser.add_argument("--node-rank", type=int, help="The node rank.")
# Model override args
parser.add_argument(
"--json-model-override-args",
type=str,
help="A dictionary in JSON string format used to override default model configurations.",
default=ServerArgs.json_model_override_args,
)
# Optimization/debug options
parser.add_argument(
"--disable-flashinfer",
@@ -459,22 +467,10 @@ class ServerArgs:
help="Turn on memory efficient weight loading with quantization (quantize per layer during loading).",
)
# Model override args
parser.add_argument(
"--json-model-override-args",
type=str,
help="A dictionary in JSON string format used to override default model configurations.",
)
@classmethod
def from_cli_args(cls, args: argparse.Namespace):
args.tp_size = args.tensor_parallel_size
args.dp_size = args.data_parallel_size
args.json_model_override_args = (
json.loads(args.json_model_override_args)
if args.json_model_override_args
else None
)
attrs = [attr.name for attr in dataclasses.fields(cls)]
return cls(**{attr: getattr(args, attr) for attr in attrs})
@@ -498,7 +494,7 @@ class ServerArgs:
self.disable_flashinfer = False
def prepare_server_args(args: argparse.Namespace) -> ServerArgs:
def prepare_server_args(argv: List[str]) -> ServerArgs:
"""
Prepare the server arguments from the command line arguments.
@@ -511,7 +507,7 @@ def prepare_server_args(args: argparse.Namespace) -> ServerArgs:
"""
parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(parser)
raw_args = parser.parse_args(args)
raw_args = parser.parse_args(argv)
server_args = ServerArgs.from_cli_args(raw_args)
return server_args