Αποτελεσματική και φορητή επικοινωνία MoE

Παρουσιάζουμε μια βιβλιοθήκη υψηλής απόδοσης, φορητή, ανοιχτού κώδικα για επικοινωνία Mixture-of-Experts (MoE) που επιτυγχάνει 10x ταχύτερη απόδοση σε σύγκριση με τα τυπικά πρωτόκολλα επικοινωνίας All-to-All.

Η υλοποίησή μας διαθέτει διάφορες βασικές τεχνικές καινοτομίες που προσφέρουν ανώτερη αποδοτικότητα επικοινωνίας MoE:

Επικοινωνία ξεκινώμενη από GPU (IBGDA): Υποστηρίζει άμεση επικοινωνία GPU-σε-NIC, μειώνοντας σημαντικά την καθυστέρηση παρακάμπτοντας τη συμμετοχή της CPU

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

Ταχύτερη απόδοση ενός κόμβου: 2.5x χαμηλότερη καθυστέρηση από την προηγουμένως ταχύτερη υλοποίηση σε διαμορφώσεις ενός κόμβου

Αποτελεσματική και φορητή απόδοση πολλαπλών κόμβων: Η υλοποίησή μας επιτυγχάνει ταχύτητες έως και 10x ταχύτερες από την τυπική επικοινωνία all-to-all. Παρόλο που είναι περίπου 2x πιο αργή από εξαιρετικά εξειδικευμένες υλοποιήσεις, η προσέγγισή μας προσφέρει καλύτερη φορητότητα σε διάφορες εκδόσεις NVSHMEM και περιβάλλοντα δικτύου (NVLink, CX-7 και EFA)

Η βιβλιοθήκη είναι πλήρως ανοιχτού κώδικα και διατίθεται στη διεύθυνση https://github.com/ppl-ai/pplx-kernels.

Σε αυτό το άρθρο, εξερευνούμε τις προκλήσεις του παραλληλισμού εμπειρογνωμόνων σε μοντέλα MoE μεγάλης κλίμακας και περιγράφουμε την προσέγγισή μας για την αποτελεσματική αποστολή και συνδυασμό token σε κατανεμημένα περιβάλλοντα GPU.

Τα μοντέλα Mixture-of-Experts (MoE), όπως τα DeepSeek R1 και Mixtral 8x7B, βελτιώνουν τα πυκνά μοντέλα περιορίζοντας τον αριθμό των βαρών που ενεργοποιούνται για κάθε token. Αντί να επιτυγχάνουν υψηλά πλήθη παραμέτρων αυξάνοντας τον αριθμό των επιπέδων ή αυξάνοντας το μέγεθος των γραμμικών επιπέδων στο Multi-Layer Perceptron (MLP) κάθε επιπέδου αποκωδικοποιητή, τα μοντέλα MoE αντικαθιστούν το παραδοσιακό MLP με πολλαπλούς εμπειρογνώμονες και έναν δρομολογητή. Για παράδειγμα, από τις 671B παραμέτρους του DeepSeek R1, μόνο οι 37B πολλαπλασιάζονται με ένα δεδομένο token κατά την εξαγωγή συμπερασμάτων. Αυτό μειώνει σημαντικά το ποσό των υπολογισμών που απαιτούνται για την αποκωδικοποίηση ενός token και βελτιώνει την καθυστέρηση σε σύγκριση με ένα πυκνό μοντέλο.

Τα μοντέλα MoE παρουσιάζουν ορισμένες πρόσθετες προκλήσεις για την εξαγωγή συμπερασμάτων, σε σύγκριση με τα πυκνά μοντέλα. Ενώ οι ίδιοι οι εμπειρογνώμονες είναι μικρά επίπεδα MLP, κάθε επίπεδο αποκωδικοποιητή περιλαμβάνει έναν δρομολογητή που αποφασίζει σε ποιούς εμπειρογνώμονες αποστέλλεται ένα token, με κάθε token να αποστέλλεται σε πολλαπλούς εμπειρογνώμονες. Ο δρομολογητής είναι συνήθως ένα μικρό γραμμικό επίπεδο που παράγει μια κατανομή πιθανότητας. Συνήθως επιλέγονται οι εμπειρογνώμονες με τα υψηλότερα σκορ K και η τελική ενεργοποίηση υπολογίζεται ως ένας σταθμισμένος μέσος όρος, αθροίζοντας τα αποτελέσματα των εμπειρογνωμόνων πολλαπλασιασμένα με ένα βάρος που προέρχεται από την κατανομή πιθανότητας του δρομολογητή.

Για να ελαχιστοποιηθεί η καθυστέρηση σε κατανεμημένα συστήματα MoE, ο παραλληλισμός μπορεί να αξιοποιηθεί σε πολλαπλές συσκευές, αλλά αυτό εισάγει προκλήσεις επικοινωνίας. Μοντέλα όπως το Mixtral 8x7B ή το Llama-70B χωρούν σε 8 συσκευές σε έναν μόνο κόμβο, επωφελούμενα από γρήγορες διασυνδέσεις NVLink χαμηλής καθυστέρησης (έως 900Gbps). Ωστόσο, τα μεγαλύτερα μοντέλα απαιτούν κατακερματισμό εμπειρογνωμόνων σε πολλαπλούς κόμβους χρησιμοποιώντας InfiniBand (με κορύφωση στα 400Gbps), γεγονός που εισάγει πρόσθετες προκλήσεις καθυστέρησης.

Δρομολόγηση token, εξειδικευμένος υπολογισμός και συλλογή εξόδου σε ένα επίπεδο mixture-of-experts.

Σε αυτό το άρθρο, εξερευνούμε το πρόβλημα του παραλληλισμού εμπειρογνωμόνων, περιγράφοντας την υλοποίησή μας πυρήνων GPU που αποστέλλουν αποτελεσματικά token σε συσκευές που φιλοξενούν τους ανατεθειμένους εμπειρογνώμονες και τα συλλέγουν πίσω υπολογίζοντας το σταθμισμένο άθροισμα των αποτελεσμάτων των εμπειρογνωμόνων. Ενώ υπάρχουν επιθετικά βελτιστοποιημένες υλοποιήσεις που στοχεύουν σε συγκεκριμένο υλικό δικτύου, παρουσιάζουμε μια εξαιρετικά φορητή υλοποίηση που βασίζεται σε ένα ελάχιστο σύνολο πρωτοκόλλων NVSHMEM, παρέχοντας παράλληλα απόδοση 10x ταχύτερη από την τυπική επικοινωνία all-to-all. Η υλοποίησή μας επιτυγχάνει κορυφαία απόδοση σε διαμορφώσεις ενός κόμβου, διατηρώντας παράλληλα εξαιρετική φορητότητα σε διάφορα περιβάλλοντα δικτύου.

Για αποδοτικό inference, τα βάρη ενός μοντέλου πρέπει να διατηρούνται στη μνήμη της συσκευής, αφήνοντας ταυτόχρονα επαρκή χώρο για activations, KV caches και άλλους buffers που απαιτούνται από το forward pass μέσω του μοντέλου. Τα πιο ικανά μοντέλα υπερβαίνουν τη χωρητικότητα ακόμη και των πιο ισχυρών GPUs, επομένως το inference πρέπει να κατανέμεται σε πολλαπλές συσκευές, οι οποίες μπορούν συλλογικά να αποθηκεύσουν τα βάρη του μοντέλου. Με βάση τα σχήματα sharding των βαρών, πρέπει να χρησιμοποιούνται διαφορετικά σχήματα επικοινωνίας και υπολογισμού για να αποφορτώνεται ο υπολογισμός και να συγχρονίζονται οι συσκευές.

Ο παραλληλισμός εμπειρογνωμόνων (Expert Parallelism - EP), που απεικονίζεται στην πρώτη εικόνα, παραλληλίζει μόνο τον υπολογισμό των εμπειρογνωμόνων. Διαφορετικοί εμπειρογνώμονες ανατίθενται σε διαφορετικές συσκευές, οι οποίες κρατούν τα βάρη τους. Μετά τη δρομολόγηση, τα token αποστέλλονται στην αντίστοιχη συσκευή, με τα αποτελέσματα να συγκεντρώνονται και να συσσωρεύονται στη συνέχεια. Η πολυπλοκότητα της δρομολόγησης εξαρτάται από τον βαθμό παραλληλισμού στα άλλα μέρη του μοντέλου: η αναπαραγωγή άλλων επιπέδων θα μπορούσε να εξαλείψει εντελώς την ανάγκη για δρομολόγηση, καθώς κάθε rank μπορεί να επιλέξει τα token από έναν τοπικά αναπαραγόμενο πίνακα δρομολόγησης. Ωστόσο, εάν μόνο ένα από τα rank εκτελεί δρομολόγηση, απαιτείται μια εκπομπή (broadcast) για την αποστολή token, δεικτών και βαρών στους αντίστοιχους εμπειρογνώμονες. Τέλος, μια εκπομπή all-gather ή all-to-all συγχρονίζει τα token εξόδου με όποιο rank συνεχίζει την εκτέλεση του μοντέλου. Μια τέτοια υλοποίηση είναι σχετικά απλή, καθώς το torch εκθέτει ήδη τα απαιτούμενα πρωτόκολλα, αν και ορισμένα οφέλη θα μπορούσαν να προκύψουν από τη συγχώνευση της συσσώρευσης με τη λειτουργία gather που συλλέγει τα token από τους εμπειρογνώμονες.

Συνδυασμός παραλληλισμού experts με παραλληλισμό tensor σε δύο GPU.

Ο παραλληλισμός μόνο για εμπειρογνώμονες δεν κλιμακώνεται ιδανικά, καθώς οι κόμβοι σε ένα σύμπλεγμα (cluster) ενδέχεται να είναι ανενεργοί ενώ το μοντέλο εκτελεί επίπεδα εκτός εμπειρογνωμόνων, όπως προσοχή (attention), κανονικοποίηση (norm) και δειγματοληψία. Ωστόσο, ο υπολογισμός αυτών των επιπέδων, κυρίως της προσοχής, μπορεί επίσης να επωφεληθεί από τον Παραλληλισμό Ταινιών (Tensor Parallelism - TP). Τα περισσότερα μοντέλα βασίζονται στην προσοχή πολλαπλών κεφαλών (multi-head attention), πράγμα που σημαίνει ότι οι κεφαλές προσοχής και οι αντίστοιχες προβολές Q, K και V μπορούν επίσης να κατακερματιστούν σε συσκευές, αναπαράγοντας ή συγκεντρώνοντας τις φέτες μεταξύ διαφόρων επιπέδων. Εάν η προσοχή κατανέμεται σε όλες τις συσκευές, ένα all-gather μπορεί να συγχρονίσει τις ενεργοποιήσεις, επιτρέποντας στη δρομολόγηση να αναπαραχθεί, απαιτώντας πρωτόκολλα συγχρονισμού παρόμοια με την περίπτωση του παραλληλισμού εμπειρογνωμόνων για μια αποτελεσματική υλοποίηση. Ωστόσο, υπάρχουν όρια στον παραλληλισμό σε αυτό το επίπεδο, καθώς η μείωση του αριθμού των κεφαλών προσοχής κάτω από ένα ορισμένο όριο θα αποφέρει φθίνουσες αποδόσεις.

Συνδυασμός παραλληλισμού experts με παραλληλισμό δεδομένων για ξεχωριστά αιτήματα σε δύο GPU.

Για την καλύτερη αξιοποίηση όλων των συσκευών και την υποστήριξη ενός πολύ υψηλού βαθμού παραλληλισμού εμπειρογνωμόνων, έως 128 ή 256 GPU, απαιτείται Παραλληλισμός Δεδομένων (Data Parallelism - DP). Σύμφωνα με αυτό το σχήμα, οι συσκευές στις οποίες χωρίζεται ένα μοντέλο ομαδοποιούνται για να χειριστούν αιτήματα ταυτόχρονα, υπολογίζοντας την προσοχή και διατηρώντας κνέσιες KV κατακερματισμένες στην τοπική τους ομάδα. Πολλαπλές περιπτώσεις αυτών των παράλληλων ομάδων συνεργάζονται στην αξιολόγηση εμπειρογνωμόνων, με καθεμία να φιλοξενεί ένα διαφορετικό υποσύνολο των βαρών των εμπειρογνωμόνων. Με βάση τον αριθμό των κεφαλών προσοχής, μια ομάδα μπορεί τυπικά να εκτείνεται έως το μέγεθος ενός ολόκληρου κόμβου, καθώς η ενδο-κόμβου επικοινωνία είναι ταχύτερη. Για παράδειγμα, στο παραπάνω σχήμα, ένα rank DP εξυπηρετεί ανεξάρτητα δύο αιτήματα, Α και Β, χειριζόμενο την προσοχή, την κανονικοποίηση και τυχόν άλλη λογιστική παρακολούθηση για τα αιτήματα. Το άλλο rank DP επεξεργάζεται ένα ξεχωριστό αίτημα C. Ωστόσο, ο πρώτος κόμβος φιλοξενεί τους μισούς εμπειρογνώμονες, ενώ ο άλλος κόμβος τους άλλους μισούς, οπότε μετά τη δρομολόγηση τα token από τα Α και Β ενδέχεται να αποσταλούν στον δεύτερο κόμβο και το αντίστροφο. Αυτό οδηγεί σε ένα πρόβλημα αραιής επικοινωνίας: κάθε συσκευή ενδέχεται να στείλει διαφορετικό αριθμό token σε οποιονδήποτε άλλο προορισμό rank. Τα υπάρχοντα πρωτόκολλα από το torch, κυρίως το all_to_all, δεν είναι ιδιαίτερα κατάλληλα, καθώς ενδέχεται να απαιτούν κάποια μορφή padding ή συγχρονισμό GPU-σε-CPU και εκπομπή μεταδεδομένων. Για την αποτελεσματική υλοποίηση της επικοινωνίας, απαιτούνται προσαρμοσμένοι πυρήνες για την επεξεργασία των πληροφοριών δρομολόγησης και την έναρξη επικοινωνίας από τις συσκευές χωρίς παρέμβαση της CPU. Μετά τη δρομολόγηση, ένας πυρήνας αποστολής (dispatch) πρέπει να στείλει τα token στα rank στα οποία δρομολογήθηκαν, ενώ στην πλευρά του συνδυασμού (combine) πρέπει να συλλεχθούν οι ενεργοποιήσεις που ανήκουν στα αιτήματα της τρέχουσας ομάδας DP. Επιπλέον, η εργασία πρέπει να εξισορροπείται εντός μιας ομάδας DP που μπορεί να έχει πολλαπλές συσκευές, προκειμένου να ελαχιστοποιηθεί η εκπομπή και να αναπαραχθούν σωστά τα token σε όλα τα rank για τα επόμενα επίπεδα.

Το NVSHMEM είναι μια υλοποίηση OpenSHMEM ειδική για NVIDIA, η οποία παρέχει φορητές διευκολύνσεις επικοινωνίας μεταξύ συσκευών που αφαιρούν την πολυπλοκότητα του υποκείμενου υλικού. Το API μπορεί να εκφράσει αναγνώσεις και εγγραφές από συσκευή σε συσκευή, οι οποίες αντιστοιχίζονται στα πρωτόκολλα επιμέρους επιπέδων μεταφοράς. Υποστηρίζονται τόσο το NVLink όσο και το RDMA, παρέχοντας έναν βαθμό φορητότητας και επιτρέποντας στους πυρήνες NVSHMEM να λειτουργούν εντός ενός κόμβου ή σε ολόκληρο ένα σύμπλεγμα πολλαπλών κόμβων. Στους πυρήνες μας, χρησιμοποιούμε απομακρυσμένες εγγραφές μνήμης και λειτουργίες σηματοδότησης (signaling) προκειμένου να μεταφέρουμε δεδομένα και να συγχρονίσουμε τις συμμετέχουσες συσκευές.

Οι λειτουργίες NVSHMEM είναι χτισμένες γύρω από την έννοια της συμμετρικής μνήμης: λειτουργούν σε buffer που έχουν δεσμευτεί σε όλες τις συσκευές που συμμετέχουν στην επικοινωνία μεταξύ συσκευών. Μια τοπική συσκευή μπορεί να χρησιμοποιήσει μια διεύθυνση που προέρχεται από το δικό της τοπικό buffer μαζί με τον δείκτη του rank προορισμού για να καθορίσει τον προορισμό μιας λειτουργίας. Η παρακάτω εικόνα απεικονίζει αυτή την έννοια: και οι δύο GPU δεσμεύουν συμμετρικά buffer του ίδιου μεγέθους, διατηρώντας δείκτες src προς αυτά. Το πρώτο rank θέλει να στείλει 3 ακέραιους αριθμούς στον δεύτερο, τοποθετώντας τους στην αρχή του buffer: η συνάρτηση nvshmem_int_put_nbi εξάγει τη διεύθυνση εκκίνησης από το τοπικό buffer, καθορίζοντας τη συσκευή προορισμού. Το δεύτερο rank εξάγει μια μετατόπιση (offset) από το δικό του buffer, στέλνοντας ένα στοιχείο στην πρώτη συσκευή, εφαρμόζοντας μετατόπιση κατά ένα. Ενώ οι διευθύνσεις προορισμού πρέπει πάντα να είναι συμμετρικά buffer που δεσμεύονται με χρήση nvshmem_alloc, τα buffer προέλευσης μπορούν να είναι αυθαίρετες περιοχές μνήμης συσκευής, υπό την προϋπόθεση ότι έχουν προ-καταχωριστεί με το NVSHMEM.

Ενώ το NVSHMEM παρέχει ένα ευρύ φάσμα πρωτοκόλλων, οι πυρήνες μας βασίζονται σε μόλις 3 συναρτήσεις, χτίζοντας όλο τον συγχρονισμό και την περίφραξη (fencing) πάνω σε αυτές.

Συμμετρική μνήμη και απομακρυσμένες εγγραφές μεταξύ δύο GPU με χρήση NVSHMEM.

nvshmemx_putmem_signal_nbi_warp: Μεταφέρει ένα μπλοκ δεδομένων από τη μία συσκευή στην άλλη, ρυθμίζοντας παράλληλα μια σημαία (flag) στην απομακρυσμένη συσκευή. Η λειτουργία είτε θέτει (NVSHMEM_SIGNAL_SET) είτε αυξάνει (NVSHMEM_SIGNAL_ADD) μια θέση 64-bit. Η σημαία ενημερώνεται αφού μεταφερθεί ολόκληρο το μπλοκ μνήμης. Εάν η απομακρυσμένη συσκευή παρατηρήσει μια αλλαγή στη σημαία, μπορεί να έχει ασφαλή πρόσβαση στο buffer στη δική της μνήμη. Αυτή η συνάρτηση είναι χρήσιμη για τη συγχώνευση μεταφοράς δεδομένων και συγχρονισμού.

Το nvshmemx_signal_op λειτουργεί σε μια ενιαία θέση μνήμης, συνήθως μια σημαία 64-bit, ρυθμίζοντάς την ή αυξάνοντάς την ατομικά (atomically). Είναι χρήσιμο για την αποστολή μεταδεδομένων και τον συγχρονισμό συσκευών.

Το nvshmem_uint64_wait_until χρησιμοποιείται στο άκρο λήψης ενός σήματος, για την ανάκληση (polling) μιας σημαίας μέχρι να την ενημερώσει ο απομακρυσμένος κόμβος.

Η διασφάλιση οποιασδήποτε μορφής σειράς μεταξύ των λειτουργιών μέσω nvshmem_quiet ή nvshmem_fence είναι δαπανηρή, επομένως η υλοποίησή μας των kernels αποφεύγει τη χρήση οποιωνδήποτε άλλων NVSHMEM primitives και ελαχιστοποιεί κάθε εξάρτηση πέρα από μεμονωμένα ζεύγη αποστολής-λήψης. Κατά την αποστολή δεδομένων, χρησιμοποιούμε πάντα τη μη μπλοκαριστική έκδοση των συναρτήσεων, χωρίς να περιμένουμε τα δεδομένα να έχουν καν σταλεί έξω από το τοπικό rank. Άλλος, έμμεσος συγχρονισμός σε όλη την ομάδα των kernels διασφαλίζει ότι αυτά τα buffers δεν αντικαθίστανται καταστροφικά ενώ τα NICs χρειάζονται πρόσβαση σε αυτά.

Το NVSHMEM χρησιμοποιεί GPUDirect RDMA έτσι ώστε η NIC να μπορεί να διαβάζει και να γράφει απευθείας στη μνήμη της GPU για ανταλλαγή δεδομένων χωρίς να εμπλέκεται η CPU. Επιπλέον, στις NIC ConnectX, το NVSHMEM υποστηρίζει επικοινωνία ξεκινώμενη από GPU (γνωστή και ως InfiniBand GPUDirect Async, ή IBGDA), η οποία επιτρέπει στην GPU να υποβάλλει λειτουργίες RDMA απευθείας στη NIC και να πραγματοποιεί polling στην ουρά ολοκλήρωσης (completion queue) της NIC απευθείας. Σε πλατφόρμες που δεν υποστηρίζουν επικοινωνία ξεκινώμενη από GPU, το NVSHMEM δημιουργεί ένα «προσωρινό» νήμα (proxy thread) στην κεντρική CPU για να ξεκινήσει την επικοινωνία και να κάνει poll την ολοκλήρωση για λογαριασμό της GPU. Το πρόγραμμα NVSHMEM είναι φορητό ανεξάρτητα από το αν υποστηρίζεται ή όχι η επικοινωνία ξεκινώμενη από GPU. Ωστόσο, η επικοινωνία ξεκινώμενη από GPU μειώνει σημαντικά την καθυστέρηση επειδή παρακάμπτει πλήρως την παράκαμψη προς την CPU.

Υλοποιούμε την επικοινωνία MoE μέσω ενός ζεύγους πυρήνων αποστολής και συνδυασμού. Οι πυρήνες αποστολής είναι υπεύθυνοι για την ανάγνωση token και πληροφοριών δρομολόγησης σε κάθε rank, αποστέλλοντάς τα στους κατάλληλους εμπειρογνώμονες. Οι πυρήνες συνδυασμού συλλέγουν τις ενεργοποιήσεις που παράγονται από τους εμπειρογνώμονες και τις στέλνουν πίσω στα rank προέλευσής τους, υπολογίζοντας παράλληλα τον σταθμισμένο μέσο όρο από τους επιλεγμένους εμπειρογνώμονες με βάση τα βάρη που υπολογίζονται από τον δρομολογητή. Οι πυρήνες χωρίζονται περαιτέρω σε ένα στοιχείο αποστολής και λήψης, προκειμένου να επιτραπεί η επικάλυψη των μεταφορών δεδομένων με τον υπολογισμό. Οι πυρήνες αποστολής είναι όλοι μη μπλοκαριζόμενοι και μη συγχρονιζόμενοι: απλώς αποστέλλουν όλες τις εγγραφές στους απομακρυσμένους κόμβους. Στην άλλη πλευρά, οι πυρήνες συνδυασμού διαβάζουν μόνο από τη μνήμη, περιμένοντας μέχρι να μεταφερθούν όλα τα απαιτούμενα δεδομένα. Μετά την αποστολή εργασίας στις NIC, ενώ τα δεδομένα μεταφέρονται ασύγχρονα μέσω του καλωδίου, οι GPU μπορούν να κάνουν άλλες χρήσιμες εργασίες τοπικά, όπως η εφαρμογή κοινών εμπειρογνωμόνων ή ο υπολογισμός προσοχής.

Κάθε ζεύγος πυρήνων έχει τα δικά του συμμετρικά buffer μνήμης, με κάθε rank να δεσμεύει ιδιωτικό αποθηκευτικό χώρο για να λαμβάνει δεδομένα από όλα τα άλλα rank. Ο μέγιστος αριθμός token που μπορούν να αποσταλούν από μια ομάδα DP είναι σταθερός σε όλα τα rank, γεγονός που θέτει επίσης το ανώτερο όριο που μπορεί να λάβει κάθε rank για κάθε τοπικό εμπειρογνώμονα από κάθε rank DP. Αυτό επιτρέπει στα rank αποστολής να εξάγουν μια μοναδική διεύθυνση στο rank προορισμού για εγγραφή, χωρίς να απαιτείται κανενός είδους συγχρονισμός μεταξύ τους. Μετά τη λήψη των δεδομένων, οι πυρήνες ανακατεύουν τα token για να τα διατάξουν στη μνήμη στη βέλτιστη μορφή που απαιτείται από τους πυρήνες λήψης. Ενώ τα buffer έχουν σημαντικές διαστάσεις, επαναχρησιμοποιούνται σε όλα τα διαδοχικά επίπεδα ενός μοντέλου.

Ξεχωριστά kernels αποστολής και λήψης για λειτουργίες dispatch και combine.

Η μόνη μορφή παγκόσμιου συγχρονισμού, όπως απεικονίζεται στην παραπάνω εικόνα, υλοποιείται στον πυρήνα συνδυασμού-λήψης (combine-receive). Μόλις τα δεδομένα αποσταλούν από την αποστολή, ένα rank επιστρέφει και συνεχίζει την εκτέλεση όταν λαμβάνει όλα τα token που χρειάζεται από τα άλλα rank και εισέρχεται στους πυρήνες αποστολής συνδυασμού μόλις υποεξεργαστούν. Το φράγμα (barrier) στο combine-receive διασφαλίζει ότι κανένα rank δεν μπορεί να προχωρήσει μπροστά και να ξεκινήσει την αποστολή-αποστολή (dispatch-send) ενώ οποιοδήποτε άλλο rank εξακολουθεί να περιμένει να λάβει δεδομένα, προκαλώντας δυνητικά καταστροφική επικάλυψη. Ο συγχρονισμός γίνεται στον πυρήνα συνδυασμού για λόγους απλότητας· ως βελτιστοποίηση, το ίδιο φράγμα θα μπορούσε να υλοποιηθεί μεταξύ των πυρήνων λήψης αποστολής και αποστολής συνδυασμού, μειώνοντας περαιτέρω την καθυστέρηση επιτρέποντας στις NIC περισσότερο χρόνο για να διευθετήσουν τις μεταφορές ενώ η GPU εκτελεί άλλους πυρήνες.

Και οι δύο πυρήνες χωρίζονται σε όλα τα διαθέσιμα SM των συσκευών: ενώ οι πυρήνες αποστολής αποστολής και λήψης συνδυασμού πρέπει να παραλληλιστούν σε έναν μέγιστο αριθμό token ανά rank (max_m), οι πυρήνες λήψης αποστολής και αποστολής συνδυασμού πρέπει να διαχειριστούν τόσα token από κάθε rank και τοπικό εμπειρογνώμονα, για μέγιστο αριθμό max_m * (num_experts // EP) * (EP // TP).

Ο πυρήνας αποστολής είναι υπεύθυνος για την αποστολή token και παραγόντων κλιμάκωσης μπλοκ (block scaling factors) σε ομότιμους κόμβους (peers) που φιλοξενούν τους εμπειρογνώμονες στους οποίους δρομολογήθηκαν τα token. Η πλευρά αποστολής βασίζεται στην εξειδίκευση warp (warp specialization) για την παραλληλοποίηση δύο εργασιών: τη συγκέντρωση πληροφοριών δρομολόγησης για τον προσδιορισμό του αριθμού των token που αποστέλλονται από κάθε rank και τη συσκευασία των token σε μηνύματα που θα αποσταλούν στον απομακρυσμένο κόμβο. Η πλευρά λήψης περιμένει πρώτα να ληφθούν όλες οι μετρήσεις token, στη συνέχεια διαβάζει και αποσυμπιέζει τα ωφέλιμα φορτία (payloads) από τη κοινόχρηστη μνήμη στα tensor που θα περάσουν στους κάτω ρεύματος πυρήνες Grouped GEMM. Αποθηκεύει επίσης πληροφορίες σε buffer που μοιράζονται με τους πυρήνες συνδυασμού, για να υποδείξει πού πρέπει να σταλεί πίσω κάθε token. Αυτές οι πληροφορίες απαιτούνται καθώς ο πυρήνας λήψης ανακατεύει τα token σε συνεχόμενα buffer σε μια μη ντετερμινιστική σειρά.

Δρομολόγηση token και σηματοδότηση της άφιξής τους μεταξύ rank προέλευσης και προορισμού.

Στο τμήμα αποστολής, ένα μεμονωμένο warp 32 νημάτων (threads) σε κάθε μπλοκ είναι υπεύθυνο για την ανάγνωση των πληροφοριών δρομολόγησης και την καταμέτρηση του αριθμού των token που αναμένεται να λάβει κάθε εμπειρογνώμονας προορισμού από αυτό το rank, αποστέλλοντας την καταμέτρηση συν ένα χρησιμοποιώντας nvshmem_signal_op. Η καταμέτρηση αυξάνεται κατά ένα, καθώς η μετάβαση από το μηδέν στο μη μηδέν στην απομακρυσμένη πλευρά σηματοδοτεί τη λήψη των μετρήσεων από ένα rank προέλευσης. Παράλληλα, τα υπόλοιπα warp συνεργάζονται για να αντιγράψουν token σε συμμετρική μνήμη σε όλα τα μπλοκ παράλληλα, συσκευάζοντας ενεργοποιήσεις, παράγοντες κλιμάκωσης και το δείκτη τους στο τοπικό rank σε ένα συνεχόμενο κομμάτι. Ο δείκτης απαιτείται από τον αποστολέα συνδυασμού για να προσδιορίσει τη διεύθυνση στην οποία θα γραφτεί το token. Στη συνέχεια, αφού διασφαλιστεί ότι όλα τα δεδομένα έχουν αντιγραφεί μέσω ενός φράγματος σε όλες τις ομάδες warp, τα warp λειτουργούν ξανά ανεξάρτητα, αποστέλλοντας το καθένα το ίδιο buffer σε διαφορετικό εμπειρογνώμονα παράλληλα. Τα token αποστέλλονται με χρήση nvshmemx_putmem_signal_nbi_warp, η οποία αυξάνει επίσης ατομικά την καταμέτρηση των απεσταλμένων token από το τοπικό rank στην απομακρυσμένη συσκευή. Εντός μιας ομάδας DP, δεδομένου ότι κάθε rank κατέχει ένα αντίγραφο του token που πρόκειται να αποσταλεί, η αποστολή εξισορροπείται ομοιόμορφα, με κάθε συσκευή να στέλνει ένα υποσύνολο των token.

Στην πλευρά λήψης, όλοι οι πυρήνες συγχρονίζονται πρώτα με τα rank αποστολής περιμένοντας να αποσταλούν οι συνολικές μετρήσεις token. Στη συνέχεια, περιμένουν όλοι να διευθετηθούν οι ατομικά αυξημένες μετρήσεις απεσταλμένων token στις συνολικές μετρήσεις, υποδεικνύοντας ότι όλα τα ωφέλιμα φορτία από τα rank προέλευσης έχουν επίσης αποσταλεί, χάρη στη σημασιολογία της κλήσης putmem. Οι πυρήνες κάνουν poll στις μετρήσεις χρησιμοποιώντας nvshmem_uint64_wait_until, παραλληλίζοντας τη λειτουργία σε όλα τα μπλοκ και νήματα. Στη συνέχεια, ένα φράγμα μεταξύ μπλοκ (cross-block barrier) διασφαλίζει ότι κανένα μπλοκ δεν διαβάζει τα buffer εκτός εάν έχουν ληφθεί σωστά όλα τα δεδομένα. Κατανεμημένα στα μπλοκ και συγχρονισμένα μέσω ενός ατομικά αυξημένου μετρητή, τα token αντιγράφονται από τη συμμετρική μνήμη σε κανονικά συνεχόμενα buffer που γίνονται τα δεδομένα εισόδου για τους κάτω ρεύματος πυρήνες που υλοποιούν τους εμπειρογνώμονες. Το rank προέλευσης, ο δείκτης εμπειρογνώμονα και ο δείκτη token αποθηκεύονται ξεχωριστά, εντοπίζοντας ακριβώς τη θέση στην οποία ο πυρήνας συνδυασμού πρέπει να στείλει τις ενεργοποιήσεις. Παρόλο που τα token από μια ομάδα DP αποστέλλονται από διαφορετικές συσκευς, ομαδοποιούνται όλα μαζί για να περάσουν στον αντίστοιχο εμπειρογνώμονα.

Οι πυρήνες συνδυασμού χωρίζονται ξανά σε μισά αποστολής και λήψης: οι αποστολείς αντιγράφουν τις μη κλασματικές (un-quantized) ενεργοποιήσεις 16-bit στις απομακρυσμένες συσκευς, με τους πυρήνες λήψης να περιμένουν να αποσταλούν τα δεδομένα και να υπολογίζουν το σταθμισμένο μέσο όρο όλων των συνεισφορών εμπειρογνωμόνων τοπικά. Επιπλέον, λειτουργούν επίσης ως φράγμα, για να συγχρονίσουν τη ακολουθία αποστολής-συνδυασμού: κάθε rank θέτει μια σημαία σε κάθε ομότιμο κόμβο κατά την είσοδο στους πυρήνες αποστολής, με τους πυρήνες λήψης να μην επιτρέπεται να επιστρέψουν εκτός εάν παρατηρήσουν ότι η σημαία έχει τεθεί. Η καθυστέρηση του συγχρονισμού είναι ελάχιστη, καθώς επικαλύπτεται με την πραγματική επικοινωνία και υπολογισμό.

Συνδυασμός των εξόδων των experts στα token σε κάθε rank παράλληλου δεδομένων.

Στην πλευρά αποστολής, οι πυρήνες διασχίζουν τη λίστα των token που έχουν ανατεθεί σε όλους τους τοπικούς εμπειρογνώμονες παράλληλα, γράφοντάς τα σε ένα buffer στο rank προορισμού. Ο rank προορισμού, ο δείκτης εμπειρογνώμονα και ο δείκτη token διαβάζονται από τα buffer ανά-token που συμπληρώνονται από τους πυρήνες scatter κατά τη λήψη των token. Κάθε αποστολέας έχει τη δική του ιδιωτική περιοχή μνήμης ανά εμπειρογνώμονα για εγγραφή, όπως υποδεικνύεται στην παραπάνω εικόνα, αποφεύγοντας την ανάγκη συγχρονισμού. Παρόμοια με την αποστολή, ο συνδυασμός αυξάνει ατομικά τους μετρητές ανά token στο rank προορισμού για να υποδείξει τη λήψη των δεδομένων: όταν ο μετρητής ταιριάζει με τον αριθμό των εμπειρογνωμόνων στους οποίους δρομολογήθηκε ένα token, είναι δυνατή η πρόσβαση στα περιεχόμενα του token.

Στον πυρήνα λήψης, η λίστα των token διασχίζεται παράλληλα σε πολλαπλά μπλοκ, περιμένοντας να φτάσουν τα περιεχόμενά τους κάνοντας poll στη σημαία που έχει τεθεί από τη λειτουργία σηματοδότησης. Κατά την άφιξη, τα ωφέλιμα φορτία διαβάζονται από τα ιδιωτικά buffer, με τον πίνακα δρομολόγησης να υποδεικνύει ποιο buffer να διαβαστεί και ποιο βάρος να ανατεθεί σε κάθε εμπειρογνώμονα. Τα αποτελέσματα γράφονται στη συνέχεια σε εξωτερικά δεσμευμένα tensor, με τον πυρήνα να ολοκληρώνει την εκτέλεση μόλις όλες οι συσκευές περάσουν το φράγμα. Σε ολόκληρη την ομάδα DP, όλα τα rank λαμβάνουν ένα αντίγραφο κάθε ενεργοποίησης εμπειρογνώμονα για να υπολογίσουν τα δικά τους αντίγραφα.

Αξιολογούμε τους πυρήνες μας σε ένα σύμπλεγμα 128 GPU κατανεμημένων σε 16 κόμβους που συνδέονται μέσω InfiniBand και NIC CX-7. Δοκιμάζουμε τόσο το GPUDirect Async (IBGDA) όσο και τη Reliable Connection (RC) με proxy CPU. Τους συγκρίνουμε με την εξαιρετικά βελτιστοποιημένη υλοποίηση DeepSeek, καθώς και με τα πυκνά πρωτόκολλα που παρέχονται από το PyTorch (μέσω NCCL) ή το NVSHMEM.

Λανθάνουσα επικοινωνία για 128 token ανά GPU σε μια πολυκόμβη διαμόρφωση EP128.

Παρόλο που στη λειτουργία πυκνού NVSHMEM all-to-all, η απόδοση των IBRC και IBGDA είναι παρόμοια (6378 µs έναντι 6180 µs), το IBGDA είναι σημαντικά ταχύτερο με τους αραιούς (sparse) πυρήνες. Προσθέτοντας το Dispatch και το Combine, το IBGDA χρησιμοποιεί 902 µs ενώ το IBRC χρειάζεται 3223 µs - μια βελτίωση 3.6x στην καθυστέρηση. Ενώ το all-to-all περιορίζεται από το εύρος ζώνης (bandwidth bound), οι αραιοί πυρήνες εκπέμπουν τάξεις μεγέθους λιγότερα δεδομένα, περιοριζόμενοι από την καθυστέρηση (latency bound). Με την ενεργοποίηση μεταφορών δικτύου απευθείας από μια GPU, χωρίς να απαιτείται proxy CPU για τον συντονισμό της GPU και της NIC, η καθυστέρηση από άκρο σε άκρο (end-to-end latency) μειώνεται σημαντικά.

Ενώ οι φορητοί πυρήνες μας είναι περίπου 2x πιο αργοί από τους εξαιρετικά βελτιστοποιημένους πυρήνες DeepSeek, βελτιώνουν την καθυστέρηση κατά 10x σε σύγκριση με τους πυκνούς πυρήνες (902 µs έναντι 9944 µs).

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

Λανθάνουσα επικοινωνία για 128 token ανά GPU σε μια μονοκόμβη διαμόρφωση EP8.

Σε διαμόρφωση ενός κόμβου (EP8), το NVSHMEM χρησιμοποιεί το NVLINK για μεταφορά, προσφέροντας χαμηλότερη καθυστέρηση και υψηλότερη απόδοση (throughput) από τη δικτύωση μεταξύ κόμβων. Οι φορητοί πυρήνες μας είναι περίπου 2.5x ταχύτεροι από το DeepEP σε διαμόρφωση ενός κόμβου (186µs έναντι 481 µs).

Οι πυρήνες που περιγράφονται εδώ ξεπερνούν τα ενσωματωμένα πρωτόκολλα των πλαισίων ML και προσφέρουν αξ Aπρεπή απόδοση χωρίς υπερβολική εξειδίκευση για συγκεκριμένες μεταφορές μεταξύ συσκευών, όπως το NVLink ή το InfiniBand. Εκτός από τις ήδη αναφερθείσες ευκαιρίες, περαιτέρω κέρδη απόδοσης είναι εφικτά αντικαθιστώντας τα πρωτόκολλα επικοινωνίας με πιο εξειδικευμένες εκδόσεις. Για παράδειγμα, σε όλο το NVLink, η χρήση συμμετρικής μνήμης θα μπορούσε να αντικατασταθεί με buffer που μοιράζονται μεταξύ των συσκευών, εξαλείφοντας ορισμένες αντιγραφές και επιτρέποντας σχήματα συγχρονισμού πιο λεπτής κοκκομέτρησης (finer-grained), σε μεμονωμένα token αντί για παρτίδες token. Σε όλο το InfiniBand, μια υλοποίηση θα μπορούσε να έχει άμεση πρόσβαση στα υποκείμενα ζεύγη ουρών (queue pairs), εξαλείφοντας την ανάγκη μεταφοράς των αναμενόμενων και απεσταλμένων μετρήσεων token και εκκένωσης (flushing) των token από τις ουρές στα buffer εξόδου καθώς φτάνουν. Ωστόσο, τέτοιες υλοποιήσεις έρχονται με το κόστος της φορητότητας, ενώ οι υλοποιήσεις μας επιτρέπουν την αξιολόγηση συμπλεγμάτων με συγκεκριμένες διαμορφώσεις πριν δεσμευτεί προσπάθεια ανάπτυξης για βελτιστοποίηση για συγκεκριμένο υλικό.

Παρουσιάσαμε μια υψηλής απόδοσης, φορητή βιβλιοθήκη για επικοινωνία MoE που επιτυγχάνει 10x ταχύτερη απόδοση σε σύγκριση με την τυπική επικοινωνία all-to-all, διατηρώντας παράλληλα συμβατότητα σε διάφορες διαμορφώσεις υλικού. Σε αναπτύξεις ενός κόμβου με NVLink, η λύση μας επιδεικνύει 2.5x χαμηλότερη καθυστέρηση από προηγούμενες υλοποιήσεις.

Η προσέγγισή μας εξισορροπεί την απόδοση με τη φορητότητα μέσω βασικών καινοτομιών που περιλαμβάνουν υποστήριξη επικοινωνίας ξεκινώμενης από GPU, αρχιτεκτονική διαχωρισμένων πυρήνων που επιτρέπει την επικάλυψη υπολογισμού-επικοινωνίας και αποτελεσματική αποστολή token με ελάχιστα πρωτόκολλα NVSHMEM. Ενώ είναι περίπου 2x πιο αργή από εξαιρετικά εξειδικευμένες υλοποιήσεις σε εγκαταστάσεις πολλαπλών κόμβων, η βιβλιοθήκη μας προσφέρει ανώτερη ευελιξία σε διάφορα περιβάλλοντα δικτύου (NVLink, CX-7 και EFA).

Καθώς τα μοντέλα MoE συνεχίζουν να κλιμακώνονται, οι αποτελεσματικές στρατηγικές επικοινωνίας όπως η δική μας θα γίνονται ολοένα και πιο σημαντικές για την πρακτική ανάπτυξη. Η πλήρως ανοιχτού κώδικα υλοποίησή μας είναι διαθέσιμη στη διεύθυνση https://github.com/ppl-ai/pplx-kernels.

DeepEP

Τεχνική Αναφορά DeepSeek-V3

NVIDIA NVSHMEM

BibTeX

@online{perplexity2025EfficientAndPortableMixtureOfExpertsCommunication,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Phi Hoang},
  title = {{Efficient and portable mixture-of-experts communication}},
  date = {2025-04-02},
  url = {https://www.perplexity.ai/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}

Εισαγωγή

Κατακερματισμός και Παραλληλισμός

NVSHMEM

Φορητοί Πυρήνες

Αποστολή (Dispatch)

Συνδυασμός (Combine)

Benchmarks

Περαιτέρω Βελτιώσεις

Συμπέρασμα

Αναφορές

Παραπομπή