55. MoE:Mistral의 초고속 추론의 비밀

필자가 오늘 이야기할 주제는 MoE다. MoE에 관심을 가지게 된 계기는 Mistral LLM의 빠른 inference 속도 때문이다. 어떻게 Mistral은 이렇게 빠르게 inference가 되는지 찾아가는 과정 중 MoE를 다시 한 번 깊게 생각하고 코드 수준까지 알아보는 시간을 갖게 되었다. 지금부터 한번 살펴보면 좋겠다. 이 글은 허깅페이스에 MoE posting을 공부하면서 이것 저것 추가하면 작성된 글입니다.

TL;DR(“Too Long; Didn’t Read”)

MoE:

동일한 수의 매개변수를 가진 모델(Dense Models)에 비해 추론 속도가 더 빠릅니다 .
모든 전문가가 메모리에 로드되므로 높은 VRAM이 필요합니다.
미세 조정에 많은 어려움이 있지만 최근 instruct모델로 인기가 많습니다.

Dense Model이라는 말이 머지? 라고 생각하는 분들을 위해 간략하게 설명을 드리자면,

인공지능에서 “dense model”이라는 용어는 특히 신경망의 맥락에서 주로 사용됩니다. 이 용어는 신경망의 특정 유형, 특히 ‘밀집층(dense layer)’ 또는 ‘완전 연결층(fully connected layer)’을 구성하는 모델을 지칭합니다. 여기서 ‘Dense’이라는 용어는 신경망의 각 뉴런(노드)이 이전 층의 모든 뉴런과 연결되어 있다는 것을 의미합니다.

Dense Model의 특징:

완전 연결: 각 뉴런은 이전 층의 모든 뉴런과 연결되어 있으며, 각 연결에는 고유한 가중치가 있습니다.
다양한 층: 밀집 모델은 여러 개의 숨겨진 층을 포함할 수 있으며, 각 층은 여러 뉴런으로 구성됩니다.
범용성: 다양한 유형의 데이터와 문제에 적용할 수 있습니다. 예를 들어, 분류, 회귀, 특징 추출 등에 사용됩니다.
매개변수의 수: 밀집층은 많은 수의 매개변수(가중치와 편향)를 가지고 있으며, 이로 인해 모델이 복잡해질 수 있습니다.

이제 다시 본론으로 돌아와서 MoE와 Mistral에 대해서 자세히 한번 살펴보겠습니다.

Mixture of Experts(MoE)란 무엇일까요?

모델의 사이즈는 더 나은 모델 품질을 위한 가장 중요한 요소 중 하나입니다. 한정된 컴퓨팅 예산이 주어진다면 모델 사이즈가 작은 모델을 많이 학습하는 것보다 큰 사이즈 모델을 조금 학습하는 것이 훨씬 경제적이라 할 수 있습니다.

Mixture of Experts을 사용하면 훨씬 적은 컴퓨팅으로 모델을 사전 학습할 수 있습니다. 즉, 일반적인 인공지능 모델과 동일한 비용으로 모델을 만든다고 했을 때 더 큰 사이즈의 모델을 만들 수 있고, 더 많은 데이터를 학습시킬 수 있습니다. 특히 MoE 모델은 사전 훈련 중에 인공지능 모델과 동일한 품질을 훨씬 더 빠르게 달성하게 됩니다.

그렇다면 MoE란 정확히 무엇일까요? Transformer 모델의 맥락에서 MoE는 Sparse MoE layers와 Gate network or router(쉽게 생각해서 softmax로 토큰을 분류하는 분류기로 생각하시면 됩니다.)로 구성됩니다.

Sparse MoE layers : 일반적인 인공지능 모델(Dense Model)에서는 FFN(feed-forward Network) layer를 사용합니다. 그러나, Mistral에서는 MoE 레이어가 사용됩니다. MoE layer에는 특정 수의 “Expert(전문가)”가 있으며 각 expert는 신경망 네트워크를 의미합니다. 코드로 살펴보면, expert는 실제로 FFN이지만 더 복잡한 네트워크이거나 심지어 MoE 자체일 수도 있어 계층적 MoE로 이어질 수도 있습니다!
쉽게 이야기를 하자면, FFN는 정보는 모든 인공신경망에 나누어 보관한다면, MoE는 특정 layer에 정보를 나누어 담을 수 있다고 생각하면 될 거 같습니다. 또는 계층적(연속적으로)으로 정보를 저장할 수도 있다! 또한, Sparse라는 단어가 궁금하실텐데요. 뒤에 별도의 캡터로 설명을 합니다. 쉽게 말하자면, 조건부 계산이라고 생각하시면 됩니다.
Gate network or router : 어떤 토큰이 어떤 expert에게 전송되는지 결정하는 게이트 네트워크 또는 라우터 .라고 생각하시면 됩니다. 예를 들어 아래 이미지에서 “More” 토큰은 두 번째 expert에게 전송되고 “Parameters” 토큰은 첫 번째 expert로 전송됩니다. 나중에 살펴보겠지만 두 명 이상의 expert에게 토큰을 보낼 수 있습니다. 신기하죠? 인공지능의 정말 무궁무진해서 재미있는거 같습니다. 다시 본론으로 돌아와서, 토큰을 expert에게 라우팅하는 방법은 MoE와 작업할 때 중요한 결정 중 하나입니다. 라우터는 학습된 매개변수로 구성되며 네트워크의 나머지 부분과 동시에 사전 훈련됩니다.

요약하면 MoE에서는 Transformer 모델의 모든 FFN 레이어를 게이트 네트워크와 특정 수의 expert로 구성된 MoE 레이어로 대체합니다.

MoE는 기존의 Dense 모델에 비해 효율적인 사전 훈련 및 더 빠른 추론과 같은 이점을 제공하지만 넘어야할 산이 아직 남아 있습니다.

Training: MoE는 훨씬 더 효율적인 컴퓨팅 효율성을 제공하지만 역사적으로 미세 조정 중에 일반화하는 데 어려움을 겪어 과적합된다는 문제점이 있습니다.
Inference: 장점으로는 MoE에는 많은 매개변수가 있을 수 있지만 추론 중에는 그 중 일부만 사용됩니다. 또한, 이는 동일한 수의 매개변수를 가진 일반적인 인공지능 모델에 비해 훨씬 더 빠른 추론으로 이어집니다. 그러나 모든 매개변수를 RAM에 로드해야 하므로 메모리 요구 사항이 높습니다.
여기에 재미있는게 있습니다. 왜 이런 현상이 발생되는지 찾느라 고생을 좀 많이 했는데요. Mistral 8x7B를 가져와서 로드를 GPU 사용량이 우리가 예상하는 것과 다릅니다. 같은 사이즈의 모델과 비교했을 때도 적습니다. 왜 이럴까 왜 이럴까? 고민하며 찾다보니, MoE 모델에서는 FFN 레이어만 개별 전문가로 처리되고 나머지 모델 매개변수는 공유되기 때문입니다. 또한 동시에, 토큰당 두 명의 전문가만 사용한다고 가정하면 추론 속도(FLOP)는 2x7B 행렬 곱셈을 계산하기 때문에 12B 모델(14B 모델과 반대)을 사용하는 것과 같지만 일부 레이어가 공유됩니다. 뒤로 읽어가시다보면 더 이해가 되실거에요.

이제 우리는 MoE가 무엇인지 대략적으로 알았으니 MoE의 시작은 어떤 연구였는지 개발 과정을 살펴보겠습니다.

MoE는 언제 만들어진걸까?

MoE의 뿌리는 1991년 논문 Adaptive Mixture of Local Experts 에서 유래되었습니다 . 앙상블 방법과 유사한 아이디어는 각각 훈련 사례의 서로 다른 하위 집합을 처리하는 별도의 네트워크로 구성된 시스템에 대한 점검하는 절차를 갖는 것이었습니다. 각각의 개별 네트워크 또는 전문가는 입력 공간의 서로 다른 영역을 전문으로 합니다. expert는 어떻게 선정될까요? 게이트 네트워크는 각 expert의 가중치를 결정합니다. 훈련 중에는 expert와 게이트는 모두 훈련됩니다.

2010년부터 2015년 사이에 두 가지 다른 연구 분야가 이후 MoE 발전에 기여했습니다.

구성 요소로서의 전문가 : 기존 MoE의 전체적인 구성은 게이팅 네트워크와 여러 전문가로 구성됩니다. MoE는 SVM, 가우스 프로세스 및 그 외 여러 연구에서 발견되었습니다. Eigen, Ranzato 및 Ilya 의 연구는 MoE를 딥러닝에 적용하였습니다. 이 연구를 통해 MoE가 multilayer network에 들어갈 수 있다는 것을 알 수 있게 되었습니다.
조건부 계산 : 기존 네트워크는 모든 계층을 통해 모든 입력 데이터를 처리합니다. 이러한 흐름이 지속되는 동안Yoshua Bengio라는 분은 입력 토큰을 기반으로 구성 요소를 동적으로 활성화하거나 비활성화하는 접근 방식을 연구했습니다.

이러한 conditional Computation 연구을 통해 NLP의 맥락에서 다양한 전문가를 연구하게 됩니다. 구체적으로 Shazeer et al. (2017, Geoffrey Hinton 및 Jeff Dean, Google의 Chuck Norris)는 spariyt를 도입하여 이 아이디어를 137B LSTM(당시 Schmidhuber가 만든 사실상의 NLP 아키텍처)으로 확장하여 매우 빠른 속도를 유지할 수 있었습니다. 대규모에서도 추론이 가능합니다. 이 논문은 번역에 중점을 두었지만 높은 커뮤니케이션 비용과 학습 불안정성 등 많은 어려움에 직면했습니다.

LSTM의 MoE 계층 — https://arxiv.org/pdf/1701.06538.pdf

MoE는 multi-trillion parameter models을 훈련시킬 수 있게 만들었습니다. 그 중 가장 대표적인 예는 1.6T parameters Switch Transformers입니다. MoE는 Computer Vision에서도 다양하게 연구되었다고 합니다. 나중에 기회가 되면 어떻게 연구가 되었는지 논문의 흐름을 따라가 보고 싶습니다. 또한, 한 가지 아쉬운 점은 논문들을 들어가보면 구글에서 발표한 논문이 많이 있다는 것입니다. 그것도 무려 몇년전에 이러한 개념을 발전시켰지만, 아쉽게도 적절한 빛을 보지 못한다는 점이 아쉽더라구요. 이러한 점은 중요한 것을 시사한다고 저는 생각합니다. 많은 천재들이 연구를 하였지만, 빛을 바라보지 못한 개념들이 많이 있다는 것이죠! 우리는 이러한 논문들을 잘 읽고 실생활에 적용하는 소규모 프로젝트를 진행할 필요가 있다고 생각합니다. 이 부분에 대해서는 나중에 별도의 포스팅을 한번 다루겠습니다.

Sparse이란 무엇입니까?

Sparse MoE layer 라는 말을 앞에서부터 사용하였습니다. 여기서 한번 자세히 살펴보고 가도록 하겠습니다. 앞에서도 말씀드렸지만, 본 포스팅의 흐름 상, Sparse는 조건부 계산이라고 생각하시면 됩니다. 일반적인 인공지능 모델에서는 모든 파라미터가 모든 입력에 사용되는 반면, Sparse를 통해 전체 시스템의 일부 부분만 실행할 수 있습니다.

번역에서 사용된 MoE Shazeer의 연구에 대해 더 자세히 살펴보겠습니다. 조건부 계산(네트워크의 일부가 예제별로 활성화됨)이라는 아이디어를 사용하면 계산을 늘리지 않고도 모델 크기를 확장할 수 있으므로 각 MoE 계층에서 수천 명의 전문가가 사용됩니다.

단점은 한쪽 전문가로 쏠려서 과적합되기 쉽다는 점이 분명히 있고, 이를 훈련시키기 어렵다는 문제점이 있습니다. 이를 좀 더 쉽게 풀어서 설명드려보겠습니다.

이 단계에는 몇 가지 문제가 여전히 남아 있습니다. 예를 들어, 일반적으로 배치 크기가 크면 성능이 더 좋지만, MoE에서는 오히려 성능이 떨어진다는 문제점이 있습니다. 이 말이 잘 이해가 안되실 수 있습니다. 입력이 들어오면 어떤 expert를 활성화시켜야하는지 결정이 됩니다. 이렇게 활성화된 expert를 active expert라고 우리는 이야기합니다. active expert는 전체 배치에서 모든 데이터 포인트가 모든 전문가에게 전달되는 것이 아니라, 오직 일부 전문가들에게만 전달됩니다. 그렇기에 10을 넣어도 10이 다 전달이 되지 않는다는 이슈가 발생되게 됩니다. 이 과정에서 각 전문가가 처리하는 데이터의 양, 즉 배치크기는 전체 배치크기보다 작아지는 현상이 발생됩니다. 이러한 MoE의 특징 때문에 MOE의 배치 크기는 줄어들게 됩니다. 그렇다면, 이러한 문제를 어떻게 해결할 수 있을까요?

학습된 게이팅 네트워크(G)는 입력의 일부분을 어떤 전문가(E)에게 보내야할지 결정합니다.

$$ y=\sum_{i=1}^n G(x)_i E_i(x) $$

위 수식은 Mixture of Experts (MoE) 모델에서의 출력을 계산하는 방법을 나타냅니다. 이 수식을 이해하기 위해서는 각 부분의 의미를 알아야 합니다.

$y$ : 이것은 최종 출력을 나타냅니다. MoE 모델의 목적은 주어진 입력 $x$에 대한 최종 출력 $y$를 계산하는 것입니다.
$\sum_{i=1}^n$ : 이것은 합계 기호입니다. MoE 모델에서는 여러 전문가의 출력을 합산하여 최종 출력을 결정합니다. 여기서 $n$은 전문가의 총 수입니다.
$G(x)_i$ : 이것은 게이팅 네트워크의 출력을 나타냅니다. 입력 $x$에 대해, 게이팅 네트워크는 각 전문가 $i$에 대한 가중치 $G(x)_i$를 계산합니다. 이 가중치는 해당 전문가가 입력 $x$에 얼마나 중요한지를 나타냅니다. 만약 $G(x)_i$가 0이면, 해당 전문가는 연산을 진행하지 않고, 최종 출력에 기여하지 않습니다.
$E_i(x)$ : 이것은 전문가 네트워크의 출력을 나타냅니다. 각 전문가 네트워크 $E_i(x)$는 입력 $x$를 받아 자신만의 출력을 생성합니다.

이를 종합해보면, MoE모델의 최종 출력 $y$는 각 전문가의 출력 $E_i(x)$에 그 전문가의 중요도를 나타내는 게이팅 가중치$G(x)_i$를 곱한 값들의 합으로 계산이 됩니다. 저는 여기서 궁금한 사항이 있었습니다. i가 0이 다면 어떻게 되는것일까? 만약 특정 전문가에 대한 G 값이 0이라면, 해당 전문가에 대한 연산을 수행할 필요가 없습니다. 이는 게이팅 네트워크가 그 전문가가 해당 입력에 대해 관련성이 없다고 판단했기 때문입니다. 결과적으로, 이러한 최적화는 불필요한 계산을 줄임으로써 전체적인 모델의 계산 효율성을 향상시킵니다.

저는 또 하나의 궁금증이 있었습니다. 일반적인 무엇이 게이팅의 역할을 하는걸까? 였습니다. 찾아보니 가장 전통적인 것은 소프트맥스 기능이 있는 간단한 네트워크를 사용합니다.

$$ G_\sigma(x)=Softmax(x \cdot W_g) $$

$x$ : 입력 벡터입니다.
$W_g$ : 게이팅 네트워크의 가중치 행렬입니다. 이 행렬은 입력 $x$와 곱해져 각 전문가에 대한 스코어를 생성합니다.
Softmax : Softmax 함수는 각 전문가에 대한 스코어를 확률 분포로 변환합니다. 이 확률 분포는 각 전문가가 해당 입력에 대해 얼마나 중요한지를 나타냅니다. Softmax 출력은 각 전문가에게 할당되는 가중치를 결정합니다. 이 가중치는 MoE 모델의 최종 출력을 결정하는데 사용됩니다. 즉, 이 수식은 MoE 모델에서 각 입력 데이터가 어떤 전문가에 의해 처리될지 결정하는 중요한 역할을 합니다. Softmax를 통해 생성된 확률 분포는 각 전문가의 중요도를 나타내며, 이를 통해 MoE 모델은 다양한 종류의 데이터에 유연하게 대응할 수 있습니다.

네트워크는 어떤 전문가가 입력을 보낼지 학습합니다. 이 외에도 Shazeer’s 논문에서는 noise가 추가된 Noisy Top-k Gating 방식이 있었으며, Noisy Top-k Gating 관련하여서는 Switch Transformers 를 참고 부탁드려요~

왜 노이즈를 추가합니까? 로드 밸런싱을 위한 것입니다!

MoE를 위한 로드 밸런싱 토큰

이번 섹션을 다루기 위해서는 로드 밸런싱이라는 개념을 알아야합니다.

일반적으로 로디 밸런싱(Load Balancing)은 네트워크 트래픽, 애플리케이션 요청 또는 컴퓨팅 작업을 여러 서버나 리소스 간에 분산시켜 처리하는 기술입니다. 이는 특히 확장성 및 서비스의 안정성을 보장하는데 중요한 역할을 합니다. 로드 밸런싱의 주요 목적은 하나의 서버에 과부하가 걸리는 것을 방지하고, 자원을 최적으로 활용하여 전체 시스템의 효율성을 높이는 것입니다.

이러한 개념을 MoE를 적용한 모델에 적용해보겠습니다.

MoE를 적용한 모델에는 많은 Experts가 있습니다. 하나의 Expert에 많은 입력이 쏠리지 않게 적절하게 배분해주는 토큰이 필요하다! 정도로 이해해볼 수 있습니다. 왜 이런 토큰이 필요한 것일까요?

이전에 논의한 것처럼, 우리의 모든 토큰이 소수의 Experts에게만 전송된다면 훈련이 비효율적이게 됩니다. 일반적인 MoE 학습에서 Gating Network는 대부분 동일한 소수의 Expert를 활성화하도록 수렴됩니다. 이는 선호하는 전문가가 더 빨리 훈련되고 더 많이 선택됨에 따라 자체적으로 강화됩니다. 이를 완화하기 위해 모든 전문가에게 동일한 중요성을 부여하도록 장려하는 Auxiliary loss가 추가됩니다. 이러한 손실을 통해 모든 전문가는 대략 동일한 수의 훈련 예제를 받을 수 있습니다.

MoE 코드로 살펴보기

import dataclasses
from typing import List

import torch
import torch.nn.functional as F
from simple_parsing.helpers import Serializable
from torch import nn


@dataclasses.dataclass
class MoeArgs(Serializable):
    num_experts: int
    num_experts_per_tok: int


class MoeLayer(nn.Module):
    def __init__(self, experts: List[nn.Module], gate: nn.Module, moe_args: MoeArgs):
        super().__init__()
        assert len(experts) > 0
        self.experts = nn.ModuleList(experts)
        self.gate = gate
        self.args = moe_args

    def forward(self, inputs: torch.Tensor):
        gate_logits = self.gate(inputs)
        weights, selected_experts = torch.topk(gate_logits, self.args.num_experts_per_tok)
        weights = F.softmax(weights, dim=1, dtype=torch.float).to(inputs.dtype)
        results = torch.zeros_like(inputs)
        for i, expert in enumerate(self.experts):
            batch_idx, nth_expert = torch.where(selected_experts == i)
            results[batch_idx] += weights[batch_idx, nth_expert, None] * expert(
                inputs[batch_idx]
            )
        return results

이 코드는 Github의 mistralai에서 mistral-src라는 곳에서 가져왔습니다.

하나하나 실펴보겠습니다.

`MoeArgs` 클래스

MoeArgs 클래스는 dataclasses 모듈을 사용하여 정의되었습니다. 이 클래스는 MoE 레이어의 설정을 저장하는 역할을 합니다.
num_experts: 전문가(experts)의 총 수를 나타냅니다.
num_experts_per_tok: 각 토큰 또는 입력 단위에 대해 선택할 전문가의 수입니다.

`MoeLayer` 클래스

MoeLayer 클래스는 nn.Module을 상속받아 PyTorch 모듈로 정의됩니다.
__init__ 함수는 MoE 레이어를 초기화합니다. 이 때, 다수의 전문가 모듈(experts), 게이트 모듈(gate), 그리고 MoeArgs 인스턴스를 인자로 받습니다.
forward 함수는 모델의 순전파를 정의합니다. 입력 텐서(inputs)가 주어지면, 게이트 모듈을 통해 각 입력에 대한 전문가의 점수(gate_logits)를 계산합니다.
torch.topk를 사용하여 가장 높은 점수를 받은 상위 num_experts_per_tok 개의 전문가를 선택합니다.
선택된 전문가에 대한 가중치는 softmax 함수를 통해 정규화됩니다.
입력 텐서에 대해 선택된 각 전문가의 출력을 가중치와 함께 합산하여 최종 결과를 생성합니다.

전문가 선택 및 출력 계산

torch.where를 사용하여 각 배치의 데이터에 대해 선택된 전문가의 인덱스를 찾습니다.
이를 통해 해당 전문가의 가중치를 적용하고, 입력에 대해 전문가 모듈을 호출합니다.
모든 선택된 전문가의 출력을 가중치와 함께 합산하여 최종 결과를 얻습니다.

이 MoE 구현은 각 입력에 대해 다수의 전문가를 선택하고, 각 전문가의 기여도를 고려하여 최종 출력을 결정합니다. 이 방식은 모델이 다양한 유형의 데이터에 대해 더 효과적으로 대응할 수 있게 해주며, 특히 대규모 모델에서 성능을 향상시키는 데 유용합니다. MoE는 각 전문가가 처리하는 작업의 특성에 따라 다양하게 구성될 수 있으며, 복잡한 문제를 해결하는 데 있어 매우 강력한 도구가 될 수 있습니다.

막상 코드로 구현을 해보면, 몇 줄 안되는데 코드는 참 많은 것들을 대변하고 있는거 같습니다.

그러기에 우리가 코드를 좋아하는게 아닐까요? 우리는 함축적으로 효율성을 적극 추구하는 사람들이니까요!

이제는 MoEs가 적용된 대표적인 예인, switch transformers에 대해서 간단하게 살펴보도록 하겠습니다.

Switch Transformers

MoE는 많은 가능성을 보여주었지만 Train과 fine-Tuning에서는 어려움을 겪고 있습니다. Switch Transformers는 이러한 주제를 심층적으로 다루는 매우 흥미로운 연구라고 저는 생각합니다. 이뿐 아니라 어려움을 해결하기 위해 최근에는 openMoE라는 깃레포도 생긴 것으로 알고 있습니다. Switch Transformers 저자는 트랜스포머로 실행할 수 있는 2048의 전문가와 함께 Hugging Face에서 1조 6천억 개의 매개변수 MoE를 공개하기도 했습니다.

저자는 FFN 레이어를 MoE 레이어로 대체했습니다. Switch Transformers 문서에서는 2개의 입력(2개의 서로 다른 토큰)을 수신하고 4명의 전문가가 있는 Switch Transformer 레이어를 제안합니다.

최소 두 명의 전문가를 사용한다는 초기 아이디어와 달리 Switch Transformers는 단순화된 단일 전문가 전략을 사용합니다. 이 접근 방식의 효과는 다음과 같습니다.

라우터 계산이 감소됩니다.
각 전문가의 배치 크기는 최소한 절반으로 줄일 수 있습니다.
통신 비용이 절감됩니다.
품질이 유지됩니다

라우터 Z-손실을 통한 안정화 훈련

앞서 설명한 균형 손실은 불안정성 문제로 이어질 수 있습니다. 품질을 희생하면서 희소 모델을 안정화하기 위해 다양한 방법을 사용할 수 있습니다. 예를 들어 드롭아웃을 도입하면 안정성이 향상되지만 모델 품질이 저하됩니다. 반면, 곱셈 구성요소를 더 추가하면 품질은 향상되지만 안정성은 감소합니다.

ST-MoE 에 도입된 라우터 z-손실은 게이팅 네트워크에 진입하는 대규모 로지트에 페널티를 적용하여 품질 저하 없이 훈련 안정성을 크게 향상시킵니다. 이러한 손실로 인해 값의 절대 크기가 더 작아지므로 반올림 오류가 줄어들고 이는 게이팅과 같은 지수 함수에 상당히 영향을 미칠 수 있습니다. 자세한 내용은 논문을 검토해 보시기 바랍니다.

MoE 미세 조정

Mixtral은 Transformers 버전 4.36.0에서 지원됩니다.

pip install "transformers==4.36.0 --upgrade

과적합 역학은 조밀한 모델과 희소 모델 간에 매우 다릅니다. 희소 모델은 과적합에 더 취약하므로 전문가 내에서 더 높은 정규화(예: 드롭아웃)를 탐색할 수 있습니다(예: 조밀한 레이어에 대해 하나의 드롭아웃 비율을 가질 수 있고 희소 레이어에 대해 다른 더 높은 드롭아웃 비율을 가질 수 있습니다).

한 가지 질문은 미세 조정을 위해 보조 손실을 사용할지 여부입니다. ST-MoE 작성자는 보조 손실을 끄는 실험을 했으며 최대 11%의 토큰이 삭제된 경우에도 품질에 큰 영향을 미치지 않았습니다. 토큰 삭제는 과적합을 방지하는 데 도움이 되는 정규화의 한 형태일 수 있습니다.

Switch Transformers는 고정된 pretrain perplexity에서 희소 모델이 다운스트림 작업, 특히 SuperGLUE와 같은 추론이 많은 작업에서 조밀한 모델보다 성능이 더 나쁘다는 것을 관찰했습니다. 반면에 TriviaQA와 같은 지식 집약적인 작업의 경우 희소 모델이 불균형적으로 잘 수행됩니다. 저자들은 또한 미세 조정에 도움을 준 전문가 수가 적다는 사실도 관찰했습니다. 일반화 문제를 확인한 또 다른 관찰은 모델이 작은 작업에서는 더 나빴지만 더 큰 작업에서는 잘 수행되었다는 것입니다.

학습 곡선 미세 조정 — 작은 작업(왼쪽)에서는 희소 모델이 검증 세트에서 훨씬 더 나쁘기 때문에 명확한 과적합을 볼 수 있습니다. 더 큰 작업(오른쪽)에서는 MoE가 잘 수행됩니다. 이 이미지는 ST-MoE 논문에서 가져온 것입니다.

전문가가 아닌 모든 가중치를 동결하는 실험을 할 수도 있습니다. 즉, MoE 레이어만 업데이트하겠습니다. 이로 인해 성능이 크게 저하됩니다. 반대의 경우도 시도해 볼 수 있습니다. 즉, MoE 레이어의 매개변수만 동결하는 것인데, 이는 거의 모든 매개변수를 업데이트하는 것과 마찬가지로 작동했습니다. 이렇게 하면 미세 조정을 위한 메모리 속도를 높이고 메모리를 줄이는 데 도움이 될 수 있습니다. 매개변수의 80%가 MoE 레이어(ST-MoE 프로젝트)에 있으므로 이는 다소 직관적이지 않을 수 있습니다. 해당 아키텍처에 대한 그들의 가설은 전문가 레이어가 1/4 레이어마다 발생하고 각 토큰이 레이어당 최대 2명의 전문가를 볼 수 있으므로 MoE 매개변수 업데이트는 다른 매개변수 업데이트보다 훨씬 적은 레이어에 영향을 미친다는 것입니다.

비 MoE 레이어 업데이트만 미세 조정에서 잘 작동합니다. — MoE 레이어만 동결하면 품질을 유지하면서 훈련 속도를 높일 수 있습니다. 이 이미지는 ST-MoE 논문에서 가져온 것입니다.

Sparse MoE를 미세 조정할 때 고려해야 할 마지막 부분은 서로 다른 미세 조정 하이퍼파라미터 설정이 있다는 것입니다. 예를 들어 희소 모델은 더 작은 배치 크기와 더 높은 학습 속도에서 더 많은 이점을 얻는 경향이 있습니다.

조밀한 모델과 희소 모델 간의 미세 조정 배치 크기와 학습률을 비교하는 표입니다. — 희소 모델의 미세 조정 품질은 학습률이 높고 배치 크기가 작을수록 향상됩니다. 이 이미지는 ST-MoE 논문에서 가져온 것입니다.

이 시점에서 사람들이 MoE를 미세 조정하는 데 어려움을 겪고 있다는 사실이 조금 슬플 수도 있습니다. 흥미롭게도 최근 논문인 MoEs Meets Instruction Tuning (2023년 7월)에서는 다음과 같은 실험을 수행합니다.

단일 작업 미세 조정
다중 작업 명령 조정
다중 작업 명령 조정 후 단일 작업 미세 조정

저자가 MoE와 T5 등가물을 미세 조정했을 때 T5 등가물이 더 좋았습니다. 저자가 Flan T5(T5 명령과 동일) MoE를 미세 조정했을 때 MoE의 성능이 훨씬 향상되었습니다. 뿐만 아니라 MoE에 대한 Flan-MoE의 개선은 T5에 비해 Flan T5보다 컸으며, 이는 MoE가 밀집 모델보다 명령어 조정에서 훨씬 더 많은 이점을 얻을 수 있음을 나타냅니다. MoE는 더 많은 작업 수를 통해 더 많은 이점을 얻습니다. 보조 손실 기능을 끄자고 제안한 이전 논의와 달리 손실은 실제로 과적합을 방지합니다.

MoE는 밀도가 높은 모델보다 지시 조정을 통해 더 많은 이점을 얻습니다. — 희소 모델은 밀집 모델에 비해 지시 조정의 이점을 더 많이 얻습니다. 이 이미지는 MoEs Meets Instruction Tuning 문서에서 가져온 것입니다.

희소 MoE와 밀도 모델을 언제 사용해야 할까요?

전문가는 많은 컴퓨터가 포함된 높은 처리량 시나리오에 유용합니다. 사전 훈련을 위한 고정된 컴퓨팅 예산이 주어지면 희소 모델이 더 최적입니다. VRAM이 거의 없는 낮은 처리량 시나리오의 경우 밀도가 높은 모델이 더 좋습니다.

참고: 희소 모델과 밀집 모델 간의 매개변수 수를 직접 비교할 수는 없습니다. 둘 다 상당히 다른 것을 나타내기 때문입니다.

MoE를 부끄럽게 만드세요

초기 MoE 작업에서는 MoE 레이어를 분기 설정으로 제시하여 GPU가 이에 맞게 설계되지 않았기 때문에 계산 속도가 느려지고 장치가 다른 사람에게 정보를 보내야 하므로 네트워크 대역폭에 병목 현상이 발생했습니다. 이 섹션에서는 이러한 모델을 사용한 사전 학습 및 추론을 보다 실용적으로 만들기 위한 일부 기존 작업에 이야기를 해볼까합니다.

병렬성

병렬성에 대해 간단히 살펴보겠습니다.

데이터 병렬성: 동일한 가중치가 모든 코어에 복제되고 데이터가 코어에 분할됩니다.
모델 병렬성: 모델은 여러 코어로 분할되고 데이터는 여러 코어로 복제됩니다.
모델 및 데이터 병렬성: 모델과 데이터를 코어 전체에 분할할 수 있습니다. 서로 다른 코어는 서로 다른 데이터 배치를 처리합니다.
전문가 병렬성 : 전문가는 서로 다른 작업자에게 배치됩니다. 데이터 병렬성과 결합하면 각 코어에는 서로 다른 전문가가 있고 데이터는 모든 코어에 걸쳐 분할됩니다.

전문가 병렬성을 사용하면 전문가가 서로 다른 작업자에게 배치되고 각 작업자는 서로 다른 훈련 샘플 배치를 사용합니다. MoE가 아닌 레이어의 경우 전문가 병렬 처리는 데이터 병렬 처리와 동일하게 작동합니다. MoE 레이어의 경우 원하는 전문가가 상주하는 작업자에게 시퀀스의 토큰이 전송됩니다.

모델, 전문가 및 데이터 병렬성을 보여주는 이미지 — 다양한 병렬 처리 기술을 사용하여 데이터와 모델이 코어를 통해 분할되는 방식을 보여주는 Switch Transformers 문서의 그림입니다.

용량 계수 및 통신 비용

용량 계수(CF)를 늘리면 품질이 향상되지만 통신 비용과 활성화 메모리가 늘어납니다. 전체 통신이 느린 경우에는 더 작은 용량 요소를 사용하는 것이 좋습니다. 좋은 출발점은 용량 계수가 1.25인 상위 2 라우팅을 사용하고 코어당 전문가 한 명을 두는 것입니다. 평가 중에 용량 요소를 변경하여 컴퓨팅을 줄일 수 있습니다.

서빙 기술

mistralai/Mixtral-8x7B-Instruct-v0.1을 추론 엔드포인트에 배포할 수 있습니다 .

MoE의 큰 단점은 매개변수의 수가 많다는 것입니다. 로컬 사용 사례의 경우 더 작은 모델을 사용하는 것이 좋습니다. 봉사에 도움이 될 수 있는 몇 가지 기술에 대해 빠르게 논의해 보겠습니다.

Switch Transformers의 저자는 초기 증류 실험을 수행했습니다. MoE를 다시 조밀한 대응물로 증류함으로써 희소성 이득의 30-40%를 유지할 수 있습니다. 따라서 증류는 더 빠른 보유 및 생산 시 더 작은 모델을 사용하는 이점을 제공합니다.
최근 접근 방식은 전체 문장이나 작업을 전문가에게 라우팅하도록 라우팅을 수정하여 서비스를 위한 하위 네트워크 추출을 허용합니다.
MoE(Aggregation of Experts): 이 기술은 전문가의 가중치를 병합하여 추론 시 매개변수 수를 줄입니다.

효율적인 훈련에 대해 자세히 알아보기

FasterMoE(2022년 3월)는 고효율 분산 시스템에서 MoE의 성능을 분석하고 다양한 병렬성 전략의 이론적 한계뿐만 아니라 전문가 인기를 왜곡하는 기술, 대기 시간을 줄이는 세분화된 통신 일정, 조정된 토폴로지 인식을 분석합니다. 가장 낮은 대기 시간을 기준으로 전문가를 선택하는 게이트로 속도가 17배 향상됩니다.

Megablocks(2022년 11월)는 MoE에 존재하는 역동성을 처리할 수 있는 새로운 GPU 커널을 제공하여 효율적인 희소 사전 학습을 탐구합니다. 그들의 제안은 토큰을 절대로 떨어뜨리지 않고 최신 하드웨어에 효율적으로 매핑하여 상당한 속도 향상을 가져옵니다. 비결은 무엇입니까? 전통적인 MoE는 모든 전문가가 동일한 모양과 동일한 수의 토큰을 가지고 있다고 가정하는 일괄 행렬 곱셈을 사용합니다. 이와 대조적으로 Megablocks는 MoE 레이어를 불균형 할당을 수용할 수 있는 블록 희소 작업으로 표현합니다.

블록 희소 연산에 최적화된 행렬 곱셈입니다. — 다양한 규모의 전문가와 토큰 수를 위한 블록 희소 행렬 곱셈([MegaBlocks](https://arxiv.org/abs/2211.15841)에서).

MoE가 무엇인지에 대해서 한번 알아보고 Mistral에서 MoE를 코드로 어떻게 구현했는지도 알아보았습니다.

또한, MoE를 적용한 Switch Transformers에 대해서도 간단하게 논문을 리뷰해보았습니다.

다음에는 좀 더 재미있는 MoE를 가져오도록 하겠습니다.

다른 재미있는 글도 있으니, 많이 놀러와주세요

긴 글 읽어주셔서 감사합니다.