Mémoire et parallélismes¶

La combinaison retenue¶

Le pré-entraînement de Kimi K3 combine six mécanismes de parallélisme et de partitionnement :

Mécanisme	Rôle
PP avec étages virtuels (VP)	Répartir les couches
EP	Répartir les experts (MoonEP)
ZeRO-1 DP	Partitionner les états de l'optimiseur
Pipeline ZeRO-2	Partitionner les gradients
CP	Partitionner la séquence (KCP)
Réplication des experts partagés	Les experts partagés sont répliqués sur les rangs EP

La communication all-to-all du dispatch et du combine des experts est recouverte par du calcul pour cacher sa latence.

Les trois problèmes critiques¶

Le rapport les énonce explicitement :

les charges de jetons sont déséquilibrées entre rangs EP → traité par MoonEP ;
activations, gradients et états d'optimiseur dépassent le budget mémoire → traité ici ;
le calcul très variable de l'encodeur visuel est exposé sur le chemin critique → traité en Encodeur multimodal.

Le gestionnaire d'activations unifié¶

L'abstraction

Chaque tenseur sauvegardé pour la passe arrière est associé à un backend de stockage enfichable.

Recalcul, quantification et déchargement ne sont plus trois techniques distinctes : ce sont trois politiques de stockage sous la même abstraction, librement composables à la granularité du tenseur.

Les politiques sont déclarées par des annotations légères sur les tenseurs, entièrement découplées du code du modèle.

C'est un choix d'architecture logicielle remarquable : la stratégie mémoire devient une préoccupation orthogonale, réglable sans toucher à la définition du modèle.

Les détails d'implémentation¶

Détail	Description
Granularité du recalcul	Par fonction, ce qui autorise le recalcul inter-couches
Allocation mémoire	Tout sur le flux de calcul principal, dans un pool unique — évite la fragmentation multi-flux et le surcoût côté hôte
Préchargement	Par couche, recouvert par le calcul — surcoût négligeable

La configuration retenue pour K3¶

La plupart des activations : quantification FP8 par blocs, combinée à du déchargement local ou distant.
Les opérateurs élément par élément : configurés en recalcul (leur recalcul est trivial, leur stockage coûteux).

Le MoE économe en mémoire¶

Deux optimisations spécifiques.

1. Réécrire la dépendance du gradient¶

Le problème

Dans l'implémentation MoE native, le calcul du gradient des probabilités permutées dépend de la sortie de la passe avant (output). Il faut donc conserver output jusqu'à la passe arrière.

Inspirés de SonicMoE, l'équipe réécrit ce gradient par une transformation mathématique en une forme qui ne dépend que de l'activation intermédiaire (act_output) et du gradient amont (doutput).

Coût : un calcul élément par élément léger supplémentaire. Gain : output n'a plus besoin d'être stocké.

2. Recalculer le dispatch¶

Dans la passe avant du GEMM groupé, seule l'entrée de l'opération de dispatch est sauvegardée. Pendant la passe arrière, l'entrée du GEMM groupé est récupérée en recalculant le dispatch.

Pourquoi c'est presque gratuit

Ce recalcul introduit de la communication. Mais cette communication peut être recouverte par une partie du calcul arrière du GEMM groupé.

Résultat : cette portion de stockage d'activation est éliminée à coût négligeable.

AttnRes économe en mémoire¶

Block AttnRes crée deux coûts : maintenir vivantes les représentations de blocs, et les transférer entre étages de pipeline. Trois optimisations les annulent.

Optimisation	Effet
Génération unique	La représentation de bloc est générée une fois à la couche frontière et partagée par toutes les couches suivantes, résidant directement sur le GPU
Checkpointing intégral	Le calcul AttnRes est entièrement encapsulé dans un checkpoint, de sorte que l'activation sauvegardée à chaque couche est identique à celle d'une architecture à résidu standard
Communication de pipeline par cache	Seuls les nouveaux blocs sont transférés incrémentalement entre étages, et libérés dès la fin du micro-lot

Le résultat revendiqué

Cette dernière optimisation atteint « la borne inférieure théorique de l'empreinte mémoire ».

Autrement dit : AttnRes est mémoire-neutre. La modification architecturale ne coûte rien en mémoire d'activation par rapport à un résidu classique. C'est une condition de son adoption à cette échelle.

Équilibrer les activations entre rangs PP¶

Le problème

Sous parallélisme de pipeline 1F1B entrelacé, les activations sont inégalement distribuées entre rangs PP à cause de l'échauffement du pipeline : le nombre d'activations résidentes décroît quand le rang PP augmente.

Le premier rang PP est donc bien plus chargé que le dernier, et risque un dépassement mémoire.

Solution : décharger à distance les activations vers la mémoire d'autres rangs PP, en utilisant le Mooncake Transfer Engine, ce qui égalise la mémoire d'activation entre rangs PP.

Une idée non triviale

On n'utilise pas la mémoire CPU ni le disque, mais la mémoire GPU inutilisée des autres étages du pipeline. Le déséquilibre structurel devient lui-même la ressource qui le compense.

Pipeline ZeRO-2 et déchargement CPU¶

Au-delà des activations, les gradients sont partitionnés entre rangs de parallélisme de données (Pipeline ZeRO-2).

De plus, les gradients partitionnés sont stockés en mémoire CPU pour réduire le pic de mémoire GPU, tout en gardant le double tampon de gradient sur le GPU.

Gradients réduits entre rangs DP
        │
        ▼
  double tampon de gradient (GPU)
        │
        ▼
  accumulation dans les partitions CPU

Une conséquence exploitée ailleurs

Ce double tampon de gradient FP32 sur GPU est réutilisé pendant le RL pour héberger temporairement les poids des modèles de référence — voir Infra RL 1 M. C'est un exemple typique de réemploi de mémoire déjà allouée.

L'orthogonalisation de Muon en P2P¶

L'optimiseur distribué partitionne uniformément les paramètres entre rangs DP. Mais l'orthogonalisation de Newton–Schulz exige la matrice complète, d'où une étape de rassemblement avant chaque mise à jour.

L'approche naïve

Effectuer un all-gather sur tout le tampon de paramètres, sur chaque rang. Cela impose une empreinte mémoire substantielle et fait de la communication le goulot principal à l'échelle.

La solution de K3 : chaque rang ne récupère que les partitions des paramètres dont il est propriétaire, par communication peer-to-peer avec les rangs détenteurs correspondants.

	`all-gather` naïf	P2P ciblé
Tampon de paramètres complet	Requis	Éliminé
Volume de communication	Élevé	Réduit
Recouvrement	Difficile	Pipeliné par tampon de morceau de modèle

Communication et calcul sont pipelinés à la granularité des tampons de morceaux de modèle, ce qui cache le surcoût de communication.

Vérification de compréhension¶

Pourquoi la quantification FP8 des activations n'est-elle pas risquée ?

Parce qu'elle est par blocs : chaque bloc de valeurs a son propre facteur d'échelle, ce qui préserve la plage dynamique locale. Et les activations n'ont besoin d'être précises que pour le calcul du gradient — une erreur de quantification s'y comporte comme du bruit, généralement toléré par la descente stochastique. Les poids, eux, sont beaucoup plus sensibles, ce qui explique le recours au QAT pour les quantifier.

Pourquoi ne pas simplement tout recalculer, plutôt que décharger ?

Le recalcul coûte du temps GPU — la ressource la plus chère. Le déchargement coûte de la bande passante PCIe/NVLink, souvent sous-utilisée et recouvrable. Le bon choix dépend du tenseur : les opérateurs élément par élément sont recalculés (trivial), les grosses activations sont déchargées. D'où l'intérêt d'une abstraction qui permet de composer les deux au cas par cas.

Quelle est la contrainte que tous ces mécanismes servent ?

Une seule : faire tenir un entraînement de 2,78 T de paramètres dans la mémoire disponible, sans sacrifier le débit. Chaque octet économisé permet soit un lot plus grand, soit une séquence plus longue, soit moins de GPU.

Chapitre précédent : MoonEP · Chapitre suivant : Encodeur multimodal