GPU/TPU의 두 병목

딥러닝 연산을 빠르게 실행하려면 두 가지 일이 모두 필요하다.

GPU/TPU 성능을 볼 때 자주 나오는 두 숫자는 이 두 일을 나타낸다.

Peak compute throughput: 초당 처리할 수 있는 FLOPs
Memory bandwidth: 초당 이동할 수 있는 bytes

연산 유닛

연산 유닛은 실제 곱셈과 덧셈을 수행한다. GPU에서는 CUDA core, tensor core 같은 단위가 있고, TPU에서는 MXU 같은 행렬 연산 유닛이 이 역할을 한다.

예를 들어 어떤 가속기의 BF16 peak compute가 C FLOP/s라면, 연산량이 F FLOPs인 작업의 순수 계산 시간 하한은:

T_\text{compute} = \frac{F}{C}

이다.

HBM은 가속기 옆에 붙어 있는 고대역폭 메모리다. 모델 가중치, activation, KV cache 같은 텐서가 여기에 저장된다.

메모리 대역폭이 BW bytes/s이고, 연산을 위해 M bytes를 이동해야 한다면 순수 메모리 이동 시간 하한은:

T_\text{memory} = \frac{M}{BW}

이다.

연산량만 보면 행렬 곱셈은 항상 빠르게 보일 수 있다. 하지만 데이터를 충분히 빠르게 가져오지 못하면 연산 유닛은 놀게 된다.

반대로 메모리 이동량이 작고 재사용이 잘 되면, 병목은 연산 유닛 쪽으로 이동한다.

그래서 성능 분석의 첫 질문은 항상 이것이다.

이 연산은 계산이 느린가, 데이터를 가져오는 것이 느린가?

이 질문을 수식으로 바꾸면 compute-memory-time으로 이어진다.