Mixture of Experts(MoE) 아키텍처 논문 리뷰와 실전 스케일링 2026

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

class Expert(nn.Module):
    """단일 전문가 FFN 모듈"""
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w2(self.dropout(F.gelu(self.w1(x))))

class MoELayer(nn.Module):
    """Mixture of Experts 레이어"""
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        num_experts: int,
        top_k: int = 2,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.experts = nn.ModuleList([
            Expert(d_model, d_ff, dropout) for _ in range(num_experts)
        ])
        # 게이팅 네트워크: 입력 -> 전문가 점수
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (batch, seq_len, d_model)
        batch_size, seq_len, d_model = x.shape
        x_flat = x.view(-1, d_model)  # (B*S, d_model)

        # 게이팅 점수 계산
        gate_logits = self.gate(x_flat)  # (B*S, num_experts)
        gate_probs = F.softmax(gate_logits, dim=-1)

        # Top-K 전문가 선택
        top_k_probs, top_k_indices = torch.topk(
            gate_probs, self.top_k, dim=-1
        )  # (B*S, top_k)
        # 선택된 전문가 가중치 재정규화
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        # 전문가별 출력 합산
        output = torch.zeros_like(x_flat)
        for k in range(self.top_k):
            expert_idx = top_k_indices[:, k]  # (B*S,)
            weight = top_k_probs[:, k].unsqueeze(-1)  # (B*S, 1)
            for i in range(self.num_experts):
                mask = (expert_idx == i)
                if mask.any():
                    expert_input = x_flat[mask]
                    expert_output = self.experts[i](expert_input)
                    output[mask] += weight[mask] * expert_output

        return output.view(batch_size, seq_len, d_model)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class TopKRouter(nn.Module):
    """Top-K 라우팅 with 노이즈 주입"""
    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.num_experts = num_experts
        self.noise_std = noise_std
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(
        self, x: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Args:
            x: (batch * seq_len, d_model)
        Returns:
            top_k_indices: (batch * seq_len, top_k)
            top_k_weights: (batch * seq_len, top_k)
            gate_logits: (batch * seq_len, num_experts)
        """
        gate_logits = self.gate(x)

        # 학습 시 노이즈 주입으로 탐색 촉진
        if self.training:
            noise = torch.randn_like(gate_logits) * self.noise_std
            gate_logits = gate_logits + noise

        gate_probs = F.softmax(gate_logits, dim=-1)
        top_k_probs, top_k_indices = torch.topk(
            gate_probs, self.top_k, dim=-1
        )
        # 재정규화
        top_k_weights = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        return top_k_indices, top_k_weights, gate_logits

def load_balancing_loss(
    gate_logits: torch.Tensor,
    top_k_indices: torch.Tensor,
    num_experts: int,
    top_k: int = 2,
) -> torch.Tensor:
    """
    Switch Transformer 스타일 로드 밸런싱 손실 계산

    각 전문가에 할당된 토큰의 비율(f_i)과
    게이트 확률의 평균(p_i)의 내적을 최소화한다.
    균등 분포일 때 최솟값을 가진다.

    Args:
        gate_logits: (num_tokens, num_experts)
        top_k_indices: (num_tokens, top_k)
        num_experts: 전문가 수
        top_k: 활성 전문가 수
    Returns:
        스칼라 로드 밸런싱 손실
    """
    num_tokens = gate_logits.shape[0]
    gate_probs = F.softmax(gate_logits, dim=-1)

    # f_i: 각 전문가에 할당된 토큰 비율
    # one-hot 인코딩으로 할당 카운트
    expert_mask = F.one_hot(
        top_k_indices, num_classes=num_experts
    ).float()  # (num_tokens, top_k, num_experts)
    expert_mask = expert_mask.sum(dim=1)  # (num_tokens, num_experts)
    tokens_per_expert = expert_mask.sum(dim=0)  # (num_experts,)
    f = tokens_per_expert / (num_tokens * top_k)

    # p_i: 각 전문가에 대한 평균 게이트 확률
    p = gate_probs.mean(dim=0)  # (num_experts,)

    # 밸런싱 손실: num_experts * sum(f_i * p_i)
    balance_loss = num_experts * (f * p).sum()

    return balance_loss

def router_z_loss(gate_logits: torch.Tensor) -> torch.Tensor:
    """
    ST-MoE Router Z-Loss 구현

    라우터 로짓의 log-sum-exp 값의 제곱 평균을 최소화.
    로짓이 과도하게 커지는 것을 방지하여 학습 안정성 확보.

    Args:
        gate_logits: (num_tokens, num_experts)
    Returns:
        스칼라 z-loss 값
    """
    # log(sum(exp(x)))의 제곱
    log_z = torch.logsumexp(gate_logits, dim=-1)  # (num_tokens,)
    z_loss = (log_z ** 2).mean()
    return z_loss


# 학습 루프에서의 통합
def compute_moe_loss(
    task_loss: torch.Tensor,
    gate_logits: torch.Tensor,
    top_k_indices: torch.Tensor,
    num_experts: int,
    alpha_balance: float = 0.01,
    alpha_z: float = 0.001,
) -> torch.Tensor:
    """
    전체 MoE 손실 = 태스크 손실 + 밸런싱 손실 + Z-Loss
    """
    bal_loss = load_balancing_loss(
        gate_logits, top_k_indices, num_experts
    )
    z_loss = router_z_loss(gate_logits)

    total_loss = task_loss + alpha_balance * bal_loss + alpha_z * z_loss
    return total_loss

class AuxLossFreeRouter(nn.Module):
    """DeepSeek-V3 스타일 Auxiliary-Loss-Free 라우터"""
    def __init__(
        self,
        d_model: int,
        num_experts: int,
        top_k: int = 8,
        bias_update_speed: float = 0.001,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.gamma = bias_update_speed
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        # 바이어스는 그래디언트 업데이트가 아닌 규칙 기반 조정
        self.register_buffer(
            "expert_bias", torch.zeros(num_experts)
        )
        self.register_buffer(
            "target_load", torch.ones(num_experts) / num_experts
        )

    def forward(self, x: torch.Tensor):
        gate_logits = self.gate(x)
        # 바이어스를 더해 라우팅 결정에 반영
        biased_logits = gate_logits + self.expert_bias.unsqueeze(0)
        gate_probs = F.softmax(gate_logits, dim=-1)

        # Top-K 선택은 바이어스 포함 로짓으로 수행
        _, top_k_indices = torch.topk(
            biased_logits, self.top_k, dim=-1
        )
        # 가중치는 원본 게이트 확률에서 추출
        top_k_probs = torch.gather(gate_probs, 1, top_k_indices)
        top_k_weights = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        # 학습 시 바이어스 업데이트 (그래디언트 불필요)
        if self.training:
            with torch.no_grad():
                expert_load = F.one_hot(
                    top_k_indices.view(-1), self.num_experts
                ).float().sum(0)
                expert_load = expert_load / expert_load.sum()
                # 과부하 전문가는 바이어스 감소, 저부하 전문가는 증가
                self.expert_bias += self.gamma * (
                    self.target_load - expert_load
                )

        return top_k_indices, top_k_weights, gate_logits

from dataclasses import dataclass
from typing import Dict, Optional
import torch

@dataclass
class ExpertCacheConfig:
    """전문가 캐시 설정"""
    gpu_cache_size: int = 4   # GPU에 상주할 전문가 수
    prefetch_count: int = 2    # 선제 로딩할 전문가 수
    device: str = "cuda:0"

class ExpertOffloadManager:
    """
    전문가 오프로딩 관리자.
    LRU 기반으로 GPU 캐시를 관리하고,
    다음 레이어 전문가를 선제 로딩한다.
    """
    def __init__(
        self,
        experts: Dict[int, torch.nn.Module],
        config: ExpertCacheConfig,
    ):
        self.experts = experts
        self.config = config
        self.gpu_cache: Dict[int, torch.nn.Module] = {}
        self.access_order: list = []

    def get_expert(self, expert_id: int) -> torch.nn.Module:
        """전문가를 GPU로 로딩 (LRU 캐시)"""
        if expert_id in self.gpu_cache:
            self.access_order.remove(expert_id)
            self.access_order.append(expert_id)
            return self.gpu_cache[expert_id]

        # GPU 캐시가 가득 찬 경우 LRU 전문가 제거
        if len(self.gpu_cache) >= self.config.gpu_cache_size:
            evict_id = self.access_order.pop(0)
            self.gpu_cache[evict_id].cpu()
            del self.gpu_cache[evict_id]

        # 전문가를 GPU로 이동
        expert = self.experts[expert_id].to(self.config.device)
        self.gpu_cache[expert_id] = expert
        self.access_order.append(expert_id)
        return expert

    def prefetch(self, predicted_expert_ids: list):
        """다음 레이어에서 사용될 전문가를 비동기 선제 로딩"""
        for eid in predicted_expert_ids[:self.config.prefetch_count]:
            if eid not in self.gpu_cache:
                self.get_expert(eid)