Snowflake Arctic 쿡 북 시리즈: MoE(전문가 혼합: Mixture-of-Expert)
이 글은 “https://medium.com/snowflake/snowflake-arctic-cookbook-series-exploring-mixture-of-experts-moe-c7d6b8f14d16” 미디엄 블로그 내용을 기반으로 작성되었습니다. 작성자의 의도와 다르게 개인적으로 이해한 내용을 바탕으로 의역되었기 때문에 원본 글도 참고해 주시기 바랍니다.
스노우플레이크는 4월 24일 업계 최고 수준의 개방성과 성능을 제공하는 기업용 LLM인 Snowflake Arctic를 출시했습니다. 스노우플레이크 AI 연구소는 Snowflake Arctic을 개발할 때 고민했던 다양한 내용에 관련된 쿡 북을 제공할 예정입니다. 여기에는 Arctic 모델을 위한 사전 학습 데이터 소싱 및 구성, MoE 아키텍처 설계, 학습과 추론을 모두 고려한 모델 설계 및 모델 파인 튜닝과 평가에 관련된 진행 과정과 방법론에 관련된 내용들이 포함됩니다.
Snowflake Arctic 쿡 북에 관련된 전체 목록은 여기에서 확인 할 수 있습니다
이번 글에서는 LLM을 구축하려는 경우 초기 설계 과정에서 결정해야 하는 LLM 아키텍처에 대한 내용을 살펴볼 예정입니다.
초창기 LLM 분야에서 연구자 또는 실무자들은 Dense 트랜스포머 아키텍처(그림1의 왼쪽)를 주로 선택했습니다. 이 아키텍처를 채택한 중요 이유는 모델 품질을 개선하기 위해 모델 크기를 확장할 수 있기 때문입니다. 그러나 임계값 이상으로 모델 크기를 확장하기 위해서는 연산 복잡도가 높아져 컴퓨팅 비용이 엄청나게 많이 듭니다. 즉, 모델 학습에 필요한 총 컴퓨팅 양은 모델 크기에 따라 선형적으로 증가하기 때문에 컴퓨팅에 대한 높은 투자 비용과 학습시간이 오래 걸리기 때문에 모델 크기를 확장하기 어려운 단점을 가지고 있었습니다.
MoE 아키텍처는 모델 학습과 추론의 컴퓨팅 비용을 증가시키지 않으면서도 모델 품질을 향상시킬 수 있습니다. 처음으로 MoE 아키텍처 아이디어는 여기에서 제안되었으며, 이후 DeepSpeed-MoE, 스위치 트랜스포머, GLaM과 같은 MoE 아키텍처와 기술들로 개선되어 왔습니다. MoE 아키텍처(그림1의 오른쪽)는 전문가라고 하는 여러 병렬 피드 포워드 네트워크(FFN)으로 구성됩니다. 전문가 선택은 개별 토근의 라우팅 동작을 결정하는 게이팅 함수에 의해 조정됩니다. 각 토근은 연산을 위해 상위 K(일반적으로 K=1 또는 k-2) 전문가를 선택합니다. 이와 같은 특성 때문에 개별 토큰에 대해 전체 네트워크의 일부만 활성화 되므로 전체 모델 크기와 전체 계산이 본질적으로 분리되는 특징을 가지고 있습니다.
MoE 아키텍처는 Dense 아키텍처와 비교하여 다음과 같은 개선점을 가지고 있습니다:
- 고정된 학습 비용으로 모델 품질 향상 가능: MoE 모델은 여러 전문가를 함께 사용함으로써 모델 크기를 효과적으로 늘릴 수 있습니다. 이를 통해 입력 데이터를 보다 전문적이고 다양하게 표현할 수 있어 궁극적으로 고 품질의 모델을 만들 수 있습니다.
- 또한, 연산을 위해 선택된 전문가 수를 고정하면(일명 top-k 고정) 각 입력 토큰에 대해 전문가의 하위 집합만 활성화 되기 때문에 전문가 수와 관계 없이 연산 비용을 통제할 수 있습니다. 따라서 MoE 모델은 모델 학습을 위한 연산에 필요한 리소스를 추가하지 않더라도 모델 품질을 개선할 수 있기 때문에 높은 확장성을 가질 수 있습니다.
경제적이고 빠른 MoE 추론: 추론 과정에서 배치 크기가 큰 경우, 추론을 위해 발생하는 비용은 모델의 총 매개변수와 관계된다기 보다는 휠씬 작은 활성 매개변수 전체 수와 관련이 있습니다. 따라서 더 높은 연산 요구사항이 있는 경우에 추론 비용이 전체 매개변수 수와 상관 관계가 있는 Dense 모델보다 MoE 모델이 추론 중에 더 경제적일 수 있습니다.
모델 품질 개선의 정량화
기존 Dense 모델에 비해 MoE 모델이 어떻게 개선 되는지 좀 더 이해하기 위해, 우리는 두 모델 성능을 직접 비교하는 테스트를 수행했습니다. 16억개의 활성 매개변수가 있는 MoE와 65개의 매개변수가 있는 Dense 모델을 동일한 토큰 수인 1조 개로 훈련한 뒤 두 결과를 비교해 봤습니다.
테스트 결과에 따르면, MoE-1.6B 모델이 더 좋은 성능과 낮은 손실으로 평가되었으며, 학습하는데 Dense 모델과 비교하여 4배 더 적은 컴퓨팅 리소스를 사용했습니다. 그림 2는 두 모델의 테스트의 손실 결과를 보여줍니다.
기술적인 세부 사항에 대해 관심이 있는 사람들을 위해 구체적인 모델 아키텍처와 핵심 하이퍼 파라미터를 표 1과 같이 공유합니다.
최고의 MoE 아키텍처를 어떻게 선택해야 하나?
MoE가 제한된 컴퓨팅 자원으로 모델 품질을 크게 향상시킬 수 있는 것은 분명하지만, 알려지지 않은 문제점도 존재할 수 있습니다. 이미 오랜 기간 동안 연구된 Dense 아키텍처와 달리 MoE 아키텍처는 아직 초기 단계이기 때문에 광범위하게 연구되지 않았습니다. 최적은 전문가의 수가 몇 개이며, 각 전문가의 크기는 얼마가 적절할 지, 한 번에 몇 개의 전문가 수(Top-K 게이팅)을 활성화하면 좋은지, 전문가 간 레이어 간격이 얼마가 적절할까와 같은 최적 구성을 파악하기 위한 연구는 매우 제한적으로 수행되어 왔습니다. 이러한 개별 선택 사항은 컴퓨팅 및 매개 변수 예산이 고정되어 있는 경우 모델 효율성에 큰 영향을 미칠 수 있습니다.
여기에서는 특히 두 가지 중요한 설계 측면에 초점을 맞춰 최적의 구성을 선택하는 방법에 대해 설명합니다:
- Top-K 선택: 각 입력에 대해 활성화된 전문가 수 또는 Top-K 선택이 고정된 컴퓨팅 및 매개변수 비용으로 모델의 품질에 어떤 영향을 미치는지에 평가
- MoE 레이어의 빈도: 여기에는 모델 내에서 MoE 레이어가 사용되는 빈도가 모델 품질에 어떤 영향을 미치는지를 평가합니다. 즉, 모든 트랜스포머 레이어에서 표준 피드포워드 네트워크(FFN)을 MoE 베리언트로 대채 하는지, 아니면 더 적은 빈도로 대체해야 하는지에 대한 인사이트를 제공합니다.
Top-k 선택: Top-1 vs. Top-2
Top-1과 Top-2 게이팅은 MoE 학습에 가장 일반적인 사용되는 두 가지 게이팅 기능입니다. 이러한 접근 방식을 공정하게 비교하기 위해 다음과 같이 조정했습니다:
- Top-2 게이팅을 위해 FFN 레이어의 크기를 절반으로 줄였습니다. 이 조정을 통해 활성화된 총 매개변수 수가 Top-1과 Top-2 설정간의 테스트 일관성을 유지하도록 했습니다.
- Top-2의 전문가 수를 두 배로 증가: 전문가당 매개변수 수가 줄어든 것을 보안하기 위해 전문가 수를 늘려 전체 모델 크기와 컴퓨팅 리소스를 유지하도록 했습니다.
컴퓨팅과 모델 크기를 모두 일정하게 유지하도록 설계된 이 테스트에서 손실 비교는 그림 3에서 보여주고 있습니다. 테스트 결과에서는 Top-2 게이팅이 통제된 조건에서도 Top-1 게이팅보다 더 효과적인 것으로 측정되었습니다. 또한 Top-3와 Top-4 게이팅 함수활성 및 총 매개변수 수라는 유사한 제약 조건에서 모델 품질을 더욱 개선할 수 있는 가능성을 보여줬습니다. 하지만, 전문가 수를 늘리면 모델 품질이 향상 되었지만, MoE 전체 통신 비용이 증가되어 높은 모델 학습 효율을 달성하기 어려웠습니다.
MoE 레어어 간 빈도 선택
MoE 모델에 대해 추가적으로 고려해야 하는 부분은 MoE 레이어간 빈도입니다.
- 모든 MoE 레이어(모든 FFN 레이어를 MoE 레이어로 총칭)
- 모든 다른 MoE 레이어(Dense FFN과 MoE 레이어 인터리브)
- 기타 등등..
테스트 인사이트
위의 두 테스트에서 얻은 주요 시사점은 연속된 레이어 간에 전문가를 결합하는 방법이 최종 모델 품질에 매우 중요하다는 것입니다. 예를 들어, E 전문가와 Top-1 게이팅의 레이어 모델을 기본 구성하는 경우 가능한 전문가의 조합은 E² 개입니다. 전문가 수를 2E 수로 늘리고 Top=2 게이팅을 적용하면(총 컴퓨팅 리소스와 매개변수는 그대로 유지) 가능한 전문가 조합의 수가 4E⁴로 증가하여 그림 3과 같이 더 많은 선택지를 제공하고 잠재적으로 모델 품질을 향상시킬 수 있습니다. 반면, Top-1을 ‘다른 모든 레이어’ MoE와 함께 사용하면, 사용 가능한 조합의 수가 2E로 감소하여 그림 4에서와 같이 모델이 약해집니다.
Arctic Dense-MoE 하이브리드 아키텍처 개요
위에서는 MoE 아키텍처 설계와 품질 트레이드 오프에 대해 살펴봤습니다. 하지만 최적의 아키텍처를 선택하기 위해서는 품질 트레이드 오프 이상을 고려해야 합니다. MoE 모델을 학습하고 배포하면 수 많은 시스템에 관련된 이슈가 발생할 수 있습니다. 이러한 문제를 전반적으로 해결할 수 있도록 시스템을 고려하여 MoE 아키텍처를 설계하는 것이 중요합니다.
Arctic은 모델 품질 개선과 아키텍처 효용성 간의 균형을 맞추기 위해 128개의 전문가와 함께 Top-2 게이팅을 사용하는 그림 5와 같이 Dense-MoE 하이브리드 아키텍처를 사용합니다. 다음 블로그에서는 Arctic과 같은 대규모 MoE 모델이 직면한 훈련 및 추론에 관련된 문제점을 자세히 살펴볼 예정입니다. 또한, 이런 문제를 효과적으로 해결하면서 높은 품질을 달성하기 위해 어떻게 Arctic 하이브 아키텍처를 설계했는지에 대한 내용도 다룰 예정입니다.
