Μεταφορά μνήμης GPU υψηλών επιδόσεων σε AWS Sagemaker Hyperpod

Η σύγχρονη υποδομή βαθιάς μάθησης απαιτεί συχνά τη μεταφορά μεγάλων ποσοτήτων δεδομένων μεταξύ GPU σε διαφορετικά μηχανήματα. Στη Perplexity, αντιμετωπίσαμε μια μοναδική τεχνική πρόκληση: την αποτελεσματική μεταφορά μη συνεχομένων περιοχών μνήμης GPU μεταξύ μηχανημάτων με τη μέγιστη δυνατή ταχύτητα. Η πλατφόρμα-στόχος μας, οι περιπτώσεις AWS p5, προσφέρει εντυπωσιακό εύρος ζώνης δικτύου 3200 Gbps μέσω 32 καρτών δικτύου. Αυτό το άρθρο μοιράζεται το ταξίδι μας για την κατασκευή μιας προσαρμοσμένης λύσης δικτύωσης υψηλών επιδόσεων που επιτυγχάνει το 97,1% αυτού του θεωρητικού εύρους ζώνης.

Η τεχνική πρόκληση

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

• Μεταφορά υψηλού εύρους ζώνης μεταξύ απομακρυσμένων GPU μη συνεχομένων τμημάτων μνήμης
• Δυνατότητα δυναμικής προσθήκης ή αφαίρεσης κόμβων από αναπτύξεις Kubernetes χωρίς διακοπή των συνεχιζόμενων λειτουργιών
• Υποστήριξη για μοτίβα επικοινωνίας ομοτίμων (peer-to-peer)

Ενώ η βιβλιοθήκη NCCL της NVIDIA είναι το de facto πρότυπο για κατανεμημένη βαθιά μάθηση, δεν ήταν ιδανική για τη δική μας περίπτωση χρήσης:

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

Σύγχρονα δίκτυα υψηλών επιδόσεων

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

Τα περισσότερα δίκτυα που χρησιμοποιούμε καθημερινά βασίζονται σε πρωτόκολλα TCP/IP, όπου οι εφαρμογές επικοινωνούν με την κάρτα δικτύου μέσω του πυρήνα του λειτουργικού συστήματος χρησιμοποιώντας sockets. Ωστόσο, τα δίκτυα υψηλών επιδόσεων χρησιμοποιούν RDMA (Remote Direct Memory Access) - μια εντελώς διαφορετική στοίβα υλικού και λογισμικού που επιτρέπει την άμεση πρόσβαση στη μνήμη μεταξύ μηχανημάτων χωρίς τη συμμετοχή της CPU.

Το AWS παρέχει Elastic Fabric Adapter (EFA), μια προσαρμοσμένη διεπαφή δικτύου που υλοποιεί το προσαρμοσμένο πρωτόκολλο της Amazon που ονομάζεται Scalable Reliable Datagram (SRD). Σε αντίθεση με την παραδοσιακή δικτύωση TCP/IP όπου τα δεδομένα πρέπει να αντιγράφονται πολλές φορές μεταξύ του χώρος χρήστη, του χώρος πυρήνα και των buffer δικτύου, το EFA με RDMA επιτρέπει την άμεση μεταφορά δεδομένων μεταξύ της μνήμης GPU και της κάρτας δικτύου, παρακάμπτοντας εντελώς την CPU.

Φιλοσοφία σχεδίασης δικτύων υψηλών επιδόσεων

Η κατασκευή συστημάτων δικτύωσης υψηλών επιδόσεων απαιτεί την εκ νέου εξέταση πολλών θεμελιωδών υποθέσεων:

1. Ιδιοκτησία Buffer: Σε αντίθεση με τα παραδοσιακά sockets όπου ο πυρήνας διαχειρίζεται τα buffer δικτύου και απαιτεί αντιγραφή μεταξύ του χώρος χρήστη και του χώρος πυρήνα, το RDMA απαιτεί από τις εφαρμογές να διαχειρίζονται τα δικά τους buffer. Όταν μια εφαρμογή ξεκινά μια λειτουργία δικτύου, μεταφέρει την κυριότητα του buffer στην κάρτα δικτύου έως ότου ολοκληρωθεί η λειτουργία, εξαλείφοντας την ανάγκη για αντιγραφή δεδομένων.
2. Καταχώριση μνήμης: Οι εφαρμογές πρέπει να καταχωρούν περιοχές μνήμης στον πυρήνα του λειτουργικού συστήματος. Ο πυρήνας ρυθμίζει αντιστοιχίσεις εικονικών διευθύνσεων που επιτρέπουν στην CPU, τις GPU και τις κάρτες δικτύου να κατανοούν όλες τις ίδιες εικονικές διευθύνσεις. Αυτή η καταχώριση είναι μια εφάπαξ λειτουργία που επιτρέπει τις επακόλουθες μεταφορές δεδομένων μηδενικής αντιγραφής.
3. Επίπεδο ελέγχου έναντι επιπέδου δεδομένων: Τα δίκτυα υψηλών επιδόσεων διαχωρίζουν τις λειτουργίες σε δύο κατηγορίες:
   1. Οι λειτουργίες επιπέδου ελέγχου (όπως η ρύθμιση σύνδεσης και η καταχώριση μνήμης) διέρχονται από τον πυρήνα για να διασφαλιστεί η ασφάλεια
   2. Οι λειτουργίες επιπέδου δεδομένων (πραγματική μεταφορά δεδομένων) παρακάμπτουν τον πυρήνα για μέγιστη απόδοση
4. Λήψη πριν από τη μετάδοση: Χωρίς buffer διαχειριζόμενα από τον πυρήνα, οι εφαρμογές πρέπει να δημοσιεύουν εκ των προτέρων λειτουργίες λήψης, καθορίζοντας πού πρέπει να τοποθετηθούν τα εισερχόμενα δεδομένα. Αυτή είναι μια θεμελιώδης αλλαγή από το μοντέλο socket όπου οι εφαρμογές μπορούν να λαμβάνουν δεδομένα ανά πάσα στιγμή.
5. Ολοκλήρωση βασισμένη σε δημοσκόπηση (poll-based): Αντί να περιμένουν ειδοποιήσεις πυρήνα μέσω μηχανισμών όπως το epoll, οι εφαρμογές ερωτούν απευθείας τις ουρές ολοκλήρωσης υλικού. Αυτό εξαλείφει την υπερφόρτωση κλήσεων συστήματος και επιτρέπει την άμεση αντίδραση σε ολοκληρωμένες λειτουργίες.
6. Επίγνωση τοπολογίας υλικού: Η κατανόηση και η βελτιστοποίηση για την τοπολογία υλικού είναι ζωτικής σημασίας για την επίτευξη μέγιστης απόδοσης.

Κατανόηση της τοπολογίας υλικού

Οι περιπτώσεις AWS p5 διαθέτουν μια εξελιγμένη εσωτερική αρχιτεκτονική. Όπως φαίνεται παρακάτω, κάθε περίπτωση περιέχει δύο υποδοχές CPU που σχηματίζουν δύο κόμβους NUMA, με κάθε κόμβο NUMA να συνδέεται σε τέσσερους διακόπτες PCIe:

Μεγέθυνση

Κάτω από κάθε διακόπτη PCIe, βρίσκουμε:

• Τέσσερις κάρτες δικτύου EFA 100 Gbps
• Μία GPU NVIDIA H100
• Ένα NVMe SSD

Οι διαδρομές δεδομένων για τις μεταφορές TCP/IP και RDMA καταδεικνύουν θεμελιώδεις διαφορές στις προσεγγίσεις τους:

Μεγέθυνση

Με το TCP/IP (αριστερή πλευρά), τα δεδομένα πρέπει να αντιγραφούν πολλές φορές:

1. Από τη μνήμη GPU στο buffer εφαρμογής στον χώρος χρήστη
2. Από το buffer χώρος χρήστη στο buffer πυρήνα
3. Από το buffer πυρήνα στην κάρτα δικτύου
4. Η αντίστροφη διαδικασία συμβαίνει στην πλευρά λήψης

Κάθε λειτουργία αντιγραφής καταναλώνει κύκλους CPU και εύρος ζώνης μνήμης. Η εφαρμογή πρέπει επίσης να κάνει εναλλαγή πλαισίου μεταξύ του χώρος χρήστη και του χώρος πυρήνα για κάθε λειτουργία δικτύου.

Αντίθετα, το RDMA (δεξιά πλευρά) παρέχει πραγματική μεταφορά δεδομένων μηδενικής αντιγραφής:

1. Η κάρτα δικτύου διαβάζει απευθείας από τη μνήμη GPU
2. Τα δεδομένα ταξιδεύουν απευθείας στην απομακρυσμένη κάρτα δικτύου
3. Η απομακρυσμένη κάρτα δικτύου γράφει απευθείας στη μνήμη GPU προορισμού
4. Η εφαρμογή χρειάζεται μόνο να ελέγξει μια ουρά ολοκλήρωσης (CQ) στον χώρος χρήστη για να μάθει πότε έχει ολοκληρωθεί η μεταφορά

Με το RDMA και τη σωστή αντιστοίχιση υλικού, η μεταφορά δεδομένων μεταξύ δύο GPU απαιτεί μόνο τη διέλευση του τοπικού διακόπτη PCIe και του δικτύου:

Μεγέθυνση

Αντίθετα, οι μεταφορές TCP/IP πρέπει να αντιγράφουν δεδομένα πολλές φορές μέσω της κύριας μνήμης, προκαλώντας σημαντική συμφόρηση στον δίαυλο PCIe:

Μεγέθυνση

Κατασκευή με το libfabric

Χρησιμοποιήσαμε το libfabric - ένα πλαίσιο που παρέχει μια γενική διεπαφή για υπηρεσίες fabric. Η υλοποίησή μας χρησιμοποιεί δύο τύπους λειτουργιών RDMA:

• RDMA δύο κατευθύνσεων (SEND/RECV) για μηνύματα ελέγχου που μεταφέρουν μεταδεδομένα σχετικά με περιοχές μνήμης (π.χ. μετατοπισμούς και μεγέθη).
• RDMA WRITE μίας κατεύθυνσης για την πραγματική μεταφορά δεδομένων, όπου κάθε λειτουργία WRITE χειρίζεται ένα συνεχόμενο τμήμα μνήμης.

Το ταξίδι προς την κορυφαία απόδοση

Η ανάπτυξή μας προχώρησε μέσω πολλών σταδίων:

1. Υλοποιήθηκε βασική μονοκατευθυντική μεταφορά μηνυμάτων με χρήση SEND/RECV
2. Επεκτάθηκε σε αμφίδρομη επικοινωνία
3. Προστέθηκε GPUDirect RDMA WRITE για απευθείας μεταφορά GPU-GPU
4. Επεκτάθηκε για να χειρίζεται πολλαπλές ταυτόχρονες μεταφορές
5. Εισήχθη η ουρά λειτουργιών για μεγαλύτερη ανθεκτικότητα
6. Επιτεύχθηκε αξιοποίηση εύρους ζώνης 97,4% σε μία μόνο κάρτα δικτύου

Κατά την κλιμάκωση σε 32 κάρτες δικτύου, υλοποιήσαμε αρκετές κρίσιμες βελτιστοποιήσεις:

1. Ουρά λειτουργιών: Αντί να υποβάλλουμε απευθείας λειτουργίες στις κάρτες δικτύου, διατηρούμε μια ουρά σε επίπεδο εφαρμογής. Αυτό παρέχει ανθεκτικότητα έναντι της συμφόρησης του δικτύου και απλοποιεί το μοντέλο προγραμματισμού.
2. Προθέρμανση δικτύου: Η εκ των προτέρων εγκαθίδρυση συνδέσεων βελτίωσε την απόδοση εκκίνησης.
3. Πολυνηματικότητα: Αφιερωμένα νήματα για τις λειτουργίες δικτύου κάθε GPU.
4. Καρφίτσωμα πυρήνων CPU: Σύνδεση νημάτων σε συγκεκριμένους πυρήνες CPU για την αποφυγή επιδράσεων NUMA και αστοχιών κρυφής μνήμης.
5. Καταμερισμός κατάστασης: Μείωση του ανταγωνισμού μεταξύ νημάτων.
6. Ομαδοποίηση λειτουργιών: Υποβολή πολλαπλών λειτουργιών μαζί.
7. Καθυστερημένη δημοσίευση λειτουργιών: Οι λειτουργίες τοποθετούνται πρώτα στην ουρά της εφαρμογής. Μετά την ανάκριση της ουράς ολοκλήρωσης, προσπαθούμε να υποβάλουμε τις εκκρεμείς λειτουργίες στην κάρτα δικτύου, εξασφαλίζοντας την αποτελεσματική χρήση των πόρων του δικτύου.
8. Κατανομή πόρων με επίγνωση NUMA: Κατανομή πόρων libfabric όπως ουρές ολοκλήρωσης στον σωστό κόμβο NUMA για την ελαχιστοποίηση της καθυστέρησης πρόσβασης στη μνήμη.

Μέσω αυτών των βελτιστοποιήσεων, πετύχαμε τελική απόδοση 3.108 Gbps - 97,1% του θεωρητικού μέγιστου εύρους ζώνης.

Το παρακάτω βίντεο δείχνει το πρόγραμμα γραμμής εντολών μας σε δράση. Μεταφέρει μη συνεχόμενα τμήματα μνήμης GPU σε έναν απομακρυσμένο κόμβο, επιτυγχάνοντας ταχύτητα μετάδοσης 3108,283 Gbps - αποδεικνύοντας την αξιοποίηση σχεδόν του θεωρητικού εύρους ζώνης της υποδομής δικτύου:

Συμπέρασμα

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

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

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

Παραπομπή

@online{perplexity2025HighPerformanceGpuMemoryTransferOnAwsSagemakerHyperpod,
  author = {{Perplexity} and Lequn Chen and Phi Hoang},
  title = {{High-performance GPU memory transfer on AWS Sagemaker Hyperpod}},
  date = {2025-02-10},
  url = {https://www.perplexity.ai/hub/blog/high-performance-gpu-memory-transfer-on-aws-sagemaker-hyperpod}
}