월말 시험 (Month 2): Attention부터 LoRA까지 백지에서 4시간 안에 직접 구현해 보기
한 달 동안 배운 Transformer, SFT, LoRA를 참고 자료 없이 처음부터 구현해 보는 시험 날이다. Multi-Head Attention과 Transformer Block, LoRALinear과 loss mask를 직접 짜고, 작은 SFT를 1 epoch 돌려 전후 출력을 비교한다. 새로 배우는 것은 없고, 내가 정말 아는지 확인하는 것이 목표
By 송민성
오늘의 목표
새로 배우는 내용은 없다. 지난 한 달간 다룬 Transformer 구조, SFT(지도 미세조정), LoRA를 책이나 이전 코드를 보지 않고 직접 구현해서 내가 실제로 이해했는지 검증한다. 제한 시간은 4시간이다. 막히는 지점이 곧 다음 주에 복습할 목록이 된다.
개념과 직관
오늘은 새 개념이 없으니, 백지에서 구현하려면 머릿속에 있어야 할 것을 점검 목록으로 정리한다. 코드를 치기 전에 종이에 아래 내용을 먼저 그려 보자.
1) Multi-Head Attention (여러 개의 시선으로 문장을 동시에 훑기)
- 각 토큰은 Q(Query, 내가 찾는 것), K(Key, 내가 가진 라벨), V(Value, 실제로 건넬 내용)로 변환된다.
- 비유: 도서관에서 Q는 검색어, K는 책 제목, V는 책 내용이다. 검색어와 제목이 잘 맞을수록 그 책 내용을 많이 가져온다.
- 헤드(head)를 여러 개 두는 것은 서로 다른 관점으로 동시에 검색하는 것과 같다.
- 텍스트 수식:
- 점수 = Q × K의 전치 / sqrt(헤드 차원) - 가중치 = softmax(점수) (마스크 적용 후) - 출력 = 가중치 × V
- 텍스트를 생성하는 모델이므로 인과 마스크(causal mask, 미래 토큰을 못 보게 가리는 장치)가 필요하다.
2) Transformer Block
- 구성: LayerNorm → Attention → 잔차 연결(residual, 입력을 출력에 더해 정보가 끊기지 않게 하는 지름길) → LayerNorm → MLP → 잔차 연결.
- 이 순서를 pre-norm이라 부른다. 이전 Day에서 쓴 구조를 따르면 된다.
3) LoRA (원래 가중치는 얼리고 작은 보정 행렬만 학습)
- 큰 가중치 W를 직접 고치는 대신, 작은 두 행렬 A, B의 곱을 옆에 붙여 학습한다.
- 텍스트 수식: 출력 = x·W + (alpha / r) · x·A·B
- r은 보정 행렬의 폭(rank)이다. r이 작을수록 학습할 파라미터가 적다.
- 초기화 규칙: B를 0으로 시작하면 학습 전에는 원래 모델과 똑같이 동작한다.
4) SFT와 loss mask
- SFT(Supervised Fine-Tuning, 정답 예시를 보여 주며 모델을 다시 학습시키는 것)는 "질문 → 답변" 쌍으로 학습한다.
- loss mask는 프롬프트 부분은 손실 계산에서 빼고 답변 부분만 채점하는 장치다. 질문을 따라 외우는 것이 아니라 답하는 법을 배워야 하기 때문이다.
- PyTorch의 cross_entropy는 라벨이 -100이면 그 위치를 무시한다. 이 규칙이 loss mask의 핵심이다.
코드로 직접 해보기
시간 배분 권장: (a) 90분, (b) 90분, SFT 1 epoch 및 비교 60분. 아래 코드는 정답 예시이니, 먼저 혼자 짜고 마지막에 비교만 하자.
(a) Multi-Head Attention + Transformer Block
import torch
import torch.nn as nn
import torch.nn.functional as F
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, n_heads):
super().__init__()
assert d_model % n_heads == 0
self.h = n_heads
self.dk = d_model // n_heads
self.qkv = nn.Linear(d_model, 3 * d_model)
self.out = nn.Linear(d_model, d_model)
def forward(self, x):
B, T, D = x.shape
q, k, v = self.qkv(x).chunk(3, dim=-1)
# (B, T, D) -> (B, h, T, dk)
q, k, v = [t.view(B, T, self.h, self.dk).transpose(1, 2) for t in (q, k, v)]
scores = q @ k.transpose(-2, -1) / (self.dk ** 0.5)
mask = torch.tril(torch.ones(T, T, device=x.device)).bool()
scores = scores.masked_fill(~mask, float("-inf"))
w = F.softmax(scores, dim=-1)
y = (w @ v).transpose(1, 2).contiguous().view(B, T, D)
return self.out(y)
class TransformerBlock(nn.Module):
def __init__(self, d_model, n_heads):
super().__init__()
self.ln1 = nn.LayerNorm(d_model)
self.attn = MultiHeadAttention(d_model, n_heads)
self.ln2 = nn.LayerNorm(d_model)
self.mlp = nn.Sequential(
nn.Linear(d_model, 4 * d_model),
nn.GELU(),
nn.Linear(4 * d_model, d_model),
)
def forward(self, x):
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
return x
# 자체 점검: 입출력 shape가 같아야 한다
x = torch.randn(2, 8, 32)
print(TransformerBlock(32, 4)(x).shape) # torch.Size([2, 8, 32])(b) LoRALinear + loss mask
class LoRALinear(nn.Module):
def __init__(self, base: nn.Linear, r=8, alpha=16):
super().__init__()
self.base = base
for p in self.base.parameters():
p.requires_grad = False # 원래 가중치는 얼린다
self.A = nn.Parameter(torch.randn(base.in_features, r) * 0.01)
self.B = nn.Parameter(torch.zeros(r, base.out_features)) # 0으로 시작
self.scale = alpha / r
def forward(self, x):
return self.base(x) + (x @ self.A @ self.B) * self.scale
# 자체 점검 1: 학습 전에는 원래 출력과 같아야 한다
base = nn.Linear(16, 16)
lora = LoRALinear(base)
x = torch.randn(4, 16)
print(torch.allclose(base(x), lora(x))) # True
# 자체 점검 2: 학습되는 파라미터 수
print(sum(p.numel() for p in lora.parameters() if p.requires_grad)) # 16*8 + 8*16 = 256def build_example(prompt_ids, answer_ids):
input_ids = prompt_ids + answer_ids
labels = [-100] * len(prompt_ids) + answer_ids # 프롬프트는 채점 제외
return input_ids, labels
def lm_loss(logits, labels):
# 위치 t의 출력은 t+1번째 토큰을 맞혀야 하므로 한 칸씩 민다
logits = logits[:, :-1, :]
labels = labels[:, 1:]
return F.cross_entropy(
logits.reshape(-1, logits.size(-1)),
labels.reshape(-1),
ignore_index=-100,
)작은 SFT 1 epoch
- 작은 모델(직접 만든 블록을 쌓은 모델이나 이전 Day에서 쓴 작은 사전학습 모델)을 준비한다.
- 질문-답변 쌍 수백 개 정도의 작은 데이터셋을 만든다. 크기와 출처는 자유다.
- 모델의 Linear 일부를 LoRALinear로 교체한다.
- 배치를 만들어 lm_loss로 1 epoch 학습한다. 옵티마이저에는
requires_grad=True인 파라미터만 넘긴다. - 학습 시작 전 loss와 끝난 뒤 loss를 둘 다 기록한다.
오늘의 실험
SFT 전후 샘플 비교 표를 만든다.
- 고정할 것: 평가용 프롬프트 5~10개(학습 데이터에 없는 것), 생성 설정(최대 길이, 샘플링 방식 등).
- 바꾸는 것: 모델 상태(SFT 전 / SFT 후).
- 관찰할 것: 답변 형식이 질문에 맞게 바뀌었는가, 엉뚱한 반복이나 끊김이 줄었는가, 학습 데이터 문장을 그대로 외운 흔적은 없는가.
기록 양식 예시:
| 프롬프트 | SFT 전 출력 | SFT 후 출력 | 한 줄 평가 |
|---|---|---|---|
| (프롬프트 1) | ... | ... | 형식은 좋아졌지만 사실은 틀림 |
표 아래에 세 가지를 덧붙인다: 학습 전/후 loss 값, 학습된 파라미터 수와 전체 대비 비율, 4시간 중 어디서 가장 오래 막혔는지. 1 epoch짜리 작은 실험이므로 품질이 낮아도 정상이다. 변화가 보이는지가 중요하고, 성능 자체는 목표가 아니다.
왜 이걸 배우나 (LLM / Post-training 연결)
Post-training 업무의 기본 반복은 "데이터를 만들고, SFT를 돌리고, 전후를 비교한다"이다. 이때 loss mask를 잘못 걸면 모델이 조용히 엉뚱하게 학습되고 에러도 나지 않으므로, 백지에서 직접 짜 본 경험이 디버깅 능력으로 이어진다. LoRA와 QLoRA는 제한된 GPU로 큰 모델을 미세조정할 때 쓰는 기본 도구라서 다음 단계에서도 계속 나온다.
초보자가 막히는 지점
- shape 에러: view/transpose 순서가 틀리는 경우가 가장 흔하다. 각 줄 옆에 주석으로 shape를 적으며 짜자.
- 마스크 방향: 인과 마스크를 반대로 걸면 loss가 비정상적으로 빨리 0에 가까워진다. 미래를 훔쳐보고 있다는 신호다.
- 라벨 시프트 누락: logits와 labels를 한 칸 밀지 않으면 모델이 "현재 토큰 복사"만 배운다.
- LoRA에서 A, B를 둘 다 0으로 초기화: 기울기가 0이 되어 학습이 시작되지 않는다. 하나는 작은 랜덤값, 하나는 0이어야 한다.
- 얼린 파라미터가 옵티마이저에 들어감 / requires_grad 누락: 파라미터 수 출력으로 항상 확인한다.
- 시간 배분 실패: 한 부분에서 40분 넘게 막히면 일단 메모하고 다음으로 넘어간다. 시험이므로 막힌 지점을 기록하는 것도 결과다.
- 막혔을 때 이전 코드를 보고 싶어지더라도 4시간 안에는 참고하지 말고, 끝난 뒤 비교하며 틀린 부분을 따로 정리한다.
알고리즘 (병행 1시간)
LeetCode Medium 3문제를 90분 안에 푼다. 구현 시험과 별도로 진행하되, 시간 배분은 본인이 조절한다. 문제마다 아래를 기록한다.
- 접근: 처음 떠올린 방법과 최종 방법을 한두 문장으로.
- 시간복잡도 / 공간복잡도: 왜 그렇게 되는지 근거 한 줄.
- 틀린 이유: 오답이나 시간 초과가 났다면 원인(경계 조건, 자료구조 선택 실수 등)을 구체적으로. 한 번에 맞혔다면 "막힐 뻔한 지점"을 적는다.
30분이 지나도 접근이 안 떠오르면 풀이를 보되, 다음 날 같은 문제를 처음부터 다시 풀어 본다.