Moonshot AI 뮤온과 문라이트 대규모 모델 훈련 최적화

대규모 언어 모델 훈련의 효율성 추구

끊임없이 더 크고 더 유능한 언어 모델을 추구하면서 효율성에 대한 절실한 필요성이 대두되었습니다. 이러한 거대 모델을 훈련하려면 엄청난 계산 능력뿐만 아니라 모든 와트와 모든 초에서 최대 성능을 끌어낼 수 있는 정교한 기술이 필요합니다. 학습 프로세스를 주도하는 엔진인 최적화 알고리즘은 절대적으로 중요합니다. 이러한 알고리즘은 수십억 또는 수조 개의 매개변수를 가진 모델이 얼마나 빠르고 효과적으로 최적 성능 상태로 수렴할 수 있는지를 결정합니다. AdamW와 같은 옵티마이저가 업계의 주력으로 자리 잡았지만, 세심한 하이퍼파라미터 튜닝에 대한 요구와 계산 자원에 대한 엄청난 소비는 보다 효율적인 대안을 찾도록 만들었습니다. 궁극적인 목표는 계산 부담을 획기적으로 줄이면서도 견고한 훈련 안정성을 제공하는 옵티마이저입니다.

기존 최적화 기술의 한계

거대한 언어 모델 훈련의 핵심 과제는 계산 요구량의 엄청난 규모에 있습니다. 모델이 커짐에 따라 각 반복마다 업데이트해야 하는 매개변수 수가 폭발적으로 증가합니다. 많은 기존 옵티마이저는 소규모 환경에서는 효과적이지만 이러한 엄청난 압력 하에서는 흔들리기 시작합니다. 효율성이 떨어지고 훈련 기간을 늘리는 지속적인 조정 및 미세 조정이 필요합니다. 또한 불안정성 문제가 발생하여 모델 성능을 저하시키는 불규칙한 업데이트로 나타날 수 있습니다. 따라서 진정으로 효과적인 솔루션은 효율성과 안정성을 모두 해결하여 과도한 계산 능력이나 끝없는 수동 매개변수 조정 없이도 원활하고 안정적인 훈련을 보장해야 합니다.

예를 들어 널리 사용되는 Adam 및 AdamW 옵티마이저는 적응형 학습률과 가중치 감소를 사용하여 모델 성능을 미세 조정합니다. 이러한 방법은 다양한 응용 분야에서 그 가치를 입증했습니다. 그러나 모델 규모가 커짐에 따라 효과가 감소합니다. 이러한 옵티마이저와 관련된 계산 오버헤드가 급격히 증가하여 진정한 대규모 훈련 노력에는 비효율적입니다. 이로 인해 대체 옵티마이저를 식별하고 개발하는 데 초점을 맞춘 활발한 연구 노력이 촉진되었습니다. 이러한 새로운 접근 방식은 우수한 성능과 효율성을 제공하는 것을 목표로 하며, 이상적으로는 힘들게 하이퍼파라미터를 튜닝할 필요 없이 안정적이고 확장 가능한 결과를 얻을 수 있습니다.

확장성을 위해 설계된 새로운 옵티마이저, Muon

Moonshot AI의 연구원들은 UCLA와 협력하여 대규모 훈련 시나리오에서 기존 방법의 한계를 극복하도록 특별히 설계된 옵티마이저인 Muon을 도입했습니다. Muon은 처음에 소규모 모델에서 인상적인 성능을 보여주었지만, 언어 모델의 거대 모델을 처리하기 위해 규모를 확장했을 때 어려움에 직면했습니다. 이러한 문제를 해결하기 위해 연구원들은 두 가지 핵심 기술을 구현했습니다.

첫째, 과적합을 방지하고 훈련 안정성을 향상시키는 정규화 기술인 가중치 감소(weight decay) 를 통합했습니다. 둘째, 일관된 RMS(root mean square) 업데이트 를 도입했습니다. 이를 통해 크기에 관계없이 모든 매개변수에 조정이 균일하게 적용됩니다. 이러한 균일성은 대규모 언어 모델의 방대한 매개변수 공간에서 균형 잡힌 학습을 유지하는 데 중요합니다. 이러한 개선을 통해 Muon은 광범위한 하이퍼파라미터 튜닝 없이도 효율적으로 작동할 수 있습니다. 이러한 “즉시 사용 가능한” 준비성은 대규모 모델 훈련을 위한 매력적인 선택이 되어 설정 및 구성 오버헤드를 크게 줄입니다.

Mixture-of-Experts 모델에서 Muon의 힘을 활용하는 Moonlight

Muon에 구현된 발전을 기반으로 연구원들은 MoE(Mixture-of-Experts) 모델인 Moonlight를 개발했습니다. Moonlight는 30억 개의 매개변수 버전과 더 큰 160억 개의 매개변수 버전의 두 가지 구성으로 제공됩니다. 둘 다 5조 7천억 개의 토큰으로 구성된 방대한 데이터 세트에서 훈련되었습니다. Moonlight는 Muon을 활용하여 성능을 최적화하는 동시에 계산 비용을 최소화합니다.

효율성을 더욱 향상시키기 위해 ZeRO-1 스타일 최적화 전략을 사용하는 Muon의 분산 버전이 개발되었습니다. 이 접근 방식은 옵티마이저 상태를 여러 장치에 분산시켜 메모리 효율성을 크게 향상시킵니다. 또한 대규모 분산 훈련에서 중요한 요소인 통신 오버헤드를 최소화합니다. 이러한 개선으로 인해 매우 안정적인 훈련 프로세스가 완성되었습니다. Moonlight는 비슷한 규모의 이전 모델에 비해 훨씬 적은 계산량으로 최첨단 성능을 달성했습니다.

성능 벤치마킹: 경쟁 모델을 능가하는 Moonlight

엄격한 성능 평가를 통해 Moonlight가 비슷한 규모의 기존 최첨단 모델보다 지속적으로 우수한 성능을 발휘하는 것으로 나타났습니다. 여기에는 LLAMA3-3B 및 Qwen2.5-3B와 같이 잘 알려진 모델이 포함됩니다. 모델 크기, 데이터 및 성능 간의 관계를 탐구하는 스케일링 법칙 실험은 Muon의 놀라운 이점을 보여주었습니다. Muon은 Adam보다 샘플 효율성이 약 두 배 입니다. 이는 경쟁력 있는 결과를 달성하면서도 훈련에 필요한 부동 소수점 연산(FLOP) 수를 크게 줄입니다.

Moonlight의 성능은 광범위한 벤치마크 작업에서 나타납니다. MMLU(Massive Multitask Language Understanding) 벤치마크에서 70.0이라는 인상적인 점수를 얻어 LLAMA3-3B(54.75) 및 Qwen2.5-3B(65.6)를 크게 능가했습니다. MMLU-pro 및 BBH(Big-Bench Hard)와 같은 보다 전문화된 벤치마크에서 Moonlight는 각각 42.4 및 65.2점을 얻어 향상된 기능을 더욱 강조했습니다. 또한 이 모델은 질문-응답 벤치마크인 TriviaQA에서 66.3점으로 강력한 성능을 보여 모든 비교 가능한 모델을 능가했습니다.

코드 생성 및 수학적 추론: 다재다능함 입증

Moonlight의 기능은 자연어 이해 및 질문 응답을 넘어섭니다. 또한 코드 관련 작업에서도 탁월합니다. 코드 생성 능력을 평가하도록 설계된 벤치마크인 HumanEval에서 48.1점을 얻었습니다. 또 다른 코드 생성 벤치마크인 MBPP(Mostly Basic Programming Problems)에서는 63.8점을 받았습니다. 이러한 결과는 기능 코드를 생성하는 데 능숙함을 보여주며 유사한 매개변수 수를 가진 다른 모델보다 성능이 뛰어납니다.

수학적 추론 영역에서 Moonlight는 뛰어난 문제 해결 능력을 보여주었습니다. 초등학교 수준의 수학 단어 문제로 구성된 벤치마크인 GSM8K(Grade School Math 8K)에서 77.4점을 얻었습니다. 고급 수학 문제에 초점을 맞춘 더 어려운 벤치마크인 MATH에서는 45.3점을 받았습니다. 이러한 결과는 Moonlight가 복잡한 수학적 추론 작업을 처리할 수 있는 능력을 강조합니다.

다국어 능력: 중국어 작업에서 탁월

Moonlight의 기능은 영어에만 국한되지 않습니다. 중국어 작업에서도 강력한 성능을 보여줍니다. 포괄적인 중국어 평가 제품군인 C-Eval에서 77.2점을 얻었습니다. 다중 작업 언어 이해에 초점을 맞춘 또 다른 중국어 벤치마크인 CMMLU에서는 78.2점을 얻었습니다. 이러한 결과는 Moonlight가 다국어 처리에서 효과적임을 입증하며 다양한 언어적 뉘앙스를 처리할 수 있는 능력을 보여줍니다. 이 모델은 이렇게 다양한 벤치마크에서 지속적으로 강력한 성능을 보여주므로강력한 일반화 능력을 입증합니다. 이전 모델에 비해 훨씬 적은 계산 비용을 유지하면서 다양한 작업에 적응하고 탁월한 성능을 발휘할 수 있습니다.

확장성 문제 해결 및 미래 연구 촉진

Muon에 구현된 혁신은 오랫동안 대규모 언어 모델 훈련을 괴롭혀 온 중요한 확장성 문제를 직접적으로 해결합니다. 연구원들은 가중치 감소와 일관된 RMS 업데이트를 통합하여 안정성과 효율성을 모두 크게 향상시켰습니다. 이를 통해 Moonlight는 성능의 한계를 뛰어넘는 동시에 훈련 비용을 절감할 수 있었습니다. 이러한 발전으로 Muon은 Adam 기반 옵티마이저에 대한 매력적인 대안으로서의 입지를 확고히 했습니다. Adam 및 그 변형과 관련된 광범위한 튜닝 없이도 우수한 샘플 효율성을 제공합니다.

또한 Muon과 Moonlight를 모두 오픈 소스로 제공하는 것은 연구 커뮤니티에 대한 중요한 기여입니다. 이러한 도구를 자유롭게 사용할 수 있도록 함으로써 연구원들은 대규모 모델을 위한 효율적인 훈련 방법에 대한 추가 탐구와 개발을 촉진하고 있습니다. 이러한 개방형 접근 방식은 협업을 장려하고 해당 분야의 발전을 가속화하여 미래에 더욱 강력하고 접근 가능한 언어 모델을 위한 길을 열어줍니다. Muon과 같은 옵티마이저의 지속적인 개선은 더 큰 모델을 구축하는 것뿐만 아니라 더 스마트하게 구축하고, 사용 가능한 리소스를 최대한 활용하고, AI 연구의 최첨단에 대한 접근을 민주화하는 것입니다.