Un modèle jouet : nano-K3¶
Pour apprendre l'architecture, il faut la faire tourner. Ce chapitre propose deux configurations réduites qui préservent toutes les propriétés structurelles de Kimi K3 tout en tenant sur un GPU grand public.
Le principe
On réduit les tailles, jamais les mécanismes. Chaque configuration conserve : l'hybride 3:1, NoPE, AttnRes par blocs, LatentMoE avec espace latent à \(0{,}5\times d\), experts partagés, SiTU-GLU, Quantile Balancing.
Si votre nano-K3 apprend, votre implémentation est probablement correcte. Si elle diverge, le bug est dans le mécanisme, pas dans l'échelle.
Les configurations¶
| Paramètre | Kimi K3 | nano-K3 | micro-K3 |
|---|---|---|---|
| Dimension cachée \(d\) | 7 168 | 512 | 256 |
| Couches | 93 | 13 | 5 |
| Composition | 69 KDA + 24 MLA | 9 KDA + 4 MLA | 3 KDA + 2 MLA |
| Motif | 3:1 + 1 MLA finale | 3:1 + 1 MLA finale | 3:1 + 1 MLA finale |
| Têtes | 96 | 8 | 4 |
| Dimension par tête | 128 | 64 | 64 |
| Latent MoE \(\ell\) | 3 584 (0,5 \(d\)) | 256 (0,5 \(d\)) | 128 (0,5 \(d\)) |
| Dim. interne expert | 3 072 | 256 | 128 |
| Experts routés | 896 | 64 | 16 |
| Experts actifs | 16 | 4 | 2 |
| Rareté | 56 | 16 | 8 |
| Experts partagés | 2 | 2 | 2 |
| Bloc AttnRes | 12 | 4 | 2 |
| Couches denses | 1 | 1 | 1 |
| Vocabulaire | 163 840 | 8 192 | 4 096 |
| Contexte | 1 048 576 | 2 048 | 512 |
| \(g_{\min}\) | −5 | −5 | −5 |
| \(\beta_1,\beta_2\) (SiTU) | 4 / 25 | 4 / 25 | 4 / 25 |
| Paramètres totaux | 2,779 T | 190,8 M | 8,5 M |
| Paramètres actifs | 104,0 G | 45,1 M | 4,7 M |
Chiffres calculés avec la même fonction que celle qui reproduit les 2,779 T de Kimi K3 — voir Retrouver les paramètres.
Quelle configuration choisir
- micro-K3 (8,5 M) : pour déboguer. Tourne sur CPU. Un entraînement sur un petit corpus de caractères prend quelques minutes.
- nano-K3 (191 M) : pour valider que l'architecture apprend réellement. Un GPU 16 Go suffit.
Le fichier de configuration¶
{
"model_type": "nano_k3",
"hidden_size": 512,
"num_hidden_layers": 13,
"num_attention_heads": 8,
"vocab_size": 8192,
"max_position_embeddings": 2048,
"first_k_dense_replace": 1,
"intermediate_size": 2048,
"hidden_act": "situ",
"activation_situ_beta": 4.0,
"activation_situ_linear_beta": 25.0,
"num_experts": 64,
"num_experts_per_token": 4,
"num_shared_experts": 2,
"routed_expert_hidden_size": 256,
"moe_intermediate_size": 256,
"moe_router_activation_func": "sigmoid",
"moe_renormalize": true,
"latent_moe_use_norm": true,
"topk_method": "noaux_tc",
"attn_res_block_size": 4,
"mla_use_nope": true,
"mla_use_output_gate": true,
"q_lora_rank": 192,
"kv_lora_rank": 64,
"linear_attn_config": {
"full_attn_layers": [4, 8, 12, 13],
"kda_layers": [1, 2, 3, 5, 6, 7, 9, 10, 11],
"head_dim": 64,
"num_heads": 8,
"short_conv_kernel_size": 4,
"gate_lower_bound": -5.0,
"use_full_rank_gate": true
},
"rms_norm_eps": 1e-5,
"tie_word_embeddings": false
}
Vérifiez le motif
full_attn_layers: [4, 8, 12, 13] — trois blocs de « 3 KDA + 1 MLA »
(couches 1–12), plus la couche 13 en MLA, exactement comme K3 place une
MLA supplémentaire en fin de pile.
Et \(13 = 3\times4 + 1\) : trois blocs AttnRes complets de 4 couches, plus un bloc partiel — la même structure que \(93 = 7\times12 + 9\).
L'ordre de construction et de test¶
Chaque étape se termine par un test qui doit passer avant de continuer.
① Tokeniseur BPE simple (ou réutiliser tiktoken/sentencepiece)
✓ test : encoder puis décoder redonne le texte exactement
② Couche KDA, forme récurrente
✓ test : identité de composition KCP (erreur < 1e-9)
✓ test : beta=0 ⟹ S ← Diag(alpha)·S
✓ test : alpha=1, beta=1 ⟹ relire avec k redonne v
③ Couche Gated MLA, NoPE
✓ test : masque causal — le jeton t ne voit jamais t+1
✓ test : permuter deux jetons change la sortie (via les couches KDA)
④ SiTU-GLU
✓ test : |sortie| ≤ β₁β₂ = 100 sur des entrées extrêmes
✓ test : écart < 1 % avec SwiGLU pour |x| < 0,5
⑤ Stable LatentMoE + routeur
✓ test : les poids de mélange somment à 1
✓ test : le biais ne change PAS les poids de mélange
✓ test : exactement k experts sont appelés par jeton
⑥ Quantile Balancing
✓ test : le déséquilibre relatif décroît sans hyperparamètre
✓ test : biais figé ⟹ routage déterministe et reproductible
⑦ Block AttnRes
✓ test : avec N=1 bloc et poids uniformes, on retrouve un résidu additif
✓ test : les poids d'attention somment à 1
✓ test : le nombre de sources vaut min(bloc_courant + 1, N+1)
⑧ Modèle complet
✓ test : la perte initiale vaut ≈ ln(V) = ln(8192) ≈ 9,01
✓ test : le modèle surapprend un lot unique jusqu'à une perte ≈ 0
✓ test : la perte décroît sur un vrai corpus
Le test le plus révélateur : la perte initiale
Un modèle non entraîné doit produire une distribution uniforme sur le vocabulaire, donc une entropie croisée de \(\ln(V)\).
- \(V = 8192 \Rightarrow \mathcal{L}_0 \approx 9{,}01\)
- \(V = 4096 \Rightarrow \mathcal{L}_0 \approx 8{,}32\)
Si votre perte initiale est très différente, vous avez un problème d'initialisation ou de normalisation — avant même d'entraîner. C'est le test le plus rentable de toute la liste.
Le second test le plus révélateur : le surapprentissage d'un lot
Entraînez sur un seul lot, répété. La perte doit descendre vers 0.
Si elle stagne, le modèle n'a pas assez de capacité ou le gradient ne circule pas correctement — typiquement un problème de résidu, de normalisation, ou de masquage.
Les pièges spécifiques à cette architecture¶
1. Les experts morts
Avec 64 experts et un corpus jouet, plusieurs experts ne recevront jamais de jeton. Instrumentez la charge par expert dès le premier jour :
loads = count_tokens_per_expert(batch)
print(f"min={min(loads)} max={max(loads)} morts={sum(1 for l in loads if l==0)}")
Si des experts sont morts et que QB ne les récupère pas, le bug est dans la mise à jour du biais.
2. La divergence des couches KDA
Symptôme : la perte explose vers NaN après quelques centaines de pas.
Causes possibles, par ordre de fréquence : - \(\mathbf{k}\) non normalisé L2 ; - décroissance non bornée (avoir oublié \(g_{\min}\)) ; - ordre inversé entre décroissance et effacement delta.
3. AttnRes qui n'apprend rien
Si les poids \(\alpha_{i\to l}\) restent quasi uniformes tout l'entraînement, vérifiez que la RMSNorm porte bien sur les clés et que les pseudo-requêtes \(\mathbf{w}_l\) sont initialisées avec une échelle raisonnable (trop petite, le softmax est plat ; trop grande, il sature).
4. La perte du canal think
Si vous implémentez le format XTML, masquez la perte sur les jetons d'entrée (utilisateur, résultats d'outils, bloc JSON de repli) et ne la calculez que sur les jetons produits par l'assistant.
Ce que nano-K3 ne permet PAS de tester¶
Les limites de l'exercice
- La stabilité à l'échelle. SiTU-GLU, la RMSNorm du LatentMoE et Per-Head Muon existent pour des problèmes qui n'apparaissent qu'à 2,8 T de paramètres. À 191 M, SwiGLU marcherait aussi bien.
- Le contexte long. Un contexte de 2 048 ne teste pas les mécanismes conçus pour 1 M.
- Les gains d'infrastructure. MoonEP, KCP, le cache de préfixe hybride n'ont aucun sens sur un seul GPU.
- La qualité. Un modèle de 191 M entraîné sur un petit corpus ne dira rien de la qualité de Kimi K3.
nano-K3 valide la correction de votre implémentation, pas les choix de conception de Kimi K3.
Chapitre précédent : Implémenter AttnRes et LatentMoE · Chapitre suivant : Feuille de route de réimplémentation