15. Mixture of Experts (MoE) — Specialized Sub-Model Routing for MangaAssist

graph TB
    subgraph "MoE Layer Processing"
        INPUT["Input tokens from attention layer<br>x ∈ ℝ^{T×d}, T=512, d=4096"]

        GATE["Gating Network G(x)<br>W_g ∈ ℝ^{8×4096}<br>Computes: softmax(W_g·x + ε)<br>Selects top-2 experts per token"]

        subgraph "8 Expert FFNs (only 2 active per token)"
            E1["Expert 1: Product<br>W_gate ∈ ℝ^{14336×4096}<br>W_up ∈ ℝ^{4096×14336}<br>117M params"]
            E2["Expert 2: Recommend<br>117M params"]
            E3["Expert 3: Order<br>117M params"]
            E4["Expert 4: Complaint<br>117M params"]
            E5["Expert 5: Knowledge<br>117M params"]
            E6["Expert 6: Chitchat<br>117M params"]
            E7["Expert 7: Multi-turn<br>117M params"]
            E8["Expert 8: Fallback<br>117M params"]
        end

        COMBINE["Weighted combination:<br>MoE(x) = G(x)₁·E₁(x) + G(x)₅·E₅(x)<br>(example: product + knowledge)"]

        OUTPUT["Output: same shape ℝ^{T×d}<br>Only 234M params activated<br>(2× expert = 2×117M)"]

        INPUT --> GATE
        GATE -->|"w=0.65"| E1
        GATE -->|"w=0"| E2
        GATE -->|"w=0"| E3
        GATE -->|"w=0"| E4
        GATE -->|"w=0.35"| E5
        GATE -->|"w=0"| E6
        GATE -->|"w=0"| E7
        GATE -->|"w=0"| E8
        E1 --> COMBINE
        E5 --> COMBINE
        COMBINE --> OUTPUT
    end

    style E1 fill:#c8e6c9
    style E5 fill:#c8e6c9
    style E2 fill:#eeeeee
    style E3 fill:#eeeeee
    style E4 fill:#eeeeee
    style E6 fill:#eeeeee
    style E7 fill:#eeeeee
    style E8 fill:#eeeeee

graph TB
    subgraph "Expert Utilization Matrix U (after training)"
        direction TB
        HEADER["Query Types →<br>Product | Recommend | Order | Complaint | Knowledge | Chitchat"]

        E1_ROW["Expert 1: 0.42 | 0.12 | 0.08 | 0.05 | 0.28 | 0.05<br>Specialization: Product + Knowledge lookups"]
        E2_ROW["Expert 2: 0.10 | 0.45 | 0.05 | 0.08 | 0.22 | 0.10<br>Specialization: Recommendations"]
        E3_ROW["Expert 3: 0.05 | 0.05 | 0.48 | 0.15 | 0.07 | 0.20<br>Specialization: Order tracking + general"]
        E4_ROW["Expert 4: 0.08 | 0.06 | 0.12 | 0.52 | 0.10 | 0.12<br>Specialization: Complaint handling"]
        E5_ROW["Expert 5: 0.15 | 0.18 | 0.05 | 0.05 | 0.47 | 0.10<br>Specialization: Manga knowledge"]
        E6_ROW["Expert 6: 0.07 | 0.10 | 0.08 | 0.12 | 0.08 | 0.55<br>Specialization: Chitchat + conversational"]
        E7_ROW["Expert 7: 0.15 | 0.20 | 0.18 | 0.15 | 0.17 | 0.15<br>Specialization: Multi-turn context (generalist)"]
        E8_ROW["Expert 8: 0.12 | 0.08 | 0.15 | 0.08 | 0.10 | 0.47<br>Specialization: Fallback / edge cases"]

        HEADER --> E1_ROW --> E2_ROW --> E3_ROW --> E4_ROW --> E5_ROW --> E6_ROW --> E7_ROW --> E8_ROW
    end

    style E1_ROW fill:#c8e6c9
    style E2_ROW fill:#c8e6c9
    style E3_ROW fill:#c8e6c9
    style E4_ROW fill:#bbdefb
    style E5_ROW fill:#c8e6c9
    style E6_ROW fill:#fff9c4
    style E7_ROW fill:#eeeeee
    style E8_ROW fill:#fff9c4

graph LR
    subgraph "Without Load Balancing Loss"
        WO_INIT["Epoch 0: Uniform<br>12.5% per expert"]
        WO_MID["Epoch 5: Collapse begins<br>Expert 1: 35%<br>Expert 3: 28%<br>Others: 5-8%"]
        WO_END["Epoch 20: Collapsed<br>Expert 1: 62%<br>Expert 3: 31%<br>Others: ~1% (dead)"]

        WO_INIT --> WO_MID --> WO_END
    end

    subgraph "With Load Balancing Loss (α=0.01)"
        W_INIT["Epoch 0: Uniform<br>12.5% per expert"]
        W_MID["Epoch 5: Slight specialization<br>Range: 9-18%"]
        W_END["Epoch 20: Balanced<br>Range: 10-16%<br>All experts active"]

        W_INIT --> W_MID --> W_END
    end

    style WO_END fill:#ffcdd2
    style W_END fill:#c8e6c9

graph TB
    subgraph "Llama 3 8B MoE Architecture (32 layers)"
        INPUT["Input embeddings"]

        L1["Layer 1-8: Dense FFN<br>(shared low-level features)<br>No MoE — pattern extraction"]

        L2["Layer 9-24: MoE FFN<br>8 experts per layer, top-2 routing<br>16 MoE layers × 939M = 15B params<br>Active: 16 × 234M = 3.7B"]

        L3["Layer 25-32: Dense FFN<br>(shared output projection)<br>No MoE — output generation"]

        HEAD["LM Head → vocabulary logits"]

        INPUT --> L1 --> L2 --> L3 --> HEAD
    end

    subgraph "Why This Split?"
        R1["Early layers learn universal<br>features (tokenization, syntax)<br>→ No need for specialization"]
        R2["Middle layers learn<br>task-specific representations<br>→ Maximum benefit from routing"]
        R3["Late layers unify<br>representations for generation<br>→ Shared output space"]
    end

    style L2 fill:#c8e6c9
    style L1 fill:#e3f2fd
    style L3 fill:#e3f2fd

sequenceDiagram
    participant T as Token "One Piece volume 107"
    participant G as Gating Network
    participant LB as Load Balancer
    participant E1 as Expert 1 (Product)
    participant E5 as Expert 5 (Knowledge)
    participant C as Combiner

    T->>G: x = hidden state (d=4096)
    G->>G: z = W_g · x + ε (8 logits)
    G->>G: p = softmax(z) = [0.38, 0.05, 0.02, 0.01, 0.42, 0.03, 0.06, 0.03]
    G->>G: top-2 = {Expert 5: 0.42, Expert 1: 0.38}

    G->>LB: Check capacity for E1, E5
    LB->>LB: E1: 58/64 slots used → OK
    LB->>LB: E5: 61/64 slots used → OK

    G->>E1: Route with weight 0.38/(0.38+0.42) = 0.475
    G->>E5: Route with weight 0.42/(0.38+0.42) = 0.525

    E1->>E1: FFN: SiLU(W_gate · x) ⊙ (W_up · x)
    E5->>E5: FFN: SiLU(W_gate · x) ⊙ (W_up · x)

    E1->>C: Product expert output
    E5->>C: Knowledge expert output
    C->>C: 0.475 × E1(x) + 0.525 × E5(x)
    C->>T: Combined MoE output

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


class ExpertFFN(nn.Module):
    """Single expert feed-forward network (SwiGLU architecture)."""

    def __init__(self, d_model: int = 4096, d_ff: int = 14336):
        super().__init__()
        self.gate_proj = nn.Linear(d_model, d_ff, bias=False)
        self.up_proj = nn.Linear(d_model, d_ff, bias=False)
        self.down_proj = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class TopKGating(nn.Module):
    """Top-K gating network with load balancing loss."""

    def __init__(
        self, d_model: int, num_experts: int, top_k: int = 2,
        noise_std: float = 0.1,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        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:
            gates: (batch, seq_len, num_experts) — routing weights
            indices: (batch, seq_len, top_k) — selected expert indices
            balance_loss: scalar load balancing loss
        """
        # Compute logits
        logits = self.gate(x)  # (B, T, E)

        # Add noise during training for exploration
        if self.training:
            noise = torch.randn_like(logits) * self.noise_std
            logits = logits + noise

        # Top-K selection
        top_k_logits, indices = torch.topk(logits, self.top_k, dim=-1)
        gates = F.softmax(top_k_logits, dim=-1)  # (B, T, K)

        # Load balancing loss
        # f_i: fraction of tokens routed to expert i
        # p_i: average routing probability for expert i
        routing_probs = F.softmax(logits, dim=-1)  # (B, T, E)
        p = routing_probs.mean(dim=[0, 1])  # (E,)

        # One-hot for selected experts
        mask = torch.zeros_like(routing_probs)
        mask.scatter_(-1, indices, 1.0)
        f = mask.mean(dim=[0, 1])  # (E,)

        balance_loss = self.num_experts * (f * p).sum()

        # Router z-loss
        z_loss = (torch.logsumexp(logits, dim=-1) ** 2).mean()

        return gates, indices, balance_loss + 0.001 * z_loss


class MoELayer(nn.Module):
    """Mixture of Experts layer replacing standard FFN."""

    def __init__(
        self,
        d_model: int = 4096,
        d_ff: int = 14336,
        num_experts: int = 8,
        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.gating = TopKGating(d_model, num_experts, top_k)
        self.experts = nn.ModuleList([
            ExpertFFN(d_model, d_ff) for _ in range(num_experts)
        ])

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            x: (batch, seq_len, d_model)
        Returns:
            output: (batch, seq_len, d_model)
            aux_loss: load balancing loss
        """
        B, T, D = x.shape
        gates, indices, aux_loss = self.gating(x)

        # Compute capacity
        capacity = int(self.capacity_factor * T / self.num_experts)

        # Gather expert outputs
        output = torch.zeros_like(x)
        expert_counts = torch.zeros(self.num_experts, device=x.device)

        for k in range(self.top_k):
            expert_idx = indices[:, :, k]  # (B, T)
            gate_weight = gates[:, :, k]   # (B, T)

            for e_idx in range(self.num_experts):
                mask = expert_idx == e_idx  # (B, T)
                if not mask.any():
                    continue

                # Capacity check
                count = mask.sum().item()
                expert_counts[e_idx] += count

                if count > capacity * B:
                    # Drop excess tokens (keep first `capacity * B`)
                    flat_mask = mask.reshape(-1)
                    indices_true = flat_mask.nonzero().squeeze(-1)
                    drop = indices_true[capacity * B:]
                    flat_mask[drop] = False
                    mask = flat_mask.reshape(B, T)

                # Route to expert
                expert_input = x[mask]  # (num_tokens, D)
                if expert_input.numel() > 0:
                    expert_output = self.experts[e_idx](expert_input)
                    output[mask] += gate_weight[mask].unsqueeze(-1) * expert_output

        return output, aux_loss

from transformers import AutoTokenizer, AutoModelForCausalLM
from torch.utils.data import DataLoader


class MoETrainer:
    """
    Train a Mixture of Experts model for MangaAssist.
    Uses LoRA on the gating network + expert selection.
    """

    def __init__(self, base_model: str = "meta-llama/Llama-3-8b-hf"):
        self.tokenizer = AutoTokenizer.from_pretrained(base_model)
        self.tokenizer.pad_token = self.tokenizer.eos_token

        # Load base model and replace FFN layers with MoE
        self.model = AutoModelForCausalLM.from_pretrained(
            base_model,
            torch_dtype=torch.bfloat16,
            device_map="auto",
        )

        # Replace middle FFN layers (9-24) with MoE
        self._replace_ffn_with_moe(layers=range(8, 24))

    def _replace_ffn_with_moe(self, layers: range):
        """Replace selected FFN layers with MoE layers."""
        for idx in layers:
            layer = self.model.model.layers[idx]
            d_model = layer.mlp.gate_proj.in_features
            d_ff = layer.mlp.gate_proj.out_features

            # Initialize first expert from the original FFN weights
            moe = MoELayer(d_model=d_model, d_ff=d_ff, num_experts=8, top_k=2)

            # Copy original FFN weights to expert 0
            with torch.no_grad():
                moe.experts[0].gate_proj.weight.copy_(layer.mlp.gate_proj.weight)
                moe.experts[0].up_proj.weight.copy_(layer.mlp.up_proj.weight)
                moe.experts[0].down_proj.weight.copy_(layer.mlp.down_proj.weight)

                # Initialize other experts with noise around expert 0
                for e in range(1, 8):
                    for param_name in ["gate_proj", "up_proj", "down_proj"]:
                        src = getattr(moe.experts[0], param_name).weight
                        tgt = getattr(moe.experts[e], param_name).weight
                        tgt.copy_(src + 0.01 * torch.randn_like(src))

            layer.mlp = moe

    def train(
        self,
        train_data: list[dict],
        epochs: int = 5,
        lr: float = 1e-5,
        balance_coeff: float = 0.01,
    ):
        """Train with combined language modeling and load balancing loss."""
        loader = DataLoader(train_data, batch_size=2, shuffle=True)
        optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr)

        for epoch in range(epochs):
            total_lm_loss = 0
            total_aux_loss = 0

            for batch in loader:
                inputs = self.tokenizer(
                    batch["text"], return_tensors="pt", padding=True,
                    truncation=True, max_length=2048,
                ).to(self.model.device)

                outputs = self.model(**inputs, labels=inputs["input_ids"])
                lm_loss = outputs.loss

                # Collect auxiliary losses from all MoE layers
                aux_loss = torch.tensor(0.0, device=self.model.device)
                for layer in self.model.model.layers:
                    if hasattr(layer.mlp, 'gating'):
                        # MoE layer stores aux_loss from last forward
                        if hasattr(layer.mlp, '_last_aux_loss'):
                            aux_loss += layer.mlp._last_aux_loss

                total_loss = lm_loss + balance_coeff * aux_loss

                total_loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                optimizer.step()
                optimizer.zero_grad()

                total_lm_loss += lm_loss.item()
                total_aux_loss += aux_loss.item()

            avg_lm = total_lm_loss / len(loader)
            avg_aux = total_aux_loss / len(loader)
            print(
                f"Epoch {epoch+1}: LM loss={avg_lm:.4f}, "
                f"Aux loss={avg_aux:.4f}"
            )

    def analyze_specialization(self, test_data: list[dict]) -> dict:
        """Analyze which experts specialize in which query types."""
        self.model.eval()
        utilization = {}  # {query_type: {expert_id: count}}

        hooks = []
        routing_log = []

        def make_hook(layer_idx):
            def hook_fn(module, input, output):
                if isinstance(output, tuple) and len(output) == 2:
                    # MoE layer returns (output, aux_loss)
                    pass
                # Log routing decisions
                _, indices, _ = module.gating(input[0])
                routing_log.append({
                    "layer": layer_idx,
                    "indices": indices.detach().cpu(),
                })
            return hook_fn

        for idx, layer in enumerate(self.model.model.layers):
            if hasattr(layer.mlp, 'gating'):
                hooks.append(
                    layer.mlp.register_forward_hook(make_hook(idx))
                )

        for ex in test_data:
            routing_log.clear()
            query_type = ex["query_type"]

            inputs = self.tokenizer(
                ex["text"], return_tensors="pt",
            ).to(self.model.device)

            with torch.no_grad():
                self.model(**inputs)

            # Aggregate routing decisions
            if query_type not in utilization:
                utilization[query_type] = {i: 0 for i in range(8)}

            for entry in routing_log:
                for expert_id in entry["indices"].flatten().tolist():
                    utilization[query_type][expert_id] += 1

        for hook in hooks:
            hook.remove()

        # Normalize to percentages
        for qt in utilization:
            total = sum(utilization[qt].values())
            if total > 0:
                utilization[qt] = {
                    k: v / total for k, v in utilization[qt].items()
                }

        return utilization

import sagemaker
from sagemaker.huggingface import HuggingFaceModel


def deploy_moe_model(model_path: str, instance_type: str = "ml.g5.2xlarge"):
    """Deploy MoE model to SageMaker endpoint."""

    # MoE requires more memory but same inference FLOPS
    # g5.2xlarge: 24GB GPU, sufficient for 32B total / 8B active
    hub = {
        "HF_MODEL_ID": model_path,
        "SM_NUM_GPUS": "1",
        "MAX_INPUT_LENGTH": "2048",
        "MAX_TOTAL_TOKENS": "4096",
    }

    model = HuggingFaceModel(
        env=hub,
        role=sagemaker.get_execution_role(),
        image_uri=sagemaker.image_uris.retrieve(
            framework="huggingface-llm",
            region="us-west-2",
            version="2.3.0",
            instance_type=instance_type,
        ),
    )

    predictor = model.deploy(
        initial_instance_count=1,
        instance_type=instance_type,
        endpoint_name="mangaassist-moe-endpoint",
    )

    return predictor