LLaMA 2 밑바닥부터 구현하기: 회전 위치 인코딩·RMSNorm·KV 캐시·그룹 쿼리 어텐션 해설

핵심 메시지

• 원조 트랜스포머가 번역 같은 과제를 위한 인코더-디코더 구조였다면, 라마는 앞의 토큰만 보고 다음 토큰을 이어 붙이는 인과적 언어 모델이다. 강의는 전체 뼈대를 먼저 세운 뒤 블록을 하나씩 채워 넣는 순서로 진행된다.
• 회전 위치 인코딩은 절대 방식과 상대 방식의 중간이다. 토큰마다 자기 위치에 비례하는 각도로 벡터를 회전시켜 두면, 어텐션의 내적이 두 토큰의 임베딩과 둘 사이의 거리에만 의존하게 된다는 것이 원 논문의 출발점이다.
• RMSNorm은 레이어 정규화에서 평균을 빼는 재중심화를 없앤 형태다. 정규화 효과의 본질이 재중심화가 아니라 재조정에 있다는 가설에 따라 평균 없이 제곱평균제곱근만으로 크기를 맞추고, 학습되는 감마 파라미터를 곱한다.
• KV 캐시는 추론에서 이미 계산한 내적을 다시 계산하지 않게 해 준다. 마지막 토큰 하나만 쿼리로 넣고 키와 값은 캐시에 쌓아 두는 방식으로, 네 번째 단계 기준 내적 횟수가 16회에서 4회로 줄어든다.
• KV 캐시가 연산량을 줄이자 이번에는 메모리 접근이 병목이 됐다. GPU는 계산보다 메모리 이동이 훨씬 느리기 때문에, 키·값의 헤드 수를 줄인 멀티 쿼리 어텐션과 그 절충안인 그룹 쿼리 어텐션이 등장했다.

쉽게 이해하기

이 영상은 라마 2의 모델 코드를 파이토치로 한 줄씩 작성하면서, 각 구성 요소가 왜 그렇게 생겼는지를 수식과 함께 설명하는 긴 실습 강의다. 진행자는 공개된 가중치를 그대로 불러 쓰기 위해 메타의 원본 코드와 변수 이름을 최대한 맞췄고, 병렬화처럼 자기 환경에서 검증할 수 없는 부분은 덜어 냈다고 밝힌다. 학습이 아니라 추론용 코드라는 점도 처음부터 분명히 한다.

가장 많은 시간을 들이는 대목은 회전 위치 인코딩이다. 원조 트랜스포머가 위치마다 정해진 벡터를 임베딩에 더했다면, 라마는 어텐션 계산 직전에 쿼리와 키에만 회전을 적용한다. 값 행렬에는 적용하지 않는다. 강의는 복소수와 오일러 공식을 동원해 회전 행렬이 왜 나오는지 유도한 뒤, 0이 대부분인 행렬을 곱하는 낭비를 피하려고 원소별 곱으로 바꾼 효율적인 구현 형태까지 따라간다. 차원이 짝수여야 한다는 제약, 기준값 10000, 위치와 각도의 모든 조합을 외적으로 미리 계산해 두는 이유도 함께 설명된다.

정규화 층에서는 레이어 정규화와 RMSNorm의 차이를 짚는다. 레이어 정규화는 평균과 분산 두 통계를 구해 값을 0 근처로 옮기고 크기를 맞추지만, RMSNorm은 평균을 아예 구하지 않는다. 원 논문이 레이어 정규화의 성공 요인을 재중심화가 아니라 재조정에서 찾았기 때문이고, 그래서 평균에 의존하지 않는 제곱평균제곱근 통계만 쓴다. 분모가 0이 되지 않도록 아주 작은 값을 더하는 처리도 코드에서 그대로 드러난다.

어텐션 부분은 추론 최적화의 역사를 압축해 보여 준다. 다음 토큰 하나만 필요한데도 매 단계마다 앞선 토큰들의 내적을 반복하는 낭비를 KV 캐시가 걷어내고, 그러자 드러난 메모리 대역폭 병목을 멀티 쿼리 어텐션이 키·값 헤드를 하나로 줄여 해결한다. 강의는 이 방식이 품질을 크게 해치지 않으면서 토큰당 처리 시간을 크게 줄였다는 논문 수치를 인용하고, 품질과 속도를 절충한 그룹 쿼리 어텐션으로 이어진다. 다만 라마의 공식 구현은 최적화된 커널 대신 부족한 키·값 헤드를 그냥 복제해 일반 멀티 헤드처럼 계산한다는 점도 솔직하게 짚는다.

마지막은 토큰을 고르는 전략이다. 매번 최고 확률만 고르는 탐욕적 방식, 여러 후보 경로를 살려 두는 빔 서치, 소프트맥스 전에 로짓을 나눠 모델의 확신 정도를 조절하는 온도, 분포에서 무작위로 뽑는 방식, 상위 K개만 남기는 방식이 차례로 비교된다. 강의는 상위 K가 분포 모양에 따라 형편없는 토큰까지 끌고 들어올 수 있다는 약점을 지적하며, 누적 확률이 기준값을 넘을 때까지만 후보를 남기는 상위 P 방식을 라마 공식 코드와 같은 선택지로 구현한다.

주요 인사이트

• 회전 위치 인코딩의 목적은 "위치를 알려 주는 것"이 아니라 어텐션의 내적이 두 토큰의 상대 거리에만 반응하도록 만드는 데 있다. 그래서 임베딩 직후가 아니라 헤드로 나눈 뒤 어텐션 직전에, 그것도 쿼리와 키에만 적용된다.
• KV 캐시를 쓰는 모델은 매 단계 입력 길이가 1이다. 강의 코드가 시퀀스 길이를 1로 못박고 학습에는 쓸 수 없다고 선을 긋는 이유가 여기에 있다. 프롬프트 토큰들도 결과를 보려고가 아니라 캐시를 채우려고 한 번씩 흘려 넣는다.
• 최적화 대상이 연산량에서 메모리 접근으로 옮겨 간 과정이 이 강의의 숨은 줄거리다. GPU 사양표에서 연산 성능과 메모리 대역폭의 격차를 직접 확인시키며, 알고리즘을 고를 때 텐서를 얼마나 옮기는지도 함께 따져야 한다고 강조한다.
• 피드포워드 층의 은닉 차원 계산식은 성능 최적화가 아니라 비교 가능성의 산물이다. 키·값 헤드를 줄여 줄어든 파라미터를 피드포워드 쪽에서 되돌려 총 파라미터 수를 맞추는 관행이며, 활성화 함수 변경에 대해 원 논문조차 "왜 더 잘 되는지 설명하지 않는다"고 적었다는 점도 인용된다.
• 상위 P가 상위 K보다 나은 이유는 후보 개수를 고정하지 않기 때문이다. 분포가 평평하면 후보를 넓게 잡아 다양성을 남기고, 한 토큰에 확률이 몰려 있으면 사실상 그 하나만 남겨 엉뚱한 선택을 막는다.

자주 묻는 질문

관련 AI 소식