server_args.py 8.01 KB
Newer Older
Lianmin Zheng's avatar
Lianmin Zheng committed
1
2
import argparse
import dataclasses
3
from typing import List, Optional, Union
Lianmin Zheng's avatar
Lianmin Zheng committed
4
5
6
7
8
9
10
11


@dataclasses.dataclass
class ServerArgs:
    model_path: str
    tokenizer_path: Optional[str] = None
    host: str = "127.0.0.1"
    port: int = 30000
12
    additional_ports: Optional[Union[List[int], int]] = None
Lianmin Zheng's avatar
Lianmin Zheng committed
13
14
    load_format: str = "auto"
    tokenizer_mode: str = "auto"
Cody Yu's avatar
Cody Yu committed
15
    chat_template: Optional[str] = None
Lianmin Zheng's avatar
Lianmin Zheng committed
16
    trust_remote_code: bool = True
Lianmin Zheng's avatar
Lianmin Zheng committed
17
    mem_fraction_static: Optional[float] = None
18
    max_prefill_num_token: Optional[int] = None
19
    context_length: Optional[int] = None
Lianmin Zheng's avatar
Lianmin Zheng committed
20
21
22
    tp_size: int = 1
    model_mode: List[str] = ()
    schedule_heuristic: str = "lpm"
23
    schedule_conservativeness: float = 1.0
Lianmin Zheng's avatar
Lianmin Zheng committed
24
    random_seed: int = 42
25
    stream_interval: int = 8
Lianmin Zheng's avatar
Lianmin Zheng committed
26
27
28
    disable_log_stats: bool = False
    log_stats_interval: int = 10
    log_level: str = "info"
29
30
    disable_regex_jump_forward: bool = False
    disable_disk_cache: bool = False
Liangsheng Yin's avatar
Liangsheng Yin committed
31
    attention_reduce_in_fp32: bool = False
Lianmin Zheng's avatar
Lianmin Zheng committed
32
33
34
35

    def __post_init__(self):
        if self.tokenizer_path is None:
            self.tokenizer_path = self.model_path
Lianmin Zheng's avatar
Lianmin Zheng committed
36
        if self.mem_fraction_static is None:
Lianmin Zheng's avatar
Lianmin Zheng committed
37
38
39
40
41
42
            if self.tp_size >= 8:
                self.mem_fraction_static = 0.80
            elif self.tp_size >= 4:
                self.mem_fraction_static = 0.82
            elif self.tp_size >= 2:
                self.mem_fraction_static = 0.85
Lianmin Zheng's avatar
Lianmin Zheng committed
43
            else:
Lianmin Zheng's avatar
Lianmin Zheng committed
44
                self.mem_fraction_static = 0.90
45
46
47
48
        if isinstance(self.additional_ports, int):
            self.additional_ports = [self.additional_ports]
        elif self.additional_ports is None:
            self.additional_ports = []
Lianmin Zheng's avatar
Lianmin Zheng committed
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

    @staticmethod
    def add_cli_args(parser: argparse.ArgumentParser):
        parser.add_argument(
            "--model-path",
            type=str,
            help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.",
            required=True,
        )
        parser.add_argument(
            "--tokenizer-path",
            type=str,
            default=ServerArgs.tokenizer_path,
            help="The path of the tokenizer.",
        )
        parser.add_argument("--host", type=str, default=ServerArgs.host)
        parser.add_argument("--port", type=int, default=ServerArgs.port)
66
67
68
69
70
71
72
73
        # we want to be able to pass a list of ports
        parser.add_argument(
            "--additional-ports",
            type=int,
            nargs="*",
            default=[],
            help="Additional ports specified for launching server.",
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
        parser.add_argument(
            "--load-format",
            type=str,
            default=ServerArgs.load_format,
            choices=["auto", "pt", "safetensors", "npcache", "dummy"],
            help="The format of the model weights to load. "
            '"auto" will try to load the weights in the safetensors format '
            "and fall back to the pytorch bin format if safetensors format "
            "is not available. "
            '"pt" will load the weights in the pytorch bin format. '
            '"safetensors" will load the weights in the safetensors format. '
            '"npcache" will load the weights in pytorch format and store '
            "a numpy cache to speed up the loading. "
            '"dummy" will initialize the weights with random values, '
            "which is mainly for profiling.",
        )
        parser.add_argument(
            "--tokenizer-mode",
            type=str,
            default=ServerArgs.tokenizer_mode,
            choices=["auto", "slow"],
            help="Tokenizer mode. 'auto' will use the fast "
            "tokenizer if available, and 'slow' will "
            "always use the slow tokenizer.",
        )
Cody Yu's avatar
Cody Yu committed
99
100
101
102
103
104
        parser.add_argument(
            "--chat-template",
            type=str,
            default=ServerArgs.chat_template,
            help="The buliltin chat template name or the path of the chat template file. This is only used for OpenAI-compatible API server",
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
105
106
107
108
109
110
111
112
113
        parser.add_argument(
            "--trust-remote-code",
            action="store_true",
            help="Whether or not to allow for custom models defined on the Hub in their own modeling files.",
        )
        parser.add_argument(
            "--mem-fraction-static",
            type=float,
            default=ServerArgs.mem_fraction_static,
114
            help="The fraction of the memory used for static allocation (model weights and KV cache memory pool). Use a smaller value if you see out-of-memory errors.",
Lianmin Zheng's avatar
Lianmin Zheng committed
115
        )
116
117
118
119
        parser.add_argument(
            "--max-prefill-num-token",
            type=int,
            default=ServerArgs.max_prefill_num_token,
120
            help="The maximum number of tokens in a prefill batch. The real bound will be the maximum of this value and the model's maximum context length.",
121
        )
122
123
124
125
126
127
        parser.add_argument(
            "--context-length",
            type=int,
            default=ServerArgs.context_length,
            help="The model's maximum context length. Use this to reduce the context length to save memory. Defaults to None (will use the value from the model's config.json instead).",
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
128
129
130
131
132
133
134
135
136
137
138
        parser.add_argument(
            "--tp-size",
            type=int,
            default=ServerArgs.tp_size,
            help="Tensor parallelism degree.",
        )
        parser.add_argument(
            "--model-mode",
            type=str,
            default=[],
            nargs="+",
Lianmin Zheng's avatar
Lianmin Zheng committed
139
140
            choices=["flashinfer", "no-cache"],
            help="Model mode: [flashinfer, no-cache]",
Lianmin Zheng's avatar
Lianmin Zheng committed
141
142
143
144
145
146
147
        )
        parser.add_argument(
            "--schedule-heuristic",
            type=str,
            default=ServerArgs.schedule_heuristic,
            help="Schudule mode: [lpm, weight, random, fcfs]",
        )
148
149
150
151
        parser.add_argument(
            "--schedule-conservativeness",
            type=float,
            default=ServerArgs.schedule_conservativeness,
152
            help="How conservative the schedule policy is. A larger value means more conservative scheduling. Use a larger value if you see requests being retracted frequently.",
153
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
154
155
156
157
158
159
        parser.add_argument(
            "--random-seed",
            type=int,
            default=ServerArgs.random_seed,
            help="Random seed.",
        )
160
161
162
        parser.add_argument(
            "--stream-interval",
            type=int,
Lianmin Zheng's avatar
Lianmin Zheng committed
163
            default=ServerArgs.stream_interval,
164
            help="The interval (or buffer size) for streaming in terms of the token length. A smaller value makes streaming smoother, while a larger value makes the throughput higher",
165
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
        parser.add_argument(
            "--log-level",
            type=str,
            default=ServerArgs.log_level,
            help="Log level",
        )
        parser.add_argument(
            "--disable-log-stats",
            action="store_true",
            help="Disable logging throughput stats.",
        )
        parser.add_argument(
            "--log-stats-interval",
            type=int,
            default=ServerArgs.log_stats_interval,
            help="Log stats interval in second.",
        )
Liangsheng Yin's avatar
Liangsheng Yin committed
183
        parser.add_argument(
184
            "--disable-regex-jump-forward",
Liangsheng Yin's avatar
Liangsheng Yin committed
185
            action="store_true",
Liangsheng Yin's avatar
Liangsheng Yin committed
186
            help="Disable regex jump-forward",
Liangsheng Yin's avatar
Liangsheng Yin committed
187
        )
188
189
190
191
192
        parser.add_argument(
            "--disable-disk-cache",
            action="store_true",
            help="Disable disk cache to avoid possible crashes related to file system or high concurrency.",
        )
Liangsheng Yin's avatar
Liangsheng Yin committed
193
194
195
196
197
        parser.add_argument(
            "--attention-reduce-in-fp32",
            action="store_true",
            help="Cast the intermidiate attention results to fp32 to avoid possible crashes related to fp16.",
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214

    @classmethod
    def from_cli_args(cls, args: argparse.Namespace):
        attrs = [attr.name for attr in dataclasses.fields(cls)]
        return cls(**{attr: getattr(args, attr) for attr in attrs})

    def url(self):
        return f"http://{self.host}:{self.port}"


@dataclasses.dataclass
class PortArgs:
    tokenizer_port: int
    router_port: int
    detokenizer_port: int
    nccl_port: int
    model_rpc_ports: List[int]