[Day3] 한권으로 끝내는 실전 LLM 파인튜닝 - 멀티헤드 어텐션과 피드포워드

다음은 [한권으로 끝내는 실전 LLM 파인튜닝] 스터디의 Day3 내용 정리입니다.

2.6 멀티헤드 어텐션과 피드포워드 

• 단일 헤드 어텐션의 한계를 극복하고 모델의 성능을 크게 향상시키는 핵심 구조. 

2.6.1 멀티헤드 어텐션 만들기

• 멀티헤드 어텐션 : 여러 개의 어텐션 메커니즘을 병렬로 사용해 다양한 관점에서 정보를 동시에 병렬처리 할 수 있도록 한다. 
• 각 헤드는 독립적으로 단일 데이터를 여러개의 다른 헤드를 통해 동시에 처리함. 그를 통해 데이터에 내재된 다양한 패턴과 관계를 더욱 정교하게 학습할 수 있게 됨. 

import torch.nn as nn
import torch.nn.functional as F

# 단일 head attention
class Head(nn.Module):
  def __init__(self, head_size):
    super().__init__()
    self.key = nn.Linear(n_embed, head_size, bias=False)
    self.query = nn.Linear(n_embed, head_size, bias=False)
    self.value = nn.Linear(n_embed, head_size, bias=False)
    self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))

  def forward(self, inputs):
    batch_size, seq_length, n_embed = inputs.shape
    keys = self.key(inputs)
    queries = self.query(inputs)
    weights = queries @ keys.transpose(-2, -1) * (embedding_dim ** -0.5)
    weights = weights.masked_fill(
        self.tril[:seq_length, :seq_length] == 0, float('-inf')
    )
    weights = F.softmax(weights, dim=-1)
    values = self.value(inputs)
    output = weights @ values
    return output

# multi head attention
class MultiHeadAttention(nn.Module):
  def __init__(self, num_heads, head_size):
    super().__init__()
    self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])

  def forward(self, inputs):
    return torch.cat([head(inputs) for head in self.heads], dim=1)

• num_heads : 입력 데이터를 동시에 처리할 방식의 수를 결정. 각 Head 인스턴스는 독립적으로 셀프 어텐션 연산을 수행.
• foward : 각 Head가 입력 데이터에 대해 독립적으로 어텐션 연산 수행, 그 결과들을 마지막 차원을 기준으로 연결함. [batch_size, sequence_length, num_heads * head_size] 형태를 띈다. 

2.6.2 피드포워드 만들기

• 어텐션 메커니즘을 통해 입력 시퀀스의 각 요소와 전체 시퀀스 간의 관계를 계산하고, 이는 입력 데이터의 전체 맥락을 파악하는데 중점을 둔다.
• 이후 피드포워드 네트워크를 통해 복잡한 패턴이나 비선형적 관계를 함께 학습시킴. 
• 각 어텐션 블록 뒤에 피드포워드 네트워크를 배치해 어텐션 메커니즘으로부터 얻은 표현을 더욱 풍부하게 만들어 모델이 더 복잡한 데이터 패턴을 학습할 수 있게 도와줌.

class FeedForward(nn.Module):
  def __init__(self, n_embed):
    super().__init__()
    self.layer = nn.Sequential(
        nn.Linear(n_embed, 4 * n_embed),
        nn.ReLU(),
        nn.Linear(4 * n_embed, n_embed)
        nn.DropOut(dropout),
    )
  
  def forward(self, input_tensor):
    return self.layer(input_tensor)

• __init__ : self.layer에 nn.Sequential을 사용해 3개의 레이어를 순차적으로 연결한 신경망을 정의함.
  첫번째 Linear()는 n_embed에서 중간차원 4*n_embed로 차원을 확장하는 선형변환.
  두번째 ReLU()는 비선형성을 도입
  세번쨰 Linear()는 4 * n_embed에서 원래 차원 n_embed로 차원을 축소하는 선형 변환.
  Dropout : 비선형 변환 이후 과적합 방지를 위해 dropout 층 추가. 
• forward : [batch_size, sequence_length, n_embed] 형태의 tensor_input을 받아 입력 데이터는 차원이 확장되고, ReLU를 거친 후 다시 원래의 차원으로 축소됨. 이렇게 입력 데이터에 비선형 변환을 적용해 모델이 복잡한 패턴과 관계를 학습할 수 있게 됨. 
• 결과적으로 FeedForward 클래스는 트랜스포머 모델의 각 어텐션 블록 뒤에 위치해 어텐션 메커니즘으로부터 얻은 정보를 추가로 처리함으로써 모델은 더 풍부하고 복잡한 표현을 학습가능하게 됨. 

2.7 _ Blocks 만들기

• Block은 모델의 설계와 구현에 중요한 구조적 단위임. 

GPT Block 구조

• 블록 구조는 모델 내 다양한 계층과 구성 요소를 하나로 묶어 모듈화, 재사용성, 확장성을 크게 향상한다. 
• 각 블록 내에서는 주로 어텐션 메커니즘과 피드포워드 네트워크가 수행되어 입력데이터로부터 더 복잡하고 추상적인 특징을 추출함. 이 구조를 통해 모델의 깊이와 복잡성을 조잘할 수 있으며 필요에 따라 블록을 추가하거나 구성을 변경해 성능 최적화가 가능함.

class Block(nn.Module):
  def __init__(self, n_embed, n_heads):
    super().__init__()
    head_size = n_embed // n_heads
    self.attention = MultiHeadAttention(n_heads, head_size)
    self.feed_forward = FeedForward(n_embed)
    self.layer_norm1 = nn.LayerNorm(n_embed)
    self.layer_norm2 = nn.LayerNorm(n_embed)

  def forward(self, input_tensor):
    input_tensor = input_tensor + self.attention(self.layer_norm1(input_tensor))
    input_tensor = input_tensor + self.feed_forward(self.layer_norm2(input_tensor))
    return input_tensor

• __init__ : head_size는 각 어텐션 헤드에서 이용할 차원의 크기로, n_embed/n_head 값으로 계산됨. 
• forward : 입력 데이터는 [batch_size, sequence_length, n_embed] 형태의 텐서
• 잔차 연결 : 깊이 신경망을 쌓을 경우 vanishing gradient 문제와 exploding gradient 문제가 발생하는 것을 해결하기 위해 제안된 기법. 각 레이어의 입력을 그 레이어의 출력에 직접 더해주는 구조. 즉, $F(x) + x$. 이 구조를 통해 신경망은 입력과 출력 사이의 잔차 함수를 학습하게 됨.
• 레이어 정규화 (Layer Norm) : 순환신경망 등의 동적 네트워크의 학습을 안정화하고 가속화하기 위해 개발됨. 각 층 내의 활성화 함수를 통과하는 모든 특성에 대해 정규화를 적용함. 이를 통하면 훈련 과정이 안정화되고, 학습 속도도 빨리지며 RNN에서 긴 시퀀스 데이터 처리의 성능이 개선됨.
  LayerNorm의 주요 장점은 입력 데이터의 배치 크기에 의존하지 않고 각 샘플 내에서 독립적으로 작동한다는 점이다. 이로 인해 동적 시퀀스 길이나 다양한 구조의 데이터 처리에 유용하게 사용됨. 

다음은 Multi Head Attention과 Feed Forward Network를 결합한 Block 구조를 가진 SemiGPT의 full code이다. 특정 모델 구조를 사용하지 않았기 때문에 성능이 그리 좋지는 않다. 

import torch
import torch.nn as nn
import torch.nn.functional as F

batch_size = 32
block_size = 8
max_iteration = 50000
eval_interval = 300
learning_rate = 1e-2
device = "cuda" if torch.cuda.is_available() else "cpu"
eval_iteration = 200
n_embed = 32
n_head = 4
n_layer = 4
dropout = 0.1

class Head(nn.Module):
    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embed, head_size, bias=False)
        self.query = nn.Linear(n_embed, head_size, bias=False)
        self.value = nn.Linear(n_embed, head_size, bias=False)
        self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))

    def forward(self, inputs):
        batch_size, sequence_length, embedding_dim = inputs.shape
        keys = self.key(inputs)
        queries = self.query(inputs)
        weights = queries @ keys.transpose(-2, -1) * (embedding_dim ** -0.5)
        weights = weights.masked_fill(self.tril[:sequence_length, :sequence_length] == 0, float("-inf"))
        weights = F.softmax(weights, dim=-1)
        values = self.value(inputs)
        output = weights @ values
        return output


def batch_function(mode):
    dataset = train_dataset if mode == "train" else test_dataset
    idx = torch.randint(len(dataset) - block_size, (batch_size,))
    x = torch.stack([dataset[index:index+block_size] for index in idx])
    y = torch.stack([dataset[index+1:index+block_size+1] for index in idx])
    x, y = x.to(device), y.to(device)
    return x, y


@torch.no_grad()
def compute_loss_metrics():
    out = {}
    model.eval()
    for mode in ["train", "eval"]:
        losses = torch.zeros(eval_iteration)
        for k in range(eval_iteration):
            inputs, targets = batch_function(mode)
            logits, loss = model(inputs, targets)
            losses[k] = loss.item()
        out[mode] = losses.mean()
    model.train()
    return out


class MultiHeadAttention(nn.Module):
    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])

    def forward(self,inputs):
        return torch.cat([head(inputs) for head in self.heads], dim=-1)


class FeedForward(nn.Module):
    def __init__(self, n_embed):
        super().__init__()
        self.layer = nn.Sequential(
            nn.Linear(n_embed, 4 * n_embed),
            nn.ReLU(),
            nn.Linear(4 * n_embed, n_embed),
            nn.Dropout(dropout),
        )

    def forward(self, input_tensor):
        return self.layer(input_tensor)


class Block(nn.Module):
    def __init__(self, n_embed, n_heads):
        super().__init__()
        head_size = n_embed // n_heads
        self.attention = MultiHeadAttention(n_heads, head_size)
        self.feed_forward = FeedForward(n_embed)
        self.layer_norm1 = nn.LayerNorm(n_embed)
        self.layer_norm2 = nn.LayerNorm(n_embed)

    def forward(self, input_tensor):
        input_tensor = input_tensor + self.attention(self.layer_norm1(input_tensor))
        input_tensor = input_tensor + self.feed_forward(self.layer_norm2(input_tensor))
        return input_tensor


class semiGPT(nn.Module):
    def __init__(self, vocab_length):
        super().__init__()
        self.embedding_token_table = nn.Embedding(vocab_length, n_embed)
        self.position_embedding_table = nn.Embedding(block_size, n_embed)
        self.blocks = nn.Sequential(*[Block(n_embed, 4) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embed)
        self.lm_head = nn.Linear(n_embed, vocab_length)

    def forward(self, inputs, targets=None):
        batch, sequence = inputs.shape

        token_embed = self.embedding_token_table(inputs) # (B, T, C)
        pos_embed = self.position_embedding_table(torch.arange(sequence, device=device)) # (T, C)
        x = token_embed + pos_embed
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)

        if targets is None:
            loss = None
        else:
            batch, sequence, embed_size = logits.shape
            logits = logits.view(batch * sequence, embed_size)
            targets = targets.view(batch * sequence)
            loss = F.cross_entropy(logits, targets)
        return logits, loss

    def generate(self, inputs, max_new_tokens):
        for _ in range(max_new_tokens):
            inputs_cond = inputs[:, -block_size:]

            logits, loss = self(inputs_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            next_inputs = torch.multinomial(probs, num_samples=1)
            inputs = torch.cat((inputs, next_inputs), dim=1)
        return inputs


model = semiGPT(ko_vocab_size).to(device)

optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)


for step in range(max_iteration):
    if step % eval_interval == 0 :
        losses = compute_loss_metrics()
        print(f'step : {step}, train loss : {losses["train"]:.4f}, val loss : {losses["eval"]:.4f}')

    example_x, example_y = batch_function("train")
    logits, loss = model(example_x, example_y)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

2.8 토크나이저 만들기

• 토크나이저 성능과 효율성은 어휘 크기에 따라 크게 달라짐. 어휘 크기가 작으면 메모리 사용량이 줄어들고 처리 속도가 빨라질 순 있지만, 복잡한 단어나 구문을 제대로 포착하지 못할 수 있음. 반대로 너무 크면 정교한 토큰화가 가능하나 계산 비용이 증가.

2.8.2 토크나이저 만들기 

import os
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import Whitespace
from datasets import load_dataset
from transformers import PreTrainedTokenizerFast

SAVE_DIR = '/content'
os.makedirs(SAVE_DIR, exist_ok=True)

VOCAB_SIZE = 10000

# 토크나이저 초기화 
tokenizer = Tokenizer(BPE(unk_token='<unk>'))
tokenizer.pre_tokenizer = Whitespace()

trainer = BpeTrainer(
    special_tokens = ['<unk>', '<s>', '</s>', '<pad>'],
    vocab_size = VOCAB_SIZE
    )

# 토크나이저 학습
def batch_iterator(batch_size=1000):
    for i in range(0, len(dataset['train']), batch_size):
        yield dataset['train'][i:i+batch_size]['document']

tokenizer.train_from_iterator(batch_iterator(), trainer=trainer)

# 토크나이저를 json 파일로 저장
tokenizer_path = os.path.join(SAVE_DIR, 'tokenizer.json')
tokenizer.save(tokenizer_path)

# 토크나이저 Huggingface 형식으로 변환
huggingface_tokenizer = PreTrainedTokenizerFast(
    tokenizer_object=tokenizer, 
    unk_token='<unk>', 
    bos_token='<s>', 
    eos_token='</s>', 
    pad_token='<pad>'
)

# Huggingface 형식의 tokenizer 저장
huggingface_path = os.path.join(SAVE_DIR, 'huggingface_tokenizer')
huggingface_tokenizer.save_pretrained(huggingface_path)

from transformers import AutoTokenizer
# huggingface 형식의 토크나이저 로드
tokenizer = AutoTokenizer.from_pretrained(huggingface_path)

# 어휘 크기 확인
print(f"Vocabulary size: {tokenizer.vocab_size}")

# 테스트
test_texts = ["안녕하세요", "자연어 처리는 매우 흥미로운 분야입니다", "인공지능과 기계학습의 발전이 놀랍습니다"]
for text in test_texts:
    encoded = tokenizer.encode(text)
    print(f"Original: {text}")
    print(f"Encoded: {encoded}")
    print(f"Decoded: {tokenizer.decode(encoded)}")
    print(f"Tokens: {tokenizer.convert_ids_to_tokens(encoded)}")
    print()