src/model/moe.py

"""Mixture-of-Experts (MoE) — écrit à la main (Phase 10).

Configuration cible imposée : **20 experts, top-2**.

Principe : chaque token est routé vers les ``top_k`` experts parmi ``n_experts``.
Le coût de calcul par token reste ~``top_k`` experts (ici 2), alors que la capacité
totale (nombre de paramètres) croît avec ``n_experts``. Un **routeur** linéaire
produit un score par expert ; on garde les ``top_k`` plus grands et on renormalise
leurs portes (gates).

Points d'implémentation :

* **Routeur** : ``Linear(n_embd, n_experts)``. On ajoute un **bruit** gaussien aux
  logits pendant l'entraînement (router jitter) pour favoriser l'exploration.
* **Loss auxiliaire d'équilibrage** (Switch-Transformer) : ``E * sum_i f_i * P_i`` où
  ``f_i`` = fraction de tokens routés vers l'expert ``i`` et ``P_i`` = probabilité
  moyenne. Elle pénalise la concentration sur peu d'experts.
* **Calcul efficace** : on ne fait passer que les tokens routés dans chaque expert
  (pas de calcul dense sur tous les experts) — c'est le point d'un MoE.

Justification : docs/architecture.md §9.
"""

from __future__ import annotations

import math
from dataclasses import dataclass

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

from src.model.attention import MultiHeadSelfAttention
from src.model.ffn import FeedForward
from src.model.text_model import TextConfig


@dataclass
class MoEConfig:
    n_experts: int = 20
    top_k: int = 2
    aux_loss_coef: float = 0.01
    jitter: float = 0.01
    ffn_mult: int = 4


class MoEFFN(nn.Module):
    """Couche MoE : routeur top-k + ``n_experts`` FFN indépendants."""

    def __init__(self, n_embd: int, config: MoEConfig, dropout: float = 0.0,
                 bias: bool = False):
        super().__init__()
        self.config = config
        self.n_experts = config.n_experts
        self.top_k = config.top_k
        self.router = nn.Linear(n_embd, config.n_experts, bias=False)
        self.experts = nn.ModuleList([
            FeedForward(n_embd, mult=config.ffn_mult, dropout=dropout, bias=bias)
            for _ in range(config.n_experts)
        ])
        nn.init.normal_(self.router.weight, std=0.02)

    def routing(self, x: torch.Tensor):
        """Retourne ``(gate, idx, probs)`` pour ``x`` de forme ``(..., n_embd)``."""
        logits = self.router(x)
        if self.training and self.config.jitter > 0:
            logits = logits + torch.randn_like(logits) * self.config.jitter
        probs = F.softmax(logits, dim=-1)
        gate, idx = probs.topk(self.top_k, dim=-1)
        gate = gate / gate.sum(dim=-1, keepdim=True).clamp_min(1e-9)
        return gate, idx, probs

    def forward(self, x: torch.Tensor):
        """``x`` : ``(B, T, n_embd)`` -> ``(out, aux_loss, load)``."""
        *lead, d = x.shape
        xf = x.reshape(-1, d)

        gate, idx, probs = self.routing(xf)          # (N, k), (N, k), (N, E)
        out = torch.zeros_like(xf)
        for e in range(self.n_experts):
            in_expert = idx == e                     # (N, k)
            token_mask = in_expert.any(dim=-1)       # (N,)
            if not bool(token_mask.any()):
                continue
            weight = (gate * in_expert).sum(dim=-1)[token_mask]  # (n_e,)
            out[token_mask] += weight.unsqueeze(-1) * self.experts[e](xf[token_mask])

        # Fraction de tokens par expert (sur les k choix) et probabilité moyenne.
        one_hot = F.one_hot(idx, self.n_experts).float().sum(dim=1)   # (N, E)
        f = one_hot.mean(dim=0)                                       # (E,)
        p = probs.mean(dim=0)                                         # (E,)
        aux_loss = self.n_experts * (f * p).sum()
        return out.reshape(*lead, d), aux_loss, f.detach()

    @property
    def active_parameters(self) -> int:
        """Paramètres mobilisés par token : routeur + ``top_k`` experts."""
        per_expert = sum(p.numel() for p in self.experts[0].parameters())
        return self.router.weight.numel() + self.top_k * per_expert


class MoETransformerBlock(nn.Module):
    """Bloc pre-norm identique à `TransformerBlock` mais avec un FFN MoE."""

    def __init__(self, n_embd: int, n_head: int, moe: MoEConfig, dropout: float = 0.0,
                 bias: bool = False, rope_theta: float = 10000.0, use_rope: bool = True):
        super().__init__()
        self.ln_1 = nn.LayerNorm(n_embd, bias=bias)
        self.attn = MultiHeadSelfAttention(n_embd, n_head, dropout=dropout, bias=bias,
                                           rope_theta=rope_theta, use_rope=use_rope,
                                           causal=True)
        self.ln_2 = nn.LayerNorm(n_embd, bias=bias)
        self.moe = MoEFFN(n_embd, moe, dropout=dropout, bias=bias)

    def forward(self, x, cos=None, sin=None):
        x = x + self.attn(self.ln_1(x), cos, sin)
        moe_out, aux, load = self.moe(self.ln_2(x))
        x = x + moe_out
        return x, aux, load


class CortexMoELM(nn.Module):
    """Transformer texte decoder-only dont les FFN sont remplacés par des MoE."""

    def __init__(self, config: TextConfig, moe: MoEConfig):
        super().__init__()
        self.config = config
        self.moe_config = moe
        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
        self.dropout = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList([
            MoETransformerBlock(config.n_embd, config.n_head, moe, dropout=config.dropout,
                                bias=config.bias, rope_theta=config.rope_theta,
                                use_rope=config.use_rope)
            for _ in range(config.n_layer)
        ])
        self.ln_f = nn.LayerNorm(config.n_embd, bias=config.bias)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        if config.tie_embeddings:
            self.lm_head.weight = self.token_embedding.weight

        from src.model.rope import build_rope_cache
        self._build_rope_cache = build_rope_cache
        self.apply(_init_weights)

    def encode(self, inputs_embeds: torch.Tensor):
        """Retourne ``(hidden, aux_loss, mean_load)``."""
        from src.model.rope import build_rope_cache

        x = self.dropout(inputs_embeds)
        t = x.size(1)
        cos, sin = build_rope_cache(t, self.config.n_embd // self.config.n_head,
                                    self.config.rope_theta, x.device, x.dtype)
        aux_total = x.new_zeros(())
        load_acc = None
        for block in self.blocks:
            x, aux, load = block(x, cos, sin)
            aux_total = aux_total + aux
            load_acc = load if load_acc is None else load_acc + load
        load_acc = load_acc / len(self.blocks)
        return self.ln_f(x), aux_total, load_acc

    def forward(self, inputs_embeds: torch.Tensor, targets: torch.Tensor | None = None):
        """Retourne ``(logits, loss, aux_loss, load)``."""
        hidden, aux, load = self.encode(inputs_embeds)
        logits = self.lm_head(hidden)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)),
                                   targets.reshape(-1), ignore_index=-100)
        return logits, loss, aux, load

    @property
    def total_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters())

    @property
    def active_parameters(self) -> int:
        """Paramètres hors experts + routeurs + top_k experts actifs."""
        expert_params = sum(p.numel() for b in self.blocks
                            for p in b.moe.experts.parameters())
        active_expert = sum(b.moe.active_parameters for b in self.blocks)
        return self.total_parameters - expert_params - sum(
            b.moe.router.weight.numel() for b in self.blocks) + active_expert


def load_balance_metrics(load: torch.Tensor) -> dict:
    """Diagnostic du routage depuis la charge par expert ``(n_experts,)``.

    ``entropy_norm`` : entropie de la distribution de charge normalisée par
    ``log(n_experts)`` (1.0 = parfaitement équilibré, 0.0 = un seul expert).
    ``max_share`` : part du token-choix reçue par l'expert le plus chargé.
    ``dead_experts`` : nombre d'experts jamais sélectionnés.
    """
    total = float(load.sum())
    p = load / max(total, 1e-9)
    ent = -(p * (p + 1e-9).log()).sum()
    return {
        "entropy_norm": float(ent / math.log(len(load))),
        "max_share": float(p.max()),
        "dead_experts": int((load == 0).sum()),
    }


def _init_weights(module: nn.Module) -> None:
    if isinstance(module, nn.Linear):
        nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Embedding):
        nn.init.normal_(module.weight, mean=0.0, std=0.02)