핵심 개념
RoPE를 사용하는 변압기 모델(LLaMA, Mistral, PaLM, Gemma 등)의 첫 번째 변압기 레이어를 사전에 계산하여 추론 속도와 비용을 향상시킬 수 있다.
초록
이 논문은 RoPE(Rotary Position Embedding)를 사용하는 변압기 모델의 첫 번째 레이어를 사전에 계산하여 추론 속도와 비용을 향상시키는 기술을 소개한다.
병렬 구조의 변압기 모델(GPT-J, Pythia, PaLM 등)의 경우, 입력 임베딩에 의존하는 Q, K, V, FFN의 출력을 사전에 계산하여 메모리에 저장할 수 있다. 이를 통해 토큰당 계산 복잡도를 낮추고 메모리 읽기 횟수를 줄일 수 있다.
직렬 구조의 변압기 모델(LLaMA 2, Mistral, Mixtral 등)의 경우, Q, K, V는 사전에 계산할 수 있지만 FFN은 그렇지 않아 상대적으로 작은 이점을 얻을 수 있다.
사전 계산으로 인한 메모리 사용량 증가는 모델 크기와 구조에 따라 다르지만, Mistral-7B의 경우 2%만 증가한다.
통계
Pythia-6.9B 모델의 첫 번째 레이어 추론 시 메모리 읽기 횟수가 11,264배 감소
Mistral-7B 모델의 첫 번째 레이어 추론 시 메모리 읽기 횟수가 2,458배 감소
가설적인 병렬 구조의 Mixtral-8x7B 모델의 첫 번째 레이어 추론 시 메모리 읽기 횟수가 140,084배 감소
인용구
"For each token, we save the operations needed for FFN and the linear layers Q, K, V. This can speed up inference if the system is limited by compute."
"This can speed up inference for systems that are memory bandwidth limited, especially during the autoregressive next-token-prediction phase."