MoE LLM 훈련 효율성 및 전문가 특화 향상을 위한 전역 배치 부하 분산 기법

MoE(Mixture-of-Experts) 기반 대규모 언어 모델(LLM) 훈련에서 부하 분산 손실(Load Balancing Loss, LBL) 계산 방식에 대한 새로운 접근 방식이 제시되었다. 이 방식은 기존의 마이크로 배치(micro-batch) 수준 부하 분산 대신 전역 배치(global-batch) 수준의 부하 분산을 도입하여 모델 성능 향상과 전문가 특화(expert specialization)를 달성한다. 구체적으로, 이 방법은 전문가 선택 빈도($f_i$)를 모든 병렬 그룹에 걸쳐 동기화하고, 각 병렬 그룹에서 부하 분산 손실을 계산한 후 모든 마이크로 배치에 걸쳐 손실을 집계하는 방식으로 작동한다. 전문가 선택 빈도는 전문가 수 차원의 벡터로, 마이크로 배치 간 동기화에 드는 비용이 거의 없다는 점이 특징이다.

MoE 아키텍처는 모델 매개변수 규모를 확장하는 데 널리 사용되는 기술이다. 일반적으로 하나의 MoE 레이어는 라우터와 여러 전문가 그룹으로 구성되며, 입력에 따라 일부 전문가만 활성화되고 라우터가 할당한 점수를 기반으로 그 출력이 집계된다. 부하 분산 손실은 MoE 기반 네트워크 훈련에서 필수적인 정규화 기법으로, 모든 전문가의 균형 잡힌 활성화를 유도한다. 그러나 Megatron-core와 같은 대부분의 기존 MoE 훈련 프레임워크는 마이크로 배치 수준에서 부하 분산 손실을 계산하고 이를 전역 배치 수준에서 평균화한다. 이러한 구현 방식은 하나의 마이크로 배치가 다양한 데이터를 포함하지 않을 경우 문제가 될 수 있다. 예를 들어, 특정 도메인(예: 코드 데이터)의 데이터만 포함된 마이크로 배치에서 라우터가 해당 토큰을 모든 전문가에게 균일하게 분배하도록 강제하면 모델 성능이 저하되고 전문가 특화가 방해될 수 있다. 특히 MoE 기반 LLM 훈련에서는 마이크로 배치 내 데이터가 동일한 도메인에서 오는 경우가 많아 이러한 문제가 더욱 두드러지며, 이는 기존 오픈소스 MoE 기반 LLM이 눈에 띄는 전문가 특화를 달성하지 못하는 부분적인 이유로 지적된다. 이러한 단점이 전역 배치 수준의 부하 분산 도입을 촉진했다.

전역 배치 부하 분산 손실은 전문가 선택 빈도($f_i^{\text{global}}$)와 전문가 $E_i$에 할당된 평균 게이팅 점수($p_i^{\text{global}}$)를 사용하여 계산된다. 여기서 $f_i^{\text{global}}$은 모든 병렬 그룹에 걸쳐 동기화된 전문가 선택 빈도이며, $p_i^{\text{global}}$은 모든 병렬 그룹의 평균 게이팅 점수이다. 실험은 세 가지 MoE 모델 구성(활성화된 매개변수 0.6B를 가진 3.4B 모델, 2.54B를 가진 15B 모델, 6.6B를 가진 43B 모델)과 두 가지 데이터 구성(120B 토큰, 400B 토큰)으로 진행되었다. 전역 배치 부하 분산 방식은 모든 모델, 데이터, 작업 설정에서 마이크로 배치 수준 손실보다 우수한 성능을 보였다. 특히, 전역 배치 부하 분산을 통해 MoE 모델은 상당한 도메인 특화를 달성했다. 기존 마이크로 배치 방식에서는 모든 전문가가 도메인과 관계없이 균일하게 활성화되었으나, 전역 배치 방식에서는 일부 전문가가 특정 도메인에 의해 자주 활성화되어 명확한 특화가 나타났다.

연구팀은 균형 배치 크기(balance BSZ)에 따른 모델 성능 변화도 비교했다. 3.4B(0.6B 활성화) 모델을 대상으로 한 실험에서 사전 훈련 PPL(Perplexity)은 BSZ 2에서 128까지 빠르게 감소했으며, 128 이후 점진적으로 포화 상태에 도달했다. 이는 균형 배치 크기가 모델 성능에 중요한 영향을 미침을 보여준다. 현재 주류 MoE 프레임워크에서 대규모 모델의 균형 배치 크기는 전문가 병렬 그룹 간 통신을 고려하더라도 일반적으로 8에서 16 사이로 설정된다. 이러한 맥락에서 전역 배치 부하 분산 기법은 더 큰 균형 배치 크기의 이점을 효과적으로 활용할 수 있게 하여, 기존 방식의 한계를 넘어설 수 있는 가능성을 제시한다.

이러한 전역 배치 부하 분산 기법의 도입은 MoE 모델을 훈련하는 개발자에게 여러 긍정적인 영향을 미친다. 개발자는 이 방식을 적용하여 MoE 모델의 전반적인 성능을 향상시킬 수 있으며, 각 전문가가 특정 데이터 도메인에 더욱 효과적으로 특화되도록 유도할 수 있다. 이는 모델의 해석 가능성을 높이고, 특정 작업이나 데이터 유형에 더 효율적이고 전문화된 모델을 구축하는 데 기여한다. 기존 MoE 훈련의 주요 한계 중 하나로 지적되던 전문가 특화 부족 문제를 해결함으로써, 이 기술은 더욱 강력하고 전문화된 MoE 모델을 다양한 도메인에서 훈련할 수 있는 기반을 마련한다. 특히, 전문가 선택 빈도 동기화에 드는 비용이 거의 없으므로, 기존 훈련 시스템에 큰 계산 오버헤드 없이 적용 가능하다는 실용적인 장점을 제공한다.

전역 배치 부하 분산 사용 시 마이크로 배치 부하 분산이 저하될 수 있으며, 이는 MoE 모델의 계산 효율성에 잠재적으로 영향을 미칠 수 있다. 이러한 가능성을 해결하기 위해, 연구팀은 전역 배치 부하 분산 손실에 마이크로 배치 부하 분산 손실을 추가하는 방법을 제안했다. 이때 마이크로 배치 손실에는 전역 배치 손실의 0.01에 해당하는 일정한 가중치를 부여한다. 이 추가 조치를 통해 모델 업데이트 단계당 시간(초)이 1.64초에서 1.59초로 개선되었으며, 모델의 효과는 거의 영향을 받지 않았다. 이는 전역 배치 부하 분산의 이점을 유지하면서 계산 효율성 저하를 방지할 수 있음을 시사한다. 연구는 주로 언어 기반 작업에 대한 실험을 진행했으나, 이 작업이 다양한 도메인에서 더 크고 전문화된 MoE 모델을 훈련하는 길을 열기를 기대한다고 밝혔다.