Transfert de poids pour le post-entraînement RL en moins de 2 s

Nous avons récemment atteint des mises à jour de paramètres intermachines de 1,3 seconde pour Kimi-K2 (1T de paramètres), transférant les poids de 256 processeurs graphiques (GPU) d'entraînement (BF16) vers 128 GPU d'inférence (FP8).

Dans le réglage fin par apprentissage par renforcement asynchrone, l'entraînement et l'inférence s'exécutent sur des GPU distincts. Après chaque étape d'entraînement, de nouveaux poids doivent être envoyés aux nœuds d'inférence. De nombreux cadres existants prennent plusieurs secondes, voire minutes, pour les modèles à mille milliards de paramètres.

En tirant parti de la communication point à point RDMA, nous parvenons à rendre le transfert de poids extrêmement rapide, sans modifier le moteur d'inférence, tout en facilitant l'écriture et la maintenance du code.

Notre solution repose sur l'ÉCRITURE RDMA, une primitive unilatérale où la source écrit directement dans la mémoire GPU de la destination.

python

def rdma_write(src_ptr, dst_ptr, size, src_mr, dst_mr):
    # Write from local [src_ptr, src_ptr+size) to remote [dst_ptr, dst_ptr+size).
    # src_mr and dst_mr contains Memory Region metadata of both sides.
    ...

Le côté destinataire ne reçoit même pas de notification pour le transfert. Cela nous procure des transferts à faible latence, à débit élevé et sans copie, pilotés par les nœuds d'entraînement sans aucune logique de contrôle sur les nœuds d'inférence.

Collecte de métadonnées – Le contrôleur rassemble les métadonnées de paramètres de tous les GPU d'entraînement et d'inférence.

Calcul du calendrier – Le contrôleur calcule un calendrier de transfert de poids statique, associant quel GPU d'entraînement envoie quel paramètre à quel GPU d'inférence, et dans quel ordre.

Distribution du calendrier – Le contrôleur envoie le calendrier à tous les GPU d'entraînement.

Exécution – Après chaque étape d'entraînement, le contrôleur signale aux GPU d'entraînement de démarrer les transferts.

Une fois le flux de travail de haut niveau défini, le principal défi consiste à exécuter efficacement les transferts de poids à l'échelle du millier de milliards de paramètres. Nous décrivons ici les détails du chemin d'exécution.

Les paramètres lors de l'entraînement sont distribués selon les placements FSDP. À l'aide de full_tensor(), tous les GPU d'un DeviceMesh peuvent reconstruire le paramètre complet et ainsi servir de source pour le transfert de poids.

Plusieurs DeviceMesh disjoints forment un groupe de maillage. Étant donné que les DeviceMesh du même groupe sont disjoints, leurs transferts n'interfèrent pas et peuvent s'exécuter entièrement en parallèle. Entre les groupes de maillage, nous insérons une barrière globale pour imposer un ordre.

Le pipeline de transfert des poids chevauche les copies hôte vers périphérique, l’assemblage des tenseurs, la fusion et la quantification, ainsi que les transferts RDMA à travers les MeshGroups.

Nous traitons le transfert de chaque tenseur de paramètre comme une tâche. Le processus de transfert de poids utilise plusieurs types de sources matérielles, c'est pourquoi nous divisons une tâche de transfert de poids en différentes étapes de pipeline qui se chevauchent dans le temps :

Copie mémoire de l'hôte vers l'appareil (memcpy) — Si FSDP décharge le poids vers le processeur (CPU)

Préparation des paramètres — Reconstruire le poids complet avec full_tensor(), appliquer la fusion de projection, quantifier si nécessaire.

Transfert RDMA — Écriture sans copie dans la mémoire distante du GPU d'inférence

Barrière globale — Une fois que tous les appels full_tensor() sont terminés, synchroniser entre les groupes de maillage à l'aide de GLOO via Ethernet.

Lors de l'implémentation, nous maintenons une file d'attente FIFO de tâches pour chaque étape du pipeline. Chaque fois que la tâche en tête de file termine l'étape, elle est déplacée vers la queue de la file de l'étape suivante.

full_tensor() et d'autres opérations GPU entraînent une utilisation supplémentaire de la mémoire GPU. Pour éviter une erreur de dépassement de mémoire (out of memory), nous lançons l'exécution d'une tâche uniquement si les tâches en cours d'exécution occupent moins de mémoire GPU temporaire qu'un seuil configurable.

Plusieurs choix de conception rendent notre système nettement plus rapide à exécuter et plus facile à maintenir que les solutions open source courantes.

Le transfert des poids basé sur Rank0 regroupe les poids avant de les distribuer au nœud d’inférence.

La communication point à point transfère les poids des deux nœuds d’entraînement vers les interfaces réseau du nœud d’inférence.

Un modèle courant consiste à acheminer tous les paramètres par l'intermédiaire des GPU de rang 0 : rassemblement sur le rang 0 d'entraînement, envoi au rang 0 d'inférence, puis nouvelle diffusion. Cela devient rapidement un goulot d'étranglement, limité par la bande passante PCIe et la carte réseau (NIC) d'un seul GPU (par exemple, 400 Gbps ≈ 50 Go/s).

En revanche, notre configuration point à point permet à chaque GPU d'entraînement d'envoyer des données directement à chaque GPU d'inférence, saturant ainsi toute la structure du réseau plutôt qu'un seul lien.

Le transfert basé sur RPC utilise un tampon côté inférence et des appels de travailleur, tandis que le transfert unilatéral écrit les poids directement sans la participation du travailleur d’inférence.

Certains systèmes reposent sur l'appel de la méthode update_weight() du moteur d'inférence pour chaque tenseur. Cela implique des modifications intrusives du code d'inférence, ainsi qu'une surcharge due aux appels de procédure distante (RPC), à la sérialisation et à la coordination du plan de contrôle.

Grâce à la primitive ÉCRITURE RDMA, nous mettons à jour les poids silencieusement dans la mémoire du GPU d'inférence, sans copies supplémentaires. Aucun message de plan de contrôle ni aucune logique de contrôle par processeur (CPU) n'est impliqué. Aucune modification du moteur d'inférence n'est requise.

Le processus de transfert de poids peut exploiter quatre types de ressources matérielles : (1) Le mouvement des données hôte-appareil (2) Le calcul GPU pour la fusion de projection et la quantification (3) Le réseau RDMA pour le plan de données (4) Ethernet pour le plan de contrôle.

Notre conception divise les tâches de transfert de poids en étapes de pipeline, ce qui permet un chevauchement facile entre différentes ressources matérielles.

Certaines implémentations recalculent un calendrier de transfert à chaque étape d'entraînement, collectant de manière répétée des métadonnées et distribuant des instructions. Cela ajoute une latence inutile au plan de contrôle.

Notre calendrier est calculé une seule fois lors de l'initialisation. Chaque itération d'entraînement reproduit simplement le plan : le contrôleur émet un signal de départ (« go ») et les GPU suivent leurs itinéraires préassignés. L'exécution est prévisible et légère.

Il est tentant de regrouper tout le processus de mise à jour des poids dans une seule fonction monolithique : collecte de métadonnées, correspondance de noms, rassemblement intra-nœud, fusion de projection, quantification, sous-découpage du monde de communication, transfert réseau inter-nœuds. C'est difficile à programmer correctement, et encore plus difficile à optimiser.

Dans notre démarche d'ingénierie, nous séparons ces étapes en composants individuels. Chaque composant peut être testé unitairement, analysé et optimisé de manière isolée.

Un transfert de poids rapide et fiable est un élément constitutif essentiel pour le réglage fin par apprentissage par renforcement à grande échelle. En combinant la primitive ÉCRITURE RDMA, un calendrier de transfert statique et une exécution en pipeline, nous avons réduit les mises à jour à l'échelle du millier de milliards de paramètres à seulement 1,3 seconde sur Kimi-K2. Cette approche est simple à appréhender, facile à maintenir et évite les goulots d'étranglement des conceptions traditionnelles.

BibTeX

@online{perplexity2025WeightTransferForRlPostTrainingInUnder2Seconds,
  author = {{Perplexity} and Lequn Chen and Erin McKnight and Vivian Wang},
  title = {{Weight transfer for RL post-training in under 2 seconds}},
  date = {2025-09-24},
  url = {https://www.perplexity.ai/hub/blog/weight-transfer-for-rl-post-training-in-under-2-seconds}
}

ÉCRITURE RDMA : transferts unilatéraux

Flux de travail de haut niveau

Exécution du transfert de poids

DeviceMesh et groupes de maillage

Pipeline de tâches

Contrôle de l'utilisation de la mémoire GPU

Pourquoi c'est rapide et simple

Communication point à point

Transfert de données unilatéral

Création de pipelines

Calendrier statique

Séparation propre

Conclusion

Citation