트랜스포머 직접 구현 — Capstone

열세 개의 lesson. 하나의 모델. 지름길은 없습니다.

유형: Build 언어: Python 선수 지식: Phase 7 · 01부터 13까지. 건너뛰지 않습니다. 예상 시간: 약 120분

문제

지금까지 모든 논문(paper)을 읽고, 어텐션(attention), 멀티 헤드 분할(multi-head split), 위치 인코딩(positional encoding), 인코더와 디코더 블록(encoder/decoder block), BERT와 GPT의 손실 함수(loss), 전문가 혼합(Mixture of Experts; MoE), KV 캐시(KV cache)까지 구현해 보았습니다. 이제 이 구성 요소들이 실제 과제(task)에서 함께 작동하도록 만들 차례입니다.

이번 capstone은 문자 수준 언어 모델링(character-level language modeling) 과제에서 작은 디코더 전용 트랜스포머(decoder-only transformer)를 처음부터 끝까지(end-to-end) 학습시키는 것입니다. 셰익스피어(Shakespeare) 텍스트를 읽어 들이고, 새로운 셰익스피어풍 텍스트를 생성합니다. 노트북에서 10분 이내에 학습할 만큼 작지만, 더 큰 데이터셋(dataset)과 더 긴 학습 시간을 주면 실제 언어 모델(LM)로 이어질 만큼 구조가 올바릅니다.

이 강의는 과정 전체의 "nanoGPT"에 해당합니다. 완전히 새로운 작업은 아닙니다. Karpathy의 2023년 nanoGPT 튜토리얼은 모든 학습자가 한 번쯤 직접 작성해 보는 표준 참조 구현(reference implementation)입니다. 여기서는 그 기본 구조를 가져와, Phase 7에서 다뤄 온 내용을 중심으로 다시 구성합니다.

개념

각 단계가 어느 lesson과 연결되는지 표시한 아키텍처는 다음과 같습니다.

input tokens (B, N)
   │
   ▼
token embedding + positional embedding  ◀── Lesson 04 (RoPE option)
   │
   ▼
┌──── block × L ────────────────────┐
│  RMSNorm                          │  ◀── Lesson 05
│  MultiHeadAttention (causal)      │  ◀── Lesson 03 + 07 (causal mask)
│  residual                         │
│  RMSNorm                          │
│  SwiGLU FFN                       │  ◀── Lesson 05
│  residual                         │
└────────────────────────────────── ┘
   │
   ▼
final RMSNorm
   │
   ▼
lm_head (tied to token embedding)
   │
   ▼
logits (B, N, V)
   │
   ▼
shift-by-one cross-entropy            ◀── Lesson 07

이번 강의에서 만들어 내는 것

GPTConfig — 모든 하이퍼파라미터(hyperparameter)를 한 곳에서 설정합니다.
MultiHeadAttention — 인과(causal) 마스크와 배치(batched) 처리를 지원하며, 선택적으로 PyTorch의 scaled_dot_product_attention을 통한 Flash-style 경로(pathway)를 사용합니다.
SwiGLUFFN — 최신 방식의 피드포워드 네트워크(FFN)입니다.
Block — 사전 정규화(pre-norm) 구조에 잔차(residual)로 감싼 어텐션과 FFN을 묶은 블록입니다.
GPT — 임베딩, 쌓아 올린 블록, LM 헤드, generate()를 포함하는 모델입니다.
AdamW, 코사인 학습률 스케줄(cosine LR), 그래디언트 클리핑(gradient clipping)을 갖춘 학습 루프(training loop).
셰익스피어 텍스트에 대한 문자 단위 토크나이저(char-level tokenizer).

이번 강의에서 다루지 않는 것

RoPE — Lesson 04에서 개념적으로 구현했습니다. 여기서는 단순화를 위해 학습된 위치 임베딩(learned positional embedding)을 사용합니다. 연습문제에서 RoPE로 교체합니다.
생성 단계에서의 KV 캐시 — 매 생성 스텝마다 전체 prefix에 대해 어텐션을 다시 계산합니다. 더 느리지만 더 단순합니다. 연습문제에서 KV 캐시를 추가합니다.
Flash Attention — PyTorch 2.0 이상에서는 입력 조건이 맞으면 자동으로 디스패치(dispatch)됩니다. 여기서는 F.scaled_dot_product_attention을 사용합니다.
MoE — 블록마다 단일 FFN을 사용합니다. MoE는 Lesson 11에서 다뤘습니다.

목표 지표(Target metrics)

Mac M2 노트북에서 4-layer, 4-head, d_model=128 GPT를 tinyshakespeare.txt에 대해 2,000 스텝 학습했을 때:

학습 손실(training loss)이 임의 초기값 약 4.2에서 약 1.5까지 약 6분 안에 내려갑니다.
생성된 출력은 완전한 셰익스피어는 아니지만, 고어체 단어, 줄 바꿈, ROMEO: 같은 고유 명사가 나타나는 "셰익스피어 같은(Shakespeare-shaped)" 텍스트가 됩니다.
텍스트의 마지막 10%를 떼어 둔 검증 손실(val loss)이 학습 손실을 비슷하게 따라가며, 이 정도 규모와 예산에서는 과적합(overfitting)이 거의 발생하지 않습니다.

직접 만들기

이 lesson은 PyTorch를 사용합니다. CPU 빌드여도 괜찮습니다. code/main.py를 참고합니다. 이 스크립트는 다음 작업을 처리합니다.

원문은 tinyshakespeare.txt가 없으면 다운로드한다고 설명합니다. 현재 저장소의 code/main.py는 네트워크 없이 실행되도록, 파일이 없을 때 내장된 셰익스피어 발췌(embedded excerpt)를 사용합니다.
바이트 수준 문자 토크나이저(byte-level char tokenizer).
90/10 비율의 학습/검증 분할(train/val split).
지원되는 하드웨어에서 bf16 자동 캐스팅(autocast)을 사용하는 학습 루프.
학습이 끝난 뒤의 샘플링(sampling).

Step 1: data

text = open("tinyshakespeare.txt").read()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
encode = lambda s: [stoi[c] for c in s]
decode = lambda xs: "".join(itos[x] for x in xs)

고유 문자(unique character)는 65개로, 어휘 사전(vocabulary)이 아주 작습니다. 4바이트 vocab_size에 충분히 들어갑니다. BPE도 필요 없고, 토크나이저 관련 골치 아픈 문제도 없습니다.

Step 2: model

code/main.py를 참고합니다. 블록은 Lesson 05에서 설명한 교과서적인 구조 그대로입니다. 사전 정규화(pre-norm), RMSNorm, SwiGLU, 인과 멀티 헤드 어텐션(causal MHA)으로 구성됩니다. 4/4/128 설정에서 파라미터 수는 약 800K입니다.

Step 3: training loop

길이 256짜리 토큰 윈도(token window)를 무작위 배치(random batch)로 가져옵니다. 순전파, 한 칸 밀린(shift-by-one) 교차 엔트로피(cross-entropy), 역전파, AdamW 스텝, 로그(log)를 반복합니다.

for step in range(max_steps):
    x, y = get_batch("train")
    logits = model(x)
    loss = F.cross_entropy(logits.view(-1, vocab_size), y.view(-1))
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    opt.zero_grad()

Step 4: sample

프롬프트(prompt)가 주어지면 순전파를 반복합니다. 매 스텝마다 top-p 로짓에서 표본을 뽑아 이어 붙이고, 500 토큰이 되면 멈춥니다.

Step 5: 출력 읽기

2,000 스텝 학습 후 다음과 같은 출력이 나옵니다.

ROMEO:
Away and mild will not thy friend, that thou shalt wit:
The chief that well shame and hath been his friends,
...

완전한 셰익스피어는 아니지만 "셰익스피어 같은" 텍스트입니다. 약 800K 파라미터와 노트북에서의 6분 학습으로 얻을 수 있는 분명한 성과입니다.

사용해보기

이 capstone은 일종의 참조 아키텍처(reference architecture)입니다. 실제 시스템에 가깝게 확장하려면 세 가지를 더합니다.

토크나이저(tokenizer)를 교체합니다. tiktoken.get_encoding("cl100k_base") 같은 BPE를 사용합니다. 어휘 크기(vocab size)는 65에서 약 50,000으로 늘어나고, 이에 맞춰 모델 용량(model capacity)도 키워야 합니다.
더 큰 말뭉치(corpus)로 학습합니다. OpenWebText나 fineweb-edu(HuggingFace)를 사용합니다. 단일 A100에서 125M 파라미터 GPT를 10B 토큰에 학습하려면 약 24시간이 걸립니다.
RoPE + KV 캐시 + Flash Attention을 추가합니다. 아래 연습문제가 각 단계를 안내합니다.

이런 확장을 거치면 유창한 영어를 생성하는 125M 파라미터 GPT가 만들어집니다. 최전선(frontier) 모델은 아니지만, 같은 코드 경로(code path)를 단지 더 크게 돌리는 방식이 2026년에 Karpathy, EleutherAI, Allen Institute가 연구용 체크포인트(research checkpoint)를 학습할 때 사용하는 방식과 같습니다.

산출물 만들기

outputs/skill-transformer-review.md를 봅니다. 이 스킬(skill)은 Phase 7의 이전 13개 lesson 기준에 비추어, 처음부터 작성한 트랜스포머 구현이 올바른지 검토(review)합니다.

연습문제

쉬움. code/main.py를 실행합니다. 학습된 모델의 최종 스텝 검증 손실(validation loss)이 2.0 아래로 내려가는지 확인합니다. max_steps를 2,000에서 5,000으로 바꾸면 검증 손실이 계속 좋아지나요?
중간. 학습된 위치 임베딩(learned positional embedding)을 RoPE로 교체합니다. MultiHeadAttention 내부에서 Q와 K에 회전(rotation)을 적용합니다. 학습 후 검증 손실이 최소한 비슷한 수준인지 확인합니다.
중간. 샘플링 루프에 KV 캐시를 구현합니다. 캐시가 있을 때와 없을 때 500 토큰을 각각 생성합니다. 노트북에서 실측 시간(wall-clock)이 5~20배 정도 빨라져야 합니다.
어려움. 그다음 다음 토큰(next-plus-one token)을 함께 예측하는 두 번째 헤드(head)를 모델에 추가합니다. DeepSeek-V3의 다중 토큰 예측(Multi-Token Prediction; MTP) 아이디어입니다. 함께(jointly) 학습했을 때 도움이 되나요?
어려움. 블록 안의 단일 FFN을 4-expert MoE로 교체합니다. 라우터(router)와 top-2 라우팅을 구현합니다. 활성 파라미터(active parameter)를 맞춘 상태에서 검증 손실이 어떻게 변하는지 확인합니다.

핵심 용어

용어	흔한 설명	실제 의미
nanoGPT	"Karpathy의 튜토리얼 저장소"	약 300줄로 구성된 최소한의 디코더 전용 트랜스포머 학습 코드이며, 사실상의 표준 참조 구현이다.
tinyshakespeare	"표준 장난감 말뭉치(toy corpus)"	약 1.1 MB 분량의 텍스트로, 2015년 이후 거의 모든 문자 단위 LM 튜토리얼이 사용해 온 말뭉치이다.
묶인 임베딩(Tied embeddings)	"입출력 행렬을 공유한다"	LM 헤드의 가중치를 토큰 임베딩 행렬의 전치(transpose)와 같게 두는 기법이다. 파라미터를 줄이면서 품질도 향상시킨다.
bf16 autocast	"학습 정밀도 트릭"	순전파와 역전파는 bf16으로 수행하고, 옵티마이저 상태는 fp32로 유지한다. 2021년 이후 표준에 가까운 설정이다.
그래디언트 클리핑(gradient clipping)	"스파이크를 막는다"	전역 그래디언트 노름(global grad norm)을 1.0으로 제한해 학습이 발산하는 것을 방지한다.
코사인 학습률 스케줄(cosine LR schedule)	"2020년 이후의 기본값"	학습률을 워밍업(warmup) 구간에서 선형으로 올린 뒤, 코사인 형태로 정점의 10% 수준까지 낮춘다.
MFU	"Model FLOP Utilization"	실제 달성한 FLOPs를 이론적 최대 FLOPs로 나눈 비율이다. 2026년 기준으로는 밀집 모델(dense) 40%, MoE 30%면 우수한 편이다.
검증 손실(val loss)	"보유 데이터에 대한 손실"	모델이 한 번도 본 적 없는 데이터에 대한 교차 엔트로피이다. 과적합 여부를 감지하는 지표다.

더 읽을거리

The Annotated Transformer (Harvard NLP) — 고전이 된 주석 달린 트랜스포머 구현입니다.
Karpathy nanoGPT — 디코더 전용 트랜스포머를 작은 규모로 끝까지 구현한 참조 코드입니다.