Qwen 글로벌-배치 로드 밸런스 MoE LLM 훈련

글로벌-배치 로드 밸런스가 MoE 성능과 전문가 특화를 향상시킴

Qwen은 Mixture-of-Experts (MoE) 대형 언어 모델(LLM) 훈련을 위한 글로벌-배치 로드 밸런스 손실 방법을 도입했습니다. 마이크로-배치 수준에서 글로벌-배치 수준으로 로드 밸런스 손실 계산을 이동시킴으로써, 모델은 다양한 작업에서 성능을 향상시키고 전문가 간의 상당한 도메인 특화를 촉진하며, 마이크로-배치 밸런스가 종종 특정 데이터 도메인에서 전문가의 특화를 방해하는 한계를 해결합니다.

마이크로-배치 로드 밸런스의 한계

Megatron-core와 같은 대부분의 기존 MoE 훈련 프레임워크는 마이크로-배치 수준에서 로드 밸런스 손실($L_{ ext{balance}}$)을 구현합니다. 이는 손실이 모든 마이크로-배치 내에서 계산되고その後 글로벌 배치에 대해 평균화된다는 것을 의미합니다.

이 접근 방식은 마이크로-배치에 다양한 데이터가 부족할 때 문제가 됩니다. 예를 들어, 마이크로-배치가 오직 코드 데이터만 포함하는 경우, 마이크로-배치 수준 손실은 라우터에게 해당 코드 토큰을 모든 전문가에 균등하게 분배하도록 강제합니다. 이는 전문가들이 특정 도메인(예: 코드)에 특화되는 것을 방해하고 전체 모델 성능을 저하시킬 수 있습니다. LLM 훈련에서 단일 마이크로-배치의 데이터는 종종 동일한 도메인에서 오기 때문에, 많은 오픈소스 MoE 모델이 눈에 띄는 전문가 특화를 갖지 못하는 이유를 설명합니다.

글로벌-배치 로드 밸런스 구현

이 문제를 해결하기 위해 Qwen은 글로벌-배치 수준에서 로드 밸런스 손실($L_{ ext{global}}$)을 계산할 것을 제안합니다. 이는 세 단계로 달성됩니다:

동기화 모든 병렬 그룹 간의 전문가 선택 빈도($f_{i}$).
계산 각 병렬 그룹에서의 로드 밸런스 손실(예: 단일 GPU에서).
집계 모든 마이크로-배치에 걸친 손실.

전문가 선택 빈도는 전문가 수와 동일한 차원을 가진 작은 벡터이므로, 이 데이터를 마이크로-배치 간에 동기화하는 것은 계산 비용이 적게 들어 구현이 "거의 무료"라고 할 수 있습니다.

성능 향상 및 전문가 특화

Qwen은 세 가지 MoE 구성(3.4B에 0.6B 활성화, 15B에 2.5B 활성화, 43B에 6.6B 활성화)과 두 가지 데이터 구성(120B 및 400B 토큰)에서 해당 방법을 테스트했습니다.

모델 효과성 향상

마이크로-배치 수준 손실과 비교했을 때, 글로벌 배치 접근 방식은 모든 테스트된 설정에서 더 나은 성능을 달성했으며, 모든 모델 크기, 데이터 양 및 작업을 포함했습니다.

도메인 특화 향상

글로벌 배치 밸런스는 전문가 특화를 가능하게 합니다. 마이크로-배치 밸런스는 도메인과 관계없이 전문가가 균등하게 활성화되는 결과를 초래하는 반면, 글로벌 배치 밸런스는 특정 도메인에 의해 특정 전문가가 자주 활성화되도록 허용하여 명확한 특화를 보여줍니다.

밸런스 배치 크기의 영향

0.6B 활성화된 3.4B 모델에서의 실험은 밸런스 배치 크기가 2에서 128로 증가함에 따라 ptr-training PPL(퍼플렉시티)이 급격히 감소하고 128 이후 포화 상태에 도달함을 보여주었습니다. 이는 글로벌 접근 방식의 중요성을 강조하는데, 주류 MoE 프레임워크는 일반적으로 더 큰 모델에 대해 밸런스 배치 크기를 8과 16 사이로 사용하기 때문입니다.

효율성과 효과성의 균형

글로벌 배치 밸런스는 모델 품질을 향상시킬 수 있지만, 마이크로-배치 밸런스의 저하를 초래할 수 있으며 이는 계산 효율성에 부정적인 영향을 미칠 수 있습니다.

이를 완화하기 위해 Qwen은 글로벌-배치 로드 밸런스 손실에 마이크로-배치 밸런스 손실을 추가하는 실험을 수행했습니다(글로벌 배치 손실의 상수 가중치 0.01 사용). 이 하이브리드 접근 방식은 업데이트 단계 속도를 개선했습니다(1.64초에서 1.59초 per 업데이트 단계)하면서 모델의 효과성에 거의 영향을 미치지 않았습니다.

결론

글로벌-배치 로드 밸런스를 구현함으로써, Qwen은 더 나은 성능과 특화를 가진 전문가를 가능하게 하는 MoE 훈련의 중요한 한계를 해결합니다. 이 방법은 특히 다양한 도메인에서 더 substantial하고 특화된 MoE 모델로 확장할 수 있는 능력을 제공하여 MoE 모델 최적화에 새로운 관점을 제시합니다.