Aller au contenu

Matériel et parallélismes

Environ 40 % du rapport Kimi K3 traite d'infrastructure. Ce chapitre donne le vocabulaire indispensable pour le lire.

Anatomie d'un GPU

Un GPU d'entraînement (NVIDIA H100/H200, ou H20 dans le cas des évaluations Kimi K3) se résume, pour notre propos, à quatre éléments :

Élément Rôle Ordre de grandeur (H100)
SM (Streaming Multiprocessor) Unité de calcul indépendante ~130 par GPU
Tensor Cores Unités dédiées aux multiplications matricielles ~1 000 TFLOPS en BF16
HBM Mémoire principale du GPU 80 Go, ~3 To/s
Mémoire partagée / SMEM Mémoire ultra-rapide locale à un SM ~228 Ko par SM

Les deux règles à retenir

  1. Les Tensor Cores sont ~10× plus rapides que le calcul générique. Une opération qui ne se formule pas comme une multiplication matricielle gaspille l'essentiel du GPU. C'est exactement l'argument de Kimi K3 pour borner la décroissance de KDA : cela permet à toutes les tuiles de passer par les Tensor Cores.
  2. Le calcul est bon marché, la mémoire est chère. Le rapport arithmétique-sur-mémoire d'un H100 dépasse 300 : lire un octet coûte autant que 300 opérations flottantes. Toute optimisation sérieuse consiste à déplacer moins de données, pas à calculer moins.

La fusion de noyaux (kernel fusion)

Un noyau (kernel) est un programme exécuté sur le GPU. Enchaîner cinq noyaux, c'est écrire et relire cinq fois les données en HBM. Les fusionner en un seul permet de garder les valeurs intermédiaires en mémoire partagée.

Kimi K3 le fait systématiquement. Exemple emblématique : le noyau de décodage KDA fusionne dans une seule boucle la convolution courte, la normalisation d'entrée, le gating, la récurrence KDA et la normalisation de sortie.

Les cinq parallélismes

Un modèle de 2,8 T de paramètres ne tient sur aucun GPU. Il faut le découper. Cinq découpes existent, et Kimi K3 les combine toutes les cinq.

1. Parallélisme de données (DP)

Chaque GPU détient une copie complète du modèle et traite un lot différent. Les gradients sont moyennés par une communication all-reduce.

  • ✅ Simple, efficace.
  • ❌ Impossible ici : une copie complète ne tient pas.

ZeRO corrige cela en partitionnant ce qui est redondant :

Étage Ce qui est partitionné Économie
ZeRO-1 États de l'optimiseur ~4×
ZeRO-2 + gradients ~8×
ZeRO-3 + paramètres ~\(N_{\text{GPU}}\)×

Kimi K3 utilise ZeRO-1 pour les états d'optimiseur et un ZeRO-2 de pipeline pour les gradients, ces derniers étant en outre stockés en mémoire CPU.

2. Parallélisme tensoriel (TP)

Chaque matrice est découpée entre plusieurs GPU. Une matrice \(7168 \times 33792\) répartie sur 8 GPU devient huit tranches \(7168 \times 4224\).

  • ✅ Réduit mémoire et calcul par GPU.
  • ❌ Nécessite une communication all-reduce à chaque couche — donc seulement viable à l'intérieur d'un nœud, sur des liens rapides (NVLink).

3. Parallélisme de pipeline (PP)

Les couches sont réparties : GPU 1 traite les couches 1–12, GPU 2 les couches 13–24, etc.

  • ✅ Communication minime : seules les activations de frontière transitent.
  • ❌ Les bulles : au démarrage, le GPU 2 attend le GPU 1.

On atténue les bulles en découpant le lot en micro-lots qui se suivent dans le pipeline (1F1B, interleaved). Kimi K3 utilise un ordonnancement 1F1B entrelacé avec des étages virtuels (VP).

Une astuce élégante de Kimi K3

Les bulles restantes ne sont pas perdues : le calcul de l'encodeur visuel y est glissé. Les passes avant du ViT des premiers micro-lots sont exécutées en amont, le reste est ordonnancé dans les bulles. Résultat : le coût de l'encodeur visuel est « largement éliminé ».

4. Parallélisme d'experts (EP)

Spécifique au MoE : les experts sont répartis entre GPU. Chaque jeton doit voyager vers les GPU hébergeant ses 16 experts, puis revenir. Deux communications all-to-all par couche.

Le problème central

Le routage est déséquilibré : certains GPU reçoivent bien plus de jetons que d'autres. Tous attendent le plus chargé. De plus, les formes des tenseurs varient d'un pas à l'autre, ce qui fragmente la mémoire et force une synchronisation hôte–GPU à chaque couche pour connaître les tailles.

C'est exactement ce que résout MoonEP, par un équilibrage parfait avec experts redondants, qui rend les formes statiques et supprime la synchronisation.

5. Parallélisme de contexte (CP)

La séquence est découpée entre GPU : le GPU 1 traite les jetons 1–250 000, le GPU 2 les jetons 250 001–500 000, etc. Indispensable pour entraîner à 1 M de jetons.

  • Pour l'attention softmax : il faut échanger des blocs clé–valeur dont la taille croît avec la séquence (méthode Ring Attention).
  • Pour l'attention linéaire : il suffit de transmettre l'état récurrent, de taille fixe. Bien meilleur.

Pourquoi KDA complique quand même les choses

Pour une attention linéaire additive, l'état d'un segment se calcule localement à partir de zéro, puis on somme. Mais KDA applique un opérateur \(\mathbf{M}_t\) à l'état entrant : l'effet d'un segment dépend de ce qui le précède.

KDA Context Parallelism (KCP) résout cela en transmettant deux quantités par rang : la transition cumulée \(\mathbf{M}\) et l'état généré depuis zéro \(\widetilde{\mathbf{S}}\). Ces deux quantités se composent de façon associative, donc un prefix scan reconstitue tous les états entrants avec un seul all-gather de taille fixe. Détail en KDA Context Parallelism.

Les communications collectives

Opération Ce qu'elle fait
all-reduce Chaque rang obtient la somme des valeurs de tous les rangs
all-gather Chaque rang obtient la concaténation des valeurs de tous
reduce-scatter Somme, puis chaque rang reçoit une tranche du résultat
all-to-all Chaque rang envoie un morceau différent à chaque autre rang
P2P Communication directe entre deux rangs

À retenir

Ces opérations sont bloquantes par nature : tout le monde attend le plus lent. D'où l'obsession du rapport pour le recouvrement (overlap) — lancer une communication puis calculer autre chose pendant qu'elle se déroule.

Kimi K3 va jusqu'à remplacer un all-gather global par des échanges P2P ciblés pour l'orthogonalisation de Muon : chaque rang ne récupère que les tranches des paramètres dont il est propriétaire, ce qui supprime le tampon de paramètres complet.

Les sandboxes : le matériel côté environnement

Le RL agentique demande un cinquième type d'infrastructure : des environnements d'exécution isolés. Un agent qui écrit du code doit pouvoir l'exécuter, échouer, planter, sans compromettre la machine hôte.

Technologie Isolation Latence de démarrage Fidélité
Conteneur (Docker) Partage le noyau Linux ~100 ms Limitée
microVM (Firecracker) Noyau séparé ~125 ms Quasi complète
VM classique Complète secondes Complète

Kimi K3 utilise AgentENV, un système à base de microVM Firecracker. Justification donnée par le rapport : avec des conteneurs classiques, l'équipe a observé « plusieurs paniques noyau et interblocages causés par des opérations non intentionnelles des agents ».

Le chiffre marquant

Au cours de l'entraînement et de l'évaluation de Kimi K3, 51 219 741 sandboxes ont été créées, à partir de 1 505 678 images distinctes. Détail en AgentENV.

Vérification de compréhension

Pourquoi combiner cinq parallélismes plutôt qu'en choisir un ?

Parce qu'aucun ne suffit seul et que chacun a un coût différent. TP est rapide mais très communicant : on le limite à l'intérieur d'un nœud. PP communique peu mais crée des bulles : on l'étale entre nœuds. EP est imposé par le MoE. CP est imposé par la longueur. DP+ZeRO absorbe le reste. Le réglage consiste à faire tenir chaque type de communication sur le lien matériel qui lui convient.

Que signifie « les formes de calcul sont statiquement connues » et pourquoi est-ce important ?

Dans un MoE classique, le nombre de jetons reçus par chaque expert varie à chaque pas. Le processeur hôte doit donc interroger le GPU avant de lancer le calcul, ce qui bloque le pipeline entre chaque couche. Avec l'équilibrage parfait de MoonEP, chaque rang reçoit exactement \(S \times K\) jetons : les tailles sont connues à l'avance, la synchronisation disparaît, et le surcoût de lancement des noyaux s'effondre.


Chapitre précédent : Inférence · Chapitre suivant : Vision et multimodalité