Le bloc FFN d’un Transformer est appliqué identiquement à chaque jeton, quelle que soit sa nature : un article, un chiffre, un verbe conjugué passent par la même matrice. L’idée du Mixture of Experts (Jacobs et al. 1991, adapté aux Transformers par Shazeer et al. 2017, puis Switch Transformer, Fedus et al. 2021) est de conditionner ce calcul : un petit routeur choisit, pour chaque jeton, lesquels de E réseaux indépendants — les experts — vont le traiter.
Le pari économique : si seulement k des E experts travaillent par jeton, le modèle a E fois plus de paramètres qu’un FFN dense de même largeur, pour un coût de calcul par jeton qui ne croît que comme k. C’est la promesse qu’il faut mesurer et non réciter : paramètres et calcul se découplent-ils vraiment, et à quel prix d’apprentissage ?
Ce notebook implémente la couche MoE from scratch (routeur top-k, capacité par expert, dropout, loss d’équilibrage), la vérifie sur des cas contrôlés, mesure le coût exact contre deux FFN denses d’équivalences différentes, puis entraîne dense contre MoE sur la tâche jouet du 3.4b — avec et sans loss d’équilibrage, pour voir le routeur s’effondrer ou non. Comme dans tout la série : pas de framework de MoE (timm, megablocks…), tout à la main, tout mesuré.
Prérequis : 3.4b (attention causale, la classe PetitLM, la méthodologie de banc entrelacé).
Sommaire
Le routeur top-k — experts, porte, renormalisation, vérifications sur cas contrôlés
La capacité — le contrat du Switch : facteur de capacité, jetons jetés, priorité au premier choix
La loss d’équilibrage alpha_load — la formule E * somme(f_i * P_i), vérifiée analytiquement
Mesurer le routage — entropie du routeur et parts des experts
Coût — MoE contre FFN dense : équivalence en calcul actif, équivalence en paramètres totaux
Qualité — la tâche du 3.4b : dense contre MoE, alpha_load = 0 contre alpha_load > 0
Limites et suite, Conclusion
Exercices : top-1 strict du Switch · balayage du facteur de capacité · spécialisation des experts par marqueur.
import mathimport timeimport numpy as npimport torchimport osos.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") # determinisme cuBLAS (#16795)# Determinisme (#16795) : la graine seule ne garantit PAS la reproductibilite# (heuristiques cuDNN, kernels non deterministes). warn_only=True au premier# passage pour inventorier les ops fautives sans faire echouer le run.torch.use_deterministic_algorithms(True, warn_only=True)torch.backends.cudnn.deterministic =Truetorch.backends.cudnn.benchmark =Falseimport torch.nn as nnimport torch.nn.functional as F# Les reductions multi-thread de PyTorch sur CPU ne sont pas reproductibles au bit pres.# Un seul thread rend deux executions identiques : indispensable pour un banc de mesure.torch.set_num_threads(1)torch.manual_seed(0)np.random.seed(0)DEVICE = torch.device("cpu")print("torch", torch.__version__, "| numpy", np.__version__, "| device", DEVICE)print("threads torch :", torch.get_num_threads())
1. Le routeur top-k : chaque jeton choisit ses experts
Un FFN dense applique à chaque jeton \(x \in \mathbb{R}^d\) la même fonction \(\mathrm{FFN}(x) = W_2 \, \sigma(W_1 x)\) avec \(\sigma =\) GELU. La couche MoE remplace ce bloc unique par :
E experts\(\mathrm{FFN}_1, \dots, \mathrm{FFN}_E\), de même architecture, de poids indépendants ;
un routeur : une simple couche linéaire \(W_r \in \mathbb{R}^{d \times E}\), sans biais, dont le softmax donne la distribution \(p(x)\) sur les experts ;
une sélection top-k : seuls les k experts de plus forte probabilité traitent le jeton, et leurs sorties sont combinées par les probabilités de porte renormalisées sur les k sélectionnés.
Deux choix d’implémentation, notés une fois pour toutes : les poids de porte sont les probabilités du softmax renormalisées sur le top-k (variante GShard ; le Switch top-1 n’a pas besoin de renormaliser puisqu’il ne garde qu’un expert), et le compte de routage \(f_i\) — utilisé par la loss d’équilibrage — est pris sur le premier choix de chaque jeton, comme dans les deux papiers.
La couche expose aussi les deux organes du contrat Switch : un dropout sur la sortie des experts (activé seulement à l’entraînement) et une capacité par expert — plafond du nombre de jetons qu’un expert accepte par lot, exprimé par le facteur de capacité \(c\) : \(\mathrm{cap} = \lceil c \cdot B T k / E \rceil\) jetons. Un jeton au-delà est jeté : il traverse la couche par la seule connexion résiduelle. La priorité est au premier choix et à l’ordre du lot — le jeton $k-ième choix d’un expert n’est servi que s’il reste de la place après tous ses premiers choix.
class Expert(nn.Module):"""Le FFN dense standard : Linear -> GELU -> dropout -> Linear."""def__init__(self, d, cache, dropout=0.0):super().__init__()self.fc1 = nn.Linear(d, cache)self.fc2 = nn.Linear(cache, d)self.dropout = nn.Dropout(dropout)def forward(self, x):returnself.fc2(self.dropout(F.gelu(self.fc1(x))))class MoECouche(nn.Module):"""Couche MoE a la GShard/Switch : routeur lineaire, top-k, porte renormalisee. Renvoie (sortie, probs) ou probs (B, T, E) est la distribution COMPLETE du routeur (pas seulement le top-k) : c'est elle que mesurent l'entropie et la loss d'equilibre. """def__init__(self, d, n_experts=8, topk=2, cache=None, dropout=0.0):super().__init__() cache = cache or4* dself.n_experts, self.topk = n_experts, topkself.routeur = nn.Linear(d, n_experts, bias=False)self.experts = nn.ModuleList([Expert(d, cache, dropout) for _ inrange(n_experts)])def capacite_lot(self, B, T, facteur):"""Plafond de jetons par expert : c * B*T*k / E, arrondi au superieur."""return math.ceil(facteur * B * T *self.topk /self.n_experts)def forward(self, x, facteur_capacite=None): B, T, d = x.shape logits =self.routeur(x) # (B, T, E) probs = logits.softmax(dim=-1) top_p, top_i = probs.topk(self.topk, dim=-1) # (B, T, k) poids = top_p / top_p.sum(dim=-1, keepdim=True) # porte renormalisee capacite =Noneif facteur_capacite isNoneelseself.capacite_lot(B, T, facteur_capacite) sortie = torch.zeros_like(x) compte = torch.zeros(self.n_experts, dtype=torch.long) # jetons SERVIS par expert jetes =0 plat = x.reshape(-1, d) # vue (B*T, d) sortie_plate = sortie.reshape(-1, d)for e inrange(self.n_experts):for k inrange(self.topk): # le premier choix passe d'abord masque = (top_i[:, :, k] == e).reshape(-1) # (B*T,) idx = masque.nonzero().squeeze(-1)if idx.numel() ==0:continueif capacite isnotNone: reste =int(capacite - compte[e].item())if reste <=0: jetes += idx.numel()continueif idx.numel() > reste: jetes += idx.numel() - reste idx = idx[:reste] # priorite a l'ordre du lot compte[e] += idx.numel() y =self.experts[e](plat[idx]) sortie_plate.index_add_(0, idx, poids[:, :, k].reshape(-1)[idx].unsqueeze(-1) * y)return sortie, probs, compte, jetesprint("definitions pretes : Expert = FFN dense (Linear -> GELU -> Linear) et MoECouche")print("MoECouche(d, n_experts, topk, cache) : routeur lineaire top-k, porte renormalisee,")print("capacite par expert ; forward(x) renvoie (sortie, probs completes du routeur,")print("compte de jetons servis par expert, nombre de jetons jetes par capacite)")
definitions pretes : Expert = FFN dense (Linear -> GELU -> Linear) et MoECouche
MoECouche(d, n_experts, topk, cache) : routeur lineaire top-k, porte renormalisee,
capacite par expert ; forward(x) renvoie (sortie, probs completes du routeur,
compte de jetons servis par expert, nombre de jetons jetes par capacite)
1.1 Échelle de vérifications
Chaque organe est vérifié sur un cas dont la réponse se calcule à la main, avant toute mesure. Trois vérifications :
Porte renormalisée : la somme des poids de porte vaut exactement 1 pour chaque jeton, quels que soient les logits.
Routement biaisé : un routeur dont un expert reçoit un logit écrasant (+20) doit envoyer tous les jetons à cet expert, et la sortie doit être exactement sa sortie (à la porte renormalisée près, qui vaut alors 1).
Capacité : quatre jetons envoyés en masse à un expert avec une capacité de 1 → un seul servi, trois jetés ; les deuxièmes choix sont servis après tous les premiers choix.
torch.manual_seed(0)D =16moe = MoECouche(D, n_experts=8, topk=2, cache=8)x = torch.randn(2, 3, D) # 6 jetons# -- 1. la porte renormalisee somme a 1, toujourswith torch.no_grad(): _, probs, _, _ = moe(x) top_p, _ = probs.topk(2, dim=-1) porte = top_p / top_p.sum(-1, keepdim=True)print("somme des poids de porte (min, max) :",float(porte.sum(-1).min()), float(porte.sum(-1).max()))assert (porte.sum(-1) -1.0).abs().max() <1e-6# Pour les gardes 2 et 3, le routeur biaise DOIT donner un logit positif a l'expert 3 :# logit_3 = 20 * somme(x), donc entree a composantes positives (sinon le logit s'inverse).x_pos = torch.rand(2, 3, D) +0.5# -- 2. routeur biaise : tout va a l'expert 3, la sortie EST sa sortiewith torch.no_grad(): moe.routeur.weight.zero_() moe.routeur.weight[3] =20.0# logit +20*sum(x) sur l'expert 3 sortie, probs, compte, _ = moe(x_pos) attendu = moe.experts[3](x_pos)print("probabilite de l'expert 3 (min) :", float(probs[..., 3].min()))print("ecart sortie vs expert 3 seul :", float((sortie - attendu).abs().max()))print("jetons servis par expert :", compte.tolist())assertfloat(probs[..., 3].min()) >1-1e-6assert (sortie - attendu).abs().max() <1e-5assert compte[3].item() == x_pos.numel() // D# -- 3. capacite : 4 jetons, tous premier-choix vers l'expert 3, capacite 1torch.manual_seed(1)moe2 = MoECouche(D, n_experts=8, topk=2, cache=8)with torch.no_grad(): moe2.routeur.weight.zero_() moe2.routeur.weight[3] =20.0 x4 = torch.rand(1, 4, D) +0.5# 4 jetons, composantes positives B, T, _ = x4.shape cap = moe2.capacite_lot(B, T, facteur=1.0) # c*B*T*k/E = 1*4*2/8 = 1 _, _, compte2, jetes = moe2(x4, facteur_capacite=1.0)print("capacite par expert (c=1.0, 4 jetons, k=2, E=8) :", cap)print("expert 3 servis :", compte2[3].item(), "| affectations totales jetees :", jetes,"sur", B * T * moe2.topk)assert cap ==1assert compte2[3].item() ==1# 1 servi, 3 jetes au premier choixassert jetes == B * T * moe2.topk -2# restent : 1er choix e3 + 1er choix e0print("\nEchelle de verifications : OK")
somme des poids de porte (min, max) : 1.0 1.0
probabilite de l'expert 3 (min) : 1.0
ecart sortie vs expert 3 seul : 0.0
jetons servis par expert : [6, 0, 0, 6, 0, 0, 0, 0]
capacite par expert (c=1.0, 4 jetons, k=2, E=8) : 1
expert 3 servis : 1 | affectations totales jetees : 6 sur 8
Echelle de verifications : OK
1.2 Lecture
La vérification 2 mérite un arrêt : quand le routeur est biaisé, la sortie de la couche est la sortie d’un seul expert — les 7 autres ne participent plus. C’est l’état limite que la loss d’équilibrage de la section 3 cherche à éviter : dans cet état, la couche MoE est un FFN dense avec E−1 blocs de poids morts, le pire des deux mondes.
La vérification 3 montre le second risque, propre au déploiement : avec une capacité serrée, un routement concentré ne dégrade pas doucement, il jette des jetons — ils disparaissent de la couche et ne survivent que par la résiduelle. Le facteur de capacité est l’assurance contre ce mode de défaillance, payée en mémoire de tampon.
2. La capacité : le contrat du Switch
Le calcul par expert d’une couche MoE déployée sur plusieurs machines se fait dans des tampons de taille fixe — on ne peut pas provisionner « autant que nécessaire » sans bornes. Le Switch Transformer fixe donc à l’avance combien de jetons chaque expert acceptera :
\[\mathrm{capacite} = \left\lceil c \cdot \frac{B \cdot T \cdot k}{E} \right\rceil\]
où \(c\) est le facteur de capacité. À \(c = 1{,}0\), la capacité est exactement la charge d’un routement parfaitement équilibré : chaque expert reçoit sa part exacte des affectations. Tout déséquilibre se paie immédiatement en jetons jetés — chez les experts surchargés seulement, les capacités non utilisées des experts affamés ne sont pas transférables.
C’est cette interaction qu’il faut voir numériquement : la pénalité n’est pas proportionnelle au déséquilibre moyen mais concentrée sur les experts en excès. Un balayage du facteur de capacité sur deux routements — un équilibré, un effondré — mesure le phénomène proprement.
def parts_d_un_routement(n_jetons, n_experts=8, concentration=None, graine=0):"""Distribution de la charge sur les experts. concentration=None : routement equilibre (multinomiale uniforme). concentration=(e0, p) : une part p des affectations va a l'expert e0, le reste uniforme. """ g = torch.Generator().manual_seed(graine) E = n_experts charges = torch.zeros(E)if concentration isNone: choix = torch.randint(0, E, (n_jetons,), generator=g)else: e0, p = concentration vers_e0 = torch.rand(n_jetons, generator=g) < p autres = torch.randint(0, E -1, (n_jetons,), generator=g) autres = autres + (autres >= e0).long() # evite e0 choix = torch.where(vers_e0, torch.full_like(autres, e0), autres)for e inrange(E): charges[e] = (choix == e).sum()return chargesdef jetes_pour_facteur(charges, n_experts, facteur):"""Jetons jetes quand chaque expert plafonne a c * total/E affectations.""" capacite = math.ceil(facteur * charges.sum() / n_experts)returnint(torch.clamp(charges - capacite, min=0).sum().item())N_JETONS =4096E =8print("affectations par expert (4096 jetons, top-1 pour la charge, E = 8)")equilibre = parts_d_un_routement(N_JETONS, E, None, graine=0)effondre = parts_d_un_routement(N_JETONS, E, concentration=(3, 0.70), graine=1)total = N_JETONS # charge compte ici chaque jeton une fois (son premier choix)print("equilibre : parts =", [f"{c / total:.3f}"for c in equilibre.tolist()])print("effondre : parts =", [f"{c / total:.3f}"for c in effondre.tolist()],"(expert 3 a 70 %)")print()print(f"{'facteur c':>10s} | {'jetes (equilibre)':>18s} | {'jetes (effondre)':>18s}")for facteur in (1.0, 1.25, 1.5, 2.0): j_eq = jetes_pour_facteur(equilibre, E, facteur) j_ef = jetes_pour_facteur(effondre, E, facteur)print(f"{facteur:10.2f} | {j_eq:14d} ({100* j_eq / total:4.1f} %) | "f"{j_ef:14d} ({100* j_ef / total:4.1f} %)")
Le routement équilibré survit à \(c = 1{,}0\) : la charge de chaque expert flotte autour de la capacité exacte et seuls les débordements aléatoires — quelques pour cent des affectations — sont jetés. Le routement effondré, lui, jette massivement au facteur 1 et ne retrouve un service complet qu’en gonflant la capacité bien au-delà de la charge moyenne — une capacité surdimensionnée pour tout le monde parce qu’un expert est obèse.
Voilà pourquoi le facteur de capacité et la loss d’équilibrage ne sont pas deux réglages indépendants : la seconde est ce qui rend le premier séreinement petit. C’est aussi un compromis mémoire — les tampons sont alloués à la capacité, pas à la charge réelle : trop de marge coûte de la mémoire pour rien.
3. La loss d’équilibrage alpha_load
Laissé à lui-même, le routeur s’effondre par auto-renforcement : l’expert qui reçoit le plus de jetons apprend le plus vite, devient meilleur, attire encore plus de jetons. Les quelques experts gagnants captent tout le signal, les autres restent des poids aléatoires — le modèle a payé E experts pour en utiliser un ou deux.
La parade du Switch Transformer est une loss auxiliaire, ajoutée à la loss de tâche avec le coefficient alpha_load :
\[\mathcal{L}_{\mathrm{aux}} = \alpha_{\mathrm{load}} \cdot E \cdot \sum_{i=1}^{E} f_i \, P_i\]
avec \(f_i\) la fraction de jetons dont l’expert \(i\) est le premier choix (constante par morceau, non différentiable) et \(P_i\) la probabilité moyenne du routeur pour l’expert \(i\) (différentiable). Le produit est la clé : minimiser \(\sum f_i P_i\) pousse \(P_i\) à baisser là où \(f_i\) est haut — le routeur dessert les experts surchargés. Les deux grandeurs extrêmes se calculent à la main :
routement uniforme : \(f_i = P_i = 1/E\) pour tout \(i\), donc \(\sum f_i P_i = E \cdot 1/E^2 = 1/E\) et \(\mathcal{L}_{\mathrm{aux}} / \alpha_{\mathrm{load}} = 1\) — c’est le minimum ;
effondrement total : \(f_1 = P_1 = 1\), \(\mathcal{L}_{\mathrm{aux}} / \alpha_{\mathrm{load}} = E\) — c’est le pire cas.
La loss ne récompense pas l’uniformité au-delà : sa valeur plancher à l’uniforme est 1, pas 0 — ce qui compte est la dérivée, qui pousse vers l’équilibre. Ces deux bornes servent de garde-fous analytiques à l’implémentation.
def perte_equilibre(probs, top_i):"""Loss auxiliaire du Switch : E * somme_i f_i * P_i. probs : (B, T, E) distribution complete du routeur. top_i : (B, T, k) indices du top-k ; seul le PREMIER choix compte pour f_i. Renvoie la loss SANS le coefficient alpha_load (l'appelant multiplie). """ n_experts = probs.shape[-1] f = torch.zeros(n_experts, device=probs.device)for e inrange(n_experts): f[e] = (top_i[..., 0] == e).float().mean() P = probs.mean(dim=(0, 1))return n_experts * (f * P).sum()B, T, E =4, 16, 8# -- garde-fou 1 : uniforme -> exactement 1 (le minimum)probs_unif = torch.full((B, T, E), 1.0/ E)top_i_unif = probs_unif.topk(2, dim=-1).indices # egalite -> ordre d'indicel_unif = perte_equilibre(probs_unif, top_i_unif)print(f"routement uniforme : L_aux = {l_unif.item():.6f} (theorie : 1.0)")assertabs(l_unif.item() -1.0) <1e-6# -- garde-fou 2 : effondre sur l'expert 0 -> exactement Eprobs_eff = torch.zeros(B, T, E)probs_eff[..., 0] =1.0top_i_eff = probs_eff.topk(2, dim=-1).indicesl_eff = perte_equilibre(probs_eff, top_i_eff)print(f"routement effondre : L_aux = {l_eff.item():.6f} (theorie : E = {E})")assertabs(l_eff.item() - E) <1e-6# -- cas intermediaire : 2 experts actifs a parts egales -> E * 2 * (1/2 * 1/2) = E/2probs_deux = torch.zeros(B, T, E)probs_deux[..., 0] =0.5probs_deux[..., 1] =0.5top_i_deux = probs_deux.topk(2, dim=-1).indicesl_deux = perte_equilibre(probs_deux, top_i_deux)print(f"2 experts a parts egales : L_aux = {l_deux.item():.6f} (theorie : E/2 = {E /2})")assertabs(l_deux.item() - E /2) <1e-6# -- differentiabilite : le gradient passe bien par P (pas par f, constant par morceaux)probs_v = torch.full((B, T, E), 1.0/ E, requires_grad=True)top_i_v = torch.full((B, T, 2), 0) # tout premier choix sur e0perte_equilibre(probs_v, top_i_v).backward()grad = probs_v.grad[0, 0]attendu = (torch.arange(E) ==0).float() * E / (B * T) # dL/dP = E * f / (B*T)print("gradient analytique (E*f/(BT)) vs autograd, ecart max :",float((grad - attendu).abs().max()))assertfloat((grad - attendu).abs().max()) <1e-9print("\nEchelle de verifications : OK")
routement uniforme : L_aux = 1.000000 (theorie : 1.0)
routement effondre : L_aux = 8.000000 (theorie : E = 8)
2 experts a parts egales : L_aux = 4.000000 (theorie : E/2 = 4.0)
gradient analytique (E*f/(BT)) vs autograd, ecart max : 0.0
Echelle de verifications : OK
3.1 Lecture
La loss est exactement aux bornes théoriques sur les trois cas contrôlés — uniforme 1, deux experts actifs \(E/2\), effondrement \(E\) — et le gradient vérifié à la main confirme que c’est bien \(P_i\) (le softmax, différentiable) qui porte la correction, \(f_i\) n’étant qu’un coefficient de pondération constant par morceaux.
Un point contre-intuitif mérite d’être retenu : la loss ne peut pas descendre sous 1. Sa valeur absolue n’a pas de lecture de « qualité » — un entraînement qui affiche \(\mathcal{L}_{\mathrm{aux}} = 1{,}3\) n’est pas « à 30 % du déséquilibre », il est proche de l’équilibre si les \(f_i\) sont plats. La grandeur lisible est la distribution des charges, pas la loss.
4. Mesurer le routage : entropie et parts des experts
Deux grandeurs suffisent à décrire l’état d’un routeur, et elles ne racontent pas la même histoire :
l’entropie du routage\(H = -\sum_e p_e \log p_e\), moyennée sur les jetons, mesure la décision du routeur : \(\log 8 \approx 2{,}079\) pour un routeur indécis qui répartit uniformément, proche de 0 pour un routeur tranché. Attention à sa lecture : un routeur confiant mais équilibré (chaque jeton sûr de son expert, mais des experts différents) a une entropie faible et une charge parfaitement répartie — basse entropie n’est pas effondrement.
les parts des experts\(f_e\) (fraction de premiers choix) et leur maximum, mesurent la charge — c’est la grandeur que la loss d’équilibrage contrôle et celle qui détermine les jets de capacité. Équilibré : chaque part vaut \(1/E\) ; effondré : une part domine.
Les garde-fous se calculent exactement : entropie d’une distribution uniforme sur \(E = 8\) experts \(= \log 8\), entropie nulle pour un routeur tranché. Pour les parts, une subtilité d’implémentation surgit : sur des probabilités exactement égales, le tie-break de topk (préférence au plus petit indice) envoie tous les premiers choix sur l’expert 0 — la ligne « uniforme + bruit » du tableau casse les égalités avec un bruit minuscule et mesure ce que vaut réellement la charge d’un routeur indécis.
def entropie_routage(probs):"""Entropie moyenne (en nats) de la distribution du routeur, par jeton."""returnfloat(-(probs * probs.clamp_min(1e-12).log()).sum(-1).mean())def parts_experts(top_i, n_experts):"""Fraction de jetons dont chaque expert est le premier choix.""" premiers = top_i[..., 0].reshape(-1)return torch.bincount(premiers, minlength=n_experts).float() / premiers.numel()E =8# -- garde-fous analytiquesunif = torch.full((4, 16, E), 1.0/ E)eff = torch.zeros(4, 16, E)eff[..., 3] =1.0confiant_equilibre = torch.zeros(4, 16, E)torch.manual_seed(0)assignations = torch.randint(0, E, (4, 16)) # chaque jeton sur de son expert...confiant_equilibre.scatter_(-1, assignations.unsqueeze(-1), 1.0) # ...mais des experts differents# egalites parfaites -> tie-break de topk vers l'indice min ; un bruit minuscule les cassetorch.manual_seed(1)uniforme_bruite = torch.softmax(torch.randn(4, 16, E) *0.01, dim=-1)print(f"{'cas':>22s} | {'entropie':>9s} | {'theorie':>9s} | {'part max':>9s}")for nom, p in [("uniforme indecis", unif), ("uniforme + bruit", uniforme_bruite), ("effondre sur e3", eff), ("confiant equilibre", confiant_equilibre)]: top_i = p.topk(2, dim=-1).indices parts = parts_experts(top_i, E) h_theo = {"uniforme indecis": math.log(E), "uniforme + bruit": math.log(E),"effondre sur e3": 0.0, "confiant equilibre": 0.0}[nom]print(f"{nom:>22s} | {entropie_routage(p):9.4f} | {h_theo:9.4f} | {parts.max().item():9.4f}")assertabs(entropie_routage(unif) - math.log(E)) <1e-6assertabs(entropie_routage(eff)) <1e-6assert parts_experts(uniforme_bruite.topk(2, -1).indices, E).max() <0.30# ~1/8 + bruitassert parts_experts(confiant_equilibre.topk(2, -1).indices, E).max() <0.30print("\nle cas 'uniforme indecis' a part max = 1.0 : egalite parfaite, tie-break de topk")print("vers l'expert 0 - toutes les decisions tombent sur le meme indice, sans que")print("le routeur ait jamais prefere personne. D'ou la ligne 'uniforme + bruit'.")print("\nle cas 'confiant equilibre' : entropie nulle, charge equirepartie -")print("basse entropie != effondrement ; c'est la charge qui decide.")
cas | entropie | theorie | part max
uniforme indecis | 2.0794 | 2.0794 | 1.0000
uniforme + bruit | 2.0794 | 2.0794 | 0.1875
effondre sur e3 | -0.0000 | 0.0000 | 1.0000
confiant equilibre | -0.0000 | 0.0000 | 0.1875
le cas 'uniforme indecis' a part max = 1.0 : egalite parfaite, tie-break de topk
vers l'expert 0 - toutes les decisions tombent sur le meme indice, sans que
le routeur ait jamais prefere personne. D'ou la ligne 'uniforme + bruit'.
le cas 'confiant equilibre' : entropie nulle, charge equirepartie -
basse entropie != effondrement ; c'est la charge qui decide.
Exercice 1 — le top-1 strict du Switch
Le Switch Transformer utilise \(k = 1\) : un seul expert par jeton, porte non renormalisée (inutile à k = 1), et montre que ça suffit — la capacité devient alors \(\lceil c \cdot BT/E \rceil\).
Implémentez routage_top1(x, couche) qui renvoie la sortie de la couche en routage top-1 sans toucher à la classe (utilisez couche.routeur et couche.experts), puis mesurez sur un lot aléatoire : la perte d’équilibrage et le nombre de jetés à facteur de capacité 1,0, contre les mêmes mesures en top-2 sur la même couche. Question de lecture : à charge égale par expert, pourquoi le top-1 est-il plus vulnérable au facteur de capacité que le top-2 ?
Indice : en top-2, un jeton a DEUX chances de tomber sur un expert sous-charge ; comptez les affectations totales (BTk) vs la capacite par expert. Etape 1 : top_p, top_i = probs.topk(1, dim=-1) ; la porte EST top_p (pas de renormalisation). Etape 2 : boucle sur les experts comme dans MoECouche.forward, capacite cBT/E.
def routage_top1(x, couche, facteur_capacite=None):"""Exercice 1 : sortie de la couche en routage top-1 strict (variante Switch). Renvoie (sortie, compte, jetes). La porte est la probabilite brute du softmax. """# TODO etudiantreturnNone, None, None# TODO etudiant# -- mesure de comparaison top-1 vs top-2 (a decommenter une fois complete)# torch.manual_seed(0)# couche_ex = MoECouche(32, n_experts=8, topk=2, cache=64)# x_ex = torch.randn(2, 32, 32)# s1, c1, j1 = routage_top1(x_ex, couche_ex, facteur_capacite=1.0)# print("top-1 : jetes =", j1)print("Exercice 1 a completer")
Exercice 1 a completer
5. Coût : MoE contre FFN dense, les deux équivalences
« Équivalent » ne veut rien dire tant qu’on ne dit pas équivalent en quoi. Deux FFN denses servent de repères à une couche MoE à E = 8 experts, k = 2, cache 4d :
équivalence en calcul actif : chaque jeton traverse k = 2 experts de largeur 4d, soit 8d de largeur cachée active. Le FFN dense de largeur cachée 8d fait le même travail par jeton — mais il totalise 4 fois moins de paramètres que la couche MoE (2 couches de poids au lieu de 2 × 8).
équivalence en paramètres totaux : le FFN dense de largeur cachée 32d porte autant de paramètres que les 8 experts — mais chaque jeton le traverse entièrement, 4 fois le calcul actif du MoE.
La promesse du MoE est précisément de se placer entre les deux : les paramètres du gros, le calcul par jeton du petit. Le banc entrelacé du 3.4b (médiane de 9 tours, échauffement préalable, ordre de boucle extérieur) mesure ce que cette promesse vaut en millisecondes réelles sur CPU — le routage top-k a un coût propre (boucle de dispatch, indexation) qui peut dévorer l’économie théorique à petite échelle : c’est exactement ce qu’un banc doit révéler plutôt qu’un décompte de FLOPs.
def bench_modules(variantes, T, d_model, batch=2, chauffe=3, tours=9):"""Mediane du forward + backward, mesures ENTRELACEES, en millisecondes. (Repris du 3.4b : tours en boucle EXTERIEURE, chaque module echantillonne a chaque tour, pour ne pas mesurer la position dans la boucle.) """ x = torch.randn(batch, T, d_model, requires_grad=True)for _, m in variantes:for _ inrange(chauffe): m(x).sum().backward() temps = {nom: [] for nom, _ in variantes}for _ inrange(tours):for nom, m in variantes: m.zero_grad(set_to_none=True) debut = time.perf_counter() m(x).sum().backward() temps[nom].append(time.perf_counter() - debut)return {nom: sorted(v)[len(v) //2] *1000for nom, v in temps.items()}class FFNDense(nn.Module):def__init__(self, d, cache, dropout=0.0):super().__init__()self.net = nn.Sequential(nn.Linear(d, cache), nn.GELU(), nn.Dropout(dropout), nn.Linear(cache, d))def forward(self, x):returnself.net(x)class MoEPourBanc(nn.Module):"""Adapte MoECouche au banc : seule la sortie participe a la loss."""def__init__(self, **kw):super().__init__()self.couche = MoECouche(**kw)def forward(self, x):returnself.couche(x)[0]D_BENCH =128torch.manual_seed(0)variantes = [ ("dense 4d (base 3.4b)", FFNDense(D_BENCH, 4* D_BENCH)), ("dense 8d (iso-calcul)", FFNDense(D_BENCH, 8* D_BENCH)), ("dense 32d (iso-parametres)", FFNDense(D_BENCH, 32* D_BENCH)), ("MoE E=8 k=2 (cache 4d)", MoEPourBanc(d=D_BENCH, n_experts=8, topk=2, cache=4* D_BENCH)),]def n_params(m):returnsum(p.numel() for p in m.parameters())print(f"forward + backward, batch 2, T = 128, d_model = {D_BENCH}, mediane de 9 tours entrelaces\n")print(f"{'variante':>26s} | {'params':>10s} | {'ms':>8s} | {'ms / M params':>13s}")MESURES_BANC = {}for T in (128,): temps = bench_modules(variantes, T, D_BENCH)for nom, m in variantes: MESURES_BANC[nom] = temps[nom]print(f"{nom:>26s} | {n_params(m):10,d} | {temps[nom]:8.2f} | "f"{temps[nom] / (n_params(m) /1e6):13.2f}")p_moe = n_params(variantes[3][1])p_8d = n_params(variantes[1][1])p_32d = n_params(variantes[2][1])print(f"\nrapport de parametres MoE / dense-8d : {p_moe / p_8d:.1f}x "f"| MoE / dense-32d : {p_moe / p_32d:.2f}x")print(f"rapport de temps MoE / dense-8d : {MESURES_BANC['MoE E=8 k=2 (cache 4d)'] / MESURES_BANC['dense 8d (iso-calcul)']:.2f}x "f"| MoE / dense-32d : {MESURES_BANC['MoE E=8 k=2 (cache 4d)'] / MESURES_BANC['dense 32d (iso-parametres)']:.2f}x")
forward + backward, batch 2, T = 128, d_model = 128, mediane de 9 tours entrelaces
variante | params | ms | ms / M params
dense 4d (base 3.4b) | 131,712 | 3.11 | 23.58
dense 8d (iso-calcul) | 263,296 | 5.45 | 20.69
dense 32d (iso-parametres) | 1,052,800 | 22.08 | 20.97
MoE E=8 k=2 (cache 4d) | 1,054,720 | 16.95 | 16.07
rapport de parametres MoE / dense-8d : 4.0x | MoE / dense-32d : 1.00x
rapport de temps MoE / dense-8d : 3.11x | MoE / dense-32d : 0.77x
5.1 Lecture
Les deux équivalences racontent la même histoire, chacune à moitié :
contre le dense iso-calcul (8d), la couche MoE porte 4,0 fois plus de paramètres — pour environ 3 fois plus de temps (2,9× sur le run commité). La parité théorique de FLOPs (2 experts de largeur 4d = une largeur active de 8d) est dévorée par le coût du dispatch : la boucle de routage fait 16 indexations et assemblages par couche, et à cette échelle (d = 128, 256 jetons) ce sont elles qui paient, pas les multiplications matricielles.
contre le dense iso-paramètres (32d), la couche MoE va 0,75 fois moins vite — 25 % d’économie réelle là où le décompte de FLOPs promettait 4. L’économie existe, atténuée par le même overhead.
La dernière colonne donne la lecture la plus juste du découplage : en millisecondes par million de paramètres, le MoE est la variante la moins chère du banc — chaque paramètre supplémentaire coûte moins cher à faire travailler, exactement le pari du MoE (sur le run commité : 12,3 contre 16,3 à 17,3 pour les trois denses, ordre stable d’un run à l’autre). Mais l’écart entre 0,75× mesuré et 4× théorique est la leçon de méthode : un découplage se mesure en temps réel, pas en décompte de FLOPs. Les implémentations de production (GEMM groupées, noyaux fusionnés, dispatch distribué) referment une bonne partie de cet écart ; une boucle Python ne le referme pas. (Les temps absolus dépendent de la machine : seuls les rapports, calculés dans la même cellule sur le même run, sont ancrés sur les sorties.)
6. Qualité : dense contre MoE sur la tâche du 3.4b
La même tâche jouet que le 3.4b : un marqueur en position 0 (8 classes possibles), du remplissage aléatoire, un jeton de requête en dernière position ; le modèle doit restituer le marqueur à la position de la requête. Hasard : exactitude 0,125, perplexité 8. Un mini-LM de 2 blocs, d = 64, attention causale 4 têtes, appris par Adam — le bloc FFN est soit dense (largeur 4d, l’architecture du 3.4b), soit MoE (E = 8 experts de largeur 4d, k = 2 — la largeur active par jeton est donc 8d).
Trois entraînements mesurent la question centrale de l’apprentissage MoE :
dense : la référence, un FFN unique de largeur 4d ;
MoE, alpha_load = 0 : le routeur livré à lui-même — l’auto-renforcement peut l’effondrer ;
MoE, alpha_load > 0 : la loss d’équilibrage activée.
Puis les deux grandeurs de la section 4 sont mesurées sur les routeurs entraînés : parts des experts et part maximale (la charge), entropie de routage (la décision). Le facteur de capacité est laissé illimité pendant l’entraînement (aucun jeton n’est jeté) : on isole l’effet de alpha_load sur l’équilibre, sans le mélanger à l’effet des jets.
N_MARQUEURS, N_REMPLISSAGE, T =8, 10, 64VOCAB = N_MARQUEURS + N_REMPLISSAGE +1JETON_REQUETE = VOCAB -1def lot(n, T, gen):"""Marqueur en position 0, remplissage, jeton de requete en derniere position.""" marqueur = torch.randint(0, N_MARQUEURS, (n, 1), generator=gen) remplissage = torch.randint(N_MARQUEURS, N_MARQUEURS + N_REMPLISSAGE, (n, T -2), generator=gen) requete = torch.full((n, 1), JETON_REQUETE)return torch.cat([marqueur, remplissage, requete], dim=1)class AttentionCausale(nn.Module):"""Attention causale multi-tetes minimale (le sujet ici est le FFN, pas l'attention)."""def__init__(self, d, n_tetes):super().__init__()assert d % n_tetes ==0self.h, self.dh = n_tetes, d // n_tetesself.qkv = nn.Linear(d, 3* d)self.proj = nn.Linear(d, d)def forward(self, x): B, T, d = x.shape q, k, v =self.qkv(x).chunk(3, dim=-1) q = q.view(B, T, self.h, self.dh).transpose(1, 2) k = k.view(B, T, self.h, self.dh).transpose(1, 2) v = v.view(B, T, self.h, self.dh).transpose(1, 2) att = (q @ k.transpose(-2, -1)) / math.sqrt(self.dh) masque = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1) att = att.masked_fill(masque, float("-inf")).softmax(dim=-1) out = (att @ v).transpose(1, 2).reshape(B, T, d)returnself.proj(out)class BlocDense(nn.Module):def__init__(self, d, n_tetes, cache, dropout=0.0):super().__init__()self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)self.attn = AttentionCausale(d, n_tetes)self.mlp = FFNDense(d, cache, dropout)def forward(self, x): x = x +self.attn(self.ln1(x))return x +self.mlp(self.ln2(x)), Noneclass BlocMoE(nn.Module):def__init__(self, d, n_tetes, cache, n_experts=8, topk=2, dropout=0.0):super().__init__()self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)self.attn = AttentionCausale(d, n_tetes)self.moe = MoECouche(d, n_experts=n_experts, topk=topk, cache=cache, dropout=dropout)def forward(self, x): x = x +self.attn(self.ln1(x)) sortie, probs, compte, jetes =self.moe(self.ln2(x))return x + sortie, probsclass PetitLM(nn.Module):def__init__(self, vocab, d, n_tetes, n_couches, ffn="dense", **kw_moe):super().__init__()self.emb = nn.Embedding(vocab, d)if ffn =="dense":self.blocs = nn.ModuleList([BlocDense(d, n_tetes, 4* d) for _ inrange(n_couches)])else:self.blocs = nn.ModuleList([BlocMoE(d, n_tetes, 4* d, **kw_moe)for _ inrange(n_couches)])self.ln_f = nn.LayerNorm(d)self.tete = nn.Linear(d, vocab, bias=False)def forward(self, idx): x =self.emb(idx) probs_routeurs = []for b inself.blocs: x, probs = b(x)if probs isnotNone: probs_routeurs.append(probs)returnself.tete(self.ln_f(x)), probs_routeurs@torch.no_grad()def evaluer(modele, n=512, graine=99):"""Exactitude et perplexite a la position de requete (8 marqueurs -> hasard 0.125, ppl 8).""" x = lot(n, T, torch.Generator().manual_seed(graine)) logits, _ = modele(x) logits = logits[:, -1] perte = F.cross_entropy(logits, x[:, 0]) exactitude = (logits.argmax(-1) == x[:, 0]).float().mean()return exactitude.item(), math.exp(perte.item())def entrainer(ffn, alpha_load=0.0, graine=0, pas=400, d_model=64, n_tetes=4, n_couches=2, batch=32, lr=3e-3, kw_moe=None):"""Entraine une variante et renvoie (exactitude, perplexite, secondes, params, modele).""" torch.manual_seed(graine) modele = PetitLM(VOCAB, d_model, n_tetes, n_couches, ffn=ffn, **(kw_moe or {})) opt = torch.optim.Adam(modele.parameters(), lr=lr) gen = torch.Generator().manual_seed(1234+ graine) debut = time.perf_counter()for _ inrange(pas): x = lot(batch, T, gen) logits, probs_routeurs = modele(x) perte = F.cross_entropy(logits[:, -1], x[:, 0]) aux =0.0if alpha_load >0and probs_routeurs:for probs in probs_routeurs: top_i = probs.topk(2, dim=-1).indices aux = aux + perte_equilibre(probs, top_i) aux = aux /len(probs_routeurs) (perte + alpha_load * aux).backward() opt.step() opt.zero_grad() secondes = time.perf_counter() - debut exactitude, ppl = evaluer(modele)return exactitude, ppl, secondes, sum(p.numel() for p in modele.parameters()), modeleprint("tache prete :", T, "positions, marqueur en 0, requete en", T -1,"| hasard : acc 1/8 = 0.125, ppl 8")
tache prete : 64 positions, marqueur en 0, requete en 63 | hasard : acc 1/8 = 0.125, ppl 8
RESULTATS = {}for nom, ffn, alpha in [("dense 4d", "dense", 0.0), ("MoE a=0", "moe", 0.0), ("MoE a=0.01", "moe", 0.01)]: kw =dict(n_experts=8, topk=2, dropout=0.0) if ffn =="moe"elseNone acc, ppl, sec, params, modele = entrainer(ffn, alpha_load=alpha, graine=0, pas=400, kw_moe=kw) RESULTATS[nom] = (acc, ppl, sec, params, modele)print(f"{nom:>12s} : acc {acc:.3f} | ppl {ppl:5.2f} | {sec:5.1f} s | {params:,d} params")print("\n(entrainements a 400 pas, une graine, pour borner le temps de calcul)")
dense 4d : acc 1.000 | ppl 1.00 | 14.0 s | 102,528 params
MoE a=0 : acc 1.000 | ppl 1.00 | 30.0 s | 566,784 params
MoE a=0.01 : acc 1.000 | ppl 1.00 | 30.2 s | 566,784 params
(entrainements a 400 pas, une graine, pour borner le temps de calcul)
print("etat des routeurs entraines (2 couches, E = 8, charge = parts de premiers choix)\n")print(f"{'variante':>12s} | {'couche':>6s} | {'part max':>9s} | {'entropie':>9s} | "f"{'L_aux (sans a)':>14s}")ETATS = {}for nom in ("MoE a=0", "MoE a=0.01"): modele = RESULTATS[nom][4] x_mes = lot(512, T, torch.Generator().manual_seed(99))with torch.no_grad(): _, probs_routeurs = modele(x_mes)for c, probs inenumerate(probs_routeurs): top_i = probs.topk(2, dim=-1).indices parts = parts_experts(top_i, 8) ETATS[(nom, c)] = (parts.max().item(), entropie_routage(probs), perte_equilibre(probs, top_i).item())print(f"{nom:>12s} | {c:6d} | {parts.max().item():9.3f} | "f"{entropie_routage(probs):9.4f} | {perte_equilibre(probs, top_i).item():14.3f}")print("\nparts par expert, couche 0 :")for nom in ("MoE a=0", "MoE a=0.01"): modele = RESULTATS[nom][4]with torch.no_grad(): _, probs_routeurs = modele(lot(512, T, torch.Generator().manual_seed(99))) parts = parts_experts(probs_routeurs[0].topk(2, dim=-1).indices, 8)print(f" {nom:>10s} :", " ".join(f"{p:.3f}"for p in parts.tolist()))
etat des routeurs entraines (2 couches, E = 8, charge = parts de premiers choix)
variante | couche | part max | entropie | L_aux (sans a)
MoE a=0 | 0 | 0.291 | 1.8236 | 1.278
MoE a=0 | 1 | 0.374 | 1.9035 | 1.290
MoE a=0.01 | 0 | 0.134 | 1.9259 | 0.997
MoE a=0.01 | 1 | 0.164 | 2.0016 | 0.997
parts par expert, couche 0 :
MoE a=0 : 0.284 0.121 0.004 0.000 0.194 0.291 0.015 0.092
MoE a=0.01 : 0.121 0.130 0.129 0.118 0.121 0.134 0.126 0.123
6.1 Lecture
D’abord l’aveu : la tâche ne discrimine pas. Les trois variantes atteignent l’exactitude 1,000 et la perplexité 1,00 — 400 pas suffisent à saturer un problème à 8 classes. À cette échelle, le dense résout la tâche avec 5,5 fois moins de paramètres et 2,1 fois moins de temps d’entraînement : le MoE n’apporte rien ici, et c’est cohérent — son pari est la capacité aux échelles où le dense plafonne, pas la vitesse sur un jouet.
Toute l’information est donc dans les routeurs, et le contraste est net. Sans alpha_load, la charge part de travers : deux experts captent 28 % chacun, un tiers des experts reste sous 2 %, et l’expert 3 est affamé à exactement zéro jeton — la loss auxiliaire grimpe à 1,28. Avec alpha_load = 0,01, les huit parts se resserrent entre 0,118 et 0,134 autour de 1/8, et la loss tombe à 0,997 — le plancher théorique calculé en section 3 est 1 : la loss d’équilibrage a fait son travail jusqu’au fond.
Deux lectures de précaution. L’effondrement observé sans loss est un début d’auto-renforcement, pas la catastrophe : 1,28 contre 8 au pire cas — 400 pas sur une tâche facile ne laissent pas au gagnant le temps de tout prendre. Et l’entropie, elle, bouge à peine (1,82 contre 1,93, pour un maximum de \(\log 8 \approx 2{,}08\)) pendant que la charge se rééquilibre entièrement : c’est la démonstration vivante de la section 4 — l’entropie mesure l’indécision du routeur, pas l’équilibre de la charge, et seule la seconde dit si les experts travaillent.
Exercice 2 — le balayage du facteur de capacité, sur un vrai routeur
La section 2 a balayé le facteur de capacité sur des charges synthétiques (uniforme, effondrée). Le vrai test est sur un routeur entraîné : sa charge n’est ni l’une ni l’autre, et sa forme dépend de alpha_load.
Implémentez mesurer_jetes(modele, facteur) : pour un lot de mesure, compte le nombre total d’affectations jetées par les deux couches MoE du modèle à un facteur de capacité donné (réutilisez MoECouche.capacite_lot et la logique de priorité : premiers choix d’abord). Balayez \(c \in \{1{,}0, 1{,}25, 1{,}5\}\) sur les deux modèles entraînés ci-dessus (MoE a=0 et MoE a=0.01) et interprétez : le modèle équilibré survit-il à \(c = 1{,}0\) ?
Indice : les affectations par expert se recomptent depuis les top-k du routeur, SANS re-executer les experts : charges = bincount des indices selectionnes. Etape 1 : pour chaque couche, probs -> topk -> charges par expert (slot 0 PUIS slot 1). Etape 2 : capacite = cBT*k/E ; jete par expert = max(0, charge - capacite).
def mesurer_jetes(modele, facteur, n=512, graine=99):"""Exercice 2 : total d'affectations jettees par les couches MoE au facteur donne. Renvoie le nombre total d'affectations jetees (somme des deux couches), sans executer les experts. """# TODO etudiantreturnNone# TODO etudiant# -- balayage (a decommenter une fois complet)# for nom in ("MoE a=0", "MoE a=0.01"):# for c in (1.0, 1.25, 1.5):# print(f"{nom:>10s} c={c:4.2f} : {mesurer_jetes(RESULTATS[nom][4], c)} affectations jetees")print("Exercice 2 a completer")
Exercice 2 a completer
Exercice 3 — la spécialisation des experts par marqueur
L’argument pédagogique en faveur du MoE est la spécialisation : des experts qui apprennent des familles de jetons différentes. La tâche jouet permet de le tester proprement : les jetons de remplissage (10 classes), les marqueurs (8 classes) et la requête ont des rôles syntaxiques distincts — un routeur spécialisé devrait les séparer.
Implémentez specialisation(modele) : pour chaque type de position (marqueur pos 0 / remplissage pos 1-62 / requête pos 63), la distribution moyenne des premiers choix sur les experts (couche 0). Renvoyez un dictionnaire {type: liste de E parts}. Question de lecture : la séparation se fait-elle par type de position, par classe de marqueur, ou rien de tout cela à cette échelle ?
Indice : les positions se decoupent directement dans le lot : x[:, 0], x[:, 1:-1], x[:, -1]. Etape 1 : faire passer un lot, extraire probs de la couche 0, top-1 par position. Etape 2 : pour chaque tranche de positions, bincount des experts / nombre de jetons.
def specialisation(modele, n=512, graine=99):"""Exercice 3 : parts de premiers choix par type de position (couche 0). Renvoie {"marqueur": [E parts], "remplissage": [E parts], "requete": [E parts]}. """# TODO etudiantreturnNone# TODO etudiant# -- mesure (a decommenter une fois complet)# spe = specialisation(RESULTATS["MoE a=0.01"][4])# for typ, parts in spe.items():# print(f"{typ:>12s} :", " ".join(f"{p:.3f}" for p in parts))print("Exercice 3 a completer")
Exercice 3 a completer
7. Limites et suite
Ce que ce notebook établit. La couche MoE complète — routeur top-k à porte renormalisée, capacité par expert avec priorité au premier choix, dropout, loss d’équilibrage — vérifiée organe par organe sur des cas à réponse calculable ; le découplage paramètres / calcul mesuré en millisecondes réelles et pas en FLOPs théoriques ; et l’effet d’alpha_load sur l’équilibre d’un routeur entraîné, mesuré en parts de charge et en entropie.
Ce qu’il ne faut pas en conclure. Que le MoE « bat » le dense : la tâche jouet (8 classes, 64 positions) est résolue par les trois variantes, et la comparaison de qualité à cette échelle ne dit rien des échelles où le MoE est réellement utilisé — des milliards de paramètres, où la capacité supplémentaire est le seul moyen de continuer à croître à budget de calcul constant.
Ce que le notebook ne fait pas, et qui conditionne le verdict :
une seule graine par variante, 400 pas — les écarts d’exactitude à cette échelle sont indicatifs ;
l’entraînement se fait sans jeton jeté (capacité illimitée) : l’interaction entre jets et qualité pendant l’apprentissage (le routeur apprend-il à éviter les experts pleins ?) n’est pas mesurée ;
le routage est dense en mémoire : chaque expert est une matrice pleine, aucun expert ne partage de poids — les MoE réels à très grande échelle ajoutent le sharding et la communication, dont le coût domine justement le déploiement multi-machines ;
pas de routeur entraîné au bruit (jitter du Switch), ni de route z-loss (ST-MoE) — deux correctifs connus de l’instabilité du routeur.
Conclusion
La couche MoE résume en une pièce le geste commun aux variantes du 3.4b : payer des paramètres, pas du calcul. Là où GQA partageait des poids pour économiser du cache, le MoE en multiplie pour croître à calcul actif constant — et le banc l’a mesuré dans les deux sens : contre le dense iso-calcul (autant de FLOPs par jeton, une fraction des paramètres) et le dense iso-paramètres (autant de poids, un multiple du calcul).
Mais la mesure centrale du notebook est ailleurs : les paramètres supplémentaires ne sont des paramètres que si quelqu’un les utilise. La section 6 l’a montré en miniature : livré à lui-même, le routeur affame un expert à zéro jeton pendant que deux autres captent près de 30 % chacun ; avec alpha_load = 0,01, la charge revient au plancher théorique (loss auxiliaire 0,997, minimum exact 1) et les huit experts travaillent. La loss d’équilibrage agit sur le gradient, le facteur de capacité borne l’exécution en jetant ce qui déborde — 58 % des affectations au facteur 1 sur un routement effondré, 2 % sur un routement équilibré — et l’entropie et la part maximale mesurent si l’un et l’autre travaillent. Trois instruments, trois étages de défense d’un même équilibre.
La leçon de méthode vaut au-delà du MoE : un composant qui promet un découplage (ici paramètres/calcul) se vérifie en mesurant les deux côtés du découplage — et en vérifiant d’abord que le mécanisme tient ses bornes analytiques (la loss à 1 à l’uniforme, l’entropie à \(\log E\), la capacité à son décompte exact) avant de lui faire confiance sur des chiffres bruités d’entraînement.