https://arxiv.org/pdf/2201.05596
거대 dense 모델의 훈련 비용과 관련하여 하드웨어 리소스의 가용성과 용량 한계를 극복할 수 있는 기술로 MoE(Mixture of Experts) 기술이 소개되었습니다. 하지만 모델 크기가 훨씬 크고, communication overhead가 추가된다는 점에서 이 MoE 모델로 추론 서비스를 제공하는 것은 상당히 어려운 문제입니다. 이 논문이 나온 것은 22년이지만, 올해 초 Deepseek R1을 출시하면서 MoE가 뜨거운 감자가 되었습니다. Deepseek R1을 알아보기 전에 MoE 관련해서 어떠한 이야기들이 있었는지 논문을 몇 편 리뷰해 보고, Deepseek R1 리뷰를 해보겠습니다. 이 글에서는 마이크로소프트사에서 publish 한 DeepSpeed-MoE 논문에서 알아봅니다.
이 논문에서 제시하는 DeepSpeed-MoE는 DeepSpeed 라이브러리의 일환으로, MoE 훈련 및 추론에 대한 end to end 솔루션입니다. MoE 모델 크기를 최대 3.7배까지 줄이는 혁신적인 MoE 아키텍처 디자인 및 모델 압축 기술과, 고도로 최적화된 추론 시스템을 포함합니다. DeepSpeed-MoE는 품질 동등 밀집 모델에 비해 최대 4.5배 더 빠르고 9배 더 저렴한 추론을 통해 대규모 MoE 모델을 제공하는 전례 없는 스케일과 효율성을 제공한다고 합니다.
메가트론-튜링 NLG 530B 모델을 예로 들어 dense(밀집) 모델이 얼마나 많은 하드웨어 리소스와 훈련 시간이 필요한지를 말하고 있는데요, 이 모델을 2000개 이상의 NVIDIA A100 GPU에서 훈련하는 데 약 3개월이 걸렸으며, 300만 GPU 시간을 소비했다고 합니다. 그리고 이 보다 더 큰 모델을 훈련할 경우, 하드웨어와 훈련 시간이 기하급수적으로 늘어날 것입니다. 그래서 이 논문에서 생각해 본 질문은 이렇습니다.
“Is it possible to make non-trivial improvement to model quality without increasing the compute cost?" or
“Is it possible to produce model with similar quality using 3 to 5 times less resources?”
모델의 품질은 희생시키지 않으면서도, 이러한 거대 모델을 훈련시키기 위한 요구 사항을 줄일 수 있는가?에 대한 답을 MoE를 통해서 찾고자 하는 것입니다.
MoE는 여러 개의 전문가(전문가 모델)로 구성된 모델 아키텍처로, 주어진 입력에 대해서 여러 전문가 중 일부만 활용하여 예측을 수행함으로써 효율성을 높입니다. 즉 전체 파라미터 수에 대비해 적은 계산 비용을 요구하게 되어, 적은 하드웨어 리소스로도 큰 모델을 훈련할 수 있게 되는 것이죠. 하지만 그것을 하기 위한 모델 구조상 더 높은 메모리 요구사항 및 훈련과 추론 시스템 모두 더 복잡해진다는 단점이 생깁니다.
이 논문이 발표 됐을 때에만 해도 NLP (Natural Language Processing) 분야에서 MoE 기반 모델의 범위는 주로 인코더-디코더 모델과 시퀀스-투-시퀀스 작업에 국한되며, 다른 도메인에서의 적용에 대한 탐구는 제한적이고, 최첨단 언어 모델 훈련의 계산 비용이 인코더-디코더 모델보다 몇 배나 더 높은 경우가 많기 때문에 MoE의 제약 사항 중 하나로 limited scope을 말하고 있습니다. 그런데 DeepSeek R1에서 기존 대비 더 적은 리소스로 훈련하고, 더 적은 비용으로 추론 서비스가 가능하다라고 보여준 것이죠. 이제 MoE의 제약 사항으로 limited scope을 말하기는 어려울 것 같습니다.
또 다른 단점으로는, MoE 모델은 dense 모델과 같은 모델 품질을 달성하기 위해 메모리 요구사항이 훨씬 크다는 것입니다. MoE 모델은 더 적은 계산을 필요로 하지만, 훨씬 더 많은 수의 파라미터가 필요합니다. 즉, MoE 기반 모델은 품질이 동등한 밀집 모델에 비해 훨씬 낮은 “파라미터 효율성”을 가지고 있습니다. 모델 크기가 크고 파라미터 효율성이 낮아지면 훈련과 추론 모두에 어려움이 있습니다. 파라미터 효율성이란, 주어진 매개변수 수에 대해 모델의 성능을 평가하는 척도로, 매개변수가 많아질수록 효율성이 떨어지는 이유는 다음과 같습니다.
- 모델 복잡성: 매개변수가 많아지면 모델이 복잡해지며, 과적합(overfitting)의 위험이 증가합니다. 즉, 훈련 데이터에 대한 성능은 좋지만, 새로운 데이터에 대한 일반화 성능은 저하될 수 있습니다. 따라서 더 많은 매개변수를 추가한다고 해서 항상 성능이 향상되지 않습니다.
- 불필요한 매개변수: 많은 모델이 불필요하게 상당수의 매개변수를 포함하고 있을 수 있고, 이들 매개변수는 유용한 정보를 전달하지 못하게 되며, 전체적인 효율성이 떨어지게 됩니다.
- 자원 소모: 매개변수가 많아지면 이를 훈련하고 추론하는 데 필요한 메모리와 계산 자원 또한 증가하게 되어, 시스템의 부하를 증가시키고, 연산 속도를 저하시켜 결국 효과적인 리소스 사용이 어려워집니다.
- 정보의 중복: 다수의 매개변수가 동일한 패턴이나 정보를 학습하게 되면, 모델이 효율적인 방법으로 정보를 포착하지 못하게 되고, 일부 매개변수는 실제로는 개선 효과가 없거나 미미하게 작용하게 됩니다. 크게 보면 불필요한 매개변수에 포함됩니다.
그리고 마지막으로, 앞에서 언급한 엄청 큰 모델 크기와 낮은 파라미터 효율성으로 인해, MoE 기반 모델의 빠른 추론이 더 어렵다는 것입니다. 더 큰 파라미터 크기 때문에 더 많은 GPU가 필요한데, 다중 GPU 추론 기술은 MoE 기반 모델을 위해 설계되지 않았다는 것이죠. MoE 모델의 경우에는 Expert가 여러 GPU에 펼쳐져 있기 때문에 여러 Expert들이 상당히 많은 all-to-all communication을 해야 하고, GPU가 많아질수록 이 cost가 급증하게 되겠죠. 또한 추론은 memory bound 연산이기 때문에 dense model 대비 10배 더 큰 파라미터가 필요한 MoE 모델은 dense 모델과 비슷한 레이턴시를 보장하기 위해서는 매우 큰 메모리 대역폭이 필요할 수 있게 되는 것입니다. 이러한 제약 사항들 때문에 소프트웨어, 하드웨어 두 측면에서 모두 추론을 잘 서비스하기 어려워지는 것입니다.
하지만 여전히 훈련 비용이 크게 줄어 들 수 있기 때문에 (이 논문에서는 GPT-3 및 MT-NLG와 같은 모델에 대해 동일한 모델 품질을 달성하기 위해 훈련 비용을 5배 감소) 포기하기에는 아쉽고 뭐 그런 것이죠. 😀
그래서 이 논문에서는 이런것을 만들어 냅니다. MoE 기반 모델의 파라미터 효율성을 향상하기 위해서, Pyramid-Residual MoE (PR-MoE)를 개발하였고, PR-MoE는 dense 모델과 MoE 모델의 혼합 형태로, expert를 가장 효과적인 곳에만 적용하여 MoE 모델의 파라미터 크기를 최대 3배까지 줄였습니다. 당연히 모델의 품질은 그대로이지만, compute 요구 사항의 변화는 최소화하고, PR-MoE의 증류 버전인 Mixture-of-Students (MoS)를 단계별 지식 증류 (knowledge distillation)를 통해 생성하였는데, MoS는 비교 가능한 모델 품질을 유지하면서 MoE 모델 크기를 최대 3.7배 줄인다고 하네요. Deepseek에서 하는 것들이 어느 순간 세상에 나온 것은 아니라는 것을 알 수 있으시겠죠? 😃
이 PR-MoE를 기반으로, 수백 개 GPU에서 확장이 가능한 고도로 최적화된 MoE 추론 시스템인 DeepSpeed-MoE 추론 시스템을 개발했고, 1TB(1000B)의 초초거대 모델의 경우에도 25ms의 레이턴시를 보여준다는 것입니다. 또한 "DeepSpeed-MoE also offers up to 4.5x faster and 9x cheaper inference for MoE models compared to quality-equivalent dense models by combining both system and model optimizations." 이렇게 표현하고 있는데요. 시스템과 모델의 최적화를 통해 추론 시 dense model 대비 4.5배 더 빠르고 비용은 9배 더 싸게 했다는 것입니다.
저는 Deepseek이 model을 open한 이유는 바로 이것이라고 생각합니다. Deepseek이 공개한 문서에서도 알 수 있는데 추론 서비스를 위해 약 350대의 GPU를 사용합니다. 한국의 H100 GPU 보유량이 약 2천대라고 하죠? 각 회사에 H100 GPU를 서비스 용도로만 400여 대를 보유한 회사가 얼마나 될까요? 그리고 이러한 모델과 시스템을 개발하고, 실험할 수 있으려면 더 많은 GPU가 필요하겠죠? 이런 것들을 모두 확보한 회사가 얼마나 될까요? 결국 model이 open 되어 있더라도, 이러한 자본(돈, GPU, 인력)이 없으면 시도조차 할 수 없고, 설사 돈으로 이러한 것들을 확보할 수 있다고 하더라도 이 MoE 기반의 서비스라는 것은 모델만으로 되는 것이 아니고, 추론 서비스 시스템(+인프라)를 고도로 최적화해서 추론 비용을 최적화할 수 있는 기술력이 있어야 한다는 것입니다. 동일한 하드웨어와 동일한 모델을 가지고 있더라도 결국 추론 서비스에 대한 기술이 없다면 서비스의 비용이 비싸지기 때문에, 이것들 중 하나라도 결여된다면 Deepseek 대비 서비스 비용을 절대 낮게 운영할 수가 없기 때문에, model을 자신 있게 open 할 수 있었을 것이라고 생각합니다. Deepseek이 등장한 이후, 사람들이 MoE의 이러한 특성 때문에 GPU가 이제 더 폭발적으로 필요하지는 않을 것이라고 생각해서 NVIDIA의 주가가 폭락했었죠. 하지만 이 기술을 자세히 살펴보았을 때, 저는 GPU가 줄어들기보다는 오히려 더 많이 필요해질 것이라고 생각합니다.
다시 논문 내용으로 돌아가서, 기존의 연구들은 MoE의 특에 의해 계산 비용을 늘리지 않고도 모델 크기를 크게 늘릴 수 있다는 것에 중점을 맞춘 반면에, 이 논문에서는 MoE 모델의 훈련, 모델 설계 및 추론에 관련된 연구를 포함하고 있다는 부분이 차이점입니다.
먼저 architecture에 대해 알아보겠습니다. 이 연구를 위해서 350M(24 layer, 1024 hidden_dim, 16 attention head), 1.3B(24 layer, 1024 hidden_dim, 32 attention head), 6.7B (32 layer, 4096 hidden_dim, 32 attention head)의 3개 모델을 선택하고, "350M+MoE-128"은 350M dense 모델을 base로 해서, 매 Feedforward 레이어마다 128개의 expert를 추가하는 MoE 모델을 나타냅니다. 350M+MoE-128과 1.3B+MoE-128 모두 총 12개의 MoE 레이어가 있다고 되어 있습니다. 350M, 1.3B 모두 24 layer인데, 기존 layer 12 + MoE layer 12개로 운영한다는 말이죠.
MoE layer에서는 각 토큰에 대해 특정 expert를 활성화 하기 위한 gating 함수를 사용하는데, 토큰 당 상위 1개의 expert만을 선택(top-1)하게 하면 1.3B+MoE-128 모델은 1.3B dense 모델과 동일한 수의 파라미터가 활성화되고, 훈련 계산량도 1.3B dense 모델과 유사합니다 top-2를 사용할 경우에는 수렴 개선 효과는 있었는데, 훈련 및 통신에 필요한 컴퓨팅 자원과 시간의 오버헤드가 크게 증가하는 문제가 있었습니다. 결국 커뮤니케이션 오버헤드를 얼마나 잘 개선하는가가 관건이 되는 것이죠.
이 그래프는 validation loss에 대해서 보여주고 있는데요, validation loss는 모델이 훈련 데이터가 아닌 검증 데이터에서 얼마나 잘 일반화되었는지를 나타내는 지표입니다. 즉 모델이 미래의 데이터에 대해 얼마나 잘 예측하는지를 평가하는데 중요한 역할을 합니다. 이 값이 낮을수록 모델이 새로운 데이터에 대해 잘 작동함을 의미하며, 모델의 성능이 좋다는 것을 나타냅니다. 반대로 높을수록 모델이 훈련 데이터에 과적합(overfitting) 되었거나, 모델이 학습을 잘하지 못하고 있음을 나타냅니다. MoE 모델과 기존의 Dense 모델을 비교해 보면, 350M+MoE-128의 검증 손실이 4배 더 큰 기초를 가진 1.3B dense모델의 검증 손실과 동등하다는 것을 알 수 있습니다. 1.3B+MoE-128과 6.7B dense 모델의 비교도 마찬가지 결과를 보이고 있죠. 즉 MoE 모델과 4-5배 더 큰 기초를 가진 dense 모델이 매우 유사한 모델 품질을 가지고 있음을 보여주는 것입니다.
이 부분의 결과에 주목해서 이 연구에서는 5배의 훈련 비용감소를 가져올 수 있다고 하는 것이고, 조 단위 매개변수 밀집 모델과 유사한 정확도를 가진 모델은 200B 매개변수(예: GPT-3) 크기의 밀집 모델 훈련 비용으로 잠재적으로 훈련될 수 있으며, 이는 수백만 달러의 훈련 비용 절감과 에너지 절약으로 이어질 수 있다고 명시하고 있습니다. DeepSeek이 이 연구 내용을 큰 모델을 기반으로 증명해 낸것이죠.
그리고 이제 MoE의 단점에 대해서 조금 더 살펴 보겠습니다. 동일한 모델 품질을 5배의 훈련 비용 절감으로 달성하기는 했는데, 파라미터 크기가 약 8배가 돼버린 것입니다. 6.7B dense 모델과 유사한 품질을 갖는 1.3B+MoE-128은 52B의 파라미터를 갖게 됩니다. 😱
이렇게 되면 메모리가 더 많이 필요해지게 되고, GPU당 메모리는 제약이 있기 때문에 결국 더 많은 GPU가 필요해지게 되겠죠? 그래서 이 논문에서는 파라미터를 더 줄이고 효율화할 수 있는 방안에 대해 연구를 하는데, 그것이 PR-MoE와 MoS(Mixture of Students)입니다.
먼저 표준 MoE 아키텍처는 모든 MoE 레이어에서 같은 수와 구조의 expert를 가지고 있습니다. 논문의 연구자들은 deep learning neural network에서 모든 레이어가 같은 표현을 학습하지 않는다는 것에서 영감을 받았다고 합니다. 이는 전이 학습(transfer learning)이라는 개념과 관련이 있습니다. 딥 러닝 모델을 학습할 때, 일반적으로 얕은 층(shallow layers)에서는 일반적인 정보(예: 에지, 텍스처)를 학습하고, 깊은 층(deep layers)에서는 특정한 문제에 대한 더 구체적인 정보를 학습한다는 관찰에 기반합니다.
그래서 2가지 half-MoE architecture의 성능에 대해 비교를 해본 것입니다. First-Half-MoE는 모델의 전반부 레이어에 MoE 레이어를 배치하고 후반부 레이어는 dense 모델과 동일하게 유지합니다. Second-Half-MoE는 MoE 레이어를 후반부로 전환하고 전반부는 dense 레이어를 사용합니다. 결과는 왼쪽에 있는 것처럼 Second-Half-MoE가 성능이 더 좋다는 것을 알 수 있습니다. 이는 모든 MoE 레이어가 동일한 수준의 표현을 학습하지 않는다는 것을 확인합니다. deep 레이어 (출력과 가까운 레이어)는 더 많은 expert의 혜택을 더 많이 누리게 된다는 것을 볼 수 있죠. (Phenomenon-I).
이 논문에서는 MoE 모델의 일반화 성능을 개선하기 위해 두 가지 일반적인 방법이 있다고 말합니다.
- expert 수를 늘리면서 expert 용량을 같게 유지하는 것
- 동일한 expert 수를 유지하면서 expert 용량을 약간 더 많은 계산(33%)의 대가로 두 배로 늘리는 것.
그러나 1 번의 경우, expert의 수가 늘어남에 따라 훈련 리소스의 메모리 요건이 증가해야 합니다 2번의 경우, 용량 증가로 인해 통신량이 두 배로 늘어나 훈련 및 추론 속도가 상당히 느려질 수 있습니다. 🤷🏻♂️ 효율성을 유지하면서도 성능을 높일 수 있는 방안이 뭘까요.
expert 용량이 더 큰 것이 정확성을 올리는 이유 중 하나는 이러한 추가 expert들이 첫 번째 expert의 “표현”을 수정하는 데 도움을 줄 수 있다는 것입니다. 이 논문에서는 이 첫 번째 전문가가 매번 변경될 필요가 있을지? 아니면 첫 번째 전문가를 고정하고 다른 전문가를 각 토큰에 할당할 수 있을지에 대해 두 가지 방법으로 비교를 수행했습니다. (Top2-MoE라고 하는 용량을 두 배로 늘리는 것과 Residual-MoE라고 하는 한 전문가는 고정하고 두 번째 전문가는 다양한 전문가 사이에서 변형하는 것)
Residual-MoE라는 것을 조금 더 살펴보면, 이 방식은 하나의 토큰은 항상 dense MLP 모듈과 MoE 모듈의 expert를 통과하고 이 두 가지 출력의 결과를 합쳐 최종 출력을 얻습니다. 여기서 중요한 직관은 MoE 모듈의 expert를 dense MLP 모듈에서 발생할 수 있는 오류를 보완할 수 있게 한 것입니다. 그래서 더 적은 통신량으로도 두 개의 전문가를 사용하는 것과 유사한 효과를 얻을 수 있게 되고, Top-1 게이팅 기능과 동일한 통신량으로 레이어당 2개의 expert를 사용하는 이점을 얻을 수 있게 된 것입니다.
위의 그림 오른쪽을 보면, 350M+MoE-32 모델에 대해서 비교를 수행한 결과입니다. Top-2 게이팅을 사용했을 경우와 Residual-MoE 방식을 비교해봤을 때 일반화 성능은 동일한 것을 알 수 있습니다. 그렇지만 Residual-MoE의 훈련 속도는 통신량이 감소하여 Top2-MoE 방식보다 10% 이상 빠르다고 합니다. (Phenomenon-II)
PR-MoE(Pyramid Residual MoE) architecture에 대한 그림입니다. 오른쪽을 보면 위에서 설명한 것처럼 swallow layer (입력과 가까운 레이어) 보다 deep layer에 더 많은 expert를 배치하였습니다. 그래서 "피라미드"라고 하는 것이고요. 마지막 레이어는 이전 레이어의 "2배"의 expert 수로 배치합니다. 그리고 residual connection 방식을 적용하여, 각 토큰은 1개의 고정 MLP와 하나의 선택된 expert를 별도로 통과하게 됩니다. 그림을 자세히 보시면, 왼쪽의 표준 MoE 경우보다 PR-MoE에서는 상위 레이어로 갈수록 연산량이 줄어들고, 메모리 요구사항이 줄어들게 될 것입니다. (Expert 개수가 더 적으니까.) 또한 PR-MoE에서는 하나의 MLP를 반드시 통과하기 때문에, 입력과 출력 path가 왼쪽 표준 MoE보다 더 쉬워질 것입니다.
표준 MoE에서는 GPU가 8대 Expert의 수가 8이라고 하면, 각 GPU당 expert 1개를 배치하면, 모든 GPU에 동일한 수의 expert가 할당되어, 부하 불균형 문제가 없습니다. 하지만 PR-MoE에서는 각 레이어의 expert 수가 다르기 때문에 단일 값으로 expert parallelism을 할 수가 없게 되는 것이죠. 선정하는 값에 따라서 어떤 레이어에서는 한 대의 GPU에서 여러 개의 expert를 담당하게 될 수도 있고, 어떤 레이어에서는 오히려 expert가 할당되지 않는 GPU가 생길 수도 있게 되는 것입니다. 즉 기존의 expert parallelism이나 data parallelism 방식이 최적화 되지 않는 문제가 발생한 것이죠.
그래서 Deepspeed-MoE에서는 다음과 같은 방식을 사용합니다.
| MoE Layer | Number of Expert | Expert parallelism | Data parallelism |
| Layer 1 | 32 | 32 | 4 |
| Layer 2 | 64 | 64 | 2 |
| Layer 3 | 128 | 128 | 1 |
이렇게 해서, 부하 불균형이나 메모리 요구량 증가 없이 효율적인 사용이 가능하다라는 것입니다.
모든 것이 다 해피해진 것 처럼 보이지만, 여기에는 문제가 하나 있습니다. 🤔 바로 커뮤니케이션 오버헤드 입니다.
기존 MoE 방식의 가장 큰 커뮤니케이션 오버헤드는, exprt 별로 토큰을 라우팅하고, 연산 후 다시 결합하는 과정에서 발생하는 과정에서 발생하는 것인데요.
입력 토큰을 gate function이 expert 들에게 분배 할 때 발생하는 all-to-all 통신.
그리고 결과를 다시 모아서 다음 레이어로 전달 할 때 발생하는 all-to-all 통신입니다.
그런데, DeepSpeed-MoE에서 처럼 MoE 레이어마다 expert 수도 다르고, expert parallelism에 data parallelism 까지 사용해야 하면 커뮤니케이션 패턴이 더 복잡해지겠죠? 위 표에서 처럼 레이어 1 -> 레이어 2로 갈 때 DP 그룹의 수가 4에서 2가 되면, GPU간 데이터 재배치가 필요하게 되면서 추가적인 커뮤니케이션 오버헤드가 발생합니다. 또한 레이어2 -> 레이어3으로 갈 때 expert의 수가 다르기 때문에 또 GPU간 데이터 이동이 필요하게 되는 것이죠.
그래서 DeepSpeed-MoE에서는 이러한 문제를 해결하기 위해 몇 가지 최적화 기법을 사용합니다.
- Hierarchical All-to-All 통신 기법
- MoE 레이어별로 최적화된 통신 패턴 적용
- Selective Activation & Load Balancing
- Fused Communication & Computation
Hierarchical All-to-All 통신 기법은, Flat All-to-All가 초래하는 무거운 통신 부담을 낮추기 위해, 먼저 같은 노드 내부에서 Expert별 토큰을 모으고, 필요한 경우에만 노드 간 통신을 수행함으로써 Cross-Node 트래픽을 최소화하는 것입니다. MoE 레이어별로 최적화된 통신 패턴은, 레이어별로 expert가 적고 data parallelism이 크면 all-reduce를 이용하고, 반대로 expert가 많고 data parallelism이 적으면 all-to-all을 사용하는 식으로 최적화를 하는 것입니다. Selective activation & load balancing은 특정 GPU의 부하가 너무 커질 경우 동적으로 로드 밸런싱을 하는 것이고요. Fused commnication & computation은 MoE에서 all-to-all 연산 후에 expert 연산이 수행되니까 GPU가 데이터를 받으면 바로 연산을 수행할 수 있도록 해서 전체적인 훈련/추론 속도를 향상 시키는 것이죠.
Mixture-of-Students에 관련된 그래프 입니다. 티처 모델에서 depth를 21로 줄이고 평가한 결과, 성능을 유사하게 가져가면서도 레이어를 12.5% 감소시킬 수 있다는 것이고요. 이를 통해 추론 시 처리량 감소 및 레이턴시 감소를 가능하게 할 수 있다는 것입니다. 이러한 연구들을 통해서 표준 MoE 대비 PR-MoE에서는 최대 3배의 메모리 소비 감소를 달성했다고 합니다.
여기까지 딱 절반 정리했습니다. 글이 너무 길어지니 2탄에서 이어서 진행하겠습니다. 남은 부분은 DeepSpeed-MoE의 추론 시스템 설계와 관련된 내용들입니다.
끝.