FLOPs

Attention에 $n^2$가 있다는 사실만으로 항상 layer 비용의 대부분이라고 판단할 수는 없다. Projection과 FFN에는 모델 폭 $d$의 제곱 항이 있고, 실제 비용은 $n$과 $d$를 함께 보아야 한다. 같은 모델에서도 context를 늘리면 어느 항이 큰지 달라진다. Parameters의 수가 아니라 이번 forward에서 수행하는 연산 수를 세어 병목이 바뀌는 지점을 찾는다.

1. Projection $8nd^{2}$ + FFN $16nd^{2}$ + Attention $4n^{2}d$
2. total $24nd^{2}$+$4n^{2}d$
3. $n$>6d crossover

Counting Operations

Forward dense vanilla full Self-Attention, $Q$/$K$/$V$/O 모두 $d\times d$, FFN4d, multiply+add=2 FLOPs 가정이다. 4 projections는 $8nd^{2}$, FFN 두 matmuls는 $16nd^{2}$, $QK^\top$/AV는 $4n^{2}d$다. Bias·Softmax·Norm·activation·position 연산은 제외하며 GQA/gated FFN이면 상수가 달라진다.

The Crossover

Attention share=$\frac{n}{6d+n}$, crossover는 n=6d다. d4096이면 24,576 tokens≈24K. n1K/8K/32K/128K의 share는 4/25/57/84%다.

128K의 logical grid는 head당약 $1.7\times 10^{10}$scores지만 실제 kernel이 grid 전체를 memory에 저장해야 한다는 뜻은 아니다.

Reading the Matrix Products

$(m\times k)(k\times r)$를 곱하면 출력은 $mr$개이고, 각 원소에 약 $k$번의 multiply–add가 필요하다. 이를 2 FLOPs로 세면 비용은 약 $2mkr$다. 위 가정을 각 component에 적용하면 다음과 같다.

Projection과 FFN은 token마다 feature를 변환하는 비용이므로 $n$에 선형, $d$에 제곱으로 커진다. Pairwise Attention은 positions 쌍을 비교하므로 $n^2$에 비례한다. $4n^2d=24nd^2$를 풀면 $n=6d$에서 두 항이 같고, 그보다 길어져야 pairwise Attention이 더 크다.

$d=4096$일 때 2K context에서는 pairwise 항이 약 8%에 불과하지만 32K에서는 약 57%다. 이 crossover는 vanilla 가정의 결과이며, 실제 GQA와 gated FFN에서는 constants가 달라진다. 모델 이름만 같게 적고 다른 architecture의 비율을 그대로 대입하지 않는다.

FLOPs Are Not Bytes

128K의 logical scores는 head 하나에 약 $1.7\times10^{10}$개다. 계산량이 크다는 사실과 이 모든 값을 GPU memory에 동시에 저장해야 한다는 주장은 다르다. 다음 section에서는 연산 수가 같은 계산도 memory 이동 방식을 바꾸어 실행 비용을 줄일 수 있음을 본다.

Parameters 개수와 forward FLOPs는 같은 양이 아니다. FLOPs와 HBM에 옮기는 bytes는 별도로 센다.

Vanilla crossover $n$>6d는 가정이 붙은 결과다.

Check

$4n^{2}d$가 $24nd^{2}$보다 커지는 조건을 유도하고 d4096의 crossover를 계산하라.

“Transformer는 항상 Attention이 연산량 대부분”이라는 명제의 반례와 조건은?

Related

• Memory Bottlenecks & FlashAttention
• Feed-Forward Network
• Vocabulary Size