Repository: vllm-project/vllm

Issue:
Maybe Wrong implementation of AttentionWithRoPE for GPTJ and GPT-NeoX?
I think there may be a wrong implementation for GPTJ and GPT-NeoX when doing 'apply rotary embedding'.

Currently implemented `PagedAttentionWithRoPE` always use the whole `query` and `key`, which is compatible with models like `llama` and `baichuan`, however for GPTJ and GPT-NeoX they may only use part of query and key when doing rope.

If current implementation does not compatible with those two models, I would suggest for those models that not using the whole query and key when applying rotary embeddings, can have another attention class that inherit from `PagedAttentionWithRoPE` and do something like:
```python
query_rot, query_pass = self._prepare_tensor_for_rope(query, self.rotary_dim)
key_rot, key_pass = self._prepare_tensor_for_rope(key, self.rotary_dim)
pos_encoding_ops.rotary_embedding_neox(
    positions,
    query_rot,
    key_rot,
    self.rotary_dim,
    self.cos_sin_cache,
)
query = self._cat_tensor_after_rope(query_rot, query_pass)
key = self._cat_tensor_after_rope(key_rot, key_pass)
```
GPTJ output not consistent with that of transformers
Hey guys,

I've been trying to get the base GPT-J model (EleutherAI/gpt-j-6b) served using vLLM, but I am running into issues with the quality of the output. Using the same offline inference code as in https://github.com/vllm-project/vllm/blob/main/examples/offline_inference.py, here is what I obtain from the model (running on an RTX A5000):


``` python3
from vllm import LLM, SamplingParams

sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=128)
llm = LLM(model="EleutherAI/gpt-j-6b")
prompts = ["AI is going to"] * 4
outputs = [t.outputs[0].text for t in llm.generate(prompts, sampling_params)]

output = [" get a lot of new stuff. A lot of other new stuff\nthat's been added and things are going to come out.\nthat you know, it's going to be implemented, that's coming out.\nin the\nand we've got like a lot of stuff coming out, and things in the like AI stuff. So we're going to be doing.\nbefore the new stuff\nand some of the AI, the future.\nthe AI\nis added in the next year and we're going to the next year.\nthat's going to be in the next gen.\nand it's going to be added",
 " change the way we work, not just the way we do business, but also how we work and we live. How we work, the way we live. In the way we do our everyday activities. It will be. We think, what we live.\n\n\n\n\n\n\n\nthe way we interact with each other, but our living.\n\nthe way we relate to each other, but how we live. AI.\n\n\nit's not even our lives. The economy, live, and build our daily life.\nhow we play, but also our lives and engage with each other, and how we",
 " be our friend.\n\n\n\n\nHere's a good reason we're gonna do this\nAnd you know, and that was interesting.\nWe'll have to be, like,\nour employees are gonna start.\nOur\nThey did you know, as you know.\nTotally, we're gonna have like you know\nAnd that's the oil.\nSo that some other question, like, and that we're gonna\nuse and and I'm going to be in the most of the middle of\nThe boss and we're going to be on TV,\nTogether. I'm gonna be\nWe're gonna",
 " save us all. Some of them, anyway. It will. AI is almost certainly. It's just a utopianism is what we can see how much about as serious, but it will be a way out of the future. I mean.\n\n\nsmart, it's going to out how it will be their business is the father than we're going to the ultimate evolution. I think we're going to be the in a perfect, the machine.\n\nnot be a bit more than 3Dot-it's just it.\n\n\n\nkill us we have long as we're going to have it will be making"]
 ```

For comparison's sake, here is what I obtain using transformers:

```python3
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-j-6b").half().eval().to(0)
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6b")
inputs = tokenizer("AI is going to", return_tensors="pt").to(0)
params = {"temperature": 0.85, "top_p": 0.94, "max_new_tokens": 128, "do_sample": True}
outputs = []
for _ in range(1,4):
    out = model.generate(inputs["input_ids"], **params)
    outputs.append(tokenizer.decode(out[0]))
outputs = [
'AI is going to be our future and this is going to be a big part of that," said Peter Lee, vice president of Google\'s self-driving cars and mapping division.\n\n"It will allow us to make the most efficient use of public roads that are more congested and have more traffic," Lee told the BBC. "It will allow us to better respond to traffic conditions and accidents."\n\nIt\'s certainly a big deal, and one that Google will have to work its way through. The company is a pioneer in the area of driverless cars, but its software has been unable to operate as a taxi driver.\n\nIts self',
'AI is going to make a comeback. It's already making a comeback. The best technology to understand AI is not the technology itself. It's human understanding of what the technology is. We're just starting to understand how human beings think, and the way that human beings think is going to revolutionize the way we think.\n\nIt's hard for us to imagine the future that AI brings. We're seeing AI in every form, and we're seeing it becoming a reality in our lives. You may not see it right away, but it's right around the corner.\n\n',
'AI is going to be a game changer and a very important factor in the success of the entire Internet. It will not be just another technology, but a game changing technology. The way we do things has been changed.\n\nAnd we can't wait to help you succeed in the online world!\n\nOur Vision\n\nOur vision is to provide the best eCommerce platform that can help your business grow. We believe in the concept of "Customer First". We will work with you and your team to deliver the best online experience.\n\nWe will continue to innovate and develop new features, functionality and services to help you',
'AI is going to change everything. The world of medicine and the world of business will be forever changed.\n\nThe world of medicine will change drastically. Doctors will be able to diagnose illnesses and provide immediate treatment without needing to send their patients to specialists. A patient's condition will be known almost instantly and he/she will be able to receive the best possible treatment.\n\nThe world of business will change drastically as well. With the help of AI, companies can easily increase their efficiency, improve their products, and ultimately increase profits.\n\nWhy is AI so important?\n\nAI has become incredibly important because it can help us solve many'
]
```

I've noticed this behaviour independent of the choice of sampling parameters.

Any help in resloving this is greatly appreciated!




Discussion:
The implementation looks correct to me. `pos_encoding_kernels.cu` uses `int rot_dim = cos_sin_cache.size(1);`
which is determined by `rotary_dim` passed from
```
rotary_dim = int(self.head_size * config.rotary_pct)
self.attn = PagedAttentionWithRoPE(self.num_heads, self.head_size, scaling, rotary_dim)
```
in GPTNeoXAttention @ vllm/model_executor/models/gpt_neox.py

Strange inference results has especially been reported for GPTJ though: https://github.com/vllm-project/vllm/issues/590
I think one of mainly problem is how to rotate tensor. Referenced from HF transformers' implementation, for GPT-J it's [`rotate_every_two`](https://github.com/huggingface/transformers/blob/fe3c8ab1af558b95f67f5fafc0c55f09fd2b09db/src/transformers/models/gptj/modeling_gptj.py#L69) and for GPT-NeoX or LLaMa it's [`rotate_half`](https://github.com/huggingface/transformers/blob/fe3c8ab1af558b95f67f5fafc0c55f09fd2b09db/src/transformers/models/gpt_neox/modeling_gpt_neox.py#L368), which will cause different results.
I haven't read the .cu code yet. But if my understanding is correct, it should always get and apply pos embeds to the top rot_dim columns in the tensor's last dimension, if so I think GPT-NeoX's implementation should be correct.
It should be addressed, since vllm output from many user's side are not aligned with hf, from the result it become more stupid than hf's output.

Addition notes, I think one can claim a vLLM model's generation quality is worse than HF's only when they doing following things:
1. loading the model from .bin or .pt file instead of .safetensors file;
2. done a ppl benchmark and confirm there is a big difference between two models from different framework
3. using beam search decode strategy rather than sampling decode strategy
@PanQiWei  

How to make sure beam search enabled in vllm?
> @PanQiWei
> 
> How to make sure beam search enabled in vllm?

Turn the flag `use_beam_search` on in `SamplingParams` or your request payload. And make sure `n` > 1(thus vLLM not support greedy search)
@PanQiWei  Does n = 2 means using 2 beam search? Does it work in stream mode?
> @PanQiWei Does n = 2 means using 2 beam search? Does it work in stream mode?

sorry my bad, in openai api it should best_of=2 aka beam_size=2; and I don't think it work in stream mode.
@PanQiWei So there still have some bias between hf and vllm.
Hi! Can you try some deterministic sampling methods like greedy decoding and compare the results? To be honest I cannot easily tell the differences between the qualities of two sampling results.
Thanks for getting back to me!

@AlekseyKorshuk and I ran the inference with deterministic sampling and got the following results:

```python3
sampling_params = SamplingParams(temperature=0, top_k=-1, max_tokens=128)
llm = LLM(model="philschmid/gpt-j-6B-fp16-sharded")
prompts = ["Obama was born in Honolulu, Hawaii. After graduating from Columbia University in 1983, he worked as a community organizer in Chicago. In 1988, he enrolled in Harvard Law School, where he was the first black president of the Harvard Law Review. After graduating,"] * 4
outputs = [t.outputs[0].text for t in llm.generate(prompts, sampling_params)]
outputs = [
' he was a year later, he was a year. He was a student of the University of Chicago, he was a community organizer for the University of Chicago, he was a community organizer, he was a year, and then he was a year later, he was a year.\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n',
 ' he was a year later, he was a year. He was a student of the University of Chicago, he was a community organizer for the University of Chicago, he was a community organizer, he was a year, and then he was a year later, he was a year.\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n',
 ' he was a year later, he was a year. He was a student of the University of Chicago, he was a community organizer for the University of Chicago, he was a community organizer, he was a year, and then he was a year later, he was a year.\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n',
 ' he was a year later, he was a year. He was a student of the University of Chicago, he was a community organizer for the University of Chicago, he was a community organizer, he was a year, and then he was a year later, he was a year.\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n']
```

Using the same prompt and model on HF transformers (using do_sample=False), we got the following completion:
`Obama served as an attorney for the Illinois State Senate and then taught constitutional law at the University of Chicago.\n`

For us, VLLM has greatly improved in terms of inference speed, but the difference in the inference results due to the different sampling methods between VLLM and Transformer is quite confusing. Transformer has a good ecosystem and there is consensus on the sampling parameters. Can it support the option of choosing between VLLM sampling and Transformer sampling methods?
VLLM's outputs are definitely (significantly) more repetitive than HF inference (as you can see in Alex's results in the first post and the greedy decoding result). I think sampling process itself has no significant delta; it might be something in cuda kernels or attention.

Related topic I posted in discussions:
https://github.com/vllm-project/vllm/discussions/471
yeah it looks to me like GPTJ responses are poor quality when I tried it out.

Other models seem to have much better inference quality, although they can be more repetitive than HF so I worry they also have regression, albeit much more minor.
Same problem here with GPT-J.
I am using the default sampling params from vLLM. I also tried playing with all the params availble, but it did not help.

I am consistently getting some gibberish in the outputs.
Here is one example.

Input:
> Let me tell you a story:

Output:
> Let me tell you a story:\n\nIt’s midnight. Your eyes begin to fall right at noon, I came out of Scotland a summer, right at the dock where this story take in the sky blue.\n\nAM,\nNever felt feverich I never...
May be it's because GPT-J's tensor rotate logic is different to other models, as described in #747

Resolving patch:
diff --git a/vllm/model_executor/layers/attention.py b/vllm/model_executor/layers/attention.py
--- a/vllm/model_executor/layers/attention.py
+++ b/vllm/model_executor/layers/attention.py
@@ -242,7 +242,7 @@ def forward(
 
 
 class PagedAttentionWithRoPE(PagedAttention):
-    """PagedAttention with GPT-NeoX style rotary embedding."""
+    """PagedAttention with rotary embedding."""
 
     def __init__(
         self,
@@ -253,8 +253,10 @@ def __init__(
         max_position: int = 8192,
         base: int = 10000,
         num_kv_heads: Optional[int] = None,
+        is_neox_style: bool = True,
     ) -> None:
         super().__init__(num_heads, head_size, scale, num_kv_heads)
+        self.is_neox_style = is_neox_style
 
         # Create the cos and sin cache.
         inv_freq = 1.0 / (base**(torch.arange(0, rotary_dim, 2) / rotary_dim))
@@ -303,12 +305,13 @@ def forward(
 
         # Apply rotary embedding to the query and key before passing them
         # to the attention op.
-        pos_encoding_ops.rotary_embedding_neox(
+        pos_encoding_ops.rotary_embedding(
             positions,
             query,
             key,
             self.head_size,
             self.cos_sin_cache,
+            self.is_neox_style,
         )
         return super().forward(
             query,
diff --git a/vllm/model_executor/models/gpt_j.py b/vllm/model_executor/models/gpt_j.py
--- a/vllm/model_executor/models/gpt_j.py
+++ b/vllm/model_executor/models/gpt_j.py
@@ -67,8 +67,11 @@ 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)
+        self.attn = PagedAttentionWithRoPE(self.num_heads,
+                                           self.head_size,
+                                           scaling,
+                                           config.rotary_dim,
+                                           is_neox_style=False)
         self.warmup = False
 
     def forward(


Test patch:
diff --git a/tests/kernels/test_pos_encoding.py b/tests/kernels/test_pos_encoding.py
--- a/tests/kernels/test_pos_encoding.py
+++ b/tests/kernels/test_pos_encoding.py
@@ -7,49 +7,64 @@
 
 from vllm import pos_encoding_ops
 
+IS_NEOX_STYLE = [True, False]
 DTYPES = [torch.half, torch.bfloat16, torch.float]
 HEAD_SIZES = [64, 80, 96, 112, 128, 256]
 ROTARY_DIMS = [None, 32]  # None means rotary dim == head size
 NUM_HEADS = [7, 12, 40, 52]  # Arbitrary values for testing
-NUM_TOKENS = [7, 83, 2048]  # Arbitrary values for testing
+NUM_TOKENS = [11, 83, 2048]  # Arbitrary values for testing
 SEEDS = [0]
 
 
-def rotate_half(x: torch.Tensor) -> torch.Tensor:
+def rotate_neox(x: torch.Tensor) -> torch.Tensor:
     x1 = x[..., :x.shape[-1] // 2]
     x2 = x[..., x.shape[-1] // 2:]
     return torch.cat((-x2, x1), dim=-1)
 
 
-def apply_rotary_pos_emb(
+def rotate_gptj(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., ::2]
+    x2 = x[..., 1::2]
+    x = torch.stack((-x2, x1), dim=-1)
+    return x.flatten(-2)
+
+
+def apply_rope(
     q: torch.Tensor,
     k: torch.Tensor,
     cos: torch.Tensor,
     sin: torch.Tensor,
+    is_neox_style: bool,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    q_embed = (q * cos) + (rotate_half(q) * sin)
-    k_embed = (k * cos) + (rotate_half(k) * sin)
+    rotate_fn = rotate_neox if is_neox_style else rotate_gptj
+    q_embed = (q * cos) + (rotate_fn(q) * sin)
+    k_embed = (k * cos) + (rotate_fn(k) * sin)
     return q_embed, k_embed
 
 
-class RefRotaryEmbeddingNeox(nn.Module):
-    """Reference implementation of the GPT-NeoX style rotary embedding."""
+class RefRotaryEmbedding(nn.Module):
+    """Reference implementation of rotary embedding."""
 
     def __init__(
         self,
         dim: int,
-        max_position_embeddings: int = 2048,
+        is_neox_style: bool,
+        max_position_embeddings: int = 8192,
         base: int = 10000,
     ) -> None:
         super().__init__()
         self.rotary_dim = dim
+        self.is_neox_style = is_neox_style
         self.max_position_embeddings = max_position_embeddings
 
         # Create cos and sin embeddings.
         inv_freq = 1.0 / (base**(torch.arange(0, dim, 2) / dim))
         t = torch.arange(max_position_embeddings).float()
         freqs = torch.einsum("i,j->ij", t, inv_freq.float())
-        emb = torch.cat((freqs, freqs), dim=-1)
+        if is_neox_style:
+            emb = torch.cat((freqs, freqs), dim=-1)
+        else:
+            emb = torch.repeat_interleave(freqs, 2, -1)
         cos = emb.cos().to(dtype=inv_freq.dtype)
         sin = emb.sin().to(dtype=inv_freq.dtype)
         self.register_buffer("cos_cached", cos, persistent=False)
@@ -61,7 +76,6 @@ def forward(
         query: torch.Tensor,  # [num_tokens, num_heads, head_size]
         key: torch.Tensor,  # [num_tokens, num_heads, head_size]
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-
         query_rot = query[..., :self.rotary_dim]
         query_pass = query[..., self.rotary_dim:]
         key_rot = key[..., :self.rotary_dim]
@@ -71,7 +85,9 @@ def forward(
         key_rot = key_rot.transpose(0, 1)
         cos = F.embedding(positions, self.cos_cached)
         sin = F.embedding(positions, self.sin_cached)
-        query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin)
+
+        query_rot, key_rot = apply_rope(query_rot, key_rot, cos, sin,
+                                        self.is_neox_style)
         query_rot = query_rot.transpose(0, 1).contiguous()
         key_rot = key_rot.transpose(0, 1).contiguous()
 
@@ -82,6 +98,7 @@ def forward(
         return query, key
 
 
+@pytest.mark.parametrize("is_neox_style", IS_NEOX_STYLE)
 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
 @pytest.mark.parametrize("num_heads", NUM_HEADS)
 @pytest.mark.parametrize("head_size", HEAD_SIZES)
@@ -89,7 +106,8 @@ def forward(
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("seed", SEEDS)
 @torch.inference_mode()
-def test_rotary_embedding_neox(
+def test_rotary_embedding(
+    is_neox_style: bool,
     num_tokens: int,
     num_heads: int,
     head_size: int,
@@ -104,15 +122,15 @@ def test_rotary_embedding_neox(
     torch.random.manual_seed(seed)
     torch.cuda.manual_seed(seed)
 
-    positions = torch.randint(0, max_position, (num_tokens, ), device='cuda')
+    positions = torch.randint(0, max_position, (num_tokens, ), device="cuda")
     query = torch.randn(num_tokens,
                         num_heads * head_size,
                         dtype=dtype,
-                        device='cuda')
+                        device="cuda")
     key = torch.randn(num_tokens,
                       num_heads * head_size,
                       dtype=dtype,
-                      device='cuda')
+                      device="cuda")
 
     # Create the rotary embedding.
     inv_freq = 1.0 / (base**(torch.arange(0, rotary_dim, 2) / rotary_dim))
@@ -126,20 +144,22 @@ def test_rotary_embedding_neox(
     # Run the kernel. The kernel is in-place, so we need to clone the inputs.
     out_query = query.clone()
     out_key = key.clone()
-    pos_encoding_ops.rotary_embedding_neox(
+    pos_encoding_ops.rotary_embedding(
         positions,
         out_query,
         out_key,
         head_size,
         cos_sin_cache,
+        is_neox_style,
     )
 
     # Run the reference implementation.
-    ref_rotary_embedding = RefRotaryEmbeddingNeox(
+    ref_rotary_embedding = RefRotaryEmbedding(
         dim=rotary_dim,
+        is_neox_style=is_neox_style,
         max_position_embeddings=max_position,
         base=base,
-    ).to(dtype=dtype, device='cuda')
+    ).to(dtype=dtype, device="cuda")
     ref_query, ref_key = ref_rotary_embedding(
         positions,
         query.view(num_tokens, num_heads, head_size),