Sparse Mixture of Experts(MoE) 아키텍처 심층 분석: 설계 원리부터 DeepSeek-V3·Qwen3까지

y = sum_{i=1}^{N} G(x)_i * E_i(x)

G(x)_i = softmax(W_g * x + noise)_i  (if i in TopK)
G(x)_i = 0                           (otherwise)

import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass

@dataclass
class MoEConfig:
    d_model: int = 1024
    d_ff: int = 4096
    num_experts: int = 8
    top_k: int = 2
    dropout: float = 0.1
    aux_loss_weight: float = 0.01

class SwiGLUExpert(nn.Module):
    """SwiGLU 활성화를 사용하는 전문가 FFN.
    LLaMA, Mistral 등 최신 모델에서 표준으로 채택된 구조."""

    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w_up = nn.Linear(d_model, d_ff, bias=False)
        self.w_down = nn.Linear(d_ff, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = F.silu(self.w_gate(x))
        up = self.w_up(x)
        return self.w_down(self.dropout(gate * up))

class TopKGating(nn.Module):
    """Top-K 게이팅 네트워크.
    Noisy Top-K Gating (Shazeer et al., 2017) 구현."""

    def __init__(self, d_model: int, num_experts: int, top_k: int = 2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        self.noise_linear = nn.Linear(d_model, num_experts, bias=False)

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

        if self.training:
            noise = F.softplus(self.noise_linear(x))
            logits = logits + noise * torch.randn_like(logits)

        probs = F.softmax(logits, dim=-1)
        top_k_probs, top_k_indices = torch.topk(probs, self.top_k, dim=-1)
        # 재정규화: 선택된 전문가 가중치 합이 1이 되도록
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        return top_k_probs, top_k_indices, probs

class ExpertChoiceGating(nn.Module):
    """Expert Choice Routing 구현.
    각 전문가가 처리할 토큰을 직접 선택하여
    로드 밸런싱을 구조적으로 보장한다."""

    def __init__(
        self,
        d_model: int,
        num_experts: int,
        capacity_factor: float = 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: torch.Tensor):
        # x: (num_tokens, d_model)
        num_tokens = x.shape[0]
        expert_capacity = int(
            num_tokens * self.capacity_factor / self.num_experts
        )

        # 게이팅 점수: (num_tokens, num_experts)
        gate_logits = self.gate(x)
        # 전문가 관점에서 점수 계산: (num_experts, num_tokens)
        gate_scores = F.softmax(gate_logits.T, dim=-1)

        # 각 전문가가 상위 capacity개 토큰 선택
        top_k_scores, top_k_indices = torch.topk(
            gate_scores, expert_capacity, dim=-1
        )  # (num_experts, capacity)

        # 디스패치 마스크 생성
        dispatch_mask = torch.zeros(
            self.num_experts, num_tokens,
            device=x.device, dtype=x.dtype,
        )
        dispatch_mask.scatter_(1, top_k_indices, top_k_scores)

        return dispatch_mask, top_k_indices, top_k_scores

def compute_load_balancing_loss(
    gate_probs: torch.Tensor,
    top_k_indices: torch.Tensor,
    num_experts: int,
    top_k: int,
) -> torch.Tensor:
    """Switch Transformer 스타일의 로드 밸런싱 보조 손실 계산.

    Args:
        gate_probs: 게이팅 확률 (num_tokens, num_experts)
        top_k_indices: 선택된 전문가 인덱스 (num_tokens, top_k)
        num_experts: 전문가 수
        top_k: 선택되는 전문가 수

    Returns:
        보조 손실 스칼라 값
    """
    num_tokens = gate_probs.shape[0]

    # f_i: 전문가 i에 배정된 토큰의 비율
    expert_mask = F.one_hot(top_k_indices, num_experts).float()
    # (num_tokens, top_k, num_experts) -> (num_tokens, num_experts)
    expert_mask = expert_mask.sum(dim=1)
    tokens_per_expert = expert_mask.sum(dim=0)  # (num_experts,)
    f = tokens_per_expert / (num_tokens * top_k)

    # P_i: 전문가 i에 대한 평균 게이팅 확률
    P = gate_probs.mean(dim=0)  # (num_experts,)

    # 보조 손실: N * sum(f_i * P_i)
    # 균등 배분일 때 최소값을 가짐
    aux_loss = num_experts * (f * P).sum()

    return aux_loss

class MoELayerWithAuxLoss(nn.Module):
    """보조 손실이 포함된 완전한 MoE 레이어 구현."""

    def __init__(self, config: MoEConfig):
        super().__init__()
        self.config = config
        self.experts = nn.ModuleList([
            SwiGLUExpert(config.d_model, config.d_ff, config.dropout)
            for _ in range(config.num_experts)
        ])
        self.gating = TopKGating(
            config.d_model, config.num_experts, config.top_k
        )
        self.aux_loss_weight = config.aux_loss_weight

    def forward(self, x: torch.Tensor):
        batch_size, seq_len, d_model = x.shape
        x_flat = x.view(-1, d_model)

        top_k_probs, top_k_indices, gate_probs = self.gating(x_flat)

        # 보조 손실 계산
        aux_loss = self.aux_loss_weight * compute_load_balancing_loss(
            gate_probs, top_k_indices,
            self.config.num_experts, self.config.top_k,
        )

        # 전문가 출력 계산
        output = torch.zeros_like(x_flat)
        for k in range(self.config.top_k):
            expert_indices = top_k_indices[:, k]  # (num_tokens,)
            expert_weights = top_k_probs[:, k]    # (num_tokens,)

            for i in range(self.config.num_experts):
                mask = (expert_indices == i)
                if mask.any():
                    expert_input = x_flat[mask]
                    expert_output = self.experts[i](expert_input)
                    output[mask] += expert_weights[mask].unsqueeze(-1) * expert_output

        output = output.view(batch_size, seq_len, d_model)
        return output, aux_loss

import logging
from collections import defaultdict

logger = logging.getLogger(__name__)

class ExpertUtilizationMonitor:
    """전문가 활용 모니터링 및 붕괴 탐지 도구.

    학습 중 각 전문가의 활용률을 추적하고,
    붕괴 징후를 조기에 감지한다.
    """

    def __init__(
        self,
        num_experts: int,
        collapse_threshold: float = 0.01,
        window_size: int = 100,
    ):
        self.num_experts = num_experts
        self.collapse_threshold = collapse_threshold
        self.window_size = window_size
        self.history: list[dict[int, float]] = []

    def record(self, expert_counts: dict[int, int], total_tokens: int):
        """배치별 전문가 활용 기록."""
        utilization = {
            i: expert_counts.get(i, 0) / max(total_tokens, 1)
            for i in range(self.num_experts)
        }
        self.history.append(utilization)

        if len(self.history) > self.window_size:
            self.history = self.history[-self.window_size:]

    def detect_collapse(self) -> list[int]:
        """전문가 붕괴 탐지. 활용률이 threshold 이하인 전문가 반환."""
        if len(self.history) < self.window_size // 2:
            return []

        collapsed = []
        for expert_id in range(self.num_experts):
            recent_util = [
                h[expert_id] for h in self.history[-self.window_size:]
            ]
            avg_util = sum(recent_util) / len(recent_util)
            if avg_util < self.collapse_threshold:
                collapsed.append(expert_id)

        if collapsed:
            logger.warning(
                f"Expert collapse detected! "
                f"Experts {collapsed} have utilization below "
                f"{self.collapse_threshold:.2%}. "
                f"Consider increasing aux_loss_weight or "
                f"reinitializing collapsed experts."
            )

        return collapsed

    def get_load_imbalance_ratio(self) -> float:
        """로드 불균형 비율 계산.
        1.0이면 완벽한 균형, 값이 클수록 불균형."""
        if not self.history:
            return 0.0

        latest = self.history[-1]
        utils = list(latest.values())
        max_util = max(utils) if utils else 0
        min_util = min(utils) if utils else 0
        avg_util = sum(utils) / len(utils) if utils else 0

        if avg_util == 0:
            return float("inf")

        return max_util / avg_util

import torch
import torch.distributed as dist
from typing import Optional

class ExpertParallelRouter:
    """Expert Parallelism을 위한 토큰 디스패치/수집 구현.

    각 GPU가 전문가 일부를 담당하고,
    All-to-All 통신으로 토큰을 라우팅한다.
    """

    def __init__(
        self,
        num_experts: int,
        ep_group: Optional[dist.ProcessGroup] = None,
    ):
        self.num_experts = num_experts
        self.ep_group = ep_group
        self.ep_size = dist.get_world_size(ep_group) if ep_group else 1
        self.ep_rank = dist.get_rank(ep_group) if ep_group else 0
        self.experts_per_rank = num_experts // self.ep_size

    def dispatch(
        self,
        tokens: torch.Tensor,
        expert_indices: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """토큰을 담당 GPU로 디스패치.

        Args:
            tokens: (num_tokens, d_model) 입력 토큰
            expert_indices: (num_tokens,) 각 토큰의 대상 전문가 인덱스

        Returns:
            dispatched_tokens: 이 GPU가 처리할 토큰
            recv_counts: 각 GPU에서 받은 토큰 수
        """
        # 각 GPU로 보낼 토큰 수 계산
        send_counts = torch.zeros(
            self.ep_size, dtype=torch.long, device=tokens.device
        )
        for rank in range(self.ep_size):
            start_expert = rank * self.experts_per_rank
            end_expert = start_expert + self.experts_per_rank
            mask = (expert_indices >= start_expert) & (
                expert_indices < end_expert
            )
            send_counts[rank] = mask.sum()

        # All-to-All로 수신 카운트 교환
        recv_counts = torch.zeros_like(send_counts)
        dist.all_to_all_single(
            recv_counts, send_counts, group=self.ep_group
        )

        # 토큰 정렬 및 All-to-All 전송
        sorted_indices = torch.argsort(expert_indices)
        sorted_tokens = tokens[sorted_indices]

        send_splits = send_counts.tolist()
        recv_splits = recv_counts.tolist()

        dispatched_tokens = torch.zeros(
            int(recv_counts.sum()), tokens.shape[1],
            dtype=tokens.dtype, device=tokens.device,
        )
        dist.all_to_all_single(
            dispatched_tokens, sorted_tokens,
            output_split_sizes=recv_splits,
            input_split_sizes=send_splits,
            group=self.ep_group,
        )

        return dispatched_tokens, recv_counts

def router_z_loss(gate_logits: torch.Tensor) -> torch.Tensor:
    """ST-MoE 스타일의 Router z-loss.
    게이팅 로짓의 크기를 제약하여 수치 안정성을 높인다.

    Args:
        gate_logits: (num_tokens, num_experts) 게이팅 로짓

    Returns:
        z_loss 스칼라
    """
    log_z = torch.logsumexp(gate_logits, dim=-1)  # (num_tokens,)
    z_loss = (log_z ** 2).mean()
    return z_loss