transformer/flops.py at main · RohanKhanBD/transformer · GitHub

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
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
# MLA flops
def mla_attention_flops(
    maxlen: int,
    embedding_dim: int,
    num_heads: int,
    qk_rope_dim: int,
    qk_nope_dim: int,
    kv_rank: int,
    v_dim: int,
):
    qk_head_dim = qk_rope_dim + qk_nope_dim
    # pre-attention
    query_linear = 2 * maxlen * embedding_dim * (qk_head_dim * num_heads)
    compressed_kv_linear = 2 * maxlen * embedding_dim * (kv_rank + qk_rope_dim)
    dcompressed_kv_linear = 2 * maxlen * kv_rank * (num_heads * (qk_nope_dim + v_dim))
    # actual attention
    atten = 2 * num_heads * (maxlen * qk_head_dim * maxlen)
    output = 2 * num_heads * (maxlen * maxlen * v_dim)
    proj = 2 * num_heads * (maxlen * v_dim * embedding_dim)
    return (
        query_linear
        + compressed_kv_linear
        + dcompressed_kv_linear
        + atten
        + output
        + proj
    )


def attention_gqa_flops(
    maxlen: int,
    embedding_dim: int,
    num_heads: int,
    kv_heads: int,
):
    head_dim = embedding_dim // num_heads
    q_proj = 2 * maxlen * embedding_dim * (num_heads * head_dim)
    k_proj = 2 * maxlen * embedding_dim * (kv_heads * head_dim)
    v_proj = k_proj
    qk = 2 * num_heads * maxlen * maxlen * head_dim
    softmax = 3 * num_heads * maxlen * maxlen
    attn_v = 2 * num_heads * maxlen * maxlen * head_dim
    out_proj = 2 * maxlen * (num_heads * head_dim) * embedding_dim
    return q_proj + k_proj + v_proj + qk + softmax + attn_v + out_proj


def ffn_flops(maxlen: int, embedding_dim: int, inter_dim: int):
    return 2 * maxlen * (2 * embedding_dim * inter_dim)


def block_flops(
    maxlen: int,
    embedding_dim: int,
    inter_dim: int,
    num_heads: int,
    kv_heads: int,
    use_mla: bool,
    qk_rope_dim: int,
    qk_nope_dim: int,
    kv_rank: int,
    v_dim: int,
):
    return (
        mla_attention_flops(
            maxlen, embedding_dim, num_heads, qk_rope_dim, qk_nope_dim, kv_rank, v_dim
        )
        if use_mla
        else attention_gqa_flops(maxlen, embedding_dim, num_heads, kv_heads)
        + ffn_flops(maxlen, embedding_dim, inter_dim)
    )


def transformer_flops(
    vocab_size: int,
    maxlen: int,
    embedding_dim: int,
    inter_dim: int,
    num_heads: int,
    kv_heads: int,
    n_layers: int,
    use_mla: bool,
    qk_rope_dim: int,
    qk_nope_dim: int,
    kv_rank: int,
    v_dim: int,
):
    block = n_layers * block_flops(
        maxlen,
        embedding_dim,
        inter_dim,
        num_heads,
        kv_heads,
        use_mla,
        qk_rope_dim,
        qk_nope_dim,
        kv_rank,
        v_dim,
    )
    logits = 2 * maxlen * embedding_dim * vocab_size
    fwf = block + logits
    bwf = 2 * fwf
    return fwf + bwf


# Fake it till you make it flops
def palm_flops(N, n_layers, num_heads, head_dim, maxlen):
    L, H, Q, T = n_layers, num_heads, head_dim, maxlen
    mf_per_token = 6 * N + 12 * L * H * Q * T
    mf = mf_per_token * maxlen
    return mf


# Total MLA params
def mla_params(
    embedding_dim: int,
    num_heads: int,
    qk_rope_dim: int,
    qk_nope_dim: int,
    kv_rank: int,
    v_dim: int,
):
    qk_head_dim = qk_rope_dim + qk_nope_dim
    query = embedding_dim * num_heads * qk_head_dim
    compressed_kv = embedding_dim * (kv_rank + qk_rope_dim)
    rms_norm = kv_rank
    dcompressed_kv = kv_rank * num_heads * (qk_nope_dim + v_dim)
    proj = num_heads * v_dim * embedding_dim
    return query + compressed_kv + rms_norm + dcompressed_kv + proj


def attention_gqa_params(embedding_dim: int, num_heads: int, kv_heads: int):
    head_dim = embedding_dim // num_heads
    query = embedding_dim * num_heads * head_dim
    key = embedding_dim * num_heads * kv_heads
    value = key
    proj = num_heads * head_dim * embedding_dim
    q_norm = k_norm = head_dim
    return query + key + value + proj + q_norm + k_norm


def ffn_params(embedding_dim: int, inter_dim: int):
    return 2 * embedding_dim * inter_dim


def block_params(
    embedding_dim: int,
    inter_dim: int,
    num_heads: int,
    kv_heads: int,
    use_mla: bool,
    qk_rope_dim: int,
    qk_nope_dim: int,
    kv_rank: int,
    v_dim: int,
):
    rms_norm = embedding_dim * 2
    return (
        mla_params(embedding_dim, num_heads, qk_rope_dim, qk_nope_dim, kv_rank, v_dim)
        if use_mla
        else attention_gqa_params(embedding_dim, num_heads, kv_heads)
        + ffn_params(embedding_dim, inter_dim)
        + rms_norm
    )


def transformer_params(
    vocab_size: int,
    embedding_dim: int,
    inter_dim: int,
    num_heads: int,
    kv_heads: int,
    n_layers: int,
    use_mla: bool,
    qk_rope_dim: int,
    qk_nope_dim: int,
    kv_rank: int,
    v_dim: int,
):
    tokemb = embedding_dim * vocab_size
    blocks = n_layers * block_params(
        embedding_dim,
        inter_dim,
        num_heads,
        kv_heads,
        use_mla,
        qk_rope_dim,
        qk_nope_dim,
        kv_rank,
        v_dim,
    )
    rms_norm = embedding_dim
    return tokemb + blocks + rms_norm