We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
CUDA error: an illegal memory acces with Falcon 40B
Hi,
I am testing different models with vllm. I see 
```CUDA error: an illegal memory access``` when I use falcon 40 b. The code I use is 
```
llm = LLM(model=ckpt_dir,tensor_parallel_size=4,trust_remote_code=True,gpu_memory_utilization=0.8)
sampling_params = SamplingParams(temperature=0, top_p=1.0, max_tokens=300)
results = llm.generate(prompts, sampling_params)
```
I am using an A100 with 4 GPUs. Please let me know if you have any questions

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `vllm/config.py`
Content:
```
1 from typing import Optional
2 
3 import torch
4 from transformers import PretrainedConfig
5 
6 from vllm.logger import init_logger
7 from vllm.transformers_utils.config import get_config
8 from vllm.utils import get_cpu_memory
9 
10 logger = init_logger(__name__)
11 
12 _GB = 1 << 30
13 
14 
15 class ModelConfig:
16     """Configuration for the model.
17 
18     Args:
19         model: Name or path of the huggingface model to use.
20         tokenizer: Name or path of the huggingface tokenizer to use.
21         tokenizer_mode: Tokenizer mode. "auto" will use the fast tokenizer if
22             available, and "slow" will always use the slow tokenizer.
23         trust_remote_code: Trust remote code (e.g., from HuggingFace) when
24             downloading the model and tokenizer.
25         download_dir: Directory to download and load the weights, default to the
26             default cache directory of huggingface.
27         load_format: The format of the model weights to load:
28             "auto" will try to load the weights in the safetensors format and
29                 fall back to the pytorch bin format if safetensors format is
30                 not available.
31             "pt" will load the weights in the pytorch bin format.
32             "safetensors" will load the weights in the safetensors format.
33             "npcache" will load the weights in pytorch format and store
34                 a numpy cache to speed up the loading.
35             "dummy" will initialize the weights with random values, which is
36                 mainly for profiling.
37         dtype: Data type for model weights and activations. The "auto" option
38             will use FP16 precision for FP32 and FP16 models, and BF16 precision
39             for BF16 models.
40         seed: Random seed for reproducibility.
41     """
42 
43     def __init__(
44         self,
45         model: str,
46         tokenizer: str,
47         tokenizer_mode: str,
48         trust_remote_code: bool,
49         download_dir: Optional[str],
50         load_format: str,
51         dtype: str,
52         seed: int,
53     ) -> None:
54         self.model = model
55         self.tokenizer = tokenizer
56         self.tokenizer_mode = tokenizer_mode
57         self.trust_remote_code = trust_remote_code
58         self.download_dir = download_dir
59         self.load_format = load_format
60         self.seed = seed
61 
62         self.hf_config = get_config(model, trust_remote_code)
63         self.dtype = _get_and_verify_dtype(self.hf_config, dtype)
64         self._verify_load_format()
65         self._verify_tokenizer_mode()
66 
67     def _verify_load_format(self) -> None:
68         load_format = self.load_format.lower()
69         if load_format not in [
70                 "auto", "pt", "safetensors", "npcache", "dummy"
71         ]:
72             raise ValueError(
73                 f"Unknown load format: {self.load_format}. Must be one of "
74                 "'auto', 'pt', 'safetensors', 'npcache', or 'dummy'.")
75         self.load_format = load_format
76 
77     def _verify_tokenizer_mode(self) -> None:
78         tokenizer_mode = self.tokenizer_mode.lower()
79         if tokenizer_mode not in ["auto", "slow"]:
80             raise ValueError(
81                 f"Unknown tokenizer mode: {self.tokenizer_mode}. Must be "
82                 "either 'auto' or 'slow'.")
83         self.tokenizer_mode = tokenizer_mode
84 
85     def verify_with_parallel_config(
86         self,
87         parallel_config: "ParallelConfig",
88     ) -> None:
89         total_num_attention_heads = self.hf_config.num_attention_heads
90         tensor_parallel_size = parallel_config.tensor_parallel_size
91         if total_num_attention_heads % tensor_parallel_size != 0:
92             raise ValueError(
93                 f"Total number of attention heads ({total_num_attention_heads})"
94                 " must be divisible by tensor parallel size "
95                 f"({tensor_parallel_size}).")
96 
97         total_num_hidden_layers = self.hf_config.num_hidden_layers
98         pipeline_parallel_size = parallel_config.pipeline_parallel_size
99         if total_num_hidden_layers % pipeline_parallel_size != 0:
100             raise ValueError(
101                 f"Total number of hidden layers ({total_num_hidden_layers}) "
102                 "must be divisible by pipeline parallel size "
103                 f"({pipeline_parallel_size}).")
104 
105     def get_hidden_size(self) -> int:
106         return self.hf_config.hidden_size
107 
108     def get_head_size(self) -> int:
109         # FIXME(woosuk): This may not be true for all models.
110         return self.hf_config.hidden_size // self.hf_config.num_attention_heads
111 
112     def get_num_heads(self, parallel_config: "ParallelConfig") -> int:
113         # For GPTBigCode & Falcon:
114         # Note: for falcon, when new_decoder_architecture is True, the
115         # multi_query flag is ignored and we use n_head_kv for the number of
116         # KV heads.
117         new_decoder_arch_falcon = (
118             self.hf_config.model_type == "falcon"
119             and getattr(self.hf_config, "new_decoder_architecture", False))
120         if not new_decoder_arch_falcon and getattr(self.hf_config,
121                                                    "multi_query", False):
122             # Multi-query attention, only one KV head.
123             return 1
124         # For Falcon:
125         if getattr(self.hf_config, "n_head_kv", None) is not None:
126             return (self.hf_config.n_head_kv //
127                     parallel_config.tensor_parallel_size)
128         # For LLaMA-2:
129         if getattr(self.hf_config, "num_key_value_heads", None) is not None:
130             return (self.hf_config.num_key_value_heads //
131                     parallel_config.tensor_parallel_size)
132         total_num_attention_heads = self.hf_config.num_attention_heads
133         return total_num_attention_heads // parallel_config.tensor_parallel_size
134 
135     def get_max_model_len(self) -> int:
136         max_model_len = float("inf")
137         possible_keys = [
138             # OPT
139             "max_position_embeddings",
140             # GPT-2
141             "n_positions",
142             # MPT
143             "max_seq_len",
144             # Others
145             "max_sequence_length",
146             "max_seq_length",
147             "seq_len",
148         ]
149         for key in possible_keys:
150             max_len_key = getattr(self.hf_config, key, None)
151             if max_len_key is not None:
152                 max_model_len = min(max_model_len, max_len_key)
153         return max_model_len
154 
155     def get_num_layers(self, parallel_config: "ParallelConfig") -> int:
156         total_num_hidden_layers = self.hf_config.num_hidden_layers
157         return total_num_hidden_layers // parallel_config.pipeline_parallel_size
158 
159 
160 class CacheConfig:
161     """Configuration for the KV cache.
162 
163     Args:
164         block_size: Size of a cache block in number of tokens.
165         gpu_memory_utilization: Fraction of GPU memory to use for the
166             vLLM execution.
167         swap_space: Size of the CPU swap space per GPU (in GiB).
168     """
169 
170     def __init__(
171         self,
172         block_size: int,
173         gpu_memory_utilization: float,
174         swap_space: int,
175     ) -> None:
176         self.block_size = block_size
177         self.gpu_memory_utilization = gpu_memory_utilization
178         self.swap_space_bytes = swap_space * _GB
179         self._verify_args()
180 
181         # Will be set after profiling.
182         self.num_gpu_blocks = None
183         self.num_cpu_blocks = None
184 
185     def _verify_args(self) -> None:
186         if self.gpu_memory_utilization > 1.0:
187             raise ValueError(
188                 "GPU memory utilization must be less than 1.0. Got "
189                 f"{self.gpu_memory_utilization}.")
190 
191     def verify_with_parallel_config(
192         self,
193         parallel_config: "ParallelConfig",
194     ) -> None:
195         total_cpu_memory = get_cpu_memory()
196         # FIXME(woosuk): Here, it is assumed that the GPUs in a tensor parallel
197         # group are in the same node. However, the GPUs may span multiple nodes.
198         num_gpus_per_node = parallel_config.tensor_parallel_size
199         cpu_memory_usage = self.swap_space_bytes * num_gpus_per_node
200 
201         msg = (f"{cpu_memory_usage / _GB:.2f} GiB out of "
202                f"the {total_cpu_memory / _GB:.2f} GiB total CPU memory is "
203                "allocated for the swap space.")
204         if cpu_memory_usage > 0.7 * total_cpu_memory:
205             raise ValueError("Too large swap space. " + msg)
206         elif cpu_memory_usage > 0.4 * total_cpu_memory:
207             logger.warning("Possibly too large swap space. " + msg)
208 
209 
210 class ParallelConfig:
211     """Configuration for the distributed execution.
212 
213     Args:
214         pipeline_parallel_size: Number of pipeline parallel groups.
215         tensor_parallel_size: Number of tensor parallel groups.
216         worker_use_ray: Whether to use Ray for model workers. Will be set to
217             True if either pipeline_parallel_size or tensor_parallel_size is
218             greater than 1.
219     """
220 
221     def __init__(
222         self,
223         pipeline_parallel_size: int,
224         tensor_parallel_size: int,
225         worker_use_ray: bool,
226     ) -> None:
227         self.pipeline_parallel_size = pipeline_parallel_size
228         self.tensor_parallel_size = tensor_parallel_size
229         self.worker_use_ray = worker_use_ray
230 
231         self.world_size = pipeline_parallel_size * tensor_parallel_size
232         if self.world_size > 1:
233             self.worker_use_ray = True
234         self._verify_args()
235 
236     def _verify_args(self) -> None:
237         if self.pipeline_parallel_size > 1:
238             raise NotImplementedError(
239                 "Pipeline parallelism is not supported yet.")
240 
241 
242 class SchedulerConfig:
243     """Scheduler configuration.
244 
245     Args:
246         max_num_batched_tokens: Maximum number of tokens to be processed in
247             a single iteration.
248         max_num_seqs: Maximum number of sequences to be processed in a single
249             iteration.
250         max_model_len: Maximum length of a sequence (including prompt
251             and generated text).
252     """
253 
254     def __init__(self, max_num_batched_tokens: int, max_num_seqs: int,
255                  max_model_len: int) -> None:
256         self.max_num_batched_tokens = max_num_batched_tokens
257         self.max_num_seqs = max_num_seqs
258         self.max_model_len = max_model_len
259 
260 
261 _STR_DTYPE_TO_TORCH_DTYPE = {
262     "half": torch.float16,
263     "float16": torch.float16,
264     "float": torch.float32,
265     "float32": torch.float32,
266     "bfloat16": torch.bfloat16,
267 }
268 
269 
270 def _get_and_verify_dtype(
271     config: PretrainedConfig,
272     dtype: str,
273 ) -> torch.dtype:
274     # NOTE: getattr(config, "torch_dtype", torch.float32) is not correct
275     # because config.torch_dtype can be None.
276     config_dtype = getattr(config, "torch_dtype", None)
277     if config_dtype is None:
278         config_dtype = torch.float32
279 
280     dtype = dtype.lower()
281     if dtype == "auto":
282         if config_dtype == torch.float32:
283             # Following the common practice, we use float16 for float32 models.
284             torch_dtype = torch.float16
285         else:
286             torch_dtype = config_dtype
287     else:
288         if dtype not in _STR_DTYPE_TO_TORCH_DTYPE:
289             raise ValueError(f"Unknown dtype: {dtype}")
290         torch_dtype = _STR_DTYPE_TO_TORCH_DTYPE[dtype]
291 
292     # Verify the dtype.
293     if torch_dtype != config_dtype:
294         if torch_dtype == torch.float32:
295             # Upcasting to float32 is allowed.
296             pass
297         elif config_dtype == torch.float32:
298             # Downcasting from float32 to float16 or bfloat16 is allowed.
299             pass
300         else:
301             # Casting between float16 and bfloat16 is allowed with a warning.
302             logger.warning(f"Casting {config_dtype} to {torch_dtype}.")
303 
304     # Check if the GPU supports the dtype.
305     if torch_dtype == torch.bfloat16:
306         compute_capability = torch.cuda.get_device_capability()
307         if compute_capability[0] < 8:
308             gpu_name = torch.cuda.get_device_name()
309             raise ValueError(
310                 "Bfloat16 is only supported on GPUs with compute capability "
311                 f"of at least 8.0. Your {gpu_name} GPU has compute capability "
312                 f"{compute_capability[0]}.{compute_capability[1]}.")
313     return torch_dtype
314
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/vllm/config.py b/vllm/config.py
--- a/vllm/config.py
+++ b/vllm/config.py
@@ -114,8 +114,9 @@
         # Note: for falcon, when new_decoder_architecture is True, the
         # multi_query flag is ignored and we use n_head_kv for the number of
         # KV heads.
+        falcon_model_types = ["falcon", "RefinedWeb", "RefinedWebModel"]
         new_decoder_arch_falcon = (
-            self.hf_config.model_type == "falcon"
+            self.hf_config.model_type in falcon_model_types
             and getattr(self.hf_config, "new_decoder_architecture", False))
         if not new_decoder_arch_falcon and getattr(self.hf_config,
                                                    "multi_query", False):

{"golden_diff": "diff --git a/vllm/config.py b/vllm/config.py\n--- a/vllm/config.py\n+++ b/vllm/config.py\n@@ -114,8 +114,9 @@\n         # Note: for falcon, when new_decoder_architecture is True, the\n         # multi_query flag is ignored and we use n_head_kv for the number of\n         # KV heads.\n+        falcon_model_types = [\"falcon\", \"RefinedWeb\", \"RefinedWebModel\"]\n         new_decoder_arch_falcon = (\n-            self.hf_config.model_type == \"falcon\"\n+            self.hf_config.model_type in falcon_model_types\n             and getattr(self.hf_config, \"new_decoder_architecture\", False))\n         if not new_decoder_arch_falcon and getattr(self.hf_config,\n                                                    \"multi_query\", False):\n", "issue": "CUDA error: an illegal memory acces with Falcon 40B\nHi,\r\nI am testing different models with vllm. I see \r\n```CUDA error: an illegal memory access``` when I use falcon 40 b. The code I use is \r\n```\r\nllm = LLM(model=ckpt_dir,tensor_parallel_size=4,trust_remote_code=True,gpu_memory_utilization=0.8)\r\nsampling_params = SamplingParams(temperature=0, top_p=1.0, max_tokens=300)\r\nresults = llm.generate(prompts, sampling_params)\r\n```\r\nI am using an A100 with 4 GPUs. Please let me know if you have any questions\n", "before_files": [{"content": "from typing import Optional\n\nimport torch\nfrom transformers import PretrainedConfig\n\nfrom vllm.logger import init_logger\nfrom vllm.transformers_utils.config import get_config\nfrom vllm.utils import get_cpu_memory\n\nlogger = init_logger(__name__)\n\n_GB = 1 << 30\n\n\nclass ModelConfig:\n    \"\"\"Configuration for the model.\n\n    Args:\n        model: Name or path of the huggingface model to use.\n        tokenizer: Name or path of the huggingface tokenizer to use.\n        tokenizer_mode: Tokenizer mode. \"auto\" will use the fast tokenizer if\n            available, and \"slow\" will always use the slow tokenizer.\n        trust_remote_code: Trust remote code (e.g., from HuggingFace) when\n            downloading the model and tokenizer.\n        download_dir: Directory to download and load the weights, default to the\n            default cache directory of huggingface.\n        load_format: The format of the model weights to load:\n            \"auto\" will try to load the weights in the safetensors format and\n                fall back to the pytorch bin format if safetensors format is\n                not available.\n            \"pt\" will load the weights in the pytorch bin format.\n            \"safetensors\" will load the weights in the safetensors format.\n            \"npcache\" will load the weights in pytorch format and store\n                a numpy cache to speed up the loading.\n            \"dummy\" will initialize the weights with random values, which is\n                mainly for profiling.\n        dtype: Data type for model weights and activations. The \"auto\" option\n            will use FP16 precision for FP32 and FP16 models, and BF16 precision\n            for BF16 models.\n        seed: Random seed for reproducibility.\n    \"\"\"\n\n    def __init__(\n        self,\n        model: str,\n        tokenizer: str,\n        tokenizer_mode: str,\n        trust_remote_code: bool,\n        download_dir: Optional[str],\n        load_format: str,\n        dtype: str,\n        seed: int,\n    ) -> None:\n        self.model = model\n        self.tokenizer = tokenizer\n        self.tokenizer_mode = tokenizer_mode\n        self.trust_remote_code = trust_remote_code\n        self.download_dir = download_dir\n        self.load_format = load_format\n        self.seed = seed\n\n        self.hf_config = get_config(model, trust_remote_code)\n        self.dtype = _get_and_verify_dtype(self.hf_config, dtype)\n        self._verify_load_format()\n        self._verify_tokenizer_mode()\n\n    def _verify_load_format(self) -> None:\n        load_format = self.load_format.lower()\n        if load_format not in [\n                \"auto\", \"pt\", \"safetensors\", \"npcache\", \"dummy\"\n        ]:\n            raise ValueError(\n                f\"Unknown load format: {self.load_format}. Must be one of \"\n                \"'auto', 'pt', 'safetensors', 'npcache', or 'dummy'.\")\n        self.load_format = load_format\n\n    def _verify_tokenizer_mode(self) -> None:\n        tokenizer_mode = self.tokenizer_mode.lower()\n        if tokenizer_mode not in [\"auto\", \"slow\"]:\n            raise ValueError(\n                f\"Unknown tokenizer mode: {self.tokenizer_mode}. Must be \"\n                \"either 'auto' or 'slow'.\")\n        self.tokenizer_mode = tokenizer_mode\n\n    def verify_with_parallel_config(\n        self,\n        parallel_config: \"ParallelConfig\",\n    ) -> None:\n        total_num_attention_heads = self.hf_config.num_attention_heads\n        tensor_parallel_size = parallel_config.tensor_parallel_size\n        if total_num_attention_heads % tensor_parallel_size != 0:\n            raise ValueError(\n                f\"Total number of attention heads ({total_num_attention_heads})\"\n                \" must be divisible by tensor parallel size \"\n                f\"({tensor_parallel_size}).\")\n\n        total_num_hidden_layers = self.hf_config.num_hidden_layers\n        pipeline_parallel_size = parallel_config.pipeline_parallel_size\n        if total_num_hidden_layers % pipeline_parallel_size != 0:\n            raise ValueError(\n                f\"Total number of hidden layers ({total_num_hidden_layers}) \"\n                \"must be divisible by pipeline parallel size \"\n                f\"({pipeline_parallel_size}).\")\n\n    def get_hidden_size(self) -> int:\n        return self.hf_config.hidden_size\n\n    def get_head_size(self) -> int:\n        # FIXME(woosuk): This may not be true for all models.\n        return self.hf_config.hidden_size // self.hf_config.num_attention_heads\n\n    def get_num_heads(self, parallel_config: \"ParallelConfig\") -> int:\n        # For GPTBigCode & Falcon:\n        # Note: for falcon, when new_decoder_architecture is True, the\n        # multi_query flag is ignored and we use n_head_kv for the number of\n        # KV heads.\n        new_decoder_arch_falcon = (\n            self.hf_config.model_type == \"falcon\"\n            and getattr(self.hf_config, \"new_decoder_architecture\", False))\n        if not new_decoder_arch_falcon and getattr(self.hf_config,\n                                                   \"multi_query\", False):\n            # Multi-query attention, only one KV head.\n            return 1\n        # For Falcon:\n        if getattr(self.hf_config, \"n_head_kv\", None) is not None:\n            return (self.hf_config.n_head_kv //\n                    parallel_config.tensor_parallel_size)\n        # For LLaMA-2:\n        if getattr(self.hf_config, \"num_key_value_heads\", None) is not None:\n            return (self.hf_config.num_key_value_heads //\n                    parallel_config.tensor_parallel_size)\n        total_num_attention_heads = self.hf_config.num_attention_heads\n        return total_num_attention_heads // parallel_config.tensor_parallel_size\n\n    def get_max_model_len(self) -> int:\n        max_model_len = float(\"inf\")\n        possible_keys = [\n            # OPT\n            \"max_position_embeddings\",\n            # GPT-2\n            \"n_positions\",\n            # MPT\n            \"max_seq_len\",\n            # Others\n            \"max_sequence_length\",\n            \"max_seq_length\",\n            \"seq_len\",\n        ]\n        for key in possible_keys:\n            max_len_key = getattr(self.hf_config, key, None)\n            if max_len_key is not None:\n                max_model_len = min(max_model_len, max_len_key)\n        return max_model_len\n\n    def get_num_layers(self, parallel_config: \"ParallelConfig\") -> int:\n        total_num_hidden_layers = self.hf_config.num_hidden_layers\n        return total_num_hidden_layers // parallel_config.pipeline_parallel_size\n\n\nclass CacheConfig:\n    \"\"\"Configuration for the KV cache.\n\n    Args:\n        block_size: Size of a cache block in number of tokens.\n        gpu_memory_utilization: Fraction of GPU memory to use for the\n            vLLM execution.\n        swap_space: Size of the CPU swap space per GPU (in GiB).\n    \"\"\"\n\n    def __init__(\n        self,\n        block_size: int,\n        gpu_memory_utilization: float,\n        swap_space: int,\n    ) -> None:\n        self.block_size = block_size\n        self.gpu_memory_utilization = gpu_memory_utilization\n        self.swap_space_bytes = swap_space * _GB\n        self._verify_args()\n\n        # Will be set after profiling.\n        self.num_gpu_blocks = None\n        self.num_cpu_blocks = None\n\n    def _verify_args(self) -> None:\n        if self.gpu_memory_utilization > 1.0:\n            raise ValueError(\n                \"GPU memory utilization must be less than 1.0. Got \"\n                f\"{self.gpu_memory_utilization}.\")\n\n    def verify_with_parallel_config(\n        self,\n        parallel_config: \"ParallelConfig\",\n    ) -> None:\n        total_cpu_memory = get_cpu_memory()\n        # FIXME(woosuk): Here, it is assumed that the GPUs in a tensor parallel\n        # group are in the same node. However, the GPUs may span multiple nodes.\n        num_gpus_per_node = parallel_config.tensor_parallel_size\n        cpu_memory_usage = self.swap_space_bytes * num_gpus_per_node\n\n        msg = (f\"{cpu_memory_usage / _GB:.2f} GiB out of \"\n               f\"the {total_cpu_memory / _GB:.2f} GiB total CPU memory is \"\n               \"allocated for the swap space.\")\n        if cpu_memory_usage > 0.7 * total_cpu_memory:\n            raise ValueError(\"Too large swap space. \" + msg)\n        elif cpu_memory_usage > 0.4 * total_cpu_memory:\n            logger.warning(\"Possibly too large swap space. \" + msg)\n\n\nclass ParallelConfig:\n    \"\"\"Configuration for the distributed execution.\n\n    Args:\n        pipeline_parallel_size: Number of pipeline parallel groups.\n        tensor_parallel_size: Number of tensor parallel groups.\n        worker_use_ray: Whether to use Ray for model workers. Will be set to\n            True if either pipeline_parallel_size or tensor_parallel_size is\n            greater than 1.\n    \"\"\"\n\n    def __init__(\n        self,\n        pipeline_parallel_size: int,\n        tensor_parallel_size: int,\n        worker_use_ray: bool,\n    ) -> None:\n        self.pipeline_parallel_size = pipeline_parallel_size\n        self.tensor_parallel_size = tensor_parallel_size\n        self.worker_use_ray = worker_use_ray\n\n        self.world_size = pipeline_parallel_size * tensor_parallel_size\n        if self.world_size > 1:\n            self.worker_use_ray = True\n        self._verify_args()\n\n    def _verify_args(self) -> None:\n        if self.pipeline_parallel_size > 1:\n            raise NotImplementedError(\n                \"Pipeline parallelism is not supported yet.\")\n\n\nclass SchedulerConfig:\n    \"\"\"Scheduler configuration.\n\n    Args:\n        max_num_batched_tokens: Maximum number of tokens to be processed in\n            a single iteration.\n        max_num_seqs: Maximum number of sequences to be processed in a single\n            iteration.\n        max_model_len: Maximum length of a sequence (including prompt\n            and generated text).\n    \"\"\"\n\n    def __init__(self, max_num_batched_tokens: int, max_num_seqs: int,\n                 max_model_len: int) -> None:\n        self.max_num_batched_tokens = max_num_batched_tokens\n        self.max_num_seqs = max_num_seqs\n        self.max_model_len = max_model_len\n\n\n_STR_DTYPE_TO_TORCH_DTYPE = {\n    \"half\": torch.float16,\n    \"float16\": torch.float16,\n    \"float\": torch.float32,\n    \"float32\": torch.float32,\n    \"bfloat16\": torch.bfloat16,\n}\n\n\ndef _get_and_verify_dtype(\n    config: PretrainedConfig,\n    dtype: str,\n) -> torch.dtype:\n    # NOTE: getattr(config, \"torch_dtype\", torch.float32) is not correct\n    # because config.torch_dtype can be None.\n    config_dtype = getattr(config, \"torch_dtype\", None)\n    if config_dtype is None:\n        config_dtype = torch.float32\n\n    dtype = dtype.lower()\n    if dtype == \"auto\":\n        if config_dtype == torch.float32:\n            # Following the common practice, we use float16 for float32 models.\n            torch_dtype = torch.float16\n        else:\n            torch_dtype = config_dtype\n    else:\n        if dtype not in _STR_DTYPE_TO_TORCH_DTYPE:\n            raise ValueError(f\"Unknown dtype: {dtype}\")\n        torch_dtype = _STR_DTYPE_TO_TORCH_DTYPE[dtype]\n\n    # Verify the dtype.\n    if torch_dtype != config_dtype:\n        if torch_dtype == torch.float32:\n            # Upcasting to float32 is allowed.\n            pass\n        elif config_dtype == torch.float32:\n            # Downcasting from float32 to float16 or bfloat16 is allowed.\n            pass\n        else:\n            # Casting between float16 and bfloat16 is allowed with a warning.\n            logger.warning(f\"Casting {config_dtype} to {torch_dtype}.\")\n\n    # Check if the GPU supports the dtype.\n    if torch_dtype == torch.bfloat16:\n        compute_capability = torch.cuda.get_device_capability()\n        if compute_capability[0] < 8:\n            gpu_name = torch.cuda.get_device_name()\n            raise ValueError(\n                \"Bfloat16 is only supported on GPUs with compute capability \"\n                f\"of at least 8.0. Your {gpu_name} GPU has compute capability \"\n                f\"{compute_capability[0]}.{compute_capability[1]}.\")\n    return torch_dtype\n", "path": "vllm/config.py"}], "after_files": [{"content": "from typing import Optional\n\nimport torch\nfrom transformers import PretrainedConfig\n\nfrom vllm.logger import init_logger\nfrom vllm.transformers_utils.config import get_config\nfrom vllm.utils import get_cpu_memory\n\nlogger = init_logger(__name__)\n\n_GB = 1 << 30\n\n\nclass ModelConfig:\n    \"\"\"Configuration for the model.\n\n    Args:\n        model: Name or path of the huggingface model to use.\n        tokenizer: Name or path of the huggingface tokenizer to use.\n        tokenizer_mode: Tokenizer mode. \"auto\" will use the fast tokenizer if\n            available, and \"slow\" will always use the slow tokenizer.\n        trust_remote_code: Trust remote code (e.g., from HuggingFace) when\n            downloading the model and tokenizer.\n        download_dir: Directory to download and load the weights, default to the\n            default cache directory of huggingface.\n        load_format: The format of the model weights to load:\n            \"auto\" will try to load the weights in the safetensors format and\n                fall back to the pytorch bin format if safetensors format is\n                not available.\n            \"pt\" will load the weights in the pytorch bin format.\n            \"safetensors\" will load the weights in the safetensors format.\n            \"npcache\" will load the weights in pytorch format and store\n                a numpy cache to speed up the loading.\n            \"dummy\" will initialize the weights with random values, which is\n                mainly for profiling.\n        dtype: Data type for model weights and activations. The \"auto\" option\n            will use FP16 precision for FP32 and FP16 models, and BF16 precision\n            for BF16 models.\n        seed: Random seed for reproducibility.\n    \"\"\"\n\n    def __init__(\n        self,\n        model: str,\n        tokenizer: str,\n        tokenizer_mode: str,\n        trust_remote_code: bool,\n        download_dir: Optional[str],\n        load_format: str,\n        dtype: str,\n        seed: int,\n    ) -> None:\n        self.model = model\n        self.tokenizer = tokenizer\n        self.tokenizer_mode = tokenizer_mode\n        self.trust_remote_code = trust_remote_code\n        self.download_dir = download_dir\n        self.load_format = load_format\n        self.seed = seed\n\n        self.hf_config = get_config(model, trust_remote_code)\n        self.dtype = _get_and_verify_dtype(self.hf_config, dtype)\n        self._verify_load_format()\n        self._verify_tokenizer_mode()\n\n    def _verify_load_format(self) -> None:\n        load_format = self.load_format.lower()\n        if load_format not in [\n                \"auto\", \"pt\", \"safetensors\", \"npcache\", \"dummy\"\n        ]:\n            raise ValueError(\n                f\"Unknown load format: {self.load_format}. Must be one of \"\n                \"'auto', 'pt', 'safetensors', 'npcache', or 'dummy'.\")\n        self.load_format = load_format\n\n    def _verify_tokenizer_mode(self) -> None:\n        tokenizer_mode = self.tokenizer_mode.lower()\n        if tokenizer_mode not in [\"auto\", \"slow\"]:\n            raise ValueError(\n                f\"Unknown tokenizer mode: {self.tokenizer_mode}. Must be \"\n                \"either 'auto' or 'slow'.\")\n        self.tokenizer_mode = tokenizer_mode\n\n    def verify_with_parallel_config(\n        self,\n        parallel_config: \"ParallelConfig\",\n    ) -> None:\n        total_num_attention_heads = self.hf_config.num_attention_heads\n        tensor_parallel_size = parallel_config.tensor_parallel_size\n        if total_num_attention_heads % tensor_parallel_size != 0:\n            raise ValueError(\n                f\"Total number of attention heads ({total_num_attention_heads})\"\n                \" must be divisible by tensor parallel size \"\n                f\"({tensor_parallel_size}).\")\n\n        total_num_hidden_layers = self.hf_config.num_hidden_layers\n        pipeline_parallel_size = parallel_config.pipeline_parallel_size\n        if total_num_hidden_layers % pipeline_parallel_size != 0:\n            raise ValueError(\n                f\"Total number of hidden layers ({total_num_hidden_layers}) \"\n                \"must be divisible by pipeline parallel size \"\n                f\"({pipeline_parallel_size}).\")\n\n    def get_hidden_size(self) -> int:\n        return self.hf_config.hidden_size\n\n    def get_head_size(self) -> int:\n        # FIXME(woosuk): This may not be true for all models.\n        return self.hf_config.hidden_size // self.hf_config.num_attention_heads\n\n    def get_num_heads(self, parallel_config: \"ParallelConfig\") -> int:\n        # For GPTBigCode & Falcon:\n        # Note: for falcon, when new_decoder_architecture is True, the\n        # multi_query flag is ignored and we use n_head_kv for the number of\n        # KV heads.\n        falcon_model_types = [\"falcon\", \"RefinedWeb\", \"RefinedWebModel\"]\n        new_decoder_arch_falcon = (\n            self.hf_config.model_type in falcon_model_types\n            and getattr(self.hf_config, \"new_decoder_architecture\", False))\n        if not new_decoder_arch_falcon and getattr(self.hf_config,\n                                                   \"multi_query\", False):\n            # Multi-query attention, only one KV head.\n            return 1\n        # For Falcon:\n        if getattr(self.hf_config, \"n_head_kv\", None) is not None:\n            return (self.hf_config.n_head_kv //\n                    parallel_config.tensor_parallel_size)\n        # For LLaMA-2:\n        if getattr(self.hf_config, \"num_key_value_heads\", None) is not None:\n            return (self.hf_config.num_key_value_heads //\n                    parallel_config.tensor_parallel_size)\n        total_num_attention_heads = self.hf_config.num_attention_heads\n        return total_num_attention_heads // parallel_config.tensor_parallel_size\n\n    def get_max_model_len(self) -> int:\n        max_model_len = float(\"inf\")\n        possible_keys = [\n            # OPT\n            \"max_position_embeddings\",\n            # GPT-2\n            \"n_positions\",\n            # MPT\n            \"max_seq_len\",\n            # Others\n            \"max_sequence_length\",\n            \"max_seq_length\",\n            \"seq_len\",\n        ]\n        for key in possible_keys:\n            max_len_key = getattr(self.hf_config, key, None)\n            if max_len_key is not None:\n                max_model_len = min(max_model_len, max_len_key)\n        return max_model_len\n\n    def get_num_layers(self, parallel_config: \"ParallelConfig\") -> int:\n        total_num_hidden_layers = self.hf_config.num_hidden_layers\n        return total_num_hidden_layers // parallel_config.pipeline_parallel_size\n\n\nclass CacheConfig:\n    \"\"\"Configuration for the KV cache.\n\n    Args:\n        block_size: Size of a cache block in number of tokens.\n        gpu_memory_utilization: Fraction of GPU memory to use for the\n            vLLM execution.\n        swap_space: Size of the CPU swap space per GPU (in GiB).\n    \"\"\"\n\n    def __init__(\n        self,\n        block_size: int,\n        gpu_memory_utilization: float,\n        swap_space: int,\n    ) -> None:\n        self.block_size = block_size\n        self.gpu_memory_utilization = gpu_memory_utilization\n        self.swap_space_bytes = swap_space * _GB\n        self._verify_args()\n\n        # Will be set after profiling.\n        self.num_gpu_blocks = None\n        self.num_cpu_blocks = None\n\n    def _verify_args(self) -> None:\n        if self.gpu_memory_utilization > 1.0:\n            raise ValueError(\n                \"GPU memory utilization must be less than 1.0. Got \"\n                f\"{self.gpu_memory_utilization}.\")\n\n    def verify_with_parallel_config(\n        self,\n        parallel_config: \"ParallelConfig\",\n    ) -> None:\n        total_cpu_memory = get_cpu_memory()\n        # FIXME(woosuk): Here, it is assumed that the GPUs in a tensor parallel\n        # group are in the same node. However, the GPUs may span multiple nodes.\n        num_gpus_per_node = parallel_config.tensor_parallel_size\n        cpu_memory_usage = self.swap_space_bytes * num_gpus_per_node\n\n        msg = (f\"{cpu_memory_usage / _GB:.2f} GiB out of \"\n               f\"the {total_cpu_memory / _GB:.2f} GiB total CPU memory is \"\n               \"allocated for the swap space.\")\n        if cpu_memory_usage > 0.7 * total_cpu_memory:\n            raise ValueError(\"Too large swap space. \" + msg)\n        elif cpu_memory_usage > 0.4 * total_cpu_memory:\n            logger.warning(\"Possibly too large swap space. \" + msg)\n\n\nclass ParallelConfig:\n    \"\"\"Configuration for the distributed execution.\n\n    Args:\n        pipeline_parallel_size: Number of pipeline parallel groups.\n        tensor_parallel_size: Number of tensor parallel groups.\n        worker_use_ray: Whether to use Ray for model workers. Will be set to\n            True if either pipeline_parallel_size or tensor_parallel_size is\n            greater than 1.\n    \"\"\"\n\n    def __init__(\n        self,\n        pipeline_parallel_size: int,\n        tensor_parallel_size: int,\n        worker_use_ray: bool,\n    ) -> None:\n        self.pipeline_parallel_size = pipeline_parallel_size\n        self.tensor_parallel_size = tensor_parallel_size\n        self.worker_use_ray = worker_use_ray\n\n        self.world_size = pipeline_parallel_size * tensor_parallel_size\n        if self.world_size > 1:\n            self.worker_use_ray = True\n        self._verify_args()\n\n    def _verify_args(self) -> None:\n        if self.pipeline_parallel_size > 1:\n            raise NotImplementedError(\n                \"Pipeline parallelism is not supported yet.\")\n\n\nclass SchedulerConfig:\n    \"\"\"Scheduler configuration.\n\n    Args:\n        max_num_batched_tokens: Maximum number of tokens to be processed in\n            a single iteration.\n        max_num_seqs: Maximum number of sequences to be processed in a single\n            iteration.\n        max_model_len: Maximum length of a sequence (including prompt\n            and generated text).\n    \"\"\"\n\n    def __init__(self, max_num_batched_tokens: int, max_num_seqs: int,\n                 max_model_len: int) -> None:\n        self.max_num_batched_tokens = max_num_batched_tokens\n        self.max_num_seqs = max_num_seqs\n        self.max_model_len = max_model_len\n\n\n_STR_DTYPE_TO_TORCH_DTYPE = {\n    \"half\": torch.float16,\n    \"float16\": torch.float16,\n    \"float\": torch.float32,\n    \"float32\": torch.float32,\n    \"bfloat16\": torch.bfloat16,\n}\n\n\ndef _get_and_verify_dtype(\n    config: PretrainedConfig,\n    dtype: str,\n) -> torch.dtype:\n    # NOTE: getattr(config, \"torch_dtype\", torch.float32) is not correct\n    # because config.torch_dtype can be None.\n    config_dtype = getattr(config, \"torch_dtype\", None)\n    if config_dtype is None:\n        config_dtype = torch.float32\n\n    dtype = dtype.lower()\n    if dtype == \"auto\":\n        if config_dtype == torch.float32:\n            # Following the common practice, we use float16 for float32 models.\n            torch_dtype = torch.float16\n        else:\n            torch_dtype = config_dtype\n    else:\n        if dtype not in _STR_DTYPE_TO_TORCH_DTYPE:\n            raise ValueError(f\"Unknown dtype: {dtype}\")\n        torch_dtype = _STR_DTYPE_TO_TORCH_DTYPE[dtype]\n\n    # Verify the dtype.\n    if torch_dtype != config_dtype:\n        if torch_dtype == torch.float32:\n            # Upcasting to float32 is allowed.\n            pass\n        elif config_dtype == torch.float32:\n            # Downcasting from float32 to float16 or bfloat16 is allowed.\n            pass\n        else:\n            # Casting between float16 and bfloat16 is allowed with a warning.\n            logger.warning(f\"Casting {config_dtype} to {torch_dtype}.\")\n\n    # Check if the GPU supports the dtype.\n    if torch_dtype == torch.bfloat16:\n        compute_capability = torch.cuda.get_device_capability()\n        if compute_capability[0] < 8:\n            gpu_name = torch.cuda.get_device_name()\n            raise ValueError(\n                \"Bfloat16 is only supported on GPUs with compute capability \"\n                f\"of at least 8.0. Your {gpu_name} GPU has compute capability \"\n                f\"{compute_capability[0]}.{compute_capability[1]}.\")\n    return torch_dtype\n", "path": "vllm/config.py"}]}