Τα Mixture of Experts, ή MoE, έχουν εξελιχθεί σε μία από τις πιο χαρακτηριστικές αρχιτεκτονικές τάσεις στην εκπαίδευση μεγάλων μοντέλων AI. Η βασική ιδέα είναι ότι, αντί όλα τα tokens να περνούν από το ίδιο πυκνό feed-forward δίκτυο, ένα σύστημα δρομολόγησης επιλέγει δυναμικά ποιοι experts θα ενεργοποιηθούν. Έτσι, η υπολογιστική ισχύς χρησιμοποιείται πιο επιλεκτικά, χωρίς να απαιτείται η πλήρης ενεργοποίηση ολόκληρου του μοντέλου σε κάθε βήμα.
Σύμφωνα με τη NVIDIA, αυτή η προσέγγιση επιτρέπει σε μοντέλα όπως τα DeepSeek, Qwen και Mixtral να ανταγωνίζονται ή και να ξεπερνούν πυκνά μοντέλα, με χαμηλότερο κόστος εκπαίδευσης. Παρ’ όλα αυτά, η θεωρητική αποδοτικότητα δεν μεταφράζεται αυτόματα σε πραγματική απόδοση στο hardware. Το νέο τεχνικό άρθρο της εταιρείας εστιάζει ακριβώς σε αυτό το κενό, δείχνοντας πώς το JAX και το NVIDIA Transformer Engine μπορούν να βελτιώσουν ουσιαστικά την εκπαίδευση dropless MoE.
Γιατί η εκπαίδευση MoE είναι τόσο απαιτητική
Σε επίπεδο παραγωγικής κλίμακας, τα MoE εισάγουν εμπόδια που δεν εμφανίζονται με την ίδια ένταση στα dense μοντέλα. Η δρομολόγηση των tokens, η αποστολή και η συλλογή των experts, η all-to-all επικοινωνία μεταξύ GPU και οι ακανόνιστοι υπολογισμοί GEMM δημιουργούν ένα σύνθετο προφίλ φόρτου. Το πρόβλημα γίνεται ακόμη πιο έντονο επειδή ο router μαθαίνει κατά την εκπαίδευση, άρα η κατανομή των tokens στους experts αλλάζει συνεχώς.
Αυτό σημαίνει ότι δύο διαφορετικά batches δεν έχουν απαραίτητα παρόμοια κατανομή φόρτου, ενώ ακόμη και μέσα στο ίδιο batch ένας expert μπορεί να φορτωθεί πολύ περισσότερο από έναν άλλο. Το αποτέλεσμα είναι ragged tensors, δηλαδή σχήματα που δεν είναι ομοιόμορφα και ορθογώνια, όπως περιμένουν οι περισσότερες βελτιστοποιημένες βιβλιοθήκες. Αν η μεταφορά και η επανασύνθεση των tokens δεν γίνουν αποδοτικά, η επικοινωνία κυριαρχεί στον χρόνο εκτέλεσης και οι GPU μένουν ανεκμετάλλευτες.
Η διαφορά ανάμεσα σε dropless και capacity-based MoE
Η NVIDIA ξεχωρίζει δύο βασικές στρατηγικές για τη διαχείριση της δρομολόγησης των tokens: το dropless και το capacity-based MoE. Στο σενάριο dropless, κάθε token επεξεργάζεται από τον expert που του έχει ανατεθεί, ακόμη κι αν η κατανομή του φόρτου είναι άνιση. Αυτή η προσέγγιση θεωρείται ελκυστική από πλευράς ποιότητας μοντέλου, επειδή αποφεύγει να απορρίπτει δεδομένα ή να τα αλλοιώνει με τεχνητό padding.
Αντίθετα, στα capacity-based frameworks κάθε expert έχει ένα σταθερό όριο tokens. Όταν αυτό ξεπερνιέται, τα επιπλέον tokens είτε αποκόπτονται είτε προστίθεται padding, ώστε η υπολογιστική μορφή να παραμένει τακτική και φιλική προς το hardware. Το τίμημα είναι ένας άμεσος συμβιβασμός ανάμεσα στην αποδοτικότητα του συστήματος και στην πληρότητα των δεδομένων που χρησιμοποιούνται για την εκπαίδευση.
Οι βελτιστοποιήσεις που φέρνει το Transformer Engine
Η προσέγγιση dropless απαιτεί kernels που να χειρίζονται αποδοτικά μεταβλητό πλήθος tokens ανά expert, χωρίς να εξαρτώνται από σταθερά σχήματα. Επιπλέον, αυτές οι δυναμικές μορφές πρέπει να λειτουργούν χωρίς αναγκαστική εμπλοκή της CPU, ώστε να μην διακόπτονται τα CUDA graphs και να αποφεύγονται οι επαναμεταγλωττίσεις. Σε αυτό το πλαίσιο, το Transformer Engine προσφέρει τρία βασικά δομικά στοιχεία: group-aware MXFP8 quantization, grouped GEMM για expert matmuls και βελτιστοποιημένες λειτουργίες expert parallelism για dispatch και combine.
Ο συνδυασμός αυτών των στοιχείων επιτρέπει στο JAX να υλοποιεί πιο πρακτικά την εκπαίδευση MoE σε μεγάλη κλίμακα. Η λογική είναι διαχωρισμένη αλλά αλληλοσυμπληρούμενη. Το grouped GEMM αναλαμβάνει ό,τι συμβαίνει μέσα σε κάθε expert, ενώ το expert parallelism βελτιστοποιεί τη μεταφορά των tokens προς και από τους experts σε πολλαπλές GPU.
Grouped GEMM και αποδοτικότεροι υπολογισμοί
Σε ένα πυκνό FFN, όλα τα tokens περνούν από το ίδιο weight matrix, άρα η πράξη GEMM έχει ένα σταθερό και προβλέψιμο σχήμα. Στο MoE αυτό καταρρέει, επειδή οι experts λαμβάνουν διαφορετικό πλήθος tokens σε κάθε βήμα. Προηγούμενες προσεγγίσεις βασίζονταν είτε σε βρόχους με πολλαπλά GEMM kernels είτε σε batched GEMM με padding, λύσεις που είτε επιβάρυναν την κρίσιμη διαδρομή είτε σπαταλούσαν υπολογιστική ισχύ.
Η λύση grouped GEMM επιτρέπει να εκτελούνται όλες οι πράξεις των experts σε μία κλήση kernel, με βάση τον πραγματικό αριθμό tokens του καθενός. Έτσι, υπολογίζονται μόνο τα έγκυρα τμήματα και αποφεύγεται η υπερκάλυψη του δυσμενέστερου σεναρίου. Η NVIDIA αναφέρει ότι το Transformer Engine αξιοποιεί τα cuBLAS και cuBLASLt για να διατηρεί υψηλή αξιοποίηση των Tensor Cores ακόμη και με ακανόνιστα σχήματα, ενώ στις NVIDIA Blackwell GPU ανοίγει και ο δρόμος για MXFP8 block scaling στα expert matmuls.
Dispatch, combine και το βάρος της επικοινωνίας
Μετά τη δρομολόγηση των tokens από τον router, το σύστημα πρέπει να τα στείλει στις σωστές συσκευές, να εκτελέσει την επεξεργασία και στη συνέχεια να επιστρέψει τα αποτελέσματα στην αρχική τους σειρά. Αυτή η διαδικασία χωρίζεται στα στάδια dispatch και combine. Σε μια αφελή υλοποίηση, τα στάδια αυτά εκτελούνται σειριακά, με αποτέλεσμα καθυστερήσεις, πολλαπλές προσπελάσεις μνήμης και κακή επικάλυψη επικοινωνίας και υπολογισμού.
Η υλοποίηση expert parallelism του Transformer Engine ενοποιεί αυτή τη διαδρομή σε μια πιο σφιχτά συγχωνευμένη ροή kernels. Σύμφωνα με τη NVIDIA, βασίζεται στο NCCL EP, ένα backend επικοινωνίας προσαρμοσμένο ειδικά στα ακανόνιστα και ανισοβαρή μοτίβα που δημιουργεί η παράλληλη εκτέλεση experts. Περιλαμβάνει επίσης μηχανισμό deduplication των tokens, ώστε όταν ένα token αποστέλλεται σε πολλούς experts στο ίδιο rank ή σε πολλαπλά ranks σε απομακρυσμένο IB node, να διασχίζει το δίκτυο μόνο μία φορά.
Τα πρόσθετα βήματα στο JAX και η μετρημένη επίδοση
Πέρα από τα kernels του Transformer Engine, η NVIDIA αναφέρει δύο ακόμη σημαντικές βελτιώσεις στο ευρύτερο stack. Η πρώτη είναι το host offloading του JAX, που επιτρέπει ορισμένες ενδιάμεσες ενεργοποιήσεις να μεταφέρονται στη μνήμη του host αντί να παραμένουν καθ’ όλη τη διάρκεια του forward pass στη μνήμη της συσκευής. Για την εκπαίδευση του DeepSeek-V3, η εταιρεία αναφέρει ότι μεταφέρει στον host τα αποτελέσματα των query και value projections για εξοικονόμηση μνήμης.
Η δεύτερη είναι τα XLA multistreaming collectives, που επιτρέπουν στον compiler να εκτελεί ανεξάρτητες συλλογικές επικοινωνίες παράλληλα σε ξεχωριστά CUDA streams. Έτσι, οι μεταφορές μέσω InfiniBand μεταξύ κόμβων μπορούν να επικαλύπτονται με την επικοινωνία μέσω NVLink μέσα στον ίδιο κόμβο. Με αυτά τα βήματα, μαζί με cuBLAS GroupedGEMM, MXFP8 GroupQuant και το βελτιστοποιημένο EP, η NVIDIA δηλώνει ότι η απόδοση στο DeepSeek-V3 671B ανέβηκε από 103 TFLOPS/GPU σε 1.068 TFLOPS/GPU σε NVIDIA GB200, δηλαδή βελτίωση 10,4x.
Κλιμάκωση σε χιλιάδες GPUs και τα επόμενα βήματα
Η πραγματική δοκιμή τέτοιων τεχνικών δεν είναι μόνο η επίδοση σε ένα rack, αλλά η συμπεριφορά όταν το σύστημα μεγαλώνει σημαντικά. Η NVIDIA υποστηρίζει ότι με το JAX MoE stack και το Transformer Engine, η επιβάρυνση από την επικοινωνία παραμένει σχετικά ελεγχόμενη ακόμη και σε multirack ανάπτυξη. Συγκεκριμένα, αναφέρει 97% αποδοτικότητα στις 1.024 GPU, ένδειξη ότι οι βελτιστοποιήσεις στη διαδρομή επικοινωνίας διατηρούν υψηλό throughput καθώς μεγαλώνει το cluster.
Για όσους θέλουν να ξεκινήσουν με εκπαίδευση dropless MoE, οι σχετικές βελτιστοποιήσεις διατίθενται στο NVIDIA NGC MaxText container με ενσωματωμένο Transformer Engine. Η προτεινόμενη πρακτική είναι να ξεκινήσει κανείς από το reference configuration, να επιβεβαιώσει την ορθότητα σε μικρότερο μοντέλο MoE και στη συνέχεια να παρακολουθήσει μετρικές όπως step time, TFLOPS/GPU, MFU, latency του grouped GEMM και καθυστερήσεις στο dispatch και στο combine. Η εταιρεία σημειώνει επίσης ότι σχεδιάζει επιπλέον προσθήκες, όπως NVFP4, quantization fused with GEMM και A2A overlap.
