Repository: vllm-project/vllm
Pull request: rope_theta and max_position_embeddings from config
This PR lets `rope_theta` and `max_position_embeddings` to be read from model configs instead of hardcoding them. Notably, this allows codellama to work without issues with longer contexts.

Fixes https://github.com/vllm-project/vllm/issues/904

Problem statement:
stuck at llm_engine.py:196
```
100%|████████████████████████████████████| 54584/54584 [00:06<00:00, 8619.58it/s]
INFO 08-29 20:48:52 llm_engine.py:70] Initializing an LLM engine with config: model='llama2-MultiTQ-20230829-01', tokenizer='llama2-MultiTQ-20230829-01', tokenizer_mode=auto, trust_remote_code=False, dtype=torch.float16, use_dummy_weights=False, download_dir=None, use_np_weights=False, tensor_parallel_size=1, seed=0)
INFO 08-29 20:48:52 tokenizer.py:29] For some LLaMA-based models, initializing the fast tokenizer may take a long time. To eliminate the initialization time, consider using 'hf-internal-testing/llama-tokenizer' instead of the original tokenizer.
INFO 08-29 20:49:14 llm_engine.py:196] # GPU blocks: 1981, # CPU blocks: 512
```
I use 54584 prompts, it stucks at llm_engine.py:196 for 30 more minutes.....did anyone  know how to solve this?

Hints from issue discussion:
Hi @cosmicexotic, thanks for trying out vLLM and reporting the error. could you share your environment? E.g., which GPU(s) are you using? How much CPU memory do you have? Are you running vLLM in a container?

Also, could you try vLLM with a small number of prompts (say 10) and see if it still hangs?
> Hi @cosmicexotic, thanks for trying out vLLM and reporting the error. could you share your environment? E.g., which GPU(s) are you using? How much CPU memory do you have? Are you running vLLM in a container?
> 
> Also, could you try vLLM with a small number of prompts (say 10) and see if it still hangs?

Thanks for replying , I use 2 Tesla V100. Total CPU memory is 131472812 kB. I tried 1,000 prompts and it works well. But when it comes to 10,000 prompts , it hangs like above.
@cosmicexotic Got it... that is weird. How long are your prompts? BTW, did you use `n` > 1 for the sampling parameter?
> @cosmicexotic Got it... that is weird. How long are your prompts? BTW, did you use `n` > 1 for the sampling parameter?

Some of my prompts are over 4096 tokens, but most of them(90%) under 4096 tokens( Llama2 max token), my sampling parameter: sampling_params = SamplingParams(temperature=0, top_k=-1) where n equals 1 by default.

@cosmicexotic, thanks for sharing it. Could you also share how you set `max_num_batched_tokens`? This should be larger than any of your prompts.
> @cosmicexotic, thanks for sharing it. Could you also share how you set `max_num_batched_tokens`? This should be larger than any of your prompts.

I set `max_num_batched_tokens` to 4096. When  I test it on 1,000 prompts, when the prompt exceed this number, I receive a warning. 
> I set max_num_batched_tokens to 4096. When I test it on 1,000 prompts, when the prompt exceed this number, I receive a warning.

@cosmicexotic Then could you set `max_num_batched_tokens` to 8192 and try again?
@WoosukKwon i have same issue when  set max_num_batched_tokens to 8192.
The more tokens permitted, the more portability error encountered in Sampler computing stage. And 4096 is fine.
i will post error message later.
@WoosukKwon there is my error message, when prompt exceed 8000(7000+ may also encounter this).
2023-08-31 17:02:19,066	ERROR worker.py:405 -- Unhandled error (suppress with 'RAY_IGNORE_UNHANDLED_ERRORS=1'): [36mray::RayWorker.execute_method()[39m (pid=2186511, ip=xxx.xxx.xxx.xxx, actor_id=114ccd0590b28ed93faa256458000000, repr=<vllm.engine.ray_utils.RayWorker object at 0x7fa619f64eb0>)
  File "/home/dell/vllm/vllm/engine/ray_utils.py", line 25, in execute_method
    return executor(*args, **kwargs)
  File "/home/dell/miniconda3/envs/ray/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/dell/vllm/vllm/worker/worker.py", line 293, in execute_model
    output = self.model(
  File "/home/dell/miniconda3/envs/ray/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/dell/vllm/vllm/model_executor/models/llama.py", line 261, in forward
    next_tokens = self.sampler(self.lm_head.weight, hidden_states,
  File "/home/dell/miniconda3/envs/ray/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/dell/vllm/vllm/model_executor/layers/sampler.py", line 89, in forward
    return _sample(probs, logprobs, input_metadata)
  File "/home/dell/vllm/vllm/model_executor/layers/sampler.py", line 391, in _sample
    next_token_ids = _sample_from_prompt(prob, sampling_params)
  File "/home/dell/vllm/vllm/model_executor/layers/sampler.py", line 305, in _sample_from_prompt
    next_token_ids = torch.multinomial(prob,
RuntimeError: probability tensor contains either `inf`, `nan` or element < 0

i see there is a PR (#936 ), although this two problems's error messages are same, but at different stage. And i try that fix, still error in sampler.
Hi @esmeetu, thanks for reporting the error. Could you share the name of your model and the exact prompt so that we can reproduce?

Commit message: Fix

Patch to review:
diff --git a/vllm/config.py b/vllm/config.py
index dd92fbccd899..4ebff0633f3a 100644
--- a/vllm/config.py
+++ b/vllm/config.py
@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Tuple
 
 import torch
 from transformers import PretrainedConfig
@@ -76,16 +76,20 @@ def __init__(
         self._verify_load_format()
         self._verify_tokenizer_mode()
         self._verify_quantization()
-        self.max_model_len = None
+        self._max_model_len = None
         if max_model_len is not None:
-            derived_max_model_len = self.get_max_model_len()
+            (derived_max_model_len,
+             derived_max_model_len_key) = self._get_max_model_len()
             if max_model_len > derived_max_model_len:
-                logger.warning(
+                raise ValueError(
                     f"User-specified max_model_len ({max_model_len}) is "
                     f"greater than the derived max_model_len "
-                    f"({derived_max_model_len}). Make sure the value is "
-                    "correct and within the model context size.")
-        self.max_model_len = max_model_len
+                    f"({derived_max_model_len_key}={derived_max_model_len} in "
+                    "model config.json). This will lead to "
+                    "incorrect model outputs or CUDA errors. "
+                    "Make sure the value is correct and within the "
+                    "model context size.")
+            self._max_model_len = max_model_len
 
     def _verify_load_format(self) -> None:
         load_format = self.load_format.lower()
@@ -167,9 +171,11 @@ def get_num_heads(self, parallel_config: "ParallelConfig") -> int:
         total_num_attention_heads = self.hf_config.num_attention_heads
         return total_num_attention_heads // parallel_config.tensor_parallel_size
 
-    def get_max_model_len(self) -> int:
-        if self.max_model_len is not None:
-            return self.max_model_len
+    def _get_max_model_len(self) -> Tuple[int, Optional[str]]:
+        """Returns the value of max_model_len and the
+        config key used to derive it."""
+        if self._max_model_len is not None:
+            return self._max_model_len, None
         max_model_len = float("inf")
         possible_keys = [
             # OPT
@@ -187,7 +193,12 @@ def get_max_model_len(self) -> int:
             max_len_key = getattr(self.hf_config, key, None)
             if max_len_key is not None:
                 max_model_len = min(max_model_len, max_len_key)
-        return max_model_len
+        # Cache the property.
+        self._max_model_len = max_model_len
+        return max_model_len, max_len_key
+
+    def get_max_model_len(self) -> int:
+        return self._get_max_model_len()[0]
 
     def get_num_layers(self, parallel_config: "ParallelConfig") -> int:
         total_num_hidden_layers = self.hf_config.num_hidden_layers
diff --git a/vllm/model_executor/models/aquila.py b/vllm/model_executor/models/aquila.py
index 1551974112b2..caf9b61ffa0e 100644
--- a/vllm/model_executor/models/aquila.py
+++ b/vllm/model_executor/models/aquila.py
@@ -105,6 +105,8 @@ def __init__(
         hidden_size: int,
         num_heads: int,
         num_kv_heads: int,
+        rope_theta: float = 10000,
+        max_position_embeddings: int = 8192,
     ):
         super().__init__()
         self.hidden_size = hidden_size
@@ -119,6 +121,8 @@ def __init__(
         self.q_size = self.num_heads * self.head_dim
         self.kv_size = self.num_kv_heads * self.head_dim
         self.scaling = self.head_dim**-0.5
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
 
         self.qkv_proj = ColumnParallelLinear(
             hidden_size,
@@ -140,6 +144,8 @@ def __init__(
             self.head_dim,
             self.scaling,
             rotary_dim=self.head_dim,
+            base=self.rope_theta,
+            max_position=self.max_position_embeddings,
         )
 
     def forward(
@@ -164,10 +170,15 @@ class AquilaDecoderLayer(nn.Module):
     def __init__(self, config: AquilaConfig):
         super().__init__()
         self.hidden_size = config.hidden_size
+        rope_theta = getattr(config, "rope_theta", 10000)
+        max_position_embeddings = getattr(config, "max_position_embeddings",
+                                          8192)
         self.self_attn = AquilaAttention(
             hidden_size=self.hidden_size,
             num_heads=config.num_attention_heads,
             num_kv_heads=config.num_attention_heads,
+            rope_theta=rope_theta,
+            max_position_embeddings=max_position_embeddings,
         )
         self.mlp = AquilaMLP(
             hidden_size=self.hidden_size,
diff --git a/vllm/model_executor/models/baichuan.py b/vllm/model_executor/models/baichuan.py
index 17e971d7bb29..277b2cc49b44 100644
--- a/vllm/model_executor/models/baichuan.py
+++ b/vllm/model_executor/models/baichuan.py
@@ -111,6 +111,8 @@ def __init__(
         hidden_size: int,
         num_heads: int,
         position_embedding: str,
+        rope_theta: float = 10000,
+        max_position_embeddings: int = 8192,
     ):
         super().__init__()
         self.hidden_size = hidden_size
@@ -122,6 +124,8 @@ def __init__(
                           tensor_model_parallel_world_size)
         self.head_dim = hidden_size // self.total_num_heads
         self.postion_embedding = position_embedding
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
 
         # pylint: disable=invalid-name
         self.W_pack = ColumnParallelLinear(
@@ -151,10 +155,13 @@ def __init__(
                                                 scaling, alibi_slopes)
         else:
             self.scaling = self.head_dim**-0.5
-            self.attn = PagedAttentionWithRoPE(self.num_heads,
-                                               self.head_dim,
-                                               self.scaling,
-                                               rotary_dim=self.head_dim)
+            self.attn = PagedAttentionWithRoPE(
+                self.num_heads,
+                self.head_dim,
+                self.scaling,
+                rotary_dim=self.head_dim,
+                base=self.rope_theta,
+                max_position=self.max_position_embeddings)
 
     def forward(
         self,
@@ -183,10 +190,15 @@ class BaiChuanDecoderLayer(nn.Module):
     def __init__(self, config: BaiChuanConfig, position_embedding: str):
         super().__init__()
         self.hidden_size = config.hidden_size
+        rope_theta = getattr(config, "rope_theta", 10000)
+        max_position_embeddings = getattr(config, "max_position_embeddings",
+                                          8192)
         self.self_attn = BaiChuanAttention(
             hidden_size=self.hidden_size,
             num_heads=config.num_attention_heads,
             position_embedding=position_embedding,
+            rope_theta=rope_theta,
+            max_position_embeddings=max_position_embeddings,
         )
         self.mlp = BaiChuanMLP(
             hidden_size=self.hidden_size,
diff --git a/vllm/model_executor/models/falcon.py b/vllm/model_executor/models/falcon.py
index dbd8a8203e4b..7aca60c30b9d 100644
--- a/vllm/model_executor/models/falcon.py
+++ b/vllm/model_executor/models/falcon.py
@@ -161,12 +161,18 @@ def __init__(self, config: FalconConfig):
             "Rotary and alibi are mutually exclusive.")
 
         if self.use_rotary:
+            rope_theta = getattr(config, "rope_theta", 10000)
+            max_position_embeddings = getattr(config,
+                                              "max_position_embeddings", 8192)
             # TODO(zhuohan): Pass in correct `max_position``
-            self.attn = PagedAttentionWithRoPE(self.num_heads,
-                                               self.head_dim,
-                                               self.inv_norm_factor,
-                                               rotary_dim=self.head_dim,
-                                               num_kv_heads=self.num_kv_heads)
+            self.attn = PagedAttentionWithRoPE(
+                self.num_heads,
+                self.head_dim,
+                self.inv_norm_factor,
+                base=rope_theta,
+                max_position=max_position_embeddings,
+                rotary_dim=self.head_dim,
+                num_kv_heads=self.num_kv_heads)
         elif self.use_alibi:
             tp_rank = get_tensor_model_parallel_rank()
             head_start = tp_rank * self.num_heads
diff --git a/vllm/model_executor/models/gpt_j.py b/vllm/model_executor/models/gpt_j.py
index c3e8da239ace..f8ffcdb7189a 100644
--- a/vllm/model_executor/models/gpt_j.py
+++ b/vllm/model_executor/models/gpt_j.py
@@ -67,11 +67,17 @@ def __init__(self, config: GPTJConfig):
         scaling = self.head_size**-0.5
         assert getattr(config, "rotary", True)
         assert config.rotary_dim % 2 == 0
-        self.attn = PagedAttentionWithRoPE(self.num_heads,
-                                           self.head_size,
-                                           scaling,
-                                           config.rotary_dim,
-                                           is_neox_style=False)
+        rope_theta = getattr(config, "rope_theta", 10000)
+        max_position_embeddings = getattr(config, "max_position_embeddings",
+                                          8192)
+        self.attn = PagedAttentionWithRoPE(
+            self.num_heads,
+            self.head_size,
+            scaling,
+            config.rotary_dim,
+            base=rope_theta,
+            max_position=max_position_embeddings,
+            is_neox_style=False)
         self.warmup = False
 
     def forward(
diff --git a/vllm/model_executor/models/gpt_neox.py b/vllm/model_executor/models/gpt_neox.py
index acbd1b47d626..225726a630cf 100644
--- a/vllm/model_executor/models/gpt_neox.py
+++ b/vllm/model_executor/models/gpt_neox.py
@@ -68,8 +68,16 @@ def __init__(self, config: GPTNeoXConfig):
         scaling = self.head_size**-0.5
         rotary_dim = int(self.head_size * config.rotary_pct)
         assert rotary_dim % 2 == 0
-        self.attn = PagedAttentionWithRoPE(self.num_heads, self.head_size,
-                                           scaling, rotary_dim)
+        rope_theta = getattr(config, "rope_theta", 10000)
+        max_position_embeddings = getattr(config, "max_position_embeddings",
+                                          8192)
+        self.attn = PagedAttentionWithRoPE(
+            self.num_heads,
+            self.head_size,
+            scaling,
+            rotary_dim,
+            base=rope_theta,
+            max_position=max_position_embeddings)
 
     def forward(
         self,
diff --git a/vllm/model_executor/models/internlm.py b/vllm/model_executor/models/internlm.py
index fdcac02a1b27..55bd76be0140 100644
--- a/vllm/model_executor/models/internlm.py
+++ b/vllm/model_executor/models/internlm.py
@@ -59,6 +59,8 @@ def __init__(
         self,
         hidden_size: int,
         num_heads: int,
+        rope_theta: float = 10000,
+        max_position_embeddings: int = 8192,
     ):
         super().__init__()
         self.hidden_size = hidden_size
@@ -70,6 +72,8 @@ def __init__(
                           tensor_model_parallel_world_size)
         self.head_dim = hidden_size // self.total_num_heads
         self.scaling = self.head_dim**-0.5
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
 
         self.qkv_proj = ColumnParallelLinear(
             hidden_size,
@@ -85,10 +89,13 @@ def __init__(
             input_is_parallel=True,
             perform_initialization=False,
         )
-        self.attn = PagedAttentionWithRoPE(self.num_heads,
-                                           self.head_dim,
-                                           self.scaling,
-                                           rotary_dim=self.head_dim)
+        self.attn = PagedAttentionWithRoPE(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            base=self.rope_theta,
+            max_position=self.max_position_embeddings,
+            rotary_dim=self.head_dim)
 
     def forward(
         self,
@@ -112,9 +119,14 @@ class InternLMDecoderLayer(nn.Module):
     def __init__(self, config: LlamaConfig):
         super().__init__()
         self.hidden_size = config.hidden_size
+        rope_theta = getattr(config, "rope_theta", 10000)
+        max_position_embeddings = getattr(config, "max_position_embeddings",
+                                          8192)
         self.self_attn = InternLMAttention(
             hidden_size=self.hidden_size,
             num_heads=config.num_attention_heads,
+            rope_theta=rope_theta,
+            max_position_embeddings=max_position_embeddings,
         )
         self.mlp = InternLMMLP(
             hidden_size=self.hidden_size,
diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py
index 0b7f4181a150..79b7bed26727 100644
--- a/vllm/model_executor/models/llama.py
+++ b/vllm/model_executor/models/llama.py
@@ -92,6 +92,7 @@ def __init__(
         num_heads: int,
         num_kv_heads: int,
         rope_theta: float = 10000,
+        max_position_embeddings: int = 8192,
         quant_config: Optional[QuantizationConfig] = None,
     ) -> None:
         super().__init__()
@@ -108,6 +109,7 @@ def __init__(
         self.kv_size = self.num_kv_heads * self.head_dim
         self.scaling = self.head_dim**-0.5
         self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
 
         self.qkv_proj = ParallelLinear.column(
             hidden_size,
@@ -126,12 +128,14 @@ def __init__(
             perform_initialization=False,
             quant_config=quant_config,
         )
-        self.attn = PagedAttentionWithRoPE(self.num_heads,
-                                           self.head_dim,
-                                           self.scaling,
-                                           base=self.rope_theta,
-                                           rotary_dim=self.head_dim,
-                                           num_kv_heads=self.num_kv_heads)
+        self.attn = PagedAttentionWithRoPE(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            base=self.rope_theta,
+            max_position=self.max_position_embeddings,
+            rotary_dim=self.head_dim,
+            num_kv_heads=self.num_kv_heads)
 
     def forward(
         self,
@@ -161,11 +165,14 @@ def __init__(
         self.hidden_size = config.hidden_size
         # Requires transformers > 4.32.0
         rope_theta = getattr(config, "rope_theta", 10000)
+        max_position_embeddings = getattr(config, "max_position_embeddings",
+                                          8192)
         self.self_attn = LlamaAttention(
             hidden_size=self.hidden_size,
             num_heads=config.num_attention_heads,
             num_kv_heads=config.num_key_value_heads,
             rope_theta=rope_theta,
+            max_position_embeddings=max_position_embeddings,
             quant_config=quant_config,
         )
         self.mlp = LlamaMLP(
diff --git a/vllm/model_executor/models/qwen.py b/vllm/model_executor/models/qwen.py
index 4ce9aea5e2c7..f572edb41db8 100644
--- a/vllm/model_executor/models/qwen.py
+++ b/vllm/model_executor/models/qwen.py
@@ -76,8 +76,13 @@ def forward(self, x):
 
 class QWenAttention(nn.Module):
 
-    def __init__(self, hidden_size: int, num_heads: int,
-                 max_position_embeddings: int):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        max_position_embeddings: int,
+        rope_theta: float = 10000,
+    ):
         super().__init__()
         self.hidden_size = hidden_size
         tensor_model_parallel_world_size = get_tensor_model_parallel_world_size(
@@ -109,6 +114,7 @@ def __init__(self, hidden_size: int, num_heads: int,
             self.head_dim,
             self.scaling,
             rotary_dim=self.head_dim,
+            base=rope_theta,
             max_position=max_position_embeddings,
         )
 
@@ -137,8 +143,11 @@ def __init__(self, config: QWenConfig):
         super().__init__()
         self.ln_1 = RMSNorm(config.n_embd, eps=config.layer_norm_epsilon)
 
-        self.attn = QWenAttention(config.n_embd, config.num_attention_heads,
-                                  config.max_position_embeddings)
+        rope_theta = getattr(config, "rope_theta", 10000)
+        self.attn = QWenAttention(config.n_embd,
+                                  config.num_attention_heads,
+                                  config.max_position_embeddings,
+                                  rope_theta=rope_theta)
 
         self.ln_2 = RMSNorm(config.n_embd, eps=config.layer_norm_epsilon)
 


Code review comments:
diff_hunk: @@ -167,9 +171,11 @@ def get_num_heads(self, parallel_config: "ParallelConfig") -> int:
         total_num_attention_heads = self.hf_config.num_attention_heads
         return total_num_attention_heads // parallel_config.tensor_parallel_size
 
-    def get_max_model_len(self) -> int:
-        if self.max_model_len is not None:
-            return self.max_model_len
+    def _get_max_model_len(self) -> Tuple[int, Optional[str]]:
+        """Returns the value of max_model_len and the
+        config key used to derive it."""
path: vllm/config.py
text: @author:
```suggestion
        """Returns the value of max_model_len and the associated key."""
```

diff_hunk: @@ -167,9 +171,11 @@ def get_num_heads(self, parallel_config: "ParallelConfig") -> int:
         total_num_attention_heads = self.hf_config.num_attention_heads
         return total_num_attention_heads // parallel_config.tensor_parallel_size
 
-    def get_max_model_len(self) -> int:
-        if self.max_model_len is not None:
-            return self.max_model_len
+    def _get_max_model_len(self) -> Tuple[int, Optional[str]]:
+        """Returns the value of max_model_len and the
+        config key used to derive it."""
path: vllm/config.py
text: @user1:
From google style guide:

> A docstring should be organized as a summary line (one physical line not exceeding 80 characters) terminated by a period, question mark, or exclamation point. When writing more (encouraged), this must be followed by a blank line, followed by the rest of the docstring starting at the same cursor position as the first quote of the first line. There are more formatting guidelines for docstrings below.

Could you make it to a single line?