[transformers] Scaled Dot Product Attention

llama구조를 참고해서 scaled dot product attention 계산 과정을 메모용으로 간단히 정리해보자.

 

먼저 필요한 라이브러리와 파라미터를 세팅하자.

import torch
from torch import nn
import transformers
import math
import torch.nn.functional as F

batch_size = 4
max_seq_len = 500
n_heads = 8
n_kv_heads = 4
embed_dim = 2048

head_dim = embed_dim // num_heads  # 256

vocab_size = 1000

 

입력 토큰에 대한 임베딩을 아래와 같이 가정하자.

wte = torch.nn.Embedding(vocab_size, embed_dim)
token_ids = torch.tensor([[4,6,10,366,10,1,4,5, 12]], dtype=torch.int64)
imput_embeds = wte(token_ids)

imput_embeds.shape
>>>>
torch.Size([1, 9, 2048])

 

 

llama3 모델(참고)에서 forward되는 부분에서 attention이 계산되는 과정을 간략히 정리해 보자.

먼저 각 token별로 q,k,v를 뽑는다.

bsz, seqlen, _ = imput_embeds.shape

# 실제 학습되는 layer
wq = nn.Linear(embed_dim, n_heads * head_dim, bias=False)
wk = nn.Linear(embed_dim, n_heads * head_dim, bias=False)
wv = nn.Linear(embed_dim, n_heads * head_dim, bias=False)

# q,k,v 계산
query = wq(imput_embeds)
key = wk(imput_embeds)
value = wv(imput_embeds)

# multi-head attention 계산을 위해 n_heads 만큼 쪼개준다.
query = query.view(bsz, seqlen, n_heads, head_dim)
key = key.view(bsz, seqlen, n_heads, head_dim)
value = value.view(bsz, seqlen, n_heads, head_dim)

query.shape, key.shape, value.shape
>>>>
(torch.Size([1, 9, 8, 256]),
 torch.Size([1, 9, 8, 256]),
 torch.Size([1, 9, 8, 256]))

 

scaled dot product attention이 계산되는 과정은 다음과 같다.

query = query.transpose(1,2)  # (bsz, n_heads, seqlen, head_dim)
key   = key.transpose(1,2)    # (bsz, n_heads, seqlen, head_dim)
value = value.transpose(1,2)  # (bsz, n_heads, seqlen, head_dim)

query.shape, key.shape, value.shape
>>>>
(torch.Size([1, 8, 9, 256]),
 torch.Size([1, 8, 9, 256]),
 torch.Size([1, 8, 9, 256]))
 
# n_heads 만큼 쪼개진 상태에서 head별로 attention 계산
attn_scores = torch.matmul(query, key.transpose(2,3)) / math.sqrt(head_dim)
attn_scores = F.softmax(attn_scores, dim=-1).type_as(query)
attn_scores.shape
>>>>
torch.Size([1, 8, 9, 9])

output = torch.matmul(attn_scores, value)  # (bsz, n_heads, seqlen, head_dim)
output.shape
>>>>
torch.Size([1, 8, 9, 256])

# 최종적으로 모든 head를 합쳐서 bsz x seqlen x embed_dim으로 변경
output = output.transpose(1,2).contiguous().view(bsz, seqlen, -1)
output.shape
>>>>
torch.Size([1, 9, 2048])