arg_utils.py 8.63 KB
Newer Older
1
import argparse
2
3
4
import dataclasses
from dataclasses import dataclass
from typing import Optional, Tuple
5

Woosuk Kwon's avatar
Woosuk Kwon committed
6
7
from vllm.config import (CacheConfig, ModelConfig, ParallelConfig,
                         SchedulerConfig)
8
9


10
@dataclass
Zhuohan Li's avatar
Zhuohan Li committed
11
class EngineArgs:
Woosuk Kwon's avatar
Woosuk Kwon committed
12
    """Arguments for vLLM engine."""
13
    model: str
14
    tokenizer: Optional[str] = None
15
    tokenizer_mode: str = 'auto'
16
    trust_remote_code: bool = False
17
    download_dir: Optional[str] = None
18
    load_format: str = 'auto'
19
    dtype: str = 'auto'
20
    seed: int = 0
21
    worker_use_ray: bool = False
22
23
24
25
    pipeline_parallel_size: int = 1
    tensor_parallel_size: int = 1
    block_size: int = 16
    swap_space: int = 4  # GiB
26
    gpu_memory_utilization: float = 0.90
27
28
29
    max_num_batched_tokens: int = 2560
    max_num_seqs: int = 256
    disable_log_stats: bool = False
30

31
    def __post_init__(self):
32
33
        if self.tokenizer is None:
            self.tokenizer = self.model
34
35
36
37
        self.max_num_seqs = min(self.max_num_seqs, self.max_num_batched_tokens)

    @staticmethod
    def add_cli_args(
38
            parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
Woosuk Kwon's avatar
Woosuk Kwon committed
39
        """Shared CLI arguments for vLLM engine."""
40
        # Model arguments
41
42
43
44
45
46
47
48
49
50
51
52
        parser.add_argument(
            '--model',
            type=str,
            default='facebook/opt-125m',
            help='name or path of the huggingface model to use')
        parser.add_argument(
            '--tokenizer',
            type=str,
            default=EngineArgs.tokenizer,
            help='name or path of the huggingface tokenizer to use')
        parser.add_argument('--tokenizer-mode',
                            type=str,
53
54
55
                            default=EngineArgs.tokenizer_mode,
                            choices=['auto', 'slow'],
                            help='tokenizer mode. "auto" will use the fast '
56
57
                            'tokenizer if available, and "slow" will '
                            'always use the slow tokenizer.')
58
59
60
        parser.add_argument('--trust-remote-code',
                            action='store_true',
                            help='trust remote code from huggingface')
61
62
        parser.add_argument('--download-dir',
                            type=str,
Zhuohan Li's avatar
Zhuohan Li committed
63
                            default=EngineArgs.download_dir,
64
                            help='directory to download and load the weights, '
65
66
                            'default to the default cache dir of '
                            'huggingface')
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
        parser.add_argument(
            '--load-format',
            type=str,
            default=EngineArgs.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.')
82
        # TODO(woosuk): Support FP32.
83
84
85
86
87
88
89
90
91
        parser.add_argument(
            '--dtype',
            type=str,
            default=EngineArgs.dtype,
            choices=['auto', 'half', 'bfloat16', 'float'],
            help='data type for model weights and activations. '
            'The "auto" option will use FP16 precision '
            'for FP32 and FP16 models, and BF16 precision '
            'for BF16 models.')
92
        # Parallel arguments
93
94
        parser.add_argument('--worker-use-ray',
                            action='store_true',
95
                            help='use Ray for distributed serving, will be '
96
97
98
99
                            'automatically set when using more than 1 GPU')
        parser.add_argument('--pipeline-parallel-size',
                            '-pp',
                            type=int,
Zhuohan Li's avatar
Zhuohan Li committed
100
                            default=EngineArgs.pipeline_parallel_size,
101
                            help='number of pipeline stages')
102
103
104
        parser.add_argument('--tensor-parallel-size',
                            '-tp',
                            type=int,
Zhuohan Li's avatar
Zhuohan Li committed
105
                            default=EngineArgs.tensor_parallel_size,
106
107
                            help='number of tensor parallel replicas')
        # KV cache arguments
108
109
        parser.add_argument('--block-size',
                            type=int,
Zhuohan Li's avatar
Zhuohan Li committed
110
                            default=EngineArgs.block_size,
Woosuk Kwon's avatar
Woosuk Kwon committed
111
                            choices=[8, 16, 32],
112
113
                            help='token block size')
        # TODO(woosuk): Support fine-grained seeds (e.g., seed per request).
114
115
116
        parser.add_argument('--seed',
                            type=int,
                            default=EngineArgs.seed,
117
                            help='random seed')
118
119
        parser.add_argument('--swap-space',
                            type=int,
Zhuohan Li's avatar
Zhuohan Li committed
120
                            default=EngineArgs.swap_space,
121
                            help='CPU swap space size (GiB) per GPU')
122
123
        parser.add_argument('--gpu-memory-utilization',
                            type=float,
Zhuohan Li's avatar
Zhuohan Li committed
124
                            default=EngineArgs.gpu_memory_utilization,
125
                            help='the percentage of GPU memory to be used for'
126
127
128
                            'the model executor')
        parser.add_argument('--max-num-batched-tokens',
                            type=int,
Zhuohan Li's avatar
Zhuohan Li committed
129
                            default=EngineArgs.max_num_batched_tokens,
130
                            help='maximum number of batched tokens per '
131
132
133
                            'iteration')
        parser.add_argument('--max-num-seqs',
                            type=int,
Zhuohan Li's avatar
Zhuohan Li committed
134
                            default=EngineArgs.max_num_seqs,
135
                            help='maximum number of sequences per iteration')
136
137
        parser.add_argument('--disable-log-stats',
                            action='store_true',
138
139
                            help='disable logging statistics')
        return parser
140
141

    @classmethod
142
    def from_cli_args(cls, args: argparse.Namespace) -> 'EngineArgs':
143
144
145
        # Get the list of attributes of this dataclass.
        attrs = [attr.name for attr in dataclasses.fields(cls)]
        # Set the attributes from the parsed arguments.
Zhuohan Li's avatar
Zhuohan Li committed
146
147
        engine_args = cls(**{attr: getattr(args, attr) for attr in attrs})
        return engine_args
148

Zhuohan Li's avatar
Zhuohan Li committed
149
    def create_engine_configs(
150
151
152
        self,
    ) -> Tuple[ModelConfig, CacheConfig, ParallelConfig, SchedulerConfig]:
        # Initialize the configs.
153
        model_config = ModelConfig(self.model, self.tokenizer,
154
                                   self.tokenizer_mode, self.trust_remote_code,
155
156
                                   self.download_dir, self.load_format,
                                   self.dtype, self.seed)
157
158
        cache_config = CacheConfig(self.block_size,
                                   self.gpu_memory_utilization,
159
160
161
                                   self.swap_space)
        parallel_config = ParallelConfig(self.pipeline_parallel_size,
                                         self.tensor_parallel_size,
162
                                         self.worker_use_ray)
163
        scheduler_config = SchedulerConfig(self.max_num_batched_tokens,
164
165
                                           self.max_num_seqs,
                                           model_config.get_max_model_len())
166
167
168
        return model_config, cache_config, parallel_config, scheduler_config


169
@dataclass
Zhuohan Li's avatar
Zhuohan Li committed
170
class AsyncEngineArgs(EngineArgs):
Woosuk Kwon's avatar
Woosuk Kwon committed
171
    """Arguments for asynchronous vLLM engine."""
Zhuohan Li's avatar
Zhuohan Li committed
172
    engine_use_ray: bool = False
173
    disable_log_requests: bool = False
174
    max_log_len: Optional[int] = None
175
176
177

    @staticmethod
    def add_cli_args(
178
            parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
Zhuohan Li's avatar
Zhuohan Li committed
179
        parser = EngineArgs.add_cli_args(parser)
180
181
        parser.add_argument('--engine-use-ray',
                            action='store_true',
Zhuohan Li's avatar
Zhuohan Li committed
182
                            help='use Ray to start the LLM engine in a '
183
184
185
                            'separate process as the server process.')
        parser.add_argument('--disable-log-requests',
                            action='store_true',
186
                            help='disable logging requests')
187
188
189
190
191
192
        parser.add_argument('--max-log-len',
                            type=int,
                            default=None,
                            help='max number of prompt characters or prompt '
                            'ID numbers being printed in log. '
                            'Default: unlimited.')
193
        return parser