GPU 구조와 플래시 어텐션 원리: 산술 강도·타일링·커널 퓨전으로 LLM 학습 속도를 끌어올리는 법

핵심 메시지

• GPU는 작은 제어 유닛으로 수많은 스레드에 같은 명령을 뿌리는 구조라, 조건 분기가 들어가면 나머지 스레드가 놀게 되어 효율이 떨어진다.
• 메모리 계층에서 공유 메모리 접근은 스무 사이클 남짓인데 글로벌 메모리는 삼백 사이클가량 걸려, 어디에 데이터를 두느냐가 곧 속도가 된다.
• 연산 성능의 증가 속도가 메모리 용량·대역폭보다 훨씬 가팔라, 요즘 최적화의 핵심은 계산을 줄이는 것이 아니라 메모리를 덜 오가게 만드는 것이다.
• 산술 강도(연산량 ÷ 통신량)가 하드웨어의 리지 포인트보다 높아야 연산 병목 상태가 되고, 낮으면 계산기가 데이터를 기다리는 통신 병목에 갇힌다.
• 저정밀도 연산·커널 퓨전·재계산·메모리 코얼레싱·타일링 다섯 가지를 결합한 결과물이 플래시 어텐션이며, 연산량은 오히려 늘었는데도 훨씬 빨라졌다.

쉽게 이해하기

강의는 최근 몇 년간 대형 언어 모델의 겉모습이 의외로 많이 바뀌지 않았다는 복습으로 시작한다. 디코더만 쓰는 구조, 왼쪽에서 오른쪽으로 토큰을 하나씩 만드는 자기회귀 방식, 위치 인코딩과 피드포워드 구성 같은 기본 골격은 사실상 표준으로 굳었다는 것이다. 대신 변한 것은 규모다. 몇 년 전 1~2조 토큰 수준이던 학습 데이터가 수십조 규모로 늘었고 파라미터도 조 단위 시대에 들어섰다. 이 정도 규모는 GPU 한 장에서 돌아갈 수 없기 때문에, 모델 이야기만으로는 설명이 끝나지 않고 인프라를 함께 봐야 한다는 것이 이날 강의의 출발점이다.

GPU와 CPU의 차이는 자동차 비유로 정리된다. CPU는 빠르지만 비싸고 수가 적은 경주용 차에 가깝고, GPU는 느리지만 아주 많이 굴릴 수 있는 오토바이 무리에 가깝다. 짧은 거리라면 빠른 차 한 대가 낫지만 물동량이 커지면 다수의 오토바이가 유리하다는 설명이다. 실제 칩 구조도 이 비유를 따라간다. CPU는 코어마다 큰 제어 유닛을 두지만 GPU는 제어 영역이 작은 대신 코어가 훨씬 많다. 그래서 GPU는 서른두 개의 스레드를 하나로 묶은 워프 단위로 같은 명령을 한꺼번에 실행하고, 조건 분기가 끼어들면 조건을 만족하지 않는 스레드들이 그냥 대기하는 손해가 생긴다.

이어지는 핵심 개념이 산술 강도다. 어떤 연산을 실행할 때 걸리는 시간은 계산 시간과 통신 시간 중 더 큰 쪽이 결정하는데, 우리가 실제로 원하는 것은 계산이고 데이터를 주고받는 시간은 버려지는 시간이다. 그래서 연산량을 데이터 이동량으로 나눈 값이 높을수록 좋다. 강의는 이 값을 하드웨어가 가진 최적점, 이른바 리지 포인트와 비교해 판단하는 방법을 보여 준다. 벡터 내적처럼 데이터만 잔뜩 옮기고 계산은 금방 끝나는 연산은 값이 0.5 수준으로 낮아 통신 병목이고, 행렬 곱은 배치 크기가 커질수록 값이 올라가 연산 병목 쪽으로 이동한다. 강의가 기준으로 든 GPU에서는 이 경계가 바이트당 295회 연산 근처였다.

다섯 가지 최적화 기법은 모두 이 값을 끌어올리는 방향으로 정렬된다. 첫째는 정밀도를 낮추는 것으로, 계산량은 그대로인데 옮기는 바이트만 절반이 되니 값이 자연스럽게 두 배가 된다. 다만 소프트맥스나 정규화처럼 표현 범위가 넓어야 하는 부분은 다시 올려서 계산하는 혼합 정밀도 전략을 쓴다. 둘째는 커널 퓨전이다. 연산을 한 줄씩 따로 쓰면 각 줄이 끝날 때마다 중간 결과를 메모리에 썼다가 다시 읽어 오는데, 여러 연산을 한 덩어리로 합치면 그 왕복이 통째로 사라진다. 셋째는 재계산으로, 중간 활성값을 저장해 두는 대신 역전파 때 다시 계산해 메모리 왕복을 줄인다. 계산은 늘지만 메모리 접근이 지배적인 연산에서는 오히려 더 빠르다. 넷째인 메모리 코얼레싱은 메모리가 어차피 인접 영역을 한 덩어리로 실어 보낸다는 성질을 이용해 워프 안의 스레드들이 필요로 하는 데이터가 그 덩어리와 겹치도록 배치하는 기법이고, 다섯째인 타일링은 큰 행렬을 작은 조각으로 잘라 조각 안에서는 빠른 공유 메모리만 쓰게 만드는 방식이다. 강의는 같은 타일링이라도 조각 크기를 어떻게 잡느냐에 따라 성능이 몇 배씩 벌어지고, 행렬 크기가 1 늘어나는 순간 필요한 타일 수가 프로세서 수를 넘겨 성능이 뚝 떨어지는 구간도 있다고 짚는다. 어휘 사전 크기를 조금 키웠더니 학습이 25% 빨라졌다는 유명한 사례도 이 정렬 문제로 설명된다.

마지막으로 플래시 어텐션이 이 다섯 가지의 종합판으로 소개된다. 표준 어텐션은 질의와 키를 불러와 곱하고, 그 결과를 다시 메모리에 썼다가 읽어 소프트맥스를 돌리고, 또 쓰고 읽는 식으로 메모리를 일곱 번이나 오간다. 문제는 소프트맥스가 전체 값을 다 봐야 분모를 구할 수 있어 타일링을 막는다는 점인데, 최댓값과 분모를 순회하면서 갱신하는 온라인 소프트맥스로 이 의존성을 풀어낸다. 그 결과 연산량은 오히려 늘었는데도 메모리 읽기·쓰기는 40기가바이트에서 4기가바이트로 열 배 가까이 줄었고, 실제 걸린 시간도 크게 짧아졌다. 성능을 볼 때 연산량만이 아니라 메모리 비용을 함께 봐야 한다는 이날 강의의 결론이 여기서 확인된다.

주요 인사이트

• AI 하드웨어 세대교체를 광고하는 성능 배수는 정밀도 축소, 명령어 개선, 공정 미세화, 희소성 등을 모두 곱한 값이라, 칩 자체가 그만큼 좋아졌다고 읽으면 과대평가가 된다.
• 최적화의 무게중심이 계산에서 데이터 이동으로 옮겨 갔다는 점이 중요하다. 같은 알고리즘을 그대로 두고 데이터를 어디에 두고 언제 옮길지만 바꿔도 실제 속도는 몇 배가 달라진다.
• 소프트맥스나 정규화 같은 기본 연산을 직접 구현하는 것은 대개 손해다. 널리 쓰이는 딥러닝 프레임워크의 구현은 이미 커널이 합쳐진 상태라 그대로 쓰는 편이 빠르다.
• 성능이 하드웨어 구조에 붙어 있는 만큼, 모델 설정값을 정할 때 프로세서 수나 타일 크기와 맞아떨어지는지 확인하는 것만으로도 공짜 속도를 얻는 경우가 있다.
• 플래시 어텐션은 새로운 수학이 아니라 기존 어텐션을 메모리 관점에서 다시 쓴 것에 가깝다. 알고리즘 혁신 못지않게 시스템 관점의 재설계가 모델 규모 확장을 떠받쳐 왔다는 뜻이다.

자주 묻는 질문

관련 AI 소식