Skip to main content
Summary

드디어 작은 GPT가 움직이기 시작한다

지금까지 우리는 조각을 하나씩 만들었습니다. 처음에는 문자 하나를 숫자로 바꾸는 토크나이저에서 출발했고, Bigram 모델로 바로 앞 글자만 보고 다음 글자를 예측해 보았습니다. 그다음에는 임베딩, 위치 임베딩, Self-Attention, Multi-Head Attention, Feed Forward Network, LayerNorm, Residual Connection을 차례로 구현했습니다.

이제 할 일은 하나입니다. 이 부품들을 하나의 모델로 조립하고, 실제 텍스트 데이터로 학습시키는 것입니다.

여기서 다시 한 번 중요한 점을 분명히 하겠습니다. 이 책에서 만드는 것은 ChatGPT가 아니라, ChatGPT의 핵심 원리를 이해하기 위한 작은 GPT입니다. 우리가 만드는 모델은 문자 단위로 동작하고, 데이터도 작고, 파라미터 수도 매우 적습니다. 하지만 구조의 뼈대는 GPT 계열 모델과 닮아 있습니다. 다음 토큰을 예측하고, causal self-attention으로 과거 문맥을 보고, Transformer Block을 여러 층 쌓아 표현을 바꿉니다.

작은 종이비행기를 직접 접어 보면 비행기의 모든 공학을 이해할 수는 없어도, 날개가 왜 필요한지, 균형이 왜 중요한지, 접는 각도가 비행에 어떤 영향을 주는지 감각을 얻을 수 있습니다. Mini GPT도 마찬가지입니다. 이 모델은 거대한 상용 LLM이 아니지만, LLM이 어떻게 학습되고 문장을 생성하는지 손끝으로 이해하게 해 주는 실험 장치입니다.

Mini GPT 학습시키기

이번 장에서 완성할 것

이번 장의 목표는 코드 조각을 더 많이 외우는 것이 아닙니다. 학습이 실제로 어떤 순서로 일어나는지 이해하는 것입니다. 모델은 텍스트를 읽고, 다음 글자를 예측하고, 틀린 정도를 손실로 계산하고, 그 손실을 줄이는 방향으로 파라미터를 조금씩 바꿉니다. 이 과정을 수천 번 반복하면 처음에는 아무 글자나 내뱉던 모델이 점차 데이터셋의 말투와 패턴을 흉내 내기 시작합니다.

이번 장에서 다룰 흐름은 다음과 같습니다.

  • 토큰 임베딩, 위치 임베딩, 여러 Transformer Block, 최종 linear head를 하나의 MiniGPT로 조립합니다.

  • AdamW optimizer와 learning_rate, batch_size, max_iters 같은 학습 설정을 정합니다.

  • forward, loss, backward, step으로 이어지는 학습 루프를 작성합니다.

  • 검증 손실을 주기적으로 측정해 모델이 정말 좋아지는지 확인합니다.

  • 체크포인트를 저장하고 불러와 실험을 이어 갈 수 있게 만듭니다.

이제 흩어져 있던 실습 코드가 하나의 훈련 프로그램으로 바뀝니다.

모델 조립하기

먼저 Mini GPT 모델을 다시 정리해 보겠습니다. 이전 장에서 만든 Block은 Transformer Block 하나입니다. 이제 이 블록을 여러 개 쌓고, 앞에는 임베딩을 붙이고, 뒤에는 다음 토큰 후보 점수를 출력하는 선형층을 붙입니다.

구조를 말로 풀면 다음과 같습니다.

토큰 ID
-> 토큰 임베딩
-> 위치 임베딩 더하기
-> Transformer Block 여러 개
-> 마지막 LayerNorm
-> Linear head
-> 각 토큰 위치의 다음 문자 점수(logits)

이 흐름을 코드로 옮기면 다음과 같습니다. 파일을 나누어 관리하고 있다면 src/model.py에 넣어도 좋고, 지금까지처럼 하나의 노트북이나 train.py 파일에서 이어서 작성해도 좋습니다.

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


class MiniGPT(nn.Module):
    def __init__(
        self,
        vocab_size,
        block_size,
        n_embd,
        n_head,
        n_layer,
        dropout,
    ):
        super().__init__()
        self.block_size = block_size

        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)

        self.blocks = nn.Sequential(*[
            Block(
                n_embd=n_embd,
                n_head=n_head,
                block_size=block_size,
                dropout=dropout,
            )
            for _ in range(n_layer)
        ])

        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape

        tok_emb = self.token_embedding_table(idx)       # (B, T, n_embd)
        pos = torch.arange(T, device=idx.device)        # (T,)
        pos_emb = self.position_embedding_table(pos)    # (T, n_embd)

        x = tok_emb + pos_emb                           # (B, T, n_embd)
        x = self.blocks(x)                              # (B, T, n_embd)
        x = self.ln_f(x)                                # (B, T, n_embd)
        logits = self.lm_head(x)                        # (B, T, vocab_size)

        loss = None
        if targets is not None:
            B, T, V = logits.shape
            logits = logits.reshape(B * T, V)
            targets = targets.reshape(B * T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.block_size:]
            logits, loss = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)

        return idx

여기서 가장 중요한 부분은 forward()입니다. 모델은 입력 idx를 받습니다. idx는 문자 ID로 이루어진 2차원 텐서입니다.

idx shape: (B, T)

B는 batch size입니다. 한 번에 몇 개의 조각을 학습할지 뜻합니다. T는 block size입니다. 한 조각 안에 들어 있는 토큰 개수입니다. 문자 단위 모델이라면 토큰 하나가 문자 하나이므로, T = 64는 한 번에 64글자의 문맥을 본다는 뜻입니다.

token_embedding_table은 각 문자 ID를 벡터로 바꿉니다. position_embedding_table은 각 위치 번호를 벡터로 바꿉니다. 두 벡터를 더하면 모델은 "어떤 문자가 어느 위치에 있는지"를 함께 알 수 있습니다.

그다음 여러 Transformer Block을 통과합니다. 이 블록들이 문맥을 섞고, 각 위치의 표현을 계속 가공합니다. 마지막 lm_head는 각 위치에서 다음 문자 후보의 점수를 만듭니다. 이 점수를 logits라고 부릅니다.

logits shape: (B, T, vocab_size)

즉 모델은 한 번의 forward에서 모든 위치에 대해 다음 문자 후보 점수를 동시에 계산합니다.

생성 함수는 학습 후의 작은 창문이다

generate() 함수도 모델 안에 함께 넣었습니다. 이 함수는 학습 자체에는 직접 필요하지 않지만, 학습된 모델이 무엇을 배웠는지 확인하는 작은 창문 역할을 합니다.

def generate(self, idx, max_new_tokens):
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -self.block_size:]
        logits, loss = self(idx_cond)
        logits = logits[:, -1, :]
        probs = F.softmax(logits, dim=-1)
        idx_next = torch.multinomial(probs, num_samples=1)
        idx = torch.cat((idx, idx_next), dim=1)

    return idx

생성 과정에서는 현재까지의 문맥을 모델에 넣고, 마지막 위치의 예측 결과만 가져옵니다. 그 예측 결과를 확률로 바꾼 뒤 하나의 다음 토큰을 뽑습니다. 그리고 그 토큰을 기존 문맥 뒤에 붙입니다. 이 일을 반복하면 글자가 하나씩 이어집니다.

아직 학습하지 않은 모델로 생성하면 거의 무작위 문자열이 나옵니다. 하지만 학습이 진행되면 점차 데이터셋에 등장하던 패턴이 보이기 시작합니다. 줄바꿈 위치가 그럴듯해지고, 자주 등장하는 단어 조각이 나타나며, 문장 부호의 사용도 조금씩 닮아 갑니다.

물론 이 작은 모델이 깊은 의미를 이해한다고 말하면 안 됩니다. 지금 모델은 문자 단위의 다음 토큰 예측을 연습하고 있습니다. 그러나 바로 이 단순한 훈련 목표가 거대한 LLM의 출발점이기도 합니다.

학습 설정 정하기

이제 모델을 실제로 학습시킬 준비를 하겠습니다. 학습에는 몇 가지 숫자를 정해야 합니다. 이 숫자들을 보통 하이퍼파라미터라고 부릅니다. 모델이 학습으로 직접 찾는 값이 아니라, 사람이 실험 전에 정해 주는 설정값입니다.

처음에는 다음 정도의 설정으로 시작해 보겠습니다.

torch.manual_seed(1337)

device = 'cuda' if torch.cuda.is_available() else 'cpu'

batch_size = 64
block_size = 64
max_iters = 5000
eval_interval = 500
eval_iters = 200
learning_rate = 3e-4
n_embd = 128
n_head = 4
n_layer = 4
dropout = 0.2

각 설정의 의미를 하나씩 살펴보겠습니다.

batch_size는 한 번의 학습 단계에서 몇 개의 텍스트 조각을 볼지 정합니다. batch_size = 64라면 모델은 매번 64개의 입력 조각과 정답 조각을 받아 손실을 계산합니다. batch size가 크면 손실 추정이 안정적일 수 있지만 메모리를 더 많이 씁니다.

block_size는 모델이 한 번에 보는 문맥 길이입니다. block_size = 64라면 모델은 최대 64개 문자까지 보고 다음 문자를 예측합니다. 이 값이 커지면 더 긴 문맥을 볼 수 있지만, Attention 계산 비용이 늘어납니다.

max_iters는 학습 단계를 몇 번 반복할지 정합니다. 여기서는 5000번 반복합니다. 작은 데이터셋이라도 모델이 한 번에 전체 데이터를 다 보는 것이 아니라, 매번 임의의 조각을 뽑아 조금씩 학습한다는 점을 기억해야 합니다.

learning_rate는 파라미터를 얼마나 크게 업데이트할지 정합니다. 너무 크면 손실이 불안정하게 튈 수 있고, 너무 작으면 학습이 지나치게 느릴 수 있습니다. 3e-4는 0.0003이라는 뜻입니다.

n_embd, n_head, n_layer는 모델 크기를 정합니다. n_embd는 각 토큰 벡터의 차원, n_head는 Attention head 개수, n_layer는 Transformer Block 개수입니다. 여기서는 작은 실습 모델이므로 모두 작게 잡았습니다.

dropout은 학습 중 일부 연결을 무작위로 꺼서 모델이 훈련 데이터에 지나치게 달라붙는 것을 완화하는 방법입니다. 단, 평가나 생성할 때는 dropout이 꺼져야 합니다. 그래서 뒤에서 model.train()과 model.eval()을 구분해서 사용할 것입니다.

모델 만들고 파라미터 수 확인하기

이제 모델 인스턴스를 만들겠습니다. 이전 장에서 만든 CharTokenizer, train_data, val_data, get_batch, vocab_size가 이미 준비되어 있다고 가정합니다.

model = MiniGPT(
    vocab_size=vocab_size,
    block_size=block_size,
    n_embd=n_embd,
    n_head=n_head,
    n_layer=n_layer,
    dropout=dropout,
)
model = model.to(device)

num_params = sum(p.numel() for p in model.parameters())
print(f'{num_params / 1e6:.3f}M parameters')

model.to(device)는 모델을 CPU나 GPU로 옮깁니다. GPU가 있다면 cuda를 사용하고, 없다면 CPU에서 실행합니다. 작은 모델은 CPU에서도 돌아가지만, 학습 시간이 더 오래 걸릴 수 있습니다.

파라미터 수는 모델이 학습해야 하는 숫자의 개수입니다. 예를 들어 출력이 다음과 비슷하게 나올 수 있습니다.

0.814M parameters

이는 약 81만 개의 숫자를 학습한다는 뜻입니다. 누군가에게는 커 보일 수 있지만, 현대 LLM에 비하면 장난감처럼 작은 규모입니다. 하지만 원리를 배우기에는 충분합니다. 중요한 것은 숫자의 크기가 아니라, 이 숫자들이 어떤 계산 구조 안에서 다음 토큰 예측 능력으로 바뀌는지 보는 것입니다.

AdamW optimizer 준비하기

모델이 손실을 계산할 수 있어도, 스스로 파라미터를 바꾸지는 않습니다. 파라미터를 어떤 방식으로 업데이트할지 정하는 도구가 optimizer입니다. 이번에는 AdamW를 사용하겠습니다.

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

AdamW는 딥러닝에서 널리 쓰이는 optimizer입니다. 단순한 경사하강법은 모든 파라미터를 같은 규칙으로 조금씩 움직입니다. 반면 Adam 계열 optimizer는 각 파라미터의 업데이트 방향과 크기를 조금 더 똑똑하게 조절합니다. 여기에 W가 붙은 AdamW는 weight decay라는 규제 방식을 더 깔끔하게 적용한 버전입니다.

너무 깊이 들어갈 필요는 없습니다. 지금은 이렇게 이해하면 충분합니다.

loss.backward()가 각 파라미터를 어느 방향으로 바꿔야 하는지 계산한다.
optimizer.step()이 그 정보를 이용해 실제 파라미터 값을 업데이트한다.

즉 optimizer는 모델을 조금씩 고쳐 가는 정비사입니다. 손실은 "지금 얼마나 틀렸는지"를 알려 주고, 그래디언트는 "어느 방향으로 고쳐야 하는지"를 알려 주며, optimizer는 "얼마나 조심스럽게 고칠지"를 결정합니다.

학습 루프의 네 단계

학습 루프는 처음 보면 기계적으로 느껴집니다. 하지만 안에서 일어나는 일을 알고 나면 아주 자연스럽습니다. 한 번의 학습 단계는 다음 네 단계로 이루어집니다.

logits, loss = model(X, Y)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()

첫 번째 줄은 forward입니다. 입력 X를 모델에 넣고 예측 결과와 손실을 얻습니다. 이때 Y는 정답입니다.

두 번째 줄은 이전 단계에서 남아 있던 그래디언트를 지웁니다. PyTorch는 기본적으로 그래디언트를 누적합니다. 그래서 매 단계마다 초기화하지 않으면 이전 계산 결과가 섞입니다.

세 번째 줄은 backward입니다. 손실을 줄이려면 각 파라미터를 어느 방향으로 바꿔야 하는지 계산합니다. 이 과정이 역전파입니다.

네 번째 줄은 step입니다. optimizer가 계산된 그래디언트를 사용해 실제 파라미터 값을 업데이트합니다.

이 네 줄을 반복하는 것이 학습입니다. 거대한 LLM도, 우리가 만든 작은 Mini GPT도 기본 리듬은 같습니다. 다만 데이터 크기, 모델 크기, 학습 인프라가 엄청나게 다를 뿐입니다.

검증 손실 측정하기

훈련 손실만 보면 위험합니다. 모델이 훈련 데이터 조각을 점점 잘 맞히는 것은 당연합니다. 하지만 우리가 원하는 것은 처음 보는 비슷한 텍스트에서도 그럴듯하게 예측하는 능력입니다. 그래서 train_data와 별도로 떼어 둔 val_data로 검증 손실을 측정합니다.

평가 함수는 다음과 같이 작성합니다.

@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()

    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)

        for k in range(eval_iters):
            X, Y = get_batch(split)
            X = X.to(device)
            Y = Y.to(device)

            logits, loss = model(X, Y)
            losses[k] = loss.item()

        out[split] = losses.mean().item()

    model.train()
    return out

@torch.no_grad()는 평가 중에는 그래디언트를 계산하지 말라는 뜻입니다. 검증 손실을 볼 때는 모델을 업데이트하지 않습니다. 예측만 하면 됩니다. 그래디언트를 계산하지 않으면 메모리를 덜 쓰고 속도도 빨라집니다.

model.eval()은 모델을 평가 모드로 바꿉니다. 특히 dropout이 꺼집니다. 학습 중에는 일부 값을 무작위로 꺼서 모델을 단련하지만, 평가 중에는 같은 입력에 대해 안정적인 출력을 얻어야 합니다.

평가가 끝난 뒤 model.train()을 호출하는 것도 중요합니다. 다시 학습 모드로 돌아가야 dropout이 학습용으로 작동합니다.

검증 손실은 한 번의 batch만 보고 판단하지 않습니다. eval_iters번 여러 batch를 뽑아 손실을 평균냅니다. 이렇게 하면 우연히 쉬운 조각이나 어려운 조각이 뽑혀 생기는 흔들림을 줄일 수 있습니다.

전체 학습 루프 작성하기

이제 실제 학습 루프를 작성하겠습니다.

history = []

for step in range(max_iters):
    if step % eval_interval == 0 or step == max_iters - 1:
        losses = estimate_loss()
        train_loss = losses['train']
        val_loss = losses['val']
        history.append((step, train_loss, val_loss))
        print(
            f'step {step}: '
            f'train loss {train_loss:.4f}, '
            f'val loss {val_loss:.4f}'
        )

    X, Y = get_batch('train')
    X = X.to(device)
    Y = Y.to(device)

    logits, loss = model(X, Y)

    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

이 코드는 매우 중요합니다. 천천히 읽어 보겠습니다.

먼저 step이 eval_interval의 배수일 때마다 검증 손실을 측정합니다. 예를 들어 eval_interval = 500이면 0, 500, 1000, 1500번째 단계에서 손실을 출력합니다. 마지막 단계에서도 한 번 더 출력합니다.

그다음 get_batch('train')으로 훈련용 입력과 정답을 가져옵니다. X는 모델에 넣을 문자 ID 조각이고, Y는 각 위치에서 맞혀야 할 다음 문자 ID입니다.

X: h e l l
Y: e l l o

문자 단위로 보면 X의 각 위치는 Y의 같은 위치에 있는 다음 문자를 예측해야 합니다. 이 단순한 이동 관계가 언어 모델 학습의 핵심입니다.

모델에 X, Y를 넣으면 logits와 loss가 나옵니다. logits는 예측 점수이고, loss는 그 예측이 정답과 얼마나 다른지 나타내는 숫자입니다.

그다음은 늘 같은 네 단계입니다.

optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()

그래디언트를 지우고, 새 그래디언트를 계산하고, 파라미터를 업데이트합니다. 이 루프가 수천 번 돌면서 Mini GPT는 조금씩 텍스트 패턴을 익힙니다.

손실 숫자 읽는 법

학습을 시작하면 다음과 비슷한 출력이 나올 수 있습니다.

step 0: train loss 4.2108, val loss 4.2135
step 500: train loss 2.4891, val loss 2.5368
step 1000: train loss 2.1914, val loss 2.3052
step 1500: train loss 2.0357, val loss 2.2196
step 2000: train loss 1.9278, val loss 2.1783
step 2500: train loss 1.8462, val loss 2.1590
step 3000: train loss 1.7815, val loss 2.1517

실제 숫자는 데이터셋과 설정에 따라 달라집니다. 중요한 것은 정확히 같은 숫자가 나오는지가 아니라 흐름입니다.

처음 손실은 보통 높습니다. 모델의 파라미터가 무작위로 초기화되어 있기 때문입니다. 이때 모델은 글자 사이의 관계를 거의 모릅니다. 그러다 학습이 진행되면 훈련 손실이 내려갑니다. 검증 손실도 함께 내려가면 모델이 훈련 데이터만 외우는 것이 아니라 일반적인 패턴을 조금씩 배우고 있다고 볼 수 있습니다.

하지만 어느 순간부터 훈련 손실은 계속 내려가는데 검증 손실은 내려가지 않거나 오히려 올라갈 수 있습니다. 이것이 과적합의 신호일 수 있습니다. 모델이 훈련 데이터의 세부 패턴을 지나치게 외워서, 새로운 조각에는 잘 맞지 않는 상태입니다.

손실을 읽을 때는 세 가지를 함께 봐야 합니다.

  • 훈련 손실이 전반적으로 내려가는가.

  • 검증 손실도 함께 내려가는가.

  • 훈련 손실과 검증 손실의 차이가 지나치게 벌어지지는 않는가.

만약 손실이 전혀 내려가지 않는다면 learning rate가 너무 크거나 작을 수 있고, 모델 코드에 오류가 있을 수도 있습니다. 입력과 정답이 제대로 한 칸씩 밀려 있는지도 확인해야 합니다.

학습 중간에 생성해 보기

손실 숫자만 보면 조금 지루합니다. 모델이 실제로 무엇을 말하는지 중간중간 확인해 보면 학습 과정이 훨씬 생생해집니다.

다음 함수를 만들어 보겠습니다.

def sample_from_model(model, tokenizer, max_new_tokens=300):
    model.eval()

    context = torch.zeros((1, 1), dtype=torch.long, device=device)
    generated = model.generate(context, max_new_tokens=max_new_tokens)
    text = tokenizer.decode(generated[0].tolist())

    model.train()
    return text

여기서는 시작 토큰으로 0번 ID 하나를 넣었습니다. 문자 단위 토크나이저에서는 0번 문자가 무엇인지는 데이터셋에 따라 다릅니다. 더 자연스러운 시작을 원한다면 특정 문자열을 인코딩해서 넣을 수도 있습니다.

print(sample_from_model(model, tokenizer, max_new_tokens=500))

학습 전 출력은 대체로 이런 느낌일 수 있습니다.

xqP$aanzz--mQ...??

학습이 진행된 뒤에는 적어도 데이터셋의 분위기를 조금씩 닮기 시작합니다. 예를 들어 줄바꿈, 공백, 자주 등장하는 단어 조각, 문장 부호가 더 그럴듯하게 나타납니다.

여기서 너무 큰 기대를 하면 실망할 수 있습니다. 작은 문자 단위 모델은 긴 논리 구조를 안정적으로 유지하기 어렵습니다. 하지만 바로 그 한계를 관찰하는 것이 중요합니다. 잘되는 부분과 안 되는 부분을 보면, 왜 더 큰 데이터, 더 큰 모델, 더 좋은 학습 방법이 필요한지 자연스럽게 이해할 수 있습니다.

체크포인트 저장하기

학습은 시간이 걸립니다. 특히 GPU가 없거나 모델 크기를 조금 키우면 몇 분, 길게는 더 오래 걸릴 수 있습니다. 학습이 끝난 뒤 모델을 저장하지 않으면 프로그램을 종료하는 순간 배운 파라미터가 사라집니다.

따라서 체크포인트를 저장해 두는 습관이 필요합니다. 체크포인트는 특정 시점의 모델 상태를 파일로 보관한 것입니다.

먼저 저장할 폴더를 만듭니다.

from pathlib import Path

checkpoint_dir = Path('checkpoints')
checkpoint_dir.mkdir(exist_ok=True)

그다음 학습이 끝난 뒤 모델과 optimizer 상태를 저장합니다.

checkpoint_path = checkpoint_dir / 'mini_gpt.pt'

torch.save({
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'step': max_iters,
    'config': {
        'vocab_size': vocab_size,
        'block_size': block_size,
        'n_embd': n_embd,
        'n_head': n_head,
        'n_layer': n_layer,
        'dropout': dropout,
        'learning_rate': learning_rate,
        'batch_size': batch_size,
    },
    'history': history,
}, checkpoint_path)

print(f'saved checkpoint to {checkpoint_path}')

여기서 model.state_dict()는 모델 파라미터를 담고 있습니다. optimizer.state_dict()는 optimizer의 내부 상태를 담고 있습니다. AdamW는 단순히 현재 파라미터만 보는 것이 아니라 이전 그래디언트 흐름에 대한 정보도 사용하므로, 학습을 이어 가려면 optimizer 상태도 함께 저장하는 편이 좋습니다.

config에는 모델을 다시 만들 때 필요한 설정을 넣었습니다. 체크포인트 파일만 남아 있고 어떤 n_embd, n_layer, block_size로 학습했는지 모르면 나중에 불러오기 어렵습니다. 실험을 반복하다 보면 이런 기록의 중요성을 금방 느끼게 됩니다.

체크포인트 불러오기

저장한 체크포인트를 다시 불러오려면 먼저 같은 구조의 모델을 만들어야 합니다. 모델 구조가 다르면 저장된 파라미터를 맞게 끼울 수 없습니다. 예를 들어 n_embd = 128로 저장한 파라미터를 n_embd = 64 모델에 넣을 수는 없습니다.

checkpoint = torch.load(checkpoint_path, map_location=device)
config = checkpoint['config']

model = MiniGPT(
    vocab_size=config['vocab_size'],
    block_size=config['block_size'],
    n_embd=config['n_embd'],
    n_head=config['n_head'],
    n_layer=config['n_layer'],
    dropout=config['dropout'],
)
model = model.to(device)
model.load_state_dict(checkpoint['model_state_dict'])

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=config['learning_rate'],
)
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

start_step = checkpoint['step']
history = checkpoint.get('history', [])

print(f'loaded checkpoint from step {start_step}')

이제 모델은 저장 당시의 파라미터를 다시 갖게 됩니다. 이어서 더 학습하고 싶다면 start_step 이후부터 루프를 돌리면 됩니다.

additional_iters = 2000

for step in range(start_step, start_step + additional_iters):
    if step % eval_interval == 0:
        losses = estimate_loss()
        train_loss = losses['train']
        val_loss = losses['val']
        history.append((step, train_loss, val_loss))
        print(
            f'step {step}: '
            f'train loss {train_loss:.4f}, '
            f'val loss {val_loss:.4f}'
        )

    X, Y = get_batch('train')
    X = X.to(device)
    Y = Y.to(device)

    logits, loss = model(X, Y)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

이렇게 하면 학습을 한 번에 끝내지 않아도 됩니다. 오늘 5000번 학습하고 저장한 뒤, 내일 다시 불러와 2000번 더 학습할 수 있습니다. 실험을 여러 번 반복하는 실제 개발에서는 이 방식이 매우 중요합니다.

좋은 실험을 위한 작은 습관

Mini GPT 학습은 작은 실습이지만, 여기서 익히는 습관은 나중에 더 큰 모델이나 LLM API 기반 서비스를 다룰 때도 도움이 됩니다. 무작정 코드를 실행하는 것보다, 설정과 결과를 함께 기록하며 실험해야 합니다.

예를 들어 다음과 같은 내용을 남겨 두면 좋습니다.

  • 데이터셋 이름과 크기

  • block_size, batch_size, n_embd, n_head, n_layer, learning_rate

  • 최종 train loss와 val loss

  • 생성 샘플 몇 개

  • 이상했던 점이나 다음에 바꾸어 볼 설정

이 기록은 단순한 메모가 아닙니다. 모델 개발은 추리 소설에 가깝습니다. 손실이 왜 내려가지 않았는지, 어느 순간부터 왜 검증 손실이 올라갔는지, block size를 늘렸을 때 왜 속도가 느려졌는지 단서를 모아야 합니다.

작은 실험을 예로 들어 보겠습니다. n_layer = 2일 때와 n_layer = 4일 때를 비교하면, 더 깊은 모델이 더 낮은 훈련 손실을 보일 수 있습니다. 하지만 검증 손실도 함께 좋아지는지는 따로 봐야 합니다. 데이터가 너무 작으면 깊은 모델이 더 빨리 외울 수도 있습니다.

learning_rate도 자주 실험하는 값입니다. 손실이 크게 출렁이거나 nan이 나오면 learning rate가 너무 클 가능성이 있습니다. 반대로 손실이 거의 움직이지 않으면 너무 작을 수 있습니다.

block_size는 문맥 길이를 바꾸는 값입니다. 길게 잡으면 더 많은 문맥을 볼 수 있지만, Attention 계산량이 늘어납니다. 특히 self-attention은 문맥 길이가 길어질수록 계산량이 빠르게 증가하므로, 작은 장비에서는 적절히 타협해야 합니다.

자주 만나는 오류들

Mini GPT를 처음 학습시키다 보면 오류를 만나는 것이 정상입니다. 오류는 실패가 아니라 모델의 내부 구조를 더 잘 이해할 기회입니다.

가장 흔한 오류는 device mismatch입니다.

Expected all tensors to be on the same device

모델은 GPU에 있는데 입력 데이터는 CPU에 있거나, 반대로 일부 텐서만 GPU에 있을 때 발생합니다. 입력 X, 정답 Y, 모델이 모두 같은 device에 있는지 확인해야 합니다.

X = X.to(device)
Y = Y.to(device)
model = model.to(device)

두 번째는 block size 오류입니다. 위치 임베딩은 block_size 길이만큼 만들어져 있습니다. 그런데 그보다 긴 입력을 넣으면 위치 인덱스가 범위를 벗어날 수 있습니다. 생성 함수에서 다음 줄을 쓰는 이유가 바로 이것입니다.

idx_cond = idx[:, -self.block_size:]

아무리 긴 문장을 생성하고 있더라도, 매번 모델에는 마지막 block_size개 토큰만 넣습니다. Mini GPT는 그 이상 과거의 문맥을 직접 볼 수 없습니다.

세 번째는 shape 오류입니다. n_embd가 n_head로 나누어떨어지지 않으면 각 head의 차원을 균등하게 나눌 수 없습니다. 예를 들어 n_embd = 128, n_head = 4는 괜찮습니다. 각 head가 32차원을 갖습니다. 하지만 n_embd = 130, n_head = 4는 깔끔하게 나누어떨어지지 않습니다.

네 번째는 손실이 nan이 되는 경우입니다. 이때는 learning rate를 낮추고, 입력 데이터에 이상한 값이 없는지 확인하고, 모델의 Attention mask가 제대로 적용되는지 살펴보는 것이 좋습니다.

전체 코드의 큰 그림

여기까지의 코드를 하나의 큰 흐름으로 압축하면 다음과 같습니다.

# 1. 설정
batch_size = 64
block_size = 64
max_iters = 5000
learning_rate = 3e-4
n_embd = 128
n_head = 4
n_layer = 4
dropout = 0.2

# 2. 모델과 optimizer
model = MiniGPT(
    vocab_size=vocab_size,
    block_size=block_size,
    n_embd=n_embd,
    n_head=n_head,
    n_layer=n_layer,
    dropout=dropout,
).to(device)

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

# 3. 학습
for step in range(max_iters):
    X, Y = get_batch('train')
    X = X.to(device)
    Y = Y.to(device)

    logits, loss = model(X, Y)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

# 4. 생성
context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(tokenizer.decode(model.generate(context, 500)[0].tolist()))

실제 코드에서는 중간중간 estimate_loss()로 검증 손실을 확인하고, 체크포인트도 저장해야 합니다. 하지만 학습의 뼈대는 위 코드에 모두 들어 있습니다.

입력을 뽑고, 예측하고, 손실을 계산하고, 역전파하고, 업데이트합니다. 이 단순한 루프가 반복되면서 파라미터 안에 데이터의 통계적 패턴이 새겨집니다.

Mini GPT가 배운다는 것

모델이 학습된다는 말은 조금 신비롭게 들릴 수 있습니다. 하지만 지금 우리가 만든 코드 안에서는 아주 구체적인 일이 벌어집니다.

처음에는 임베딩 테이블도, Attention의 query/key/value 가중치도, Feed Forward Network의 가중치도 거의 무작위입니다. 그래서 어떤 문자가 다음에 올지 제대로 맞히지 못합니다. 손실은 높습니다.

학습 루프가 한 번 돌 때마다 모델은 자신이 얼마나 틀렸는지 계산합니다. 그리고 그 틀림을 줄이려면 각 숫자를 어느 방향으로 바꿔야 하는지 역전파로 알아냅니다. AdamW는 그 방향으로 파라미터를 조금씩 움직입니다.

이 과정이 반복되면 어떤 문자 뒤에 어떤 문자가 자주 오는지, 줄바꿈은 어떤 상황에서 나오는지, 공백과 문장 부호가 어떤 패턴을 이루는지 모델의 파라미터에 반영됩니다. 문자 단위 모델이므로 단어를 명시적으로 아는 것은 아니지만, 자주 나타나는 문자 배열을 통해 단어처럼 보이는 조각을 만들어 냅니다.

이것이 다음 토큰 예측의 힘입니다. 모델은 "문장을 이해하라"는 직접 명령을 받지 않았습니다. 오직 다음 문자를 맞히는 연습만 했습니다. 그런데 그 연습을 충분히 하면 텍스트의 구조를 어느 정도 흉내 내게 됩니다.

물론 작은 모델의 한계도 분명합니다. 긴 문맥을 오래 기억하지 못하고, 의미가 흔들리며, 그럴듯한 단어 조각을 만들다가 갑자기 이상한 흐름으로 빠질 수 있습니다. 하지만 이 한계야말로 다음 장에서 관찰할 중요한 주제입니다.

다음 장으로 가기 전에

이번 장에서 우리는 Mini GPT를 실제로 학습시켰습니다. 이제 이 모델은 더 이상 빈 구조물이 아닙니다. 데이터셋을 보고, 손실을 줄이고, 파라미터를 바꾸며, 자신만의 방식으로 텍스트 패턴을 품게 되었습니다.

지금까지 배운 내용을 한 문장으로 정리하면 이렇습니다.

Mini GPT는 토큰 임베딩과 위치 임베딩으로 입력을 만들고, 여러 Transformer Block으로 문맥 표현을 갱신한 뒤, 다음 토큰을 맞히도록 손실을 줄이며 학습된다.

우리는 이제 모델을 만들었고, 학습도 시켰고, 저장하고 불러오는 방법까지 배웠습니다. 하지만 아직 가장 재미있는 질문이 남아 있습니다.

학습된 모델은 글자를 어떻게 하나씩 이어 붙일까요? 왜 같은 시작 문장을 넣어도 매번 다른 문장이 나올 수 있을까요? 확률이 가장 높은 토큰만 고르면 좋은 결과가 나올까요, 아니면 약간의 무작위성이 필요할까요?

다음 장에서는 모델이 문장을 생성하는 과정을 더 가까이에서 들여다보겠습니다. 이제 Mini GPT의 입에서 첫 문장이 나올 차례입니다.

책 목록

처음 만나는 LLM: 원리부터 Mini GPT 구현, RAG 서비스 개발까지

전체 목차 보기

이 노트는 보도·비평·교육·연구 목적으로 작성되었습니다. 저작권 문의가 있으시면 에서 알려주세요.

틸노트는 사용자가 AI를 활용해 노트를 쉽게 작성하거나 편집할 수 있는 서비스입니다.

AI 스퀘어 단톡방에 함께해요

AI 활용과 틸노트 사용 팁을 나누는 오픈채팅방입니다. 입장 비밀번호는 tilnote입니다.

참여하기