Mixture of Experts(MoE) 아키텍처 심층 분석: Switch Transformer에서 Mixtral까지의 발전과 효율적 스케일링 전략

import torch
import torch.nn as nn
import torch.nn.functional as F

class GatingNetwork(nn.Module):
    """MoE 기본 게이팅 네트워크"""
    def __init__(self, d_model, num_experts, top_k=2):
        super().__init__()
        self.top_k = top_k
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x):
        # x: (batch_size, seq_len, d_model)
        logits = self.gate(x)  # (batch_size, seq_len, num_experts)

        # Top-k 전문가 선택
        top_k_logits, top_k_indices = torch.topk(logits, self.top_k, dim=-1)
        top_k_gates = F.softmax(top_k_logits, dim=-1)

        return top_k_gates, top_k_indices

class MoELayer(nn.Module):
    """기본 Sparse MoE 레이어"""
    def __init__(self, d_model, d_ff, num_experts, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.gate = GatingNetwork(d_model, num_experts, top_k)

        # 각 전문가는 독립적인 FFN
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff),
                nn.GELU(),
                nn.Linear(d_ff, d_model)
            )
            for _ in range(num_experts)
        ])

    def forward(self, x):
        B, T, D = x.shape
        gates, indices = self.gate(x)  # gates: (B,T,k), indices: (B,T,k)

        # 출력 초기화
        output = torch.zeros_like(x)

        # 각 전문가별로 배치 처리
        for i in range(self.num_experts):
            # i번째 전문가에 할당된 토큰 마스크
            expert_mask = (indices == i).any(dim=-1)  # (B, T)
            if not expert_mask.any():
                continue

            # 해당 토큰 추출 및 전문가 연산
            expert_input = x[expert_mask]
            expert_output = self.experts[i](expert_input)

            # 게이트 가중치 적용
            gate_values = gates[indices == i]
            output[expert_mask] += gate_values.unsqueeze(-1) * expert_output

        return output

class SwitchRouter(nn.Module):
    """Switch Transformer 라우터 (Top-1)"""
    def __init__(self, d_model, num_experts, capacity_factor=1.25):
        super().__init__()
        self.num_experts = num_experts
        self.capacity_factor = capacity_factor
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x):
        B, T, D = x.shape
        # 게이팅 로짓
        logits = self.gate(x)  # (B, T, num_experts)
        probs = F.softmax(logits, dim=-1)

        # Top-1 선택
        gate_values, expert_indices = probs.max(dim=-1)  # (B, T)

        # 전문가 용량 계산
        capacity = int(self.capacity_factor * T / self.num_experts)

        # 용량 초과 토큰 드롭 처리
        dispatch_mask = torch.zeros(B, T, self.num_experts, dtype=torch.bool)
        for i in range(self.num_experts):
            expert_mask = (expert_indices == i)
            # 용량 초과 시 초과분 드롭
            positions = expert_mask.nonzero(as_tuple=True)
            if len(positions[1]) > capacity:
                drop_indices = positions[1][capacity:]
                expert_mask[positions[0][capacity:], drop_indices] = False
            dispatch_mask[:, :, i] = expert_mask

        return gate_values, expert_indices, dispatch_mask

def load_balancing_loss(gates, expert_indices, num_experts, alpha=0.01):
    """Switch Transformer 로드 밸런싱 보조 손실"""
    B, T = expert_indices.shape

    # f_i: 각 전문가에 배정된 토큰 비율
    f = torch.zeros(num_experts, device=gates.device)
    for i in range(num_experts):
        f[i] = (expert_indices == i).float().sum() / (B * T)

    # P_i: 각 전문가의 평균 라우팅 확률
    probs = F.softmax(gates, dim=-1)  # 전체 전문가에 대한 확률
    P = probs.mean(dim=[0, 1])  # (num_experts,)

    # 밸런싱 손실
    loss = alpha * num_experts * (f * P).sum()
    return loss

class MixtralMoELayer(nn.Module):
    """Mixtral 8x7B 스타일 MoE 레이어"""
    def __init__(self, d_model=4096, d_ff=14336, num_experts=8, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k

        # 라우터: 선형 게이팅
        self.gate = nn.Linear(d_model, num_experts, bias=False)

        # 8개의 독립적인 SwiGLU FFN 전문가
        self.experts = nn.ModuleList([
            SwiGLUExpert(d_model, d_ff) for _ in range(num_experts)
        ])

    def forward(self, x):
        B, T, D = x.shape
        x_flat = x.view(-1, D)  # (B*T, D)

        # 라우팅 확률 계산
        logits = self.gate(x_flat)  # (B*T, 8)
        weights, indices = torch.topk(logits, self.top_k, dim=-1)
        weights = F.softmax(weights, dim=-1)  # Top-k 내 정규화

        # 전문가 출력 가중 합산
        output = torch.zeros_like(x_flat)
        for i in range(self.top_k):
            expert_idx = indices[:, i]  # (B*T,)
            gate_weight = weights[:, i]  # (B*T,)

            for j in range(self.num_experts):
                mask = (expert_idx == j)
                if mask.any():
                    expert_out = self.experts[j](x_flat[mask])
                    output[mask] += gate_weight[mask].unsqueeze(-1) * expert_out

        return output.view(B, T, D)


class SwiGLUExpert(nn.Module):
    """Mixtral에서 사용하는 SwiGLU 기반 FFN 전문가"""
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)
        self.w3 = nn.Linear(d_model, d_ff, bias=False)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

class DeepSeekMoELayer(nn.Module):
    """DeepSeek-MoE 스타일 레이어: 세분화된 전문가 + 공유 전문가"""
    def __init__(
        self,
        d_model=2048,
        d_ff=10944,
        num_routed_experts=64,
        num_shared_experts=2,
        top_k=6,
    ):
        super().__init__()
        self.top_k = top_k

        # 세분화된 라우팅 전문가 (작은 FFN)
        expert_d_ff = d_ff // 4  # 기존 대비 1/4 크기
        self.routed_experts = nn.ModuleList([
            SwiGLUExpert(d_model, expert_d_ff)
            for _ in range(num_routed_experts)
        ])

        # 공유 전문가 (항상 활성화)
        self.shared_experts = nn.ModuleList([
            SwiGLUExpert(d_model, d_ff)
            for _ in range(num_shared_experts)
        ])

        # 라우터
        self.gate = nn.Linear(d_model, num_routed_experts, bias=False)

    def forward(self, x):
        B, T, D = x.shape
        x_flat = x.view(-1, D)

        # 공유 전문가 출력 (항상 활성)
        shared_out = sum(expert(x_flat) for expert in self.shared_experts)

        # 라우팅 전문가 출력
        logits = self.gate(x_flat)
        weights, indices = torch.topk(logits, self.top_k, dim=-1)
        weights = F.softmax(weights, dim=-1)

        routed_out = torch.zeros_like(x_flat)
        for i in range(self.top_k):
            expert_idx = indices[:, i]
            gate_weight = weights[:, i]
            for j in range(len(self.routed_experts)):
                mask = (expert_idx == j)
                if mask.any():
                    out = self.routed_experts[j](x_flat[mask])
                    routed_out[mask] += gate_weight[mask].unsqueeze(-1) * out

        return (shared_out + routed_out).view(B, T, D)

class ExpertChoiceRouter(nn.Module):
    """Expert Choice 라우팅: 전문가가 토큰을 선택"""
    def __init__(self, d_model, num_experts, capacity_factor=1.0):
        super().__init__()
        self.num_experts = num_experts
        self.capacity_factor = capacity_factor
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x):
        B, T, D = x.shape
        x_flat = x.view(-1, D)  # (N, D) where N = B*T
        N = x_flat.shape[0]

        logits = self.gate(x_flat)  # (N, num_experts)
        scores = F.softmax(logits, dim=0)  # 토큰 차원으로 softmax

        # 각 전문가가 처리할 토큰 수
        k = int(self.capacity_factor * N / self.num_experts)

        # 전문가별 Top-k 토큰 선택
        top_k_scores, top_k_indices = torch.topk(
            scores.t(), k, dim=-1
        )  # (num_experts, k)

        return top_k_scores, top_k_indices

# DeepSpeed MoE 설정 예시
deepspeed_config = {
    "train_batch_size": 256,
    "fp16": {"enabled": True},
    "zero_optimization": {"stage": 2},
    "moe": {
        "enabled": True,
        "ep_size": 8,            # Expert Parallelism: 8개 GPU에 전문가 분산
        "num_experts": 64,
        "top_k": 2,
        "capacity_factor": 1.25,
        "min_capacity": 4,
        "use_residual": True,    # 잔차 MoE
        "moe_param_group": True  # 전문가 파라미터 별도 그룹
    }
}