Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_fused_rotary_emb.py 4.51 KB
Newer Older
zhouxiang's avatar
0.2.6版本新增文件补充  
zhouxiang committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142 
import pytest
import torch
from torch import nn

from lmdeploy.pytorch.kernels.fused_rotary_emb import fused_rotary_emb


class DummyRotaryEmbedding(nn.Module):

    def __init__(self,
                 dim,
                 max_position_embeddings=2048,
                 base=10000,
                 device=None):
        super().__init__()
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        inv_freq = 1.0 / (self.base**(torch.arange(
            0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
        self.register_buffer('inv_freq', inv_freq, persistent=False)

    def forward(self, x, position_ids, seq_len=None):
        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(
            position_ids.shape[0], -1, 1)
        position_ids_expanded = position_ids[:, None, :].float()
        freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos().to(dtype=x.dtype)
        sin = emb.sin().to(dtype=x.dtype)
        # backwards compatibility
        return cos, sin


class DummyLinearScalingRotaryEmbedding(DummyRotaryEmbedding):

    def __init__(self,
                 dim,
                 max_position_embeddings=2048,
                 base=10000,
                 device=None,
                 scaling_factor=1.0):
        self.scaling_factor = scaling_factor
        super().__init__(dim, max_position_embeddings, base, device)

    def forward(self, x, position_ids, seq_len=None):
        position_ids = position_ids.float() / self.scaling_factor
        cos, sin = super().forward(x, position_ids, seq_len)
        return cos, sin


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)


def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=2):
    """Applies Rotary Position Embedding to the query and key tensors."""
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


class TestFusedRotaryEmb:

    @pytest.fixture
    def dtype(self):
        yield torch.float16

    @pytest.fixture
    def batch_size(self):
        yield 2

    @pytest.fixture
    def head_dim(self):
        yield 64

    @pytest.fixture
    def q_num_heads(self):
        yield 4

    @pytest.fixture
    def k_num_heads(self):
        yield 2

    @pytest.fixture
    def seq_len(self):
        yield 100

    @pytest.fixture
    def q(self, batch_size, seq_len, q_num_heads, head_dim, dtype):
        yield torch.rand(batch_size,
                         seq_len,
                         q_num_heads,
                         head_dim,
                         dtype=dtype).to('cuda')

    @pytest.fixture
    def k(self, batch_size, seq_len, k_num_heads, head_dim, dtype):
        yield torch.rand(batch_size,
                         seq_len,
                         k_num_heads,
                         head_dim,
                         dtype=dtype).to('cuda')

    @pytest.fixture
    def position_ids(self, batch_size, seq_len):
        yield torch.randint(0, seq_len + 100, (batch_size, seq_len)).cuda()

    @pytest.fixture
    def rotary_emb(self, head_dim):
        yield DummyLinearScalingRotaryEmbedding(head_dim,
                                                scaling_factor=1.0).to('cuda')

    @pytest.fixture
    def gt(self, q, k, position_ids, rotary_emb):
        with torch.inference_mode():
            cos, sin = rotary_emb(q, position_ids)
            yield apply_rotary_pos_emb(q,
                                       k,
                                       cos,
                                       sin,
                                       position_ids=position_ids)

    def test_fused_rotary_emb(self, q, k, position_ids, rotary_emb, gt):
        inv_freq = rotary_emb.inv_freq
        scaling_factor = rotary_emb.scaling_factor

        with torch.inference_mode():
            outq, outk = fused_rotary_emb(q,
                                          k,
                                          position_ids,
                                          inv_freq,
                                          scaling_factor=scaling_factor)

        gtq, gtk = gt
        torch.testing.assert_close(outq, gtq, atol=1e-3, rtol=1e-5)
        torch.testing.assert_close(outk, gtk, atol=1e-3, rtol=1e-5)