MoE (Mixture of Experts) 아키텍처 완전 정복: Mixtral부터 DeepSeek-MoE까지 구현 원리와 실전 최적화 기법

import torch
import torch.nn as nn

class MoELayer(nn.Module):
    def __init__(self, hidden_size, num_experts=8, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k

        # 게이팅 네트워크
        self.gate = nn.Linear(hidden_size, num_experts)

        # 전문가 네트워크들
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(hidden_size, hidden_size * 4),
                nn.ReLU(),
                nn.Linear(hidden_size * 4, hidden_size)
            ) for _ in range(num_experts)
        ])

    def forward(self, x):
        # 게이팅 점수 계산
        gate_scores = self.gate(x)  # [batch, seq_len, num_experts]

        # Top-K 전문가 선택
        top_k_scores, top_k_indices = torch.topk(gate_scores, self.top_k, dim=-1)
        top_k_weights = torch.softmax(top_k_scores, dim=-1)

        # 선택된 전문가들의 출력 결합
        output = torch.zeros_like(x)
        for i in range(self.top_k):
            expert_idx = top_k_indices[:, :, i]
            expert_weight = top_k_weights[:, :, i:i+1]

            # 각 전문가 실행
            for expert_id in range(self.num_experts):
                mask = (expert_idx == expert_id)
                if mask.any():
                    expert_out = self.experts[expert_id](x[mask])
                    output[mask] += expert_weight[mask] * expert_out

        return output

class DeepSeekMoE(nn.Module):
    def __init__(self, hidden_size, num_experts=64, experts_per_token=6, shared_expert_size=2):
        super().__init__()

        # 공유 전문가 (항상 활성화)
        self.shared_experts = nn.ModuleList([
            FFN(hidden_size) for _ in range(shared_expert_size)
        ])

        # 라우팅 전문가 (선택적 활성화)
        self.routed_experts = nn.ModuleList([
            # 더 작은 크기의 전문가들
            FFN(hidden_size, intermediate_size=hidden_size * 2)
            for _ in range(num_experts)
        ])

        self.gate = nn.Linear(hidden_size, num_experts)
        self.experts_per_token = experts_per_token

# 각 전문가가 처리할 수 있는 최대 토큰 수 제한
capacity_factor = 1.25
expert_capacity = (batch_size * seq_len * top_k / num_experts) * capacity_factor

def load_balancing_loss(gate_logits, top_k_indices, num_experts):
    # 각 전문가의 사용 빈도 계산
    expert_counts = torch.zeros(num_experts)
    for idx in top_k_indices.flatten():
        expert_counts[idx] += 1

    # 균등 분포와의 차이를 패널티로 부여
    target_distribution = torch.ones(num_experts) / num_experts
    actual_distribution = expert_counts / expert_counts.sum()

    return torch.sum((actual_distribution - target_distribution) ** 2)