Infrastructure pour le RL à un million de jetons¶

Le cadre : la colocalisation¶

Kimi K3 adopte un entraînement RL colocalisé : le déploiement (rollout) et l'entraînement partagent les mêmes GPU.

Le résultat annoncé

Cela garde chaque expérience RL à contexte de 1 M dans quelques centaines de GPU.

C'est un chiffre remarquable. Une architecture désagrégée — flotte d'inférence séparée de la flotte d'entraînement — aurait exigé bien davantage de matériel, dont une partie inactive à tout instant.

À cela s'ajoutent les partial rollouts pour réduire la latence de queue des trajectoires ultra-longues.

Le problème que la colocalisation crée

Une contention mémoire entre :

le cache KV du rollout, qui doit être persisté pour l'itération suivante (à cause du partial rollout) ;
la mémoire nécessaire à l'entraînement.

Ce conflit devient sévère en RL à contexte long : le cache KV d'une seule trajectoire de 1 M de jetons se compte en dizaines de gigaoctets.

1. Le pool de cache KV externe¶

Pourquoi les échecs de cache sont catastrophiques ici¶

Trois facteurs aggravants

À 1 M de contexte et en rollout multi-étapes, un échec de cache de préfixe est extrêmement coûteux — il faut re-prefiller un contexte énorme.
Le partial rollout aggrave le problème au début de chaque itération : de nombreuses requêtes de prefill longues et inachevées de l'itération précédente arrivent toutes en même temps.
Le décodage spéculatif accélère le renouvellement des requêtes dans des intervalles d'appels d'outils relativement fixes, ce qui augmente le brassage des blocs de préfixe.

Ces effets peuvent déclencher des préemptions et abaisser le taux de succès du cache — critique pour le RL à contexte long.

La solution : découpler la rétention de la résidence GPU¶

Une conception en write-back (réécriture différée) :

État du préfixe	Emplacement
Blocs de décodage actifs	Cache KV GPU
Préfixes réutilisables inactifs	Pool de cache KV externe en DRAM CPU — écrits seulement au moment de leur éviction du GPU, et préchargés avant leur prochaine réutilisation

Write-back contre write-through

Une stratégie write-through copierait chaque bloc en CPU dès sa création. La stratégie write-back n'engage DRAM CPU et bande passante de transfert que pour les préfixes qui quittent le chemin de décodage actif.

Gain : on évite les copies CPU redondantes de blocs encore résidents et actifs sur GPU.

Les états KDA suivent

KDA states are offloaded and prefetched together with the corresponding MLA KV cache blocks, keeping their lifecycles aligned.

C'est la conséquence directe de l'architecture hybride : un préfixe n'est réutilisable que si les deux caches sont restaurables ensemble. Les séparer briserait la cohérence. Voir Cache de préfixe hybride.

D'où vient la DRAM¶

Pour fournir assez de DRAM au pool externe, les états d'entraînement (poids du modèle et états de l'optimiseur) sont déchargés sur NVMe après la fin d'une itération d'entraînement. Après une itération de rollout, le pool est libéré pour éviter la contention avec les charges d'entraînement.

Itération d'entraînement terminée
        │
        ▼  poids + états d'optimiseur → NVMe
   DRAM libérée → pool de cache KV
        │
        ▼  phase de rollout
        │
        ▼  rollout terminé → pool libéré
   DRAM rendue → états d'entraînement rechargés

L'idée

La DRAM CPU est traitée comme une ressource alternée entre deux consommateurs qui n'en ont jamais besoin en même temps. C'est du multiplexage temporel appliqué à la mémoire.

2. L'ordonnanceur à auto-throttling¶

Le problème¶

La concurrence fixe ne marche pas

En rollout multi-étapes, les contextes croissent progressivement à mesure que la trajectoire avance.

Une concurrence fixe fondée sur la longueur moyenne de trajectoire complète est difficile à estimer et trop conservatrice au début (quand les contextes sont courts, on sous-utilise le matériel).
Une concurrence trop élevée crée une pression sur le cache KV dans les phases tardives et peut déclencher des préemptions.

La solution¶

Un mécanisme d'auto-throttling à la couche d'ordonnancement des requêtes LLM, qui utilise des signaux d'exécution :

nombre de requêtes actives ;
nombre de requêtes en file ;
taux d'utilisation du cache KV.

pour contrôler dynamiquement combien de requêtes sont envoyées au moteur d'inférence.

Le comportement obtenu

Rollout bien utilisé au début, concurrence réduite à mesure que la pression sur le cache KV monte. Ni sous-saturation, ni surcharge — sans réglage manuel.

3. La réutilisation des tampons de gradient¶

Le problème¶

Le calcul de la perte RL nécessite des passes avant uniquement sur des modèles non-politique, notamment les modèles de référence. Leurs poids sont trop volumineux pour rester résidents sur GPU.

La solution¶

L'astuce

Garder ces poids en mémoire CPU, et les matérialiser seulement quand on en a besoin, en adossant leurs tenseurs de paramètres au stockage du tampon de gradient FP32 du modèle de politique.

Cela réutilise de la mémoire GPU déjà allouée, sans allocation supplémentaire ni fragmentation. Et c'est sûr, parce que ces tampons seront de toute façon écrasés quand les vrais gradients seront calculés plus tard.

Le streaming par morceaux¶

Avec le partitionnement et déchargement ZeRO-2, chaque GPU ne conserve les tampons de gradient que pour deux morceaux VPP en RL.

Les poids de référence sont donc diffusés morceau par morceau dans ces emplacements :

emplacement A : calcul avant en cours sur le morceau n
emplacement B : préchargement du morceau n+1
        ↓ (échange)
emplacement B : calcul avant sur le morceau n+1
emplacement A : préchargement du morceau n+2

Un double tampon classique, qui cache le coût de copie sans augmenter la mémoire GPU.

Récapitulatif¶

Problème	Mécanisme	Ressource échangée
Cache KV persistant entre itérations	Pool externe en DRAM CPU, write-back	DRAM contre re-prefill
DRAM insuffisante	Déchargement des états d'entraînement sur NVMe	Bande passante NVMe contre DRAM
Concurrence de rollout mal calibrée	Auto-throttling par signaux d'exécution	Réactivité contre débit
Modèles de référence trop gros	Adossement aux tampons de gradient	Rien — réemploi pur

Le fil rouge

Aucun de ces mécanismes n'ajoute de matériel. Tous consistent à mieux utiliser ce qui existe déjà : de la DRAM inutilisée, des tampons GPU momentanément vides, des signaux d'exécution ignorés. C'est de l'optimisation d'occupation, pas de la mise à l'échelle.

Vérification de compréhension¶

Pourquoi ne pas simplement recalculer le préfixe au lieu de le mettre en cache ?

Parce qu'un prefill de 1 M de jetons est plusieurs ordres de grandeur plus coûteux qu'un transfert DRAM. Le rapport donne un ordre de grandeur analogue côté production : un préfixe de 400 K jetons pour un incrément de 4 K. Le recalcul serait un gaspillage de 99 % du travail.

Pourquoi les états KDA doivent-ils suivre exactement les blocs MLA ?

Parce qu'un préfixe n'est réutilisable qu'à une frontière où les deux caches sont restaurables. Si l'état KDA était évincé sans son cache MLA correspondant, ou l'inverse, le préfixe deviendrait inutilisable et il faudrait tout recalculer. Aligner les cycles de vie est la seule façon de garantir la cohérence.

Le déchargement des poids sur NVMe ne ralentit-il pas trop ?

Il s'exécute entre les phases, pas pendant. L'alternance entraînement/rollout du RL colocalisé crée des fenêtres où les états d'entraînement ne servent à rien. Le transfert est amorti sur la durée d'une phase de rollout entière, qui se compte en minutes ou en heures.

Chapitre précédent : Encodeur multimodal · Chapitre suivant : Sandboxes AgentENV