TPU에서 LLaMA 3 서빙하기

LLaMA 3-70B는 KV head가 8개이므로 토큰당 크기는 2 * K * H * L = 2 * 8 * 128 * 80 = 160kB이다.

이게 얼마나 큰지 눈여겨보라! 시퀀스 길이가 (흔히 그렇듯) 32k 토큰이면 시퀀스당 160e3 * 32,768 = 5.3GB를 쓴다. BS=240이면 1.3TB다! TPU v5e는 개당 16GB뿐이므로 이만큼의 메모리를 담는 데만 약 (70e9 + 1.3e12) / 16e9 = 86개의 TPU v5e 칩이 필요하다. 모델 파라미터 70GB에 비해 얼마나 큰지도 눈여겨보라.

KV가 int8에서 160e3 바이트이므로 총 KV 메모리는 160e3 * 8192 * 32 = 41.9e9 바이트다. 파라미터는 파라미터당 1바이트이므로 70e9 바이트다. 따라서 총 메모리 사용량은 41.9e9 + 70e9 = 112GB다.

쓸 수 있는 가장 작은 slice는 112e9 / 16e9 = 7 TPU, 즉 (짝수 크기로 반올림하면) TPU v5e 4x2다. 꽤 빠듯하게 들어가는 수준이라 다른 오버헤드까지 감안하면 다 안 들어갈 수도 있으므로, 최소 4x4가 필요할 수도 있다(아니면 batch size를 낮추거나).

이전 장의 공식을 그대로 가져오면 된다:

여기서는 파라미터가 int8이고 FLOPs는 bfloat16이므로 임계 batch size는 약 120이다. 오른쪽 항의 max를 직접 계산할 수도 있지만, 그건 이미 여러 번 해 본 계산이다. matmul과 FLOPs 양쪽 모두 memory-bound 영역 깊숙이 들어와 있다.

그러면 메모리 bandwidth만 엄밀히 봤을 때 step time은 기본적으로 (KV size + param size) / (8 * HBM bandwidth) = 112e9 / (8 * 8.2e11) = 17ms다. 따라서 이론적인 step time은 약 17ms다. throughput은 32 / .017 = 1882 tokens / sec, 칩당으로는 1882 / 8 = 235 tokens / sec / chip이 된다.

한 가지 확인해 둘 것은 matmul이 ICI bound일 가능성이다. 여기서는 2개 axis를 여기에 쓸 수 있으므로, 이론상 ICI bound가 되는 조건은 $Y > 2 * F / 2200 = 2 * 28672 / 2200 = 26$이다. 그러니 문제없다!

4x4에서 돌려도 ICI 측면에서는 여전히 괜찮으므로 latency는 17 / 2 = 8.5ms로 떨어지지만, 칩당 throughput은 그대로다.

7장에서 논의했듯이 $B \ll D, F$인 bfloat16 matmul에서는

이다. weight가 int8이면 분모에서 2배가 사라지므로 $2BDF / DF = 2B > 240$, 즉 $B > 120$이 되어 임계 batch size가 이전의 절반이 된다. 우리에게 정말 유리하다! int8 weight에 int8 FLOPs를 쓰면 TPU FLOPs/s로 int8 값을 써야 하는데, 이 값은 bfloat16의 1.97e14에서 3.94e14로 거의 두 배가 된다. 그러면 다시 원점인 약 $B > 240$으로 돌아온다.

int8 weight + bfloat16 FLOPs 조합은 꽤 흔하다. 파라미터를 손실 없이 quantize하는 것이 저정밀 산술 연산을 수행하는 것보다 쉬운 경우가 많기 때문이다.

쉽다! 아주 작은 batch size라도 괜찮다면 유일한 제약은 파라미터 메모리를 HBM에 담는 것이다. ceil(num_params * sizeof(dtype) / HBM per TPU), 다시 말해 ceil(70e9 * sizeof(dtype) / 16e9)를 가장 가까운 합리적인 topology(2의 배수)로 올림하면 된다:

꽤 멋지다! 원한다면 LLaMA 70B를 TPU v5e 2x2에 넣을 수 있다는 뜻이니까. 다만 KV cache 수가 아주 작다는 점이 눈에 띌 것이다. 그게 곧 우리의 batch size다! FLOPs 활용률이 끔찍하게 낮을 것이라는 뜻이다. batch size를 240까지 끌어올릴 수 있다면 우리는 기꺼이 더 큰 topology를 쓸 것이다.

이것도 쉽다. HBM을 꽉 채우도록 batch size를 골랐기 때문이다! TPU v5e 한 대 분량의 바이트를 MXU로 로드하는 데 얼마나 걸리는지의 문제일 뿐이다. 그냥 v5e HBM / v5e HBM memory bandwidth = 16GB / 8.2e11 = 19ms, 즉 step당 19ms다. 생성 길이의 중앙값이 512 토큰이라고 가정하면 decode 한 번에 약 9초다. batch size를 줄이면 latency가 조금은 나아질 수 있다는 점도 알아두자. 예를 들어 int4에서 모델 파라미터만 따진다면 HBM이 더 이상 꽉 차 있지 않으므로 최소 latency는 약 10ms / step이다.

토큰당 비용과 정확히 상관되는 값이므로 중요한 질문이다.

decode 길이 중앙값 가정을 쓰면 throughput은 그냥 $B / (\text{per-step latency} \cdot \text{median steps} \cdot N) \approx 43 / (0.019 * 512 * N)$이다. 대략 $(4.42 / N)$ QPS가 나오므로, $N$을 대입하면:

forward pass의 작업 메모리(activation과 attention에 할당되는 메모리)를 완전히 무시했으므로 상당히 낙관적인 수치다. Flash Attention이라면 터무니없는 가정은 아니지만 현실적이지도 않다. 실제 수치는 이것의 절반 정도일 가능성이 크다. 절대적인 최대 throughput을 원한다면 아마 칩 수를 두 배 이상으로 늘리고 batch size도 크게 키워야 할 것이다.

bfloat16으로 4x8 slice를 쓰면 KV cache용으로 372GB가 남아 batch size를 140까지 키울 수 있다. 그러면 step time은 그대로이므로 throughput은 14.39 / num_chips가 되어:

더 키우면 이득이 훨씬 더 커진다! 여기서 큰 교훈은, KV cache 크기에 제약을 받는 상황이라면 모든 경우에 가장 작은 topology가 가장 성능 좋은 topology는 아니라는 것이다.

이전 장에서 논의했듯이 generation 중의 sharding 선택지는 사실상 하나뿐이다: model parallelism. communication-bound가 되기 전까지 얼마나 할 수 있을까? 역시 이전 장에서 봤듯이, 모델은 대략 다음 조건에서 communication-bound가 된다:

LLaMA 3-70B는 F = 28,672이므로, model sharding에 2개 axis를 쓰면 대략 $Y = 28672 \cdot 2 / 2200 = 26$이 나온다. 즉 일반적으로 약 16개 칩까지는 communication-bound 없이 확장할 수 있고 4x4는 쓸 수 있지만 4x8은 안 된다. 게다가 연산을 완벽히 겹치지는 못하므로 이 추정치조차 지나치게 낙관적이다.

요점: 순수 model parallelism으로는 실제로 4x8에서 서빙할 수 없다. 여기서 할 수 있는 최선은 4x2이고, 어쩌면 4x4까지다.

하지만 앞서 논의했듯이 batch size가 작을 때는 모델이 FLOPs-bound가 아니라 memory-bandwidth-bound이므로, 흔히 throughput을 크게 해치지 않으면서 model parallelism을 더 할 수 있다. 앞에서 이 값이 대략 $Y=F / (8\cdot B)$라고 했으므로, batch size 64라면 이론상 ICI-bound가 되기 전까지 Y = 28,672 / (8 * 64) = 56-way model parallelism까지 갈 수 있다. 검산을 위해 단일 matmul의 $T_\text{ici comms}$, $T_\text{hbm comms}$, $T_\text{math}$를 보자. 명확하게:

4x8에서는 $T_\text{ici comms}$ = (2 * 64 * 8192) / 9e10 = 11us, $T_\text{hbm comms}$ = (2 * 8192 * 28,672) / (32 * 8.2e11) = 18us, $T_\text{math}$ = (2 * 64 * 8192 * 28,672) / (32 * 1.97e14) = 4us가 되어, 이론상 여전히 HBM bandwidth bound다. 아주 좋다! 다만 4x4에서 4x8로 키우는 것은 throughput 관점에서는 아마 별 도움이 안 되고, 대신 latency를 줄여 준다!

int8과 int4 config를 보면, 그쪽은 순수 model parallelism으로 할 수 있다. 이제 quantization이 더 빠른 FLOPs를 넘어 실질적인 이점을 주는 지점에 도달한 것이다: comms-bound가 되기 전에 더 큰 batch size를 쓸 수 있게 해 준다. 이 이야기의 결말은, 4x8에서 peak throughput은 달성할 수 없지만 int8과 int4 config라면 순수 model parallelism이 가능하다는 것이다.

8k 토큰이면 확실히 compute-bound이므로 FLOPs만 따지면 된다. 모델 파라미터가 70e9개이므로 forward pass 한 번에 2 * 70e9 * B FLOPs를 쓴다. 40% MFU(FLOPs 활용률)를 가정하면 실행 시간은 약 2 * 70e9 * 8192 / (16 * 1.97e14 * 0.4) = 0.91s다. 지금까지 봐 온 수치들과 비교하면 사실 꽤 긴 시간이다!

꽤 단순한 문제다. decode 길이 중앙값이 4096 토큰이므로 시퀀스는 대략 1 / 4096 토큰마다 하나씩 끝난다. batch size가 32이므로 step당 32 / 4096개의 시퀀스가 evict된다. KV cache 길이가 대략 8192 + 4096이므로 step당 evict되는 토큰은 32 * (8192 + 4096) / 4096 = 96개다. 일반 공식은 $B * (P + G) / G$이며, 여기서 $P$와 $G$는 각각 prefill과 generate 길이다.

꽤 재미있는 문제다. prefill 서버 수를 $P$, generate 서버 수를 $G$라고 하자. 일반적으로 P / prefill_latency의 속도로 시퀀스를 밀어 넣고 B * G / (generate_latency * median_decode_length)의 속도로 소비하는 pipeline 문제다. 앞에서 prefill step당 910ms, batch size 43(그냥 32라고 치자)에서 decode step당 19ms를 계산했다. 따라서 P / 0.91 = 32 * G / (0.019 * 512), 즉 P = 3G가 필요하다. generation 서버보다 prefill 서버가 약 3배 더 필요하다!

이 roofline들을 계산하는 코드는 다음과 같다:

import numpy as np

num_chips = 16  # we fix 16 as the amount of total model parallelism we do
bytes_per_param = 1  # int8 means 1 byte per param
param_count = 70e9
param_size = bytes_per_param * param_count
sequence_length = 8192  # can vary this

hbm_bandwidth = 8.20E+11  # v5e
flops = 1.97E+14  # v5e

def kv_cache_size(bs):
    return 2 * bs * 128 * 8 * 80

def min_topology(bytes):
    return 2 ** np.ceil(np.log2(bytes / 16e9))

def get_max_batch_size(
    num_chips: int,
    sequence_length: int,
    param_size: float,
) -> int:
  batch_sizes = np.arange(1, 1024, 4)
  kv_sizes = kv_cache_size(sequence_length * batch_sizes)
  required_chips = min_topology(kv_sizes + param_size)
  max_idx = np.where(required_chips <= num_chips)[0][-1]
  return max_idx

max_idx = get_max_batch_size(
    num_chips=num_chips,
    sequence_length=sequence_length,
    param_size=param_size,
)  # get the largest batch size that can fit
batch_sizes = np.arange(1, 512, 1)[:max_idx]
kv_sizes = kv_cache_size(sequence_length * batch_sizes)

kv_comms_time = kv_sizes / (num_chips * hbm_bandwidth)

param_comms_time = param_size / (num_chips * hbm_bandwidth)
param_comms_time = np.asarray([param_comms_time] * batch_sizes.shape[0])

flops_time = 2 * param_size * batch_sizes / (num_chips * flops)  # roughly true in a 2ND sense

mlp_time = np.maximum(flops_time, param_comms_time)
attn_time = kv_comms_time  # always bandwidth-bound for generate

latency = 1000 * (mlp_time + attn_time)
throughput = batch_sizes / (latency * num_chips)

latency를 KV 로딩과 파라미터 로딩이라는 두 원천으로 아주 명시적으로 분해했다는 점, 그리고 FLOPs와 comms 중 더 큰 쪽이 latency를 결정한다는 점을 눈여겨보라.