Generative AI / LLMs

LLM 아키텍처에 Mixture of Experts(MoE)를 활용하기

Reading Time: 7 minutes

Mixture of Experts(MoE) 거대 언어 모델(LLM) 아키텍처는 최근 GPT-4와 같은 독점 LLM은 물론 Mixtral 8x7B의 오픈 소스 출시와 함께 커뮤니티 모델에서도 등장하고 있습니다. Mixtral 모델의 강력한 상대적 성능으로 인해 MoE와 LLM 아키텍처에서의 사용에 대한 많은 관심과 수많은 질문이 제기되었습니다. 그렇다면 MoE란 무엇이며 왜 중요한 것일까요?

MoE는 하나의 레이어 또는 연산(예: 선형 레이어, MLP 또는 attention projection)의 계산을 여러 개의 “전문가(expert)” 하위 네트워크로 분할하는 신경망의 아키텍처 패턴입니다. 이러한 하위 네트워크는 각각 독립적으로 자체 연산을 수행하며, 그 결과를 결합하여 MoE 레이어의 최종 출력을 생성합니다. MoE 아키텍처는 모든 입력에 대해 모든 전문가가 사용되는 고밀도(dense) 또는 모든 입력에 대해 전문가 하위 집합이 사용되는 스파스(sparse) 아키텍처가 될 수 있습니다.

이 게시물은 주로 LLM 아키텍처에서의 MoE 적용에 초점을 맞추고 있습니다. 다른 영역에서 MoE를 적용하는 방법에 대한 자세한 내용은 Sparse Mixture of Experts로 비전 확장하기, 다국어 ASR 스트리밍을 위한 Mixture of Expert Conformer 및 FEDformer: 롱텀 시리즈 예측을 위한 빈도 강화 분해 트랜스포머(Frequency Enhanced Decomposed Transformer)를 참조하세요.

LLM 아키텍처의 Mixture of Experts

이 섹션에서는 몇 가지 배경 정보를 제공하고 LLM 아키텍처에서 MoE를 사용할 때의 이점을 강조합니다.

모델 용량

모델 용량은 모델이 이해하거나 표현할 수 있는 복잡성 수준으로 정의할 수 있습니다. 일반적으로 매개변수 수가 많은 모델(충분히 훈련된 경우)이 역사적으로 더 큰 용량을 갖는 것으로 입증되었습니다.

MoE는 용량에 어떻게 반영되나요? 일반적으로 매개변수가 많은 모델은 더 큰 용량을 가지며, MoE 모델은 모델의 레이어를 전문가 하위 네트워크가 원래 레이어와 동일한 크기의 MoE 레이어로 대체함으로써 기본 모델에 비해 용량을 효과적으로 늘릴 수 있습니다.

연구원들은 동일한 양의 토큰으로 훈련된 비슷한 크기의 완전 밀도(fully dense) 모델과 비교하여 MoE 모델의 정확도를 조사했습니다(완전 밀도 크기: EP 매개변수와 비교한 MoE 크기: E*P 매개변수). 완전 밀도 모델이 일반적으로 더 나은 성능을 보이지만, 이는 여전히 활발한 연구 분야입니다. 자세한 내용은 라우팅된 언어 모델의 통합 스케일링 법칙을 참조하세요.

그렇다면 왜 그냥 밀집(dense) 모델을 사용하지 않는가 하는 의문이 생깁니다. 여기에 대한 답은 스파스(sparse) MoE, 특히 스파스 MoE가 사용되는 매개변수당 플롭 효율이 더 높다는 사실에 있습니다.

각 토큰에 대해 두 명의 전문가(expert)만 활용되는 8-전문가(expert) MoE를 사용하는 모델인 Mixtral 8x7B를 예로 들어보겠습니다. 이 경우, 모델 내에서 단일 토큰의 포워드 패스에서 특정 토큰에 사용되는 매개변수의 수가 훨씬 적습니다(총 460억 개의 매개변수 중 120억 개의 매개변수 사용). 따라서 8명의 전문가를 모두 사용하거나 비슷한 크기의 완전 밀도 모델을 사용할 때보다 컴퓨팅이 덜 필요합니다. 트레이닝에서 토큰이 일괄 처리된다는 점을 감안할 때, 모든 전문가가 사용되지는 않더라도 대부분 사용됩니다. 즉, 이 체제에서 스파스 MoE는 같은 크기의 고밀도 모델에 비해 컴퓨팅과 메모리 용량을 더 적게 사용합니다.

GPU 시간이 매우 귀한 자원인 오늘날, 대규모로 완전 고밀도 모델을 훈련하는 것은 시간과 비용 측면에서 엄청난 비용이 소요됩니다. Meta가 훈련한 Llama 2 모델 세트(완전 밀도)는 사전 훈련에 330만 NVIDIA A100 GPU 시간을 사용한 것으로 알려졌습니다. 이를 맥락에 맞게 설명하자면, 1,024개의 GPU를 다운타임 없이 최대 성능으로 330만 GPU 시간을 사용하려면 약 134일이 소요됩니다. 이는 실험, 하이퍼파라미터 스윕(sweeps) 또는 훈련 중단을 고려하지 않은 수치입니다.

비용을 줄이면서 더 큰 모델을 훈련하는 MoE

MoE 모델은 가중치당 플롭 효율이 높아 비용을 절감하는 데 도움이 되며, 이는 시간이나 컴퓨팅 비용에 제약이 있는 체제에서 더 많은 토큰을 처리하고 모델을 더 많이 훈련할 수 있다는 의미입니다. 매개변수가 많은 모델이 완전히 수렴하려면 더 많은 샘플이 필요하다는 점을 감안할 때, 이는 기본적으로 고정된 예산으로 고밀도 모델보다 더 나은 MoE 모델을 훈련할 수 있다는 것을 의미합니다.

지연 시간 단축을 제공하는 MoE

계산이 병목 현상을 일으키는 대규모 프롬프트와 배치에서 MoE 아키텍처를 사용하면 첫 번째 토큰 제공 지연 시간을 단축할 수 있습니다. 검색 증강 생성(RAG) 및 자율 에이전트와 같은 사용 사례에서는 모델에 대한 많은 호출이 필요하여 단일 호출 지연 시간이 늘어날 수 있으므로 지연 시간 단축이 더욱 중요합니다.

MoE 아키텍처는 어떻게 작동하나요?

MoE 모델에는 두 가지 주요 구성 요소가 있습니다. 첫째, 혼합을 구성하는 “전문가(expert)” 하위 네트워크로, 밀집형(dense)과 희소형(sparse) MoE 모두에 사용됩니다. 둘째, 희소성 모델에서 어떤 전문가가 어떤 토큰을 처리할지 결정하기 위해 사용하는 라우팅 알고리즘입니다. 밀집형 및 희소형 MoE의 일부 공식에서는 전문가 출력의 가중 평균을 수행하는 데 사용되는 가중치 메커니즘이 MoE에 포함될 수 있습니다. 희소형 MoE의 경우에 초점을 맞추겠습니다.

발표된 많은 논문에서 MoE 기법은 트랜스포머 블록 내의 다층 퍼셉트론(MLP)에 적용됩니다. 이 경우 트랜스포머 블록 내의 MLP는 일반적으로 일련의 전문가 MLP 서브 네트워크로 대체되며, 그 결과는 평균 또는 합계를 사용하여 MLP MoE 출력을 생성하기 위해 결합됩니다.

연구에 따르면 MoE의 개념을 트랜스포머 아키텍처의 다른 부분에도 적용할 수 있다고 합니다. 최근 논문 SwitchHead: 전문가 혼합 어텐션으로 트랜스포머 가속화하기에서는 어텐션 연산에서 사용할 입력을 Q, K, V 행렬로 변환하는 투영 레이어에도 MoE를 적용할 수 있다고 제안합니다. 다른 논문에서는 조건부 실행 MoE 개념을 주의력 헤드 자체에 적용할 것을 제안하기도 했습니다.

라우팅 네트워크(또는 알고리즘)는 특정 입력이 있을 때 어떤 전문가가 활성화될지 결정하는 데 사용됩니다. 라우팅 알고리즘은 전문가 선택 라우팅을 사용한 Mixture-of-Experts에서 설명한 것처럼 단순(텐서 평균값에 대한 균일한 선택 또는 비닝)부터 복잡한 것까지 다양할 수 있습니다.

특정 라우팅 알고리즘의 문제 적용 가능성을 결정하는 많은 요소 중에서 특정 라우팅 체제에서의 모델 정확도와 특정 체제에서의 부하 분산이라는 두 가지 핵심 요소가 자주 논의됩니다. 올바른 라우팅 알고리즘을 선택하는 것은 정확도와 플롭 효율성 사이의 절충안일 수 있습니다. 완벽하게 부하가 분산된 라우팅 알고리즘은 토큰당 정확도를 낮출 수 있는 반면, 가장 정확한 라우팅 알고리즘은 전문가 간에 토큰을 불균등하게 분배할 수 있습니다.

제안된 많은 라우팅 알고리즘은 특정 전문가에 의해 발생하는 병목 현상을 최소화하면서 모델 정확도를 최대화하도록 설계되었습니다. Mixtral 8x7B는 Top-K 알고리즘을 사용하여 토큰을 라우팅하지만, 전문가 선택 라우팅을 사용한 Mixture-of-Expert와 같은 논문에서는 특정 전문가에 과도하게 라우팅되지 않도록 하는 개념을 소개합니다. 이를 통해 병목 현상을 방지할 수 있습니다.

Mixtral 모델 실험하기

실제로 각 전문가는 무엇을 배울까요? 구두점, 동사, 형용사 등과 같은 낮은 수준의 언어 구조에 특화되어 있을까요, 아니면 코딩, 수학, 생물학, 법률과 같은 높은 수준의 개념과 영역의 달인일까요?

이를 위해 Mixtral 8x7B 모델을 사용하여 실험을 설계했습니다. 이 모델에는 32개의 순차적 트랜스포머 블록이 있으며, 각 MLP 레이어는 8명의 전문가가 포함된 희소성 MoE 블록으로 대체되고, 이 중 각 토큰에 대해 두 개만 활성화됩니다. 셀프 어텐션 및 정규화 레이어(normalization layer)를 포함한 나머지 레이어는 모든 토큰이 공유합니다.

8x7B라는 이름을 읽을 때 전문가가 각각 70억 개의 파라미터로 구성된 8개의 개별 전체 네트워크라고 상상할 수 있으며, 각 토큰은 이 8개의 전체 네트워크 중 하나에서 엔드투엔드로 완전히 처리된다는 점에 주목할 필요가 있습니다(그림 1). 이 설계는 8x7B = 56B 모델을 산출합니다.

그림 1. Mixtral 8x7B 모델에 대한 가능한 해석

이것은 확실히 그럴듯한 설계이기는 하지만, Mixtral 8x7B에 사용된 설계는 아닙니다. 실제 설계는 그림 2에 나와 있으며, 각 토큰은 70억 개의 파라미터로 처리됩니다. 토큰과 그 사본(각 계층에서 두 번째로 선택된 전문가가 처리)은 2x7B = 14B가 아니라 총 129억 개의 매개변수로만 처리된다는 점에 유의하세요. 그리고 전체 네트워크의 파라미터는 공유 레이어로 인해 8x7B = 56B가 아닌 470억 개에 불과합니다.

그림 2. 단순화된 Mixtral 8x7B 모델 아키텍처

따라서 네트워크를 통과하는 각 토큰은 격자와 같은 구조를 거쳐야 하며, 각 레이어에서 두 전문가의 가능한 조합은 \binom{8}{2} = 28개입니다. Mixtral 8x7B에는 32개의 트랜스포머 블록이 있으므로 네트워크의 가능한 인스턴스화는 총 28^{32}$입니다.

이러한 인스턴스화 각각을 “풀스택 전문가”(토큰을 종단간 처리하는 전문가)로 간주한다면, 이들이 제공하는 전문성을 알아낼 수 있을까요? 안타깝게도 28^{32}는 매우 큰 숫자(~2×10^{46})로, LLM을 훈련하는 데 사용되는 모든 데이터(대부분의 LLM의 경우 ~3T~10T 토큰)보다 훨씬 크기 때문에 동일한 인스턴스화에서 두 토큰을 처리하는 경우는 거의 없을 것입니다. 따라서 전체 전문가 조합이 아닌 각 계층 전문가가 전문적으로 다루는 분야를 연구할 것입니다.

실험 결과

모델을 통해 대규모 멀티태스크 언어 이해(MMLU) 벤치마크의 모든 샘플을 실행했습니다. 여기에는 추상 대수학, 세계 종교, 전문 법률, 해부학, 천문학, 기업 윤리 등 다양한 57개 주제에 대한 객관식 문제가 포함되어 있습니다. 저희는 8명의 전문가 각각에 대한 토큰-전문가 과제를 레이어 1, 16, 32에 기록했습니다. 데이터를 분석한 결과, 몇 가지 주목할 만한 관찰 결과가 나왔습니다.

로드 밸런싱

로드 밸런싱 덕분에 전문가들은 균등한 부하를 받지만, 가장 활성화가 많이 일어나는 전문가(expert)가 가장 활성화가 적게 일어나는 전문가보다 최대 40~60% 더 많은 토큰을 받습니다.

그림 3. 모든 MMLU 주제에 대한 전문가 로딩 분포

도메인-전문가 할당

일부 도메인은 다른 도메인보다 일부 전문가를 더 많이 활성화합니다.

레이어 32에서는 전문가 3과 8을 다른 전문가보다 훨씬 더 많이 사용하는 추상 대수학을 예로 들 수 있습니다.

그림 4. 추상 대수학의 레이어 32 전문가에 대한 토큰 분포

반면에 전문법 영역에서는 전문가 4가 대부분 활성화되고 상대적으로 전문가 3과 8이 비활성화됩니다.

그림 5. 전문 법률에 대한 레이어 32의 전문가에 대한 토큰 분포

세계 종교는 또 다른 흥미로운 예로, 전문가 7이 전문가 8보다 5배 이상 적은 토큰을 받는다는 점이 특징입니다.

그림 6. 세계 종교에 대한 레이어 32의 전문가에 대한 토큰 분포

그림 6. 세계 종교에 대한 레이어 32의 전문가에 대한 토큰 분포
이러한 실험 결과는 전문가들의 부하 분포가 다양한 주제에 걸쳐 균일한 경향을 보인다는 것을 보여줍니다. 그러나 샘플이 모두 특정 주제에만 속할 경우, 분포의 불균형이 크게 나타날 수 있습니다.

전문가(expert)들이 가장 선호하는 토큰

그림 7의 워드 클라우드는 각 전문가가 가장 자주 처리한 토큰을 보여줍니다.

그림 7. 전문가가 처리한 가장 일반적인 토큰

토큰별로 가장 선호하는 전문가

각 토큰마다 선호하는 전문가가 있나요? 다음 예에서 볼 수 있듯이 각 토큰에는 선호하는 전문가 세트가 있는 것 같습니다.

토큰 “:”에 대한 전문가 할당 , 모든 “:” 토큰은 계층 1에서 전문가 1과 7이 처리하고 계층 32에서 전문가 3과 8이 처리합니다(그림 8). 그림 9, 10, 11은 다양한 토큰에 대한 전문가 할당을 보여줍니다.

그림 8. 토큰 “:”에 대한 전문가 할당
그림 9. 토큰 “:”에 대한 전문가 할당
그림 10. 토큰 “무엇(what)”에 대한 전문가 할당
그림 11. 토큰 “누구(who)”에 대한 전문가 할당

요약

MoE 모델은 모델 사전 학습 처리량에 확실한 이점을 제공하므로 고밀도 모델과 동일한 양의 컴퓨팅으로 보다 표현력이 풍부한 희소 MoE 모델을 학습할 수 있습니다. 따라서 동일한 컴퓨팅 예산으로 더 경쟁력 있는 모델을 만들 수 있습니다. MoE 모델은 전체 네트워크 또는 기존 네트워크 내의 특정 레이어를 대상으로 할 수 있습니다. 일반적으로 라우팅이 포함된 희소 MoE는 일부 전문가만 사용하도록 하기 위해 적용됩니다.

이번 실험에서는 토큰이 할당되는 방식과 전문가 간의 상대적인 부하의 균형도를 살펴보았습니다. 이 실험은 로드 밸런싱 알고리즘에도 불구하고 여전히 큰 분포 불균형이 존재할 수 있으며, 일부 전문가는 작업을 일찍 완료하는 반면 다른 전문가는 과부하가 걸리기 때문에 추론의 비효율성에 영향을 미칠 수 있음을 보여줍니다. 이는 현재 활발히 연구 중인 흥미로운 분야입니다.

NVIDIA NGC 카탈로그를 통해 Mixtral 8x7B Instruct와 다른 AI 기반 모델을 사용해 볼 수 있습니다.

더 자세히 알아보고 싶으신가요? NVIDIA GTC 2024의 패널 세션인 Mistral AI: 손안의 프론티어 AI를 놓치지 마세요.

관련 리소스

GTC 세션: LLM 인프라 구축, 트레이닝 속도 가속화, 생성형 AI 혁신을 위한 엔드투엔드 솔루션 설계(Aivres 발표)
GTC 세션 RAG 시스템의 효율성 향상을 위한 기술(제공: Aivres)
GTC 세션 새로운 언어 모델 스택을 위한 아키텍팅
NGC 컨테이너: genai-llm-playground
웨비나 거대 언어 모델 구현하기
웨비나 클라우드 지원 AI 추론 솔루션의 성능을 활용하고 클라우드에서 LLM 추론 배포의 단계별 데모 체험하기

Discuss (0)

Tags