"""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)