[DeepSeek-V3] Different rotary embedding implementation between DeepSeek-AI and Transformers
Author: wwwjn
Pull request: no
Labels: name: ROPE

Counterpart issue : https://github.com/deepseek-ai/DeepSeek-V3/issues/938

**Describe the issue**
Hi team, I'm working on reproducing the great deepseek-v3 model on [torchtitan](https://github.com/pytorch/torchtitan/tree/refs/heads/main) . While I'm trying to run numerical verification, I noticed the rotary embedding in HF and this repo is different.

HF: https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/modeling_deepseek.py#L339
- In HF rotary embedding implementation, they explicitly permute (interleave the odd column and even columns) the q_pe / k_pe.

Deepseek-AI: https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/modeling_deepseek.py#L339
- In DeepSeek-AI's implementation, they didn't permute the q_pe / k_pe.
- In the `convert.py`, the script load HF weights as it's original order, and this script didn't permute the weights, to accommodate the ordering difference in `apply_rotary_embedding()`.

And this discrepancy will result in different mathematical results after the attention module of the first dense layer.

I want to double-check with the team if I missed something here. Thank you for your help in advance! cc @tianyu-l

**To Reproduce**
Environment: transformers ==4.54.0

I run following 2 runs with code: https://github.com/wwwjn/DeepSeek-V3 . I randomized the same inputs for both runs.
1.  a single forward pass using HuggingFace `transformers`, with weights from:https://huggingface.co/deepseek-ai/DeepSeek-V3-0324/tree/main .
    ```python hf_implementation/hf_implementation.py --num_layers 5 > hf_outputs.txt 2>&1```

2. Using this repo's model implementation and run a single forward pass (To simplify, I didn't use the     - First,  I used the `convert.py` to convert HF checkpoint weights. 
     ```python convert.py --hf-ckpt-path /data/users/jianiw/dsv3-weights/ --save-path /data/users/jianiw/dsv3-weights-5-layer/ --n-experts 256 --model-parallel 8```
    - Second, I run a single forward pass using
      ```torchrun --nnodes 1 --nproc-per-node 8 inference/run_single_forward.py --config inference/configs/config_671B.json > dsv3-output.txt 2>&1```


Here's the detailed numerical comparison I've seen:

<img width="917" height="488" alt="Image" src="https://github.com/user-attachments/assets/a1b3f865-d0ec-4ed1-a5c0-f039fe6879f2" />

<img width="917" height="628" alt="Image" src="https://github.com/user-attachments/assets/73cc92a1-eca0-4cd0-8543-fb049107b66b" />


**Expected behavior**
Expected behavior: After first dense layer's attention layer, the output should be almost the same (There might be slightly difference because of fp8 vs other pericision). 


**Additional context**
We observed the same issue for llama3 model before #30872, and get better understanding with @ArthurZucker 's help - The weights for llama3 on HuggingFace is permuted compared to original Meta's weights . so we need to manually permute the weights back to accommodate the rotary embedding implementation difference. Reference: https://github.com/pytorch/torchtitan/issues/335, https://github.com/pytorch/torchtitan/issues/1291#issuecomment-2997077080

Comments:
Hey! You had me worried for a while there! 
We have this https://github.com/huggingface/transformers/blob/main/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py#L401-L402 

which we added exactly for that! 
The documentation should probably be updated and if it is not a default, that should be the default! 

Hey @ArthurZucker Thank so much for your reply!  

Let use following terminology to different 2 rotary embedding implementation: 
-  Implementation No.1: `apply_rotary_embedding_interleaved()` ([Code](https://github.com/huggingface/transformers/blob/main/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py#L276)) : Permute q_pe and k_pe. 
- Implementation No.2: `apply_rotary_embedding()` ([code](https://github.com/huggingface/transformers/blob/main/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py#L211)): No permutation of q_pe, k_pe.  This rotary embedding is used in DeepSeek-AI github [repo](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/model.py).

Implementation No.1 and No.2 are **mathematically different** given the same k_pe/q_pe inputs, and should correspond to different wq_a / wq_b weights to achieve the same output result. 

In my test,  I used`config = AutoConfig.from_pretrained(local_hf_path, trust_remote_code=True); model = AutoModelForCausalLM.from_pretrained(local_hf_path, config=config, yyy)`. `transformers` leads to the same function([link](https://huggingface.co/deepseek-ai/DeepSeek-V3-Base/blob/main/modeling_deepseek.py#L339)) as Implementation No.1.

I have few follow up questions:
1. In reality, `transformers` will go to `if self.config.rope_interleave` branch and call implementation No.1, and  `rope_interleave` default should be True. Is this correct understanding? I checked the configuration [here](https://github.com/huggingface/transformers/blob/main/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py#L189)

2. Is the checkpoint weights from HF [repo](https://huggingface.co/deepseek-ai/DeepSeek-V3-Base/tree/main) corresponding to Implementation No.1 or Implementation No.2? If it corresponds to implementation No.1, is there any way I could permute back the weights to make it work with implementation No.2 (Deepseek-AI original implementation)?

Much appreciated!

Hey @ArthurZucker , thanks for your help and any updates here? :)


Hey! Sorry forgot to answer

1. Deepseek uses the [`complex`](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/model.py#L375) formulation, which used to be licenced by meta, which is why transformers went with Eleuther's equivalent implementation that does permute the q and k. BUT this is a bit less efficient, so for most of the recent models (like llama) we permuted q and k in the weights, this way we have a more efficient (but equivalently efficient to the complex formulation) implementation.
2. The original weights can be permutted for your needs of course. We have no power on the organization's checkpoints, my recommendation is to use a branch (you open a PR to that repo with permuted weights) and use it!

I am sorry for the confusion it created and very happy if you want to submit a PR to update our doc about that!

TLDR:
- if you use rotated weights, you don't need the permute
- deepseek original codebase does not need the permute because they use the complex formulation
- if you use original weights, but `transformers` eleuther based implementation of rope, you need to permute the weights or permute q, k on the fly (as it is done today, which is not super efficient 😉 )

@ArthurZucker  Thanks! We don't have access to deepseek's original weights as well , and we only have access to the weights on HF. 

So, to simplify, we should use one of the following options:
1. rotated weights <-> Don't permute q, k in RoPE (deepseek-AI codebase)
2. non-rotated weights <-> permute q, k in RoPE  (HF `transofmers` codebase)

So the deepseek-v3's weights on HF ([here](https://huggingface.co/deepseek-ai/DeepSeek-V3-Base/)) has NOT been  rotated, is this a correct understanding?

> deepseek original codebase does not need the permute because they use the complex formulation

Also have another question about this conclusion. I think the permutation is not cause by complex formulation or Eleuther's [formulation](https://blog.eleuther.ai/rotary-embeddings/). Both formulation are equivalent if they take the same q, k input.

The q, k need to be permute is because the wq / wk that used to generate q, k is different (the columns has been rearranged).

I also want to check if my understanding about this conclusion is correct.
 

Hi @ArthurZucker great Monday! Thank you so much for your constantly help and reply,  I want to quickly follow up and double confirm this question

> [@ArthurZucker](https://github.com/ArthurZucker) Thanks! We don't have access to deepseek's original weights as well , and we only have access to the weights on HF.
> 
> So, to simplify, we should use one of the following options:
> 
> 1. rotated weights <-> Don't permute q, k in RoPE (deepseek-AI codebase)
> 2. non-rotated weights <-> permute q, k in RoPE  (HF `transofmers` codebase)
> 
> So the deepseek-v3's weights on HF ([here](https://huggingface.co/deepseek-ai/DeepSeek-V3-Base/)) has NOT been rotated, is this a correct understanding?



Yes sorry! 
> We don't have access to deepseek's original weights as well , and we only have access to the weights on HF.

The original weights, AFAIK are the ones published on `huggingface` no? 

> So, to simplify, we should use one of the following options:
> rotated weights <-> Don't permute q, k in RoPE (deepseek-AI codebase)
> non-rotated weights <-> permute q, k in RoPE (HF transofmers codebase)

Absolutely! The only **thing** is https://huggingface.co/deepseek-ai/DeepSeek-R1/blob/main/modeling_deepseek.py#L364 the original code DOES permute! On the fly at that! 

Regarding: 
> Also have another question about this conclusion. I think the permutation is not cause by complex formulation or Eleuther's [formulation](https://blog.eleuther.ai/rotary-embeddings/). Both formulation are equivalent if they take the same q, k input.
>
>The q, k need to be permute is because the wq / wk that used to generate q, k is different (the columns has been rearranged).

Yes, but we HAVE to re-arrange them in the weights, otherwise we need to do it in the forward to have correct outputs.

TLDR of the entire issue:
- rotated weights <-> Eleuther formulation -> Don't permute q, k in RoPE (Llama in `transformers`)
- non rotated-weights <-> Llama licensed formulation (complex) -> Don't permute in ROPE (deepseek original codebase) 
- non-rotated weights <-> Eleuther formulation -> permute q, k in RoPE (Deepseek remote code codebase)

@wwwjn mega sorry for answering slowly, it's a very important issue but I am just completely under water! 


Thank you so much for the help and clarification! That helps a lot!!

I tried to write a minimal reproduce of what I observed before:  https://gist.github.com/wwwjn/381771de10d1a30b9f7ca2df433e04a3. The script is comparing the following 2 paths:
1. `transformers` and Deepseek remote code codebase: 
    - non-rotated weights <-> Eleuther formulation -> permute q, k in RoPE (Deepseek remote code codebase)
2. Deepseek original codebase:  
    - non rotated-weights <-> Llama licensed formulation (complex) -> Don't permute in ROPE 

By running the script, I could see these 2 implementation are mathematical NOT equal with some dimension (eg, seq_len=2048, num_heads=128, rope_head_dim=64).  Could you help quickly checking if this is expected?


Yes! 
`Implementations produce slightly different results.` this is absolutely expected!

This does not induce an error that propagates, so it's safe. 
One thing, when there is a layer norm after the rope, it can get messy (depending on the type, layer-wise etc you need to shift the distribution of the weights as well to match the inputs q and k that were shifted) 

Do you think we can close this issue? 

> Implementations produce slightly different results

Thank you so much for your help and answer. From the script above, I see there's huge differences in the 2 results, and this will directly influence the output of attention module, the output matrices are very different for the attention module

Ah then no 😅 there should not be huge difference


> Ah then no 😅 there should not be huge difference

Thanks! One step back: 

1. Is this code originally from Deepseek-AI or HuggingFace? https://huggingface.co/deepseek-ai/DeepSeek-R1/blob/main/modeling_deepseek.py#L364

2. Is the weights on HF rotated or non-rotated?  Based on our previous discuss, it should be non-rotated and originally from DeepSeek-AI, right?

1. The code on the hub is provided by deepseek AI yes. If you check it out, it does do the rotation, like we do in `transformers`
2. Weights are NOT rotated! 
🤗 


The code you are using is not the one we use in `transformers` for the cos and sin 

Also you cannot compare the way you do, q and k are bound to be different, what you need to compare is once you combine them for the final results, because then the rotation is negated

+ you need to transpose q and k to get the script right:
```python
from dataclasses import dataclass

import numpy as np

import torch
from torch import nn


def print_tensor_stats(name, tensor):
    mean = tensor.mean().item()
    std = tensor.std().item()
    min_val = tensor.min().item()
    max_val = tensor.max().item()
    print(
        f"{name} - Shape: {tensor.shape} Mean: {mean:.6f}, Min: {min_val:.6f}, Max: {max_val:.6f}, Std: {std:.6f}, First 10 values: {tensor.flatten()[:10].tolist()}"
    )


# ============================================= HuggingFace implementation =============================================
from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding, LlamaConfig


# Copied from transformers.models.llama.modeling_llama.rotate_half
def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
def apply_rotary_pos_emb_interleaved(q, k, cos, sin, position_ids, unsqueeze_dim=1):
    """Applies Rotary Position Embedding to the query and key tensors.

    Args:
        q (`torch.Tensor`): The query tensor.
        k (`torch.Tensor`): The key tensor.
        cos (`torch.Tensor`): The cosine part of the rotary embedding.
        sin (`torch.Tensor`): The sine part of the rotary embedding.
        position_ids (`torch.Tensor`):
            The position indices of the tokens corresponding to the query and key tensors. For example, this can be
            used to pass offsetted position ids when working with a KV-cache.
        unsqueeze_dim (`int`, *optional*, defaults to 1):
            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
    Returns:
        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
    #"""
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)

    b, h, s, d = q.shape
    q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
    k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)

    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


# ============================================= Complex implementation =============================================
def precompute_freqs_cis(rope_head_dim, seqlen, base=10000) -> torch.Tensor:
    """
    Precomputes frequency-based complex exponential values for rotary positional embeddings.

    Args:
        args (DeepSeekV3ModelArgs): Model arguments containing positional embedding parameters.

    Returns:
        torch.Tensor: Precomputed complex exponential values for positional embeddings.
    """
    # Use real deepseek-v3 args
    dim = rope_head_dim

    # Basic RoPE frequency calculation
    freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))

    # Create position indices
    t = torch.arange(seqlen)

    # Outer product: [positions] × [frequencies]
    freqs = torch.outer(t, freqs)

    # Convert to complex exponentials: e^(i*freq*pos)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis


def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
    """
    Applies rotary positional embeddings to the input tensor using complex number operations.

    Args:
        x (torch.Tensor): Input tensor with positional embeddings to be applied.
        freqs_cis (torch.Tensor): Precomputed complex exponential values for positional embeddings.

    Returns:
        torch.Tensor: Tensor with rotary embeddings applied.
    """
    dtype = x.dtype
    x = torch.view_as_complex(x.float().view(*x.shape[:-1], -1, 2))
    freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1))
    y = torch.view_as_real(x * freqs_cis).flatten(3)
    return y.to(dtype)


def compare_implementations(
    batch_size=2, seq_len=2048, num_heads=128, rope_head_dim=64
):
    """
    Compare the EleutherAI and complex number implementations of rotary position embeddings.

    Head_dim:  "qk_rope_head_dim": 64,
    """
    print(
        f"Comparing implementations with shape: batch={batch_size}, seq_len={seq_len}, heads={num_heads}, qk_rope_head_dim={rope_head_dim}"
    )

    # Create random q and k tensors [batch, seq_len, num_heads, head_dim] as input
    torch.manual_seed(42)
    q = torch.randn(batch_size, seq_len, num_heads, rope_head_dim)
    k = torch.randn(batch_size, seq_len, num_heads, rope_head_dim)
    v = torch.randn(batch_size, num_heads, seq_len, rope_head_dim)
    config = LlamaConfig(
        head_dim=rope_head_dim,
        num_attention_heads=num_heads,
        # rope_type="yarn",
        max_position_embeddings=seq_len,
    )
    # 1) HuggingFace implementation
    hf_rope = LlamaRotaryEmbedding(config)
    position_ids = torch.arange(
        0,
        seq_len,
        dtype=torch.long,
        device=q.device,
    )
    position_ids = position_ids.unsqueeze(0)
    cos, sin = hf_rope(q, position_ids)
    hf_q, hf_k = apply_rotary_pos_emb_interleaved(
        q.transpose(1, 2),
        k.transpose(1, 2),
        cos,
        sin,
        position_ids,
    )
    print(hf_q.shape, hf_k.shape, v.shape)

    qk = torch.matmul(hf_q, hf_k.transpose(2, 3))
    print(f"qk shape: {qk.shape}")
    # 2) complex implementation: DeepSeek-AI github version
    freqs_cis = precompute_freqs_cis(
        rope_head_dim=rope_head_dim,
        seqlen=seq_len,
    )
    complex_q = apply_rotary_emb(q, freqs_cis)
    complex_k = apply_rotary_emb(k, freqs_cis)
    print(complex_q.shape, complex_k.shape, v.shape)
    complex_qk = torch.matmul(complex_q.transpose(1, 2), complex_k.permute(0, 2, 3, 1))
    print(torch.testing.assert_close(qk, complex_qk, rtol=1e-5, atol=1e-5))


if __name__ == "__main__":
    print("Rotary Position Embedding Comparison")

    # The following combination didn't work
    # Compare implementations: original DeepSeek-V3 671B model parameters
    compare_implementations(batch_size=1, seq_len=2048, num_heads=128, rope_head_dim=64)
    compare_implementations(batch_size=1, seq_len=16, num_heads=8, rope_head_dim=4)

    # This combination works
    compare_implementations(batch_size=1, seq_len=16, num_heads=4, rope_head_dim=2)
```

> Also you cannot compare the way you do, q and k are bound to be different, what you need to compare is once you combine them for the final results, because then the rotation is negated

Thank you so much for your help! That's much clear now.  Where should I compare the results (if not comparing q,k ), and where does the rotation negated later in the code? Eg, at the end of attention module? 

@ArthurZucker  Thank you so much for detailed explanation! After update the script, I got the exact same results. Now I understand more about the difference between RoPE implementation. 

To summary,  we should not compare k_pe / q_pe, we should compare q @ k.T or q_pe @ k_pe.T

And also highlight @ArthurZucker 's insight again as a conclusion:

- rotated weights <-> Eleuther formulation -> Don't permute q, k in RoPE (Llama in transformers)
- non rotated-weights <-> Llama licensed formulation (complex) -> Don't permute in ROPE (deepseek original codebase)
- non-rotated weights <-> Eleuther formulation -> permute q, k in RoPE (Deepseek remote code codebase)


Nice, very happy that its clearer now! I think this was also something our community had trouble with, a deep answer was long due! 🤗