ThAME: 3D Memory-Enabled Heterogeneous Accelerator for LLM Mixture of Experts

인터넷이 매일 더 똑똑해지는 책들이 가득한 거대한 도서관이라고 상상해 보세요. 이 "똑똑한 책들"은 거대 언어 모델(LLM)이라 불리며, 이야기를 쓰고, 수학 문제를 풀고, 인간처럼 대화할 수 있습니다. 하지만 이 모델들이 정말 똑똑해지려면 수십억 개의 사실을 기억할 수 있는 거대한 용량이 필요합니다. 최근 과학자들은 "전문가 혼합(Mixture of Experts, MoE)"이라는 영리한 기술을 발견했습니다. 하나의 거대한 뇌가 모든 것을 기억하려고 노력하는 대신, MoE는 소규모 전문가 팀을 구축합니다. 질문을 던지면 시스템은 전체 팀을 깨우는 대신, 정답을 알고 있는 몇 명의 전문가만을 호출합니다. 이는 마치 사서에게 책을 가져다 달라고 요청하는 것과 같습니다. 도서관의 모든 직원을 깨우는 대신, 역사 섹션의 벨을 누르는 것과 같죠.

하지만 여기에는 함정이 있습니다. 현실 세계에서 이 전문가들을 호출하는 과정은 매우 무질서합니다. 전문가들은 사방에 흩어져 있고, 사서는 적절한 책을 가져오기 위해 앞뒤로 뛰어다녀야 하며, 그 후 답변을 모으기 위해 다시 달려와야 합니다. 이 "앞뒤로 뛰어다니는 과정"은 교통 체증을 유발합니다. 컴퓨터는 실제로 생각하는 시간보다 데이터가 도착하기를 기다리는 데 더 많은 시간을 보냅니다. 이것이 바로 "메모리 벽(memory wall)"이며, 가장 빠른 슈퍼컴퓨터조차 느려지게 만드는 병목 현상입니다. 문제는, 어떻게 하면 독자 바로 옆에 책이 있고, 러너(runners)들이 교통 체증을 피할 수 있는 완벽한 고속도로 시스템을 갖춘 도서관을 만들 수 있느냐는 것입니다.

---

문제점: 뇌 속의 교통 체증

이 논문은 바로 이 문제를 해결하기 위해 ThAME(3D 메모리 기반 이종 가속기)이라는 새로운 아키텍처를 소개합니다. 저자들은 현재의 컴퓨터가 "생각하는" 부분(CPU 또는 GPU)과 "데이터가 사는" 부분(메모리)이 서로 멀리 떨어져 있는 바쁜 도시와 같다고 설명합니다. 거대 언어 모델이 전문가 혼합(MoE) 기술을 사용할 때, 이는 혼란스러운 상황을 초래합니다.

학교 급식실을 상상해 보세요. 학생들(토큰)이 각자의 음식을 먹기 위해 서로 다른 배식대(전문가)로 가야 합니다. 일반적인 급식실에서는 모두가 같은 줄에 섭�니다. 하지만 MoE 시스템에서는 학생들이 원하는 음식에 따라 각기 다른 무작위의 배식대로 보내집니다. 어떤 배식대에는 엄청난 인파가 몰리는 반면, 어떤 곳은 아무도 오지 않습니다. 학생들은 음식을 얻기 위해 급식실을 가로질러 달려가야 하고, 그 후 식사를 섞기 위해 중앙 테이블로 다시 달려와야 합니다. 이는 "스캐터-개더(scatter-gather, 흩뿌리고 모으기)" 교통 패턴을 만듭니다. 즉, 학생들이 사방으로 달려가며 충돌하고 대기 시간이 길어지는 혼란스러운 상황입니다. 논문은 표준 컴퓨터 칩이 이런 예측 불가능하고 무질서한 교통량을 처리하도록 설계되지 않았다고 주장합니다. 칩은 정체 상태에 빠지게 되며, 시스템 전체가 느려집니다.

해결책: 특별한 고속도로를 갖춘 3D 도시

이를 해결하기 위해 저자들은 이 AI 모델들을 위한 완전히 새로운 하이테크 도시를 짓는 것과 같은 ThAME를 제안합니다. 그들은 세 가지 주요 아이디어를 사용하여 이를 구현합니다.

1. 적재적소의 도구 (이종 메모리):
   논문은 모든 메모리가 동일하지 않다는 점을 지적합니다. AI의 어떤 부분은 매우 자주 쓰여야 하는 반면(화이트보드처럼), 다른 부분은 거대한 책장에서 읽기만 하면 됩니다.

• "화이트보드" 부분(어텐션 메커니즘)을 위해서는 빠르고 쓰기 쉽지만 공간을 많이 차지하는 DRAM을 사용합니다.
• "거대한 책장" 부분(전문가 가중치)을 위해서는 매우 밀도가 높고 자주 다시 쓸 필요가 없는 3D 메모리인 FeFET-NAND를 사용합니다.
• 이것은 참고 서적이 거대하고 고밀도인 타워(FeFET)에 쌓여 있는 한편, 활발한 작업대에는 다른 종류의 더 빠른 책상(DRAM)을 사용하는 도서관과 같습니다. 이 층들을 수직(3D)으로 쌓음으로써, 프로세서 바로 옆에 방대한 양의 데이터를 배치하여 "러너"들이 멀리 이동할 필요가 없도록 만듭니다.

2. 특별한 고속도로 (NoC):
   이것이 논문의 가장 큰 혁신입니다. 메모리가 가까이 있더라도, "러너"(데이터)들은 여전히 서로 다른 전문가 배식대 사이를 이동해야 합니다. 저자들은 표준 도로망(단순한 그리드나 링 형태)이 예측 불가능한 교통량 앞에서 실패한다는 것을 깨달았습니다.

• 그들은 **계층적 네트워크 온 칩(Hierarchical Network-on-Chip, NoC)**을 설계했습니다. 지역 사회가 로컬 트래픽을 빠르게 처리하기 위해 작은 전용 도로(crossbars)를 가진 도시를 상상해 보세요. 그런 다음, 이 지역들은 서로 다른 구역 간의 큰 흐름을 처리할 수 있는 스마트한 트리 구조의 고속도로 시스템으로 연결됩니다.
• 이 시스템은 어떤 전문가가 바쁘든 상관없이(학생들이 어떻게 분포되든) 교통 체증을 최소화하도록 복잡한 수학적 방법을 통해 최적화되었습니다. 이는 사고가 발생하기 전에 자동으로 차량을 우회시키는 교통 제어 시스템과 같습니다.

3. 스마트 디스패처 (Smart Dispatcher):
   시스템에는 스마트한 교통 경찰 역할을 하는 스케줄러가 포함되어 있습니다. 이 스케줄러는 군중을 관찰하고, 모든 사람이 거의 동시에 작업을 마칠 수 있도록 각 전문가 배식대에 정확히 몇 명의 "러너"(컴퓨터 코어)를 할당할지 결정합니다. 이는 한 명의 느린 전문가가 전체 그룹을 지연시키는 것을 방지합니다.

연구 결과

저자들은 단순히 종이 위에서만 이를 구축한 것이 아니라, 성능을 확인하기 위해 상세한 시뮬레이션을 실행했습니다. 그들은 ThAME를 기존의 최고 시스템들(Stratum 및 H3D-T 등) 및 표준 GPU와 비교했습니다.

• 속도: 시뮬레이션 결과, ThAME는 믿기 힘들 정도로 빨랐습니다. 텍스트를 생성할 때 기존 최고의 하드웨어보다 최대 15.7배 더 빨랐습니다. 이는 표준 컴퓨터가 문단을 쓰는 데 10초가 걸린다면, ThAME는 1초 미만으로 끝낼 수 있음을 의미합니다.
• 에너지: 또한 훨씬 효율적이어서, 동일한 작업에 대해 최대 9.8배 적은 에너지를 사용했습니다. 이는 거대 모델을 실행하는 데 엄청난 전력을 소비한다는 점에서 매우 중요한 성과입니다.
• 견고성: 논문은 메모리 기술이 완벽하지 않더라도(예를 들어, 읽기 속도가 예상보다 느린 경우), ThAME가 충분한 대역폭을 내장하고 있기 때문에 여전히 잘 작동한다는 것을 보여줍니다.

제한 사항

이 논문이 주장하지 않는 점을 명시하는 것도 중요합니다. 저자들은 공장에서 물리적인 칩을 직접 제작하여 실제 서버 팜에서 테스트한 것이 아닙/ 모든 결과는 하드웨어가 어떻게 작동하는지 모사하는 매우 상세한 컴퓨터 모델인 **사이클 단위 정확도 시뮬레이션(cycle-accurate simulations)**을 통해 도출되었습니다. 또한 모든 AI 문제를 해결했다고 주장하지도 않았습니다. 그들은 특히 현재의 컴퓨터가 가장 어려워하는 부분인 "디코드(decode)" 단계(텍, 한 단어씩 생성하는 단계)에 집중했습니다. 그들은 "프리필(prefill)" 단계(초기 프롬프트를 읽는 단계)는 표준적이고 강력한 프로세서(TPU)에 의해 처리된다는 점을 인정했으며, 그들의 혁신은 그 뒤에 이어지는 무질서하고 메모리 집약적인 부분에 특화되어 있습니다.

결론

이 논문은 서로 다른 유형의 메모리를 혼합하고 칩 내부에 교통량에 최적화된 스마트 고속도로를 구축함으로써, 드디어 전문가 혼합(MoE) 모델을 빠르고 효율적으로 실행할 수 있다는 점을 시사합니다. 비록 현재는 시뮬레이션 단계이지만, 이 "스마트 고속도로를 갖춘 3D 도시" 접근 방식이 전력망을 낭비하지 않으면서 차세대 초지능 AI를 여는 열쇠가 될 수 있다는 유망한 결과를 보여줍니다.