같은 계산을 n 번 반복하고 있었다

1편의 생성 코드는 이렇게 생겼다.

for _ in range(n):
    logits = model(idx[:, -CTX:])[:, -1, :]      # 문맥 전체를 넣는다
    idx = torch.cat([idx, sample(logits)], 1)

글자를 하나 붙이고 나서 처음부터 다시 넣는다. 100번째 글자를 뽑을 때 앞의 99개는 이미 지난 걸음에서 계산했던 것인데 또 계산한다.

그래도 답은 맞다. 인과 마스크가 있어서 앞쪽 토큰의 출력은 뒤에 무엇이 붙든 변하지 않기 때문이다. 13편에서 마스크를 뺐을 때 망가졌던 그 성질이 여기서는 다시 계산할 필요가 없다는 보장으로 쓰인다.

무엇을 저장하면 되나

블록 안에서 이전 토큰들에 대해 필요한 값은 딱 두 개다. 어텐션의 K 와 V. 새 토큰의 Q 는 그 둘과 만나기만 하면 된다.

q, k, v = qkv(norm(x)).split(D, dim=2)     # x 는 새 토큰 하나뿐
k = torch.cat([cache_k, k], dim=2)         # 지난 걸음의 K 를 이어 붙인다
v = torch.cat([cache_v, v], dim=2)
cache_k, cache_v = k, v
a = softmax(q @ k.transpose(-1, -2) / sqrt(dh))   # 마스크가 필요 없다

마스크가 사라진 것도 눈여겨볼 만하다. 새 토큰의 Q 는 자기보다 앞에 있는 K 만 보고 있으므로, 가릴 것이 애초에 없다.

같은 답이 나오는가

속도보다 이것이 먼저다. 같은 씨앗으로 128글자를 뽑아 비교한다.

글자열 일치            True
마지막 로짓 최대 차이   9.06e-06

글자열은 완전히 같다. 로짓에 9e-06 이 남는 것은 계산 순서가 달라서다. 한쪽은 128개를 한 번에 행렬곱하고 다른 쪽은 한 개씩 128번 쌓으니, 부동소수점 덧셈의 결합 순서가 다르다. 어느 쪽도 틀린 게 아니고 둘 다 같은 값의 다른 반올림이다.

여기서 정직하게 짚을 것이 있다. 비트 단위로 같지는 않다. 확률이 아주 가까운 두 후보가 붙어 있는 자리에서는 이 9e-06 이 다른 글자를 뽑게 할 수 있다. 위 실험에서는 안 그랬을 뿐이다.

산술은 130분의 1로 줄어든다

곱셈 횟수를 세어 보자. 캐시가 없으면 t 번째 걸음에서 t 개 토큰을 통과시킨다.

캐시 없음, t 번째 걸음   선형층 12·t·d²      어텐션 2·t²·d
캐시 있음, t 번째 걸음   선형층 12·d²        어텐션 2·t·d

n 걸음을 다 더하면 캐시 없는 쪽은 d²n² 과 dn³ 이 되고 있는 쪽은 d²n 과 dn² 이 된다. 각 항이 n 배씩 줄어든다.

d=128, 블록 3개, n=256 으로 세면 비가 134.5 다. 산술만 보면 130배 빨라져야 한다.

실제로는 3배다

같은 모형으로 길이만 바꿔 가며 재면 이렇다. 3회 이상 돌려 최솟값을 취했다.

     n   캐시 없음   캐시 있음   실측 배수   산술 배수
    64      71 ms      42 ms       1.70       32.9
   128     192 ms      78 ms       2.45       66.1
   256     582 ms     198 ms       2.94      134.5
   512    2031 ms     400 ms       5.08      277.5
  1024   11304 ms    1001 ms      11.29      579.9

n=256 에서 산술은 134.5 배를 약속하는데 실제는 2.94 배다. 약속의 2.2% 만 받았다.

이유는 이 모형이 너무 작아서다. d=128 짜리 행렬곱 하나는 CPU 에게 일도 아니고, 시간은 파이썬 한 줄 한 줄과 커널 실행 준비에 들어간다. 블록 3개 × 연산 여덟 개 × n 걸음이면 그 준비 횟수가 양쪽 다 같다. 줄인 것은 곱셈이고, 값을 치르고 있던 것은 곱셈이 아니었다.

그래도 캐시가 이긴다

그런데 표의 마지막 열이 아니라 실측 배수 열이 커지고 있다. 1.70 에서 11.29 로, 길이가 16배 늘어나는 동안 6.6배 벌어졌다.

기울기로 보면 분명해진다. 두 축을 로그로 놓고 적합하면

캐시 없음   시간 ∝ n^1.80
캐시 있음   시간 ∝ n^1.15

캐시가 바꾸는 것은 상수가 아니라 차수다. 상수 배는 오버헤드에 먹혀 안 보이지만 차수는 안 먹힌다. 그래서 짧을 때는 별것 없어 보이고 길어지면 반드시 이긴다.

1.80 이라는 값에는 단서를 붙여야 한다. 캐시 없는 쪽은 n² 항과 n³ 항이 섞여 있어 순수한 거듭제곱이 아니고, 적합 오차가 최대 23.4% 다. 캐시 있는 쪽은 10.4% 다. 정확한 지수라기보다 기울기가 확실히 다르다 는 표시로 읽어야 한다.

모형을 키우면 약속에 가까워진다

오버헤드가 문제라면 곱셈을 크게 만들면 된다. n=256 을 고정하고 폭만 키워 본다.

      d   캐시 없음   캐시 있음   실측 배수   산술 배수   받아낸 비율
    128     583 ms     195 ms       2.99      134.5        2.2%
    256    1203 ms     265 ms       4.54      131.7        3.4%
    512    3163 ms     415 ms       7.63      130.2        5.9%
   1024   10918 ms    1048 ms      10.42      129.4        8.1%

산술 배수는 130 근처에서 거의 안 움직이는데 실측은 2.99 에서 10.42 로 오른다. 받아낸 비율이 2.2% 에서 8.1% 로 는다. 곱셈이 커질수록 곱셈을 아낀 값이 실제 시간으로 돌아온다.

실제 모형은 d 가 4096쯤이고 블록이 수십 개다. 그 규모에서는 오버헤드가 무시할 수 있는 몫이 되고, 캐시의 이득이 산술 계산에 붙는다. 이 실험이 보여 주는 것은 작은 모형에서 잰 속도 이득을 큰 모형에 그대로 옮길 수 없다는 것이다.

값은 메모리로 치른다

캐시는 공짜가 아니다. 블록마다 K 와 V 를 토큰 수만큼 들고 있어야 한다.

캐시 크기 = 2 × 블록 수 × 토큰 수 × d × 바이트

이 모형에서 float32 로 재면

n =  128    0.38 MB
n =  512    1.50 MB
n = 1024    3.00 MB
가중치      2.43 MB   (637,156개 × 4바이트)

n=1024 에서 캐시가 가중치보다 1.23배 크다. 63만 개짜리 모형을 돌리는데 저장해야 할 것이 모형 자체보다 많아진다.

실제 규모에서는 이게 더 심해진다. d=4096, 블록 32개, n=8192, float16 이면 4.0 GB 다. 요즘 모형이 헤드마다 따로 두던 K, V 를 여러 헤드가 나눠 쓰게 바꾸는(GQA, MQA) 이유가 이것이고, 10편에서 헤드를 나눈 값이 여기서 다시 계산된다.

그래서

• 매 걸음 문맥을 다시 넣는 것은 인과 마스크 덕분에 답은 맞지만 낭비다. 필요한 것은 블록마다 K 와 V 두 개뿐이다
• 결과는 같다. 글자열은 완전히 일치하고 로짓 차이는 9.06e-06, 덧셈 순서에서 오는 반올림이다. 비트 단위로 같지는 않다
• 산술은 n=256 에서 134.5 배 줄어드는데 실제 시간은 2.94 배만 빨라진다. 받아낸 것이 2.2% 다. 작은 모형에서는 곱셈이 아니라 오버헤드가 값이다
• 그래도 캐시가 바꾸는 것은 차수다. n^1.80 대 n^1.15. 길어지면 반드시 이긴다
• 폭을 128 에서 1024 로 키우면 받아낸 비율이 2.2% 에서 8.1% 로 오른다
• 값은 메모리다. n=1024 에서 캐시 3.00 MB 대 가중치 2.43 MB. 실제 규모에서는 4.0 GB

다음 편에서는 그 메모리를 직접 줄여 본다. 가중치를 float32 에서 int8 로 바꾸면 무엇이 남고 무엇이 사라지는지 재 본다.