Μεταφορά βαρών για RL Post-Training σε λιγότερο από 2 δευτερόλεπτα

Πρόσφατα πετύχαμε ενημερώσεις παραμέτρων μεταξύ μηχανών σε 1,3 δευτερόλεπτα για το Kimi-K2 (1T παράμετροι), μεταφέροντας βάρη από 256 GPU εκπαίδευσης (BF16) σε 128 GPU συμπερασμού (FP8).

Στη λεπτή ρύθμιση ασύγχρονης ενισχυτικής μάθησης, η εκπαίδευση και η συμπερασμός εκτελούνται σε ξεχωριστές GPU. Μετά από κάθε βήμα εκπαίδευσης, τα νέα βάρη πρέπει να προωθηθούν στους κόμβους συμπερασμού. Πολλοί υπάρχοντες σκελετοί χρειάζονται αρκετά δευτερόλεπτα —ή ακόμα και λεπτά— για μοντέλα τρισεκατομμυρίων παραμέτρων.

Αξιοποιώντας την επικοινωνία σημείου-σε-σημείο RDMA, είμαστε σε θέση να κάνουμε τη μεταφορά βαρών εξαιρετικά γρήγορη, χωρίς να αλλάξουμε τη μηχανή συμπερασμού, και να κάνουμε τον κώδικα ευκολότερο στη συγγραφή και τη συντήρηση.

Η λύση μας βασίζεται στο RDMA WRITE, μια πρωτόγονη λειτουργία μιας κατεύθυνσης όπου η πηγή γράφει απευθείας στη μνήμη GPU του προορισμού.

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.
    ...

Η πλευρά του προορισμού δεν θα ειδοποιηθεί καν για τη μεταφορά. Αυτό μας προσφέρει μεταφορές χαμηλής καθυστέρησης, υψηλής απόδοσης και μηδενικής αντιγραφής που οδηγούνται από τους κόμβους εκπαίδευσης χωρίς καμία λογική ελέγχου στους κόμβους συμπερασμού.

Συλλογή μεταδεδομένων – Ο ελεγκτής συγκεντρώνει τα μεταδεδομένα παραμέτρων από όλες τις GPU εκπαίδευσης και συμπερασμού.

Υπολογισμός χρονοδιαγράμματος – Ο ελεγκτής υπολογίζει ένα στατικό πρόγραμμα μεταφοράς βαρών, αντιστοιχίζοντας ποια GPU εκπαίδευσης στέλνει ποια παράμετρο σε ποια GPU συμπερασμού και με ποια σειρά.

Διανομή χρονοδιαγράμματος – Ο ελεγκτής στέλνει το χρονοδιάγραμμα σε όλες τις GPU εκπαίδευσης.

Εκτέλεση – Μετά από κάθε βήμα εκπαίδευσης, ο ελεγκτής δίνει σήμα στις GPU εκπαίδευσης να ξεκινήσουν τις μεταφορές.

Με καθορισμένη τη ροή εργασίας υψηλού επιπέδου, η βασική πρόκληση είναι πώς να εκτελεστούν οι μεταφορές βαρών αποτελεσματικά σε κλίμακα τρισεκατομμυρίων παραμέτρων. Εδώ περιγράφουμε τις λεπτομέρειες της διαδρομής εκτέλεσης.

Οι παράμετροι στην εκπαίδευση κατανέμονται σύμφωνα με τις τοποθετήσεις FSDP. Χρησιμοποιώντας το full_tensor(), όλες οι GPU σε ένα DeviceMesh μπορούν να αναδημιουργήσουν την πλήρη παράμετρο, επομένως όλες μπορούν να λειτουργήσουν ως πηγή για μεταφορά βαρών.

Πολλαπλά αποσυνδεδεμένα DeviceMeshes σχηματίζουν μια ομάδα mesh. Επειδή τα DeviceMeshes στην ίδια ομάδα είναι αποσυνδεδεμένα, οι μεταφορές τους δεν αλληλεπιδρούν και μπορούν να εκτελεστούν πλήρως παράλληλα. Μεταξύ των ομάδων mesh, εισάγουμε ένα καθολικό εμπόδιο για την επιβολή της σειράς.

Ο αγωγός μεταφοράς βαρών επικαλύπτει τις αντιγραφές από τον κεντρικό επεξεργαστή στη συσκευή, τη συναρμολόγηση των tensor, τη συγχώνευση και την ποσοτικοποίηση, καθώς και τις μεταφορές RDMA μεταξύ των MeshGroups.

Αντιμετωπίζουμε τη μεταφορά κάθε τανυστή παραμέτρων ως εργασία. Η διαδικασία μεταφοράς βαρών αξιοποιεί πολλαπλούς τύπους πηγών υλικού, επομένως χωρίζουμε μια εργασία μεταφοράς βαρών σε διαφορετικά στάδια pipeline που επικαλύπτονται χρονικά:

Memcpy από κεντρικό υπολογιστή σε συσκευή — Εάν το FSDP εκφορτώνει βάρος στη CPU

Προετοιμασία παραμέτρων — Αναδημιουργία πλήρους βάρους με full_tensor(), εφαρμογή σύντηξης προβολής, κβαντισμός εάν χρειάζεται.

Μεταφορά RDMA — Εγγραφή μηδενικής αντιγραφής σε απομακρυσμένη μνήμη GPU συμπερασμού

Καθολικό εμπόδιο — Αφού ολοκληρωθούν όλες οι κλήσεις full_tensor(), συγχρονισμός σε ομάδες mesh χρησιμοποιώντας GLOO μέσω Ethernet.

Στην υλοποίηση, διατηρούμε μια ουρά FIFO εργασιών για κάθε στάδιο pipeline. Κάθε φορά που η κεφαλή της εργασίας ουράς ολοκληρώνει το στάδιο, μετακινείται στην ουρά του επόμενου σταδίου.

Το full_tensor() και άλλες λειτουργίες GPU εισάγουν επιπλέον χρήση μνήμης GPU. Για να αποφύγουμε σφάλμα εξάντλησης μνήμης (OOM), ξεκινάμε την εκτέλεση μιας εργασίας μόνο εάν οι τρέχουσες εν εξελίξει εργασίες καταλαμβάνουν λιγότερη προσωρινή μνήμη GPU από ένα ρυθμιζόμενο όριο (watermark).

Καθαρός διαχωρισμός

Η μεταφορά βαρών με βάση το Rank0 συγκεντρώνει τα βάρη πριν τα διανείμει στον κόμβο inference.

Η επικοινωνία σημείο προς σημείο μεταφέρει βάρη από και τους δύο κόμβους εκπαίδευσης στις διεπαφές δικτύου του κόμβου inference.

Ένα κοινό μοτίβο είναι η διοχέτευση όλων των παραμέτρων μέσω GPU rank-0: συγκέντρωση στο rank-0 εκπαίδευσης, αποστολή στο rank-0 συμπερασμού και στη συνέχεια εκ νέου διασπορά. Αυτό γίνεται γρήγορα σημείο συμφόρησης, περιοριζόμενο από το εύρος ζώνης PCIe και το NIC μιας μόνο GPU (π.χ., 400 Gbps ≈ 50 GB/s).

Αντίθετα, η ρύθμιση σημείου-σε-σημείο επιτρέπει σε κάθε GPU εκπαίδευσης να στέλνει απευθείας σε κάθε GPU συμπερασμού, κορεσμό ολόκληρου του δικτύου αντί για έναν μόνο σύνδεσμο.

Η μεταφορά με βάση το RPC χρησιμοποιεί ένα buffer στην πλευρά του inference και κλήσεις worker, ενώ η μονόπλευρη μεταφορά γράφει τα βάρη απευθείας χωρίς τη συμμετοχή του worker του inference.

Ορισμένα συστήματα βασίζονται στην κλήση της μεθόδου update_weight() της μηχανής συμπερασμού για κάθε τανυστή. Αυτό σημαίνει επεμβατικές αλλαγές στον κώδικα συμπερασμού, συν επιπλέον κόστος από RPC, σειριοποίηση και συντονισμό επιπέδου ελέγχου.

Με το πρωτόγονο RDMA WRITE, ενημερώνουμε τα βάρη σιωπηρά στη μνήμη GPU συμπερασμού, χωρίς επιπλέον αντίγραφα. Δεν εμπλέκεται μήνυμα επιπέδου ελέγχου ούτε λογική ελέγχου CPU. Δεν απαιτείται καμία τροποποίηση στη μηχανή συμπερασμού.

Η διαδικασία μεταφοράς βαρών μπορεί να αξιοποιήσει τέσσερις τύπους πόρων υλικού: (1) Κίνηση δεδομένων κεντρικού υπολογιστή-συσκευής (2) Υπολογισμός GPU για σύντηξη προβολής και κβαντισμό (3) Δίκτυο RDMA για επίπεδο δεδομένων (4) Ethernet για επίπεδο ελέγχου.

Ο σχεδιασμός μας χωρίζει τις εργασίες μεταφοράς βαρών σε στάδια pipeline, επιτρέποντας την εύκολη επικάλυψη σε διαφορετικούς πόρους υλικού.

Ορισμένες υλοποιήσεις επανυπολογίζουν ένα πρόγραμμα μεταφοράς σε κάθε βήμα εκπαίδευσης, συλλέγοντας επανειλημμένα μεταδεδομένα και διανέμοντας οδηγίες. Αυτό προσθέτει περιττή καθυστέρηση επιπέδου ελέγχου.

Το χρονοδιάγραμμά μας υπολογίζεται μία φορά κατά την προετοιμασία. Κάθε επανάληψη εκπαίδευσης απλώς επαναλαμβάνει το σχέδιο: ο ελεγκτής εκδίδει ένα σήμα "go" και οι GPU ακολουθούν τις προκαθορισμένες διαδρομές τους. Η εκτέλεση είναι προβλέψιμη και ελαφριά.

Είναι δελεαστικό να μπλέξετε ολόκληρη τη διαδικασία ενημέρωσης βαρών σε μία μονολιθική συνάρτηση: συλλογή μεταδεδομένων, αντιστοίχιση ονομάτων, συγκέντρωση εντός κόμβου, σύντηξη προβολής, κβαντισμός, υποδιαίρεση κόσμου επικοινωνίας, μεταφορά δικτύου μεταξύ κόμβων. Είναι δύσκολο να προγραμματιστεί σωστά και ακόμη πιο δύσκολο να βελτιστοποιηθεί.

Στη μηχανική μας, διαχωρίζουμε αυτά τα βήματα ως μεμονωμένα στοιχεία. Κάθε στοιχείο μπορεί να ελεγχθεί μεμονωμένα, να αξιολογηθεί και να βελτιστοποιηθεί απομονωμένα.

Η γρήγορη και αξιόπιστη μεταφορά βαρών είναι ένα κρίσιμο δομικό στοιχεία για λεπτή ρύθμιση RL μεγάλης κλίμακας. Συνδυάζοντας την πρωτόγονη λειτουργία RDMA WRITE, ένα στατικό πρόγραμμα μεταφοράς και την εκτέλεση pipeline, μειώσαμε τις ενημερώσεις τρισεκατομμυρίων παραμέτρων σε μόλις 1,3 δευτερόλεπτα στο Kimi-K2. Η προσέγγιση είναι απλή στην κατανόηση, εύκολη στη συντήρηση και αποφεύγει τα σημεία συμφόρησης των παραδοσιακών σχεδίων.

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: μεταφορές μιας κατεύθυνσης

Ροή εργασίας υψηλού επιπέδου

Εκτέλεση μεταφοράς βαρών

DeviceMesh και ομάδες Mesh

Pipeline εργασιών

Έλεγχος χρήσης μνήμης GPU

Γιατί είναι γρήγορο και απλό

Επικοινωνία σημείου-σε-σημείο

Μεταφορά δεδομένων μιας κατεύθυνσης

Pipelining

Στατικό χρονοδιάγραμμα

Καθαρός διαχωρισμός

Συμπέρασμα

Παραπομπή