Η AWS βελτιώνει κατά 40% την απόδοση εκπαίδευσης MoE reinforcement learning στο Amazon EKS

Η νέα αρχιτεκτονική συνδυάζει Amazon EKS, EFA και DeepEP για μεγαλύτερη αποδοτικότητα μεγάλων μοντέλων AI

Η AWS ανακοίνωσε μια νέα τεχνική αρχιτεκτονική για την κλιμάκωση της εκπαίδευσης μεγάλων μοντέλων reinforcement learning τύπου Mixture-of-Experts (MoE) πάνω στην υπηρεσία Amazon Elastic Kubernetes Service (EKS). Η καινοτομία αυτή βασίζεται στη χρήση της Elastic Fabric Adapter (EFA) και του DeepEP, προσφέροντας σύμφωνα με την εταιρεία έως και 40% αυξημένη απόδοση στην παραγωγή rollout.

Η συγκεκριμένη λύση αντιμετωπίζει τις προκλήσεις της ανισοκατανομής φορτίου και της υπερβολικής επικοινωνίας που δημιουργούνται από το super-sparse token routing στα MoE μοντέλα. Για να το πετύχει, το DeepEP αντικαθιστά το κλασικό πρωτόκολλο all-to-all επικοινωνίας (μέσω NCCL) με ειδικούς πυρήνες GPU που αξιοποιούν NVLink για εσωτερική επικοινωνία κόμβου και libfabric/EFA για συνδέσεις ανάμεσα σε διαφορετικούς κόμβους. Έτσι μειώνεται σημαντικά ο χρόνος αναμονής λόγω επικοινωνιών, που παραδοσιακά επιβραδύνει την εκπαίδευση.

Η δομή της υποδομής συνδυάζει πολλαπλές ομάδες κόμβων μέσα στο Amazon EKS, οι οποίες είναι βελτιστοποιημένες αναλόγως των καθηκόντων τους. Υπάρχουν ξεχωριστές ομάδες για την εκπαίδευση πολιτικών (policy training), για παραγωγή rollout, για επεξεργασία περιβαλλόντων σε CPU και για αποθήκευση δεδομένων εμπειρίας και checkpoint caches σε κόμβους με μεγάλη μνήμη. Επιπλέον, η χρήση Spot instances για την παραγωγή rollout επιτρέπει εξοικονόμηση κόστους, καθώς η διαδικασία αυτή μπορεί να αντεπεξέλθει σε προσωρινές απώλειες υπολογιστικού δυναμικού.

Στο πλαίσιο δοκιμών που περιγράφει η AWS, σε σύστημα με 48 GPU P5en instances (16 για εκπαίδευση και 32 για rollout), η ενεργοποίηση του DeepEP πάνω στην EFA οδήγησε σε αύξηση 40% στην συνολική απόδοση rollout. Η αρχιτεκτονική μπορεί να επεκταθεί για την υποστήριξη χιλιάδων επιταχυντών, προσφέροντας έτσι προσαρμοστικότητα σε πολύ μεγάλα clusters εκπαίδευσης μοντέλων MoE.

Το σύστημα απαιτεί συγκεκριμένες εκδόσεις λογισμικού και υποδομών, όπως Amazon EKS version 1.31 ή νεότερη, EFA installer 1.49 με το AWS OFI NCCL plugin, DeepEP 2.0.0, NCCL 2.31.2 και PyTorch 2.12.1 με CUDA 13.0. Υποστηρίζεται η χρήση GPU instances τύπου p5.48xlarge, p5e.48xlarge και p6-b200.48xlarge.

Η υποστήριξη για EFA και placement groups, που είναι βασικά συστατικά αυτής της λύσης, είναι πλέον διαθέσιμη σε όλες τις περιοχές που προσφέρει η AWS τις υπηρεσίες Amazon EKS, σύμφωνα με προηγούμενη ανακοίνωση της ίδιας της εταιρείας.

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

Ενώ η AWS αναφέρει τα αποτελέσματα βάσει εσωτερικών δοκιμών της, δεν υπάρχουν ακόμη ανεξάρτητα benchmarks για επιβεβαίωση. Παρόλα αυτά, η αρχιτεκτονική αυτή φαίνεται να τοποθετεί την AWS σε ισχυρή θέση στον ανταγωνισμό για υποδομές cloud που υποστηρίζουν απαιτητικά AI workloads σε μεγάλη κλίμακα.