[CloudNeta] Hands-On LLM Serving 1주차 - GPT-2 해부도: state_dict 이름과 구조의 자리

#AI #LLM #MLOps #cloudneta #transformers

이 문서는 1주차 part 2 - LLM과 트랜스포머 기본의 트랜스포머 절에서 갈라져 나온 실물 편입니다. 거기서 구조를 개념으로 훑었다면, 여기서는 그 구조가 실제 파일 안에서 어떤 이름으로 앉아 있는지 봅니다.

용어가 낯설다면 LLM 기초 For Dummies를 먼저 보시면 수월합니다.

state_dict를 열면 h.0.attn.c_attn.weight (768, 2304) 같은 이름이 쏟아집니다. 아래 숫자는 전부 openai-community/gpt2를 실제로 올려서 잰 값입니다.

먼저 숫자

항목	값	어디서 오는가
블록 수	12	`h.0` 부터 `h.11`
hidden 차원	768	토큰 하나를 나타내는 벡터 길이
어텐션 헤드	12	768을 12로 나눠 헤드당 64
어휘 크기	50,257	`wte`의 행 개수
최대 위치	1,024	`wpe`의 행 개수
총 파라미터	124,439,808	흔히 "124M"이라 부르는 그 숫자

헤드당 차원이 64라는 점을 기억해 두시면 좋습니다. 어텐션 수식의 에서 그 가 768이 아니라 64입니다. 점수를 나눠 주는 값이 8이라는 뜻이죠.

골격: 이름과 자리를 맞춰 보기

state_dict의 키를 위에서 아래로 늘어놓으면 그대로 모델의 통과 순서가 됩니다. 가운데 블록만 열두 번 반복되고, 나머지는 앞뒤에 한 번씩 붙습니다.

h.N의 N이 0부터 11까지 돕니다. 열두 블록은 이름만 다를 뿐 모양이 똑같습니다. 입출력 크기가 같아서 레고처럼 쌓을 수 있습니다.

회색으로 적은 attn.bias만 성격이 다릅니다. 크기가 백만 개를 넘어서 처음 열면 제일 큰 항목처럼 보이는데, 미래 토큰을 가리는 삼각형 마스크라 학습으로 바뀌지 않고 그래서 파라미터로 세지 않습니다. part 2에서 마스크드 어텐션이라 불렀던 그 가림막입니다.

입구와 출구가 같은 행렬입니다

lm_head는 따로 파라미터를 갖지 않습니다. wte를 전치해서 재사용합니다. 실제로 두 텐서의 메모리 주소가 같습니다. 이것을 weight tying이라고 부릅니다.

입구에서 "번호 → 벡터"로 쓰던 표를 출구에서 "벡터 → 번호별 점수"로 뒤집어 쓰는 셈이라, 3,860만 개를 한 벌만 들고 두 군데에 씁니다.

블록 하나: Pre-LN과 잔차 연결

블록 안에서 어텐션과 FFN은 결과를 그대로 넘기지 않습니다. 원래 입력을 결과에 다시 더해서 넘깁니다. 이것이 잔차 연결입니다. 그리고 GPT-2는 각 층에 들어가기 전에 정규화합니다.

들어온 모양과 나가는 모양이 [T, 768]로 같습니다. 그래서 같은 블록을 열두 번 쌓을 수 있습니다.

두 구조의 이유는 part 2의 해당 절에서 다뤘습니다. 요약하면 잔차 연결은 기울기가 블록 열두 개를 지나도 지름길로 흐르게 하고, Pre-LN은 깊은 네트워크에서 그 기울기를 더 안정적으로 만들어 줍니다.

c_attn: 2304은 어디서 온 숫자인가

h.0.attn.c_attn.weight의 모양이 (768, 2304)인 게 처음엔 뜬금없어 보입니다. 2304를 3으로 나누면 768입니다. Q, K, V 세 벌을 한 행렬에 붙여 놓은 것입니다. 행렬곱을 세 번 하는 대신 한 번에 끝내려는 최적화죠.

2304 = 3 × 768이고 768 = 12 × 64입니다. 이 두 나눗셈이 c_attn의 정체 전부입니다.

헤드마다 다른 관계에 집중하도록 학습됩니다. 어떤 헤드는 문법을, 어떤 헤드는 의미를 봅니다. 그래서 c_proj (768, 768)이 필요합니다. 열두 헤드가 각자 본 것을 한 번 더 섞어 주는 자리입니다.

파라미터 배분: 1억 2천만 개가 어디에 쓰였나

작은 모델일수록 임베딩 비중이 큽니다. 어휘 크기는 모델이 작다고 줄지 않으니, 124M짜리에서는 그 표 하나가 3분의 1을 먹습니다.

그리고 블록 안에서는 어텐션이 주인공처럼 보이지만 무게는 FFN이 두 배입니다. c_fc가 768을 3072로 네 배 부풀렸다가 c_proj가 되돌리기 때문입니다.

텐서	모양	개수	하는 일
`wte.weight`	`(50257, 768)`	38,597,376	번호를 벡터로. 출구에서 재사용
`wpe.weight`	`(1024, 768)`	786,432	몇 번째 자리인지 더해 줌
`h.N.ln_1`	`(768,)` ×2	1,536	어텐션 앞 정규화
`h.N.attn.c_attn`	`(768, 2304)`	1,771,776	Q, K, V를 한 번에 생성
`h.N.attn.c_proj`	`(768, 768)`	590,592	헤드 12개 결과를 섞음
`h.N.ln_2`	`(768,)` ×2	1,536	FFN 앞 정규화
`h.N.mlp.c_fc`	`(768, 3072)`	2,362,368	네 배로 부풀림
`h.N.mlp.c_proj`	`(3072, 768)`	2,360,064	다시 되돌림
`ln_f`	`(768,)` ×2	1,536	마지막 정규화
`h.N.attn.bias`	`(1,1,1024,1024)`	0	인과 마스크. 학습 대상이 아닌 버퍼

hidden을 두 배로 키우면 블록 안 행렬이 대부분 hidden × hidden이라 파라미터는 네 배가 됩니다. GPT-3이 768에서 12,288로 16배 키웠으니 블록 하나의 무게는 256배가 된 셈입니다. 거기에 층까지 12에서 96으로 늘렸습니다. 1.25억이 1,750억이 되는 길입니다.

이름을 읽는 법

h.7.attn.c_attn.weight
└─┬┘
  여덟 번째 블록 (0부터 세니까)

h.7.attn.c_attn.weight
    └─┬┘
      어텐션 쪽. mlp 면 FFN 쪽

h.7.attn.c_attn.weight
         └──┬─┘
            Q, K, V 를 만드는 행렬
            c_proj 면 헤드를 다시 섞는 행렬

한 줄로

GPT-2는 같은 블록 열두 개를 쌓고 앞뒤에 임베딩과 정규화를 붙인 것입니다. 블록 하나는 다시 어텐션과 FFN 두 덩어리이고, 각각 앞에 LayerNorm이 있고 뒤에 잔차 연결이 있습니다. state_dict에 보이는 이름은 전부 이 구조 위의 자리 하나씩을 가리킵니다.

직접 확인해보기

from transformers import AutoModelForCausalLM, AutoConfig
c = AutoConfig.from_pretrained("gpt2")
m = AutoModelForCausalLM.from_pretrained("gpt2")
sd = m.state_dict()

print(c.n_layer, c.n_embd, c.n_head, c.vocab_size)     # 12 768 12 50257
print(sum(p.numel() for p in m.parameters()))          # 124439808

for k, v in sd.items():
    if k.startswith("transformer.h.0."):
        print(f"{k:44} {tuple(v.shape)}")

# lm_head 와 wte 가 같은 메모리를 쓰는지 확인
print(m.lm_head.weight.data_ptr() == sd["transformer.wte.weight"].data_ptr())