최적의 배치 크기(Batch size)를 갖는 것은 추론 시 모델의 토큰당 비용을 줄일 수 있습니다.

이 글은 배치 크기가 추론 비용에 어떻게 영향을 미치는지에 대한 설명이 될 것입니다. 우리는 깊이 있게 파고들어 기초부터 프레임워크를 구축해 나갈 것입니다. 이 개념에 대한 비형식적인 설명이므로 격식 있는 언어를 기대하지 마세요.

이 글에서는 plot과 graph라는 단어가 혼용되어 사용되므로 혼동하지 마세요.

종속성 (Dependencies)

가장 먼저 우리가 이해해야 할 것은 LLM 추론이 실제로 무엇에 의존하는가입니다.

위 방정식에서 우리는 추론을 수행하는 데 걸리는 실제 시간이 두 가지 주요 요인에 달려 있음을 알 수 있습니다.

1 -> 메모리 (Memory)

KV cache 및 모델 가중치(Model weights) 가져오기 시간

2 -> 연산 (Compute)

수학적 연산을 수행하는 데 걸리는 시간

방정식에 따르면, 추론에 걸리는 시간은 연산에 걸리는 시간과 메모리를 가져오는 데 걸리는 시간 중 최댓값보다 크거나 같습니다.

메모리를 가져오는 데 걸리는 시간 (t_memory)

위 방정식에서 메모리를 가져오는 데 걸리는 시간은 모델의 바이트(bytes) 크기와 KV 캐시(kv cache)의 바이트 합계를 메모리 대역폭(memory bandwidth)으로 나눈 값입니다.

여기서 KV 캐시는 다음과 같습니다.

B × len_ctx × (bytes/token)

연산에 걸리는 시간 (t_compute)

t_compute 방정식은 배치 크기(batch size)에 활성 파라미터(active parameter) 수를 곱한 값을 FLOPS로 나눈 값입니다.

어텐션 연산 (attention computation)에 소요되는 시간은 가중치 행렬 곱셈 (weight matrix multiplication) 비용과 비교했을 때 사소하기 때문에 고려하지 않습니다.

그래프 그리기

여기에 있는 것은 지연 시간 (latency) 그래프입니다. 우리는 기본적으로 앞서 다루었던 모든 값들을 이 그래프에 나타내고 있습니다.

plot 1.1 

• t_compute는 배치 크기 (batch size)가 증가함에 따라 선형적으로 증가합니다.
• t_kv-fetch 또한 선형적으로 증가합니다.
• t_weight-fetch는 전혀 증가하지 않고 일정하게 유지됩니다.

plot 1.2 

• t_weight-fetch와 t_kv-fetch의 결합을 t_memory로 플로팅합니다.
• t_memory의 y 절편이 정확히 t_weight-fetch의 절편에서 시작하는 것을 확인할 수 있습니다.

plot 1.3

• 여기서는 t_compute와 t_memory 중 최댓값을 취합니다.
• 이제 이 그래프는 배치 크기 (batch size)가 증가함에 따라 지연 시간 (latency)이 어떻게 증가하는지를 알려줍니다.
• 여기서 주목할 점은 지연 시간에 하한선 (lower bound)이 존재한다는 것입니다. 모델 파라미터 (model parameters)를 로드하는 데 걸리는 시간은 일정하기 때문입니다.

우리는 배치 크기 (batch size)의 함수로서 지연 시간 (t)을 도식화해 왔지만, 토큰당 비용 (cost per token)을 구하기 위해서는 (t/B) 비용을 배치 크기의 함수로 도식화해야 합니다.

plot 2.1

그래프 2.1은 (B)의 함수로서 (t/B)를 도식화한 것이며, 변수들이 상당히 많이 변합니다.

• t_compute는 선형적(linear)이었으나, 이를 (B)로 나눈 후에는 상수가 됩니다.
• t_KV-fetch 또한 선형적이었으나, 이를 (B)로 나눈 후에는 상수가 됩니다.
• t_weight-fetch는 상수였으나, 이를 (B)로 나눈 후에는 쌍곡선(hyperbolic curve)이 됩니다.

plot 2.2

• t_weight-fetch와 t_kv-fetch의 결합을 t_memory로 도식화합니다.
• 이 결합이 쌍곡선을 위로 이동시키는 현상이 발생합니다.

plot 2.3

• 그래프 2.2의 최댓값(maximum)을 취하면 그래프 2.3을 얻게 됩니다.
• 이제 토큰당 비용 (cost per token) 도식을 확인할 수 있습니다.

관찰 사항 (Observations)

• 그래프 2.3을 보면 배치 크기(Batch Size)를 증가시킴에 따라 비용이 하한선(lower bound)에 도달할 때까지 감소하는 것을 볼 수 있습니다.
• 원한다면 비용과 추론 시간(inference time) 사이의 트레이드오프(tradeoff)를 수행할 수 있습니다. 즉, 배치 크기를 줄이면 비용은 올라가지만 추론 시간은 감소합니다.
• 비용에는 $t_{compute}$라는 하한선이 존재합니다.
• 여기서 주의할 점(caveat)이 있습니다. 배치 크기를 증가시키다 보면 어느 시점에서 $t_{kv-fetch}$가 $t_{compute}$를 초과하게 되며, 그 시점부터는 연산(compute)이 더 이상 하한선이 아니게 됩니다. 대신 메모리 대역폭(memory bandwidth)이 병목이 되며 MFU (Model Flop Utilization, 모델 연산 효율)가 크게 감소합니다. 따라서 배치 크기가 너무 크면 토큰당 비용(cost/token)이 증가합니다.

읽어주셔서 감사합니다 :)