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