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

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

Lors du 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, les nouveaux poids doivent être transférés vers les nœuds d'inférence. De nombreux frameworks existants nécessitent plusieurs secondes, voire plusieurs minutes, pour les modèles à mille milliards de paramètres.

En tirant parti de la communication point à point RDMA, nous sommes en mesure d'accélérer considérablement le transfert de poids, sans modifier le moteur d'inférence, et de rendre le code plus facile à écrire et à maintenir.

Notre solution repose sur RDMA WRITE, une primitive unilatérale par laquelle la source écrit directement dans la mémoire GPU de 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é destination ne sera même pas notifié du transfert. Cela nous permet d'obtenir des transferts à faible latence, haut débit et zéro 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.

Planification du calcul – Le contrôleur calcule un calendrier de transfert de poids statique, en mappant 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 commencer les transferts.

Avec le flux de travail de haut niveau défini, le défi clé est de savoir comment 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. En utilisant full_tensor(), tous les GPU dans un DeviceMesh peuvent reconstruire le paramètre complet, et donc tous peuvent servir de source pour le transfert de poids.

Plusieurs DeviceMeshes disjoints forment un groupe de maillage. Parce que les DeviceMeshes dans le 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 de l’hôte vers le périphérique, l’assemblage des tenseurs, la fusion et la quantification, ainsi que les transferts RDMA à travers MeshGroups.

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

Memcpy hôte vers périphérique — Si FSDP décharge le poids vers le 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 zéro copie vers la mémoire GPU d'inférence distante

Barrière globale — Une fois tous les appels full_tensor() terminés, synchroniser à travers les groupes de maillage en utilisant GLOO via Ethernet.

Dans la mise en œuvre, 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 d'attente de l'étape suivante.

full_tensor() et d'autres opérations GPU introduisent une utilisation supplémentaire de la mémoire GPU. Pour éviter une erreur de mémoire insuffisante, nous ne commençons l'exécution d'une tâche que si les tâches en cours 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 de poids basé sur Rank0 collecte les poids avant de les distribuer au nœud d’inférence.

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

Un modèle courant consiste à canaliser tous les paramètres via des GPU de rang 0 : rassembler sur le rang 0 d'entraînement, envoyer vers le rang 0 d'inférence, puis redistribuer. Cela devient rapidement un goulot d'étranglement, limité par la bande passante PCIe d'un seul GPU et la carte réseau (par exemple, 400 Gbps ≈ 50 Go/s).

En revanche, notre configuration point à point permet à chaque GPU d'entraînement d'envoyer directement vers chaque GPU d'inférence, saturant ainsi toute la structure réseau plutôt qu'une seule liaison.

Le transfert basé sur RPC utilise un tampon côté inférence et des appels de workers, tandis que le transfert unilatéral écrit les poids directement, sans participation des workers 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 signifie des changements intrusifs dans le code d'inférence, plus la surcharge des RPC, de la sérialisation et de la coordination du plan de contrôle.

Avec la primitive RDMA WRITE, nous mettons à jour les poids silencieusement dans la mémoire GPU d'inférence, sans copie supplémentaire. Aucun message du plan de contrôle et aucune logique de contrôle CPU n'est impliquée. Aucune modification du moteur d'inférence n'est requise.

Le processus de transfert de poids peut tirer parti de quatre types de ressources matérielles : (1) Mouvement de données hôte-périphérique (2) Calcul GPU pour la fusion de projection et la quantification (3) 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, permettant un chevauchement facile entre différentes ressources matérielles.

Certaines implémentations recalculent un calendrier de transfert à chaque étape d'entraînement, collectant et distribuant à plusieurs reprises les instructions. Cela ajoute une latence inutile au plan de contrôle.

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

Il est tentant d'entremêler tout le processus de mise à jour des poids dans une 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 ingénierie, nous séparons ces étapes en tant que composants individuels. Chaque composant peut être testé unitairement, raisonné et optimisé isolément.

Un transfert de poids rapide et fiable est un élément constitutif essentiel pour le réglage fin RL à grande échelle. En combinant la primitive RDMA WRITE, un calendrier de transfert statique et une exécution en pipeline, nous avons réduit les mises à jour de mille milliards de paramètres à seulement 1,3 seconde sur Kimi-K2. L'approche est simple à comprendre, 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}
}

RDMA WRITE : 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

Pipelining

Calendrier statique

Séparation propre

Conclusion

Citation