Mixture of Experts(MoE) 아키텍처 논문 심층 분석: GShard에서 DeepSeek-MoE까지

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

class BasicMoEGating(nn.Module):
    """Sparsely-Gated MoE의 기본 게이팅 네트워크"""
    def __init__(self, d_model: int, num_experts: int, top_k: int = 2,
                 noise_std: float = 1.0):
        super().__init__()
        self.top_k = top_k
        self.noise_std = noise_std
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        self.noise_gate = nn.Linear(d_model, num_experts, bias=False)

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

        # 학습 시 노이즈 추가 (탐색 촉진)
        if self.training:
            noise = F.softplus(self.noise_gate(x))
            logits = logits + torch.randn_like(logits) * noise * self.noise_std

        # 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 레이어 (Capacity Factor 포함)"""
    def __init__(self, d_model: int, d_ff: int, num_experts: int,
                 top_k: int = 2, capacity_factor: float = 1.25):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.capacity_factor = capacity_factor
        self.gate = BasicMoEGating(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: torch.Tensor):
        B, T, D = x.shape
        gates, indices = self.gate(x)
        expert_capacity = int(self.capacity_factor * T / self.num_experts)

        output = torch.zeros_like(x)
        tokens_dropped = 0

        for i in range(self.num_experts):
            # i번째 전문가에 할당된 토큰 마스크
            expert_mask = (indices == i).any(dim=-1)  # (B, T)
            # 용량 초과 토큰 드롭
            token_count = expert_mask.sum(dim=-1)
            for b in range(B):
                if token_count[b] > expert_capacity:
                    overflow = token_count[b] - expert_capacity
                    tokens_dropped += overflow.item()

            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, tokens_dropped

class GShardRouter(nn.Module):
    """GShard Top-2 라우터 (랜덤 라우팅 포함)"""
    def __init__(self, d_model: int, num_experts: int,
                 capacity_factor: float = 2.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: torch.Tensor):
        # x: (batch_size, seq_len, d_model)
        B, T, D = x.shape
        logits = self.gate(x)
        probs = F.softmax(logits, dim=-1)

        # Top-1: 결정론적 선택
        top1_probs, top1_indices = probs.max(dim=-1)

        # Top-2: 확률적 선택 (나머지 전문가에서 샘플링)
        mask = torch.zeros_like(probs)
        mask.scatter_(-1, top1_indices.unsqueeze(-1), 1.0)
        remaining_probs = probs * (1 - mask)
        remaining_probs = remaining_probs / (remaining_probs.sum(dim=-1, keepdim=True) + 1e-8)

        # 확률에 비례하여 두 번째 전문가 샘플링
        top2_indices = torch.multinomial(
            remaining_probs.view(-1, self.num_experts), 1
        ).view(B, T)
        top2_probs = probs.gather(-1, top2_indices.unsqueeze(-1)).squeeze(-1)

        # 게이트 값 재정규화
        total = top1_probs + top2_probs
        top1_gates = top1_probs / total
        top2_gates = top2_probs / total

        return (top1_gates, top1_indices, top2_gates, top2_indices)

class SwitchRouter(nn.Module):
    """Switch Transformer 라우터 (Top-1 단일 전문가)"""
    def __init__(self, d_model: int, num_experts: int,
                 capacity_factor: float = 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: torch.Tensor):
        B, T, D = x.shape
        logits = self.gate(x)
        probs = F.softmax(logits, dim=-1)

        # Top-1 선택 (단일 전문가)
        gate_values, expert_indices = probs.max(dim=-1)  # (B, T)

        # Capacity 기반 토큰 드롭
        capacity = int(self.capacity_factor * T / self.num_experts)

        # 각 전문가별 할당 토큰 수 추적
        expert_counts = torch.zeros(B, self.num_experts, device=x.device)
        drop_mask = torch.zeros(B, T, dtype=torch.bool, device=x.device)

        for t in range(T):
            for b in range(B):
                eid = expert_indices[b, t].item()
                if expert_counts[b, eid] >= capacity:
                    drop_mask[b, t] = True  # 용량 초과 토큰 드롭
                else:
                    expert_counts[b, eid] += 1

        # 드롭된 토큰의 게이트 값을 0으로
        gate_values = gate_values * (~drop_mask).float()

        return gate_values, expert_indices, drop_mask

def switch_load_balancing_loss(router_probs: torch.Tensor,
                                expert_indices: torch.Tensor,
                                num_experts: int,
                                alpha: float = 0.01) -> torch.Tensor:
    """Switch Transformer의 Simplified Load Balancing Loss

    Args:
        router_probs: 라우터 확률 (B, T, num_experts)
        expert_indices: 선택된 전문가 인덱스 (B, T)
        num_experts: 전문가 수
        alpha: 밸런싱 계수
    """
    # f_i: 각 전문가에 배정된 토큰 비율
    one_hot = F.one_hot(expert_indices, num_experts).float()
    f = one_hot.mean(dim=[0, 1])  # (num_experts,)

    # P_i: 각 전문가에 대한 평균 라우터 확률
    P = router_probs.mean(dim=[0, 1])  # (num_experts,)

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

class DeepSeekMoELayer(nn.Module):
    """DeepSeek-MoE 레이어 (Fine-grained Segmentation + Shared Expert)"""
    def __init__(self, d_model: int, d_ff: int, num_routed_experts: int,
                 num_shared_experts: int, top_k: int, segmentation_factor: int = 4):
        super().__init__()
        self.num_routed = num_routed_experts
        self.num_shared = num_shared_experts
        self.top_k = top_k

        # Fine-grained: FFN 중간 차원을 segmentation_factor로 나눔
        fine_d_ff = d_ff // segmentation_factor

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

        # 라우팅 전문가 (세분화된 크기)
        self.routed_experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, fine_d_ff),
                nn.SiLU(),
                nn.Linear(fine_d_ff, d_model)
            ) for _ in range(num_routed_experts)
        ])

        # 게이팅 네트워크 (라우팅 전문가용)
        self.gate = nn.Linear(d_model, num_routed_experts, bias=False)

    def forward(self, x: torch.Tensor):
        B, T, D = x.shape

        # 1. 공유 전문가 출력 (모든 토큰에 적용)
        shared_output = sum(expert(x) for expert in self.shared_experts)

        # 2. 라우팅 전문가 선택 및 출력
        logits = self.gate(x)
        probs = F.softmax(logits, dim=-1)
        top_k_probs, top_k_indices = torch.topk(probs, self.top_k, dim=-1)
        top_k_gates = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        routed_output = torch.zeros_like(x)
        for k in range(self.top_k):
            expert_idx = top_k_indices[:, :, k]  # (B, T)
            gate_val = top_k_gates[:, :, k]       # (B, T)
            for i in range(self.num_routed):
                mask = (expert_idx == i)
                if mask.any():
                    inp = x[mask]
                    out = self.routed_experts[i](inp)
                    routed_output[mask] += gate_val[mask].unsqueeze(-1) * out

        return shared_output + routed_output

def compute_expert_utilization(expert_indices: torch.Tensor,
                                num_experts: int) -> dict:
    """전문가 활용도 모니터링 유틸리티

    Args:
        expert_indices: 선택된 전문가 인덱스 (B, T) 또는 (B, T, K)
        num_experts: 총 전문가 수

    Returns:
        활용도 메트릭 딕셔너리
    """
    flat = expert_indices.flatten()
    counts = torch.bincount(flat, minlength=num_experts).float()
    total = flat.numel()

    # 각 전문가의 토큰 비율
    fractions = counts / total

    # 활용도 메트릭
    max_fraction = fractions.max().item()
    min_fraction = fractions.min().item()
    ideal_fraction = 1.0 / num_experts

    # 불균형 비율 (1에 가까울수록 균등)
    balance_ratio = min_fraction / (max_fraction + 1e-8)

    # 비활성 전문가 수
    dead_experts = (counts == 0).sum().item()

    return {
        "balance_ratio": balance_ratio,
        "max_load": max_fraction,
        "min_load": min_fraction,
        "ideal_load": ideal_fraction,
        "dead_experts": dead_experts,
        "load_std": fractions.std().item(),
    }

def router_z_loss(router_logits: torch.Tensor) -> torch.Tensor:
    """Router z-loss: 라우터 로짓의 크기를 억제하여 안정성 향상

    Args:
        router_logits: 라우터의 원시 로짓 (B, T, num_experts)
    """
    # log-sum-exp의 제곱
    log_z = torch.logsumexp(router_logits, dim=-1)  # (B, T)
    z_loss = (log_z ** 2).mean()
    return z_loss


class StableMoETrainer:
    """MoE 학습 안정화를 위한 트레이너 예시"""
    def __init__(self, model, optimizer, alpha_balance=0.01,
                 alpha_z=0.001):
        self.model = model
        self.optimizer = optimizer
        self.alpha_balance = alpha_balance
        self.alpha_z = alpha_z

    def train_step(self, input_ids, labels):
        self.optimizer.zero_grad()

        # Forward pass
        logits, router_logits, expert_indices = self.model(input_ids)

        # Task loss
        task_loss = F.cross_entropy(
            logits.view(-1, logits.size(-1)), labels.view(-1)
        )

        # Load balancing loss
        router_probs = F.softmax(router_logits, dim=-1)
        balance_loss = switch_load_balancing_loss(
            router_probs, expert_indices, self.model.num_experts
        )

        # Router z-loss
        z_loss = router_z_loss(router_logits)

        # Total loss
        total_loss = task_loss + self.alpha_balance * balance_loss + self.alpha_z * z_loss

        total_loss.backward()

        # 그래디언트 클리핑 (MoE 학습에서 특히 중요)
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)

        self.optimizer.step()

        return {
            "task_loss": task_loss.item(),
            "balance_loss": balance_loss.item(),
            "z_loss": z_loss.item(),
            "total_loss": total_loss.item(),
        }