Wysokowydajny transfer pamięci GPU na AWS Sagemaker Hyperpod

Nowoczesna infrastruktura deep learning często wymaga przesyłania dużych ilości danych między procesorami GPU w różnych maszynach. W Perplexity napotkaliśmy unikalne wyzwanie techniczne: wydajne przesyłanie nieciągłych obszarów pamięci GPU między maszynami z maksymalną możliwą prędkością. Nasza docelowa platforma, instancje AWS p5, oferuje imponującą przepustowość sieciową 3200 Gb/s za pośrednictwem 32 kart sieciowych. Niniejszy artykuł opisuje naszą podróż w kierunku zbudowania niestandardowego rozwiązania sieciowego o wysokiej wydajności, które osiąga 97,1% teoretycznej przepustowości.

Wyzwanie techniczne

Nasz przypadek użycia wiązał się z kilkoma kluczowymi wymaganiami technicznymi:

• Wysokoprzepustowy transfer między zdalnymi procesorami GPU nieciągłych fragmentów pamięci
• Możliwość dynamicznego dodawania lub usuwania węzłów z wdrożeń Kubernetes bez zakłócania bieżących operacji
• Wsparcie dla wzorców komunikacji typu peer-to-peer

Chociaż biblioteka NVIDIA NCCL jest standardem de facto dla rozproszonego deep learningu, nie była ona idealna dla naszego przypadku użycia:

• NCCL świetnie sprawdza się w komunikacji zbiorowej, ale wymaga ustanowienia statycznego "świata", co wiąże się z koniecznością restartu całego klastra przy zmianie uczestniczących węzłów.
• Synchroniczny model komunikacji NCCL zwiększa złożoność w przypadku naszego asynchronicznego obciążenia
• Chcieliśmy mieć bezpośrednią kontrolę nad naszymi wzorcami transferu pamięci w celu optymalizacji
• Budowa własnego rozwiązania zapewniła cenne możliwości edukacyjne

Nowoczesne sieci o wysokiej wydajności

Aby zrozumieć nasze rozwiązanie, najpierw przyjrzyjmy się, czym nowoczesne sieci o wysokiej wydajności różnią się od tradycyjnych rozwiązań sieciowych.

Większość sieci, z których korzystamy na co dzień, opiera się na protokołach TCP/IP, w których aplikacje komunikują się z kartą sieciową za pośrednictwem jądra systemu operacyjnego przy użyciu gniazd. Jednak sieci o wysokiej wydajności wykorzystują RDMA (Remote Direct Memory Access) – zupełnie inny stos sprzętowy i programowy, który umożliwia bezpośredni dostęp do pamięci między maszynami bez angażowania procesora CPU.

AWS udostępnia Elastic Fabric Adapter (EFA), niestandardowy interfejs sieciowy, który implementuje własny protokół Amazon o nazwie Scalable Reliable Datagram (SRD). W przeciwieństwie do tradycyjnej sieci TCP/IP, gdzie dane muszą być wielokrotnie kopiowane między przestrzenią użytkownika, przestrzenią jądra i buforami sieciowymi, EFA z RDMA pozwala na bezpośredni transfer danych między pamięcią GPU a kartą sieciową, całkowicie pomijając procesor CPU.

Filozofia projektowania sieci o wysokiej wydajności

Budowa systemów sieciowych o wysokiej wydajności wymaga ponownego przemyślenia kilku fundamentalnych założeń:

1. Własność bufora: W przeciwieństwie do tradycyjnych gniazd, gdzie jądro zarządza buforami sieciowymi i wymaga kopiowania między przestrzenią użytkownika a przestrzenią jądra, RDMA wymaga od aplikacji samodzielnego zarządzania buforami. Gdy aplikacja inicjuje operację sieciową, przenosi własność bufora na kartę sieciową do momentu zakończenia operacji, eliminując potrzebę kopiowania danych.
2. Rejestracja pamięci: Aplikacje muszą rejestrować obszary pamięci w jądrze systemu operacyjnego. Jądro ustanawia mapowania adresów wirtualnych, które pozwalają procesorowi, procesorom GPU i kartom sieciowym rozumieć te same adresy wirtualne. Rejestracja ta jest jednorazową operacją, która umożliwia późniejsze transfery danych typu zero-copy.
3. Płaszczyzna sterowania a płaszczyzna danych: Sieci o wysokiej wydajności dzielą operacje na dwie kategorie:
   1. Operacje płaszczyzny sterowania (takie jak konfiguracja połączenia i rejestracja pamięci) przechodzą przez jądro, aby zapewnić bezpieczeństwo
   2. Operacje płaszczyzny danych (faktyczny transfer danych) omijają jądro w celu uzyskania maksymalnej wydajności
4. Odbiór przed transmisją: Bez buforów zarządzanych przez jądro, aplikacje muszą wstępnie publikować operacje odbioru, określając miejsce umieszczenia danych przychodzących. Jest to fundamentalna zmiana w stosunku do modelu gniazd, w którym aplikacje mogą odbierać dane w dowolnym momencie.
5. Zakończenie oparte na odpytywaniu: Zamiast czekać na powiadomienia jądra za pomocą mechanizmów takich jak epoll, aplikacje bezpośrednio odpytują kolejki zakończeń sprzętowych. Eliminuje to narzut wywołań systemowych i pozwala na natychmiastową reakcję na zakończone operacje.
6. Świadomość topologii sprzętowej: Zrozumienie i optymalizacja pod kątem topologii sprzętowej ma kluczowe znaczenie dla osiągnięcia maksymalnej wydajności.

Zrozumienie topologii sprzętowej

Instancje AWS p5 mają wyrafinowaną architekturę wewnętrzną. Jak pokazano poniżej, każda instancja zawiera dwa gniazda procesora tworzące dwa węzły NUMA, z których każdy węzeł NUMA łączy się z czterema przełącznikami PCIe:

Powiększ

Pod każdym przełącznikiem PCIe znajdziemy:

• Cztery karty sieciowe EFA o przepustowości 100 Gb/s
• Jeden procesor NVIDIA H100 GPU
• Jeden dysk NVMe SSD

Ścieżki danych dla transferów TCP/IP i RDMA wykazują fundamentalne różnice w podejściu:

Powiększ

W przypadku TCP/IP (po lewej stronie) dane muszą być wielokrotnie kopiowane:

1. Z pamięci GPU do bufora aplikacji w przestrzeni użytkownika
2. Z bufora przestrzeni użytkownika do bufora jądra
3. Z bufora jądra do karty sieciowej
4. Odwrotny proces zachodzi po stronie odbiorczej

Każda operacja kopiowania zużywa cykle procesora i przepustowość pamięci. Aplikacja musi również przełączać kontekst między przestrzenią użytkownika a przestrzenią jądra dla każdej operacji sieciowej.

W przeciwieństwie do tego, RDMA (po prawej stronie) zapewnia prawdziwy transfer danych typu zero-copy:

1. Karta sieciowa odczytuje dane bezpośrednio z pamięci GPU
2. Dane przesyłane są bezpośrednio do zdalnej karty sieciowej
3. Zdalna karta sieciowa zapisuje dane bezpośrednio w docelowej pamięci GPU
4. Aplikacja musi jedynie sprawdzić kolejkę zakończeń (CQ) w przestrzeni użytkownika, aby dowiedzieć się, kiedy transfer został zakończony

Dzięki RDMA i odpowiedniemu parowaniu sprzętowemu, przesyłanie danych między dwoma procesorami GPU wymaga jedynie przejścia przez lokalny przełącznik PCIe i sieć:

Powiększ

W przeciwieństwie do tego, transfery TCP/IP muszą kopiować dane wielokrotnie przez pamięć główną, powodując znaczne przeciążenie magistrali PCIe:

Powiększ

Budowanie z wykorzystaniem libfabric

Użyliśmy libfabric – platformy, która zapewnia ogólny interfejs dla usług sieciowych. Nasza implementacja wykorzystuje dwa rodzaje operacji RDMA:

• Dwustronne RDMA (SEND/RECV) dla wiadomości sterujących, które przenoszą metadane o obszarach pamięci (np. przesunięcia i rozmiary).
• Jednostronne RDMA WRITE dla faktycznego przesyłania danych, gdzie każda operacja WRITE obsługuje jeden ciągły fragment pamięci.

Droga do szczytowej wydajności

Nasz rozwój przebiegał przez kilka etapów:

1. Zaimplementowano podstawowy jednokierunkowy transfer wiadomości przy użyciu SEND/RECV
2. Rozszerzono o komunikację dwukierunkową
3. Dodano GPUDirect RDMA WRITE dla bezpośredniego transferu GPU-GPU
4. Rozszerzono, aby obsługiwać wiele jednoczesnych transferów
5. Wprowadzono kolejkowanie operacji dla zwiększenia solidności
6. Osiągnięto wykorzystanie przepustowości na poziomie 97,4% na pojedynczej karcie sieciowej

Podczas skalowania do 32 kart sieciowych zaimplementowaliśmy kilka kluczowych optymalizacji:

1. Kolejkowanie operacji: Zamiast bezpośredniego przesyłania operacji do kart sieciowych, utrzymujemy kolejkę na poziomie aplikacji. Zapewnia to odporność na przeciążenie sieci i upraszcza model programistyczny.
2. Rozgrzewka sieci: Uprzednie ustanawianie połączeń poprawiło wydajność przy uruchamianiu.
3. Wielowątkowość: Dedykowane wątki dla operacji sieciowych każdego procesora GPU.
4. Przypinanie rdzeni procesora: Powiązanie wątków z konkretnymi rdzeniami procesora w celu uniknięcia efektów NUMA i chybień pamięci podręcznej.
5. Sharding stanu: Redukcja rywalizacji między wątkami.
6. Grupowanie operacji: Przesyłanie wielu operacji jednocześnie.
7. Leniwe publikowanie operacji: Operacje są najpierw kolejkowane w aplikacji. Po odpytaniu kolejki zakończeń próbujemy przesłać oczekujące operacje do karty sieciowej, zapewniając wydajne wykorzystanie zasobów sieciowych.
8. Alokacja zasobów z uwzględnieniem NUMA: Przydzielanie zasobów libfabric, takich jak kolejki zakończeń, na właściwym węźle NUMA w celu zminimalizowania opóźnień dostępu do pamięci.

Dzięki tym optymalizacjom osiągnęliśmy końcową wydajność na poziomie 3108 Gb/s – 97,1% teoretycznej maksymalnej przepustowości.

Poniższy film pokazuje nasz program wiersza poleceń w działaniu. Przesyła on nieciągłe fragmenty pamięci GPU do zdalnego węzła, osiągając prędkość transmisji 3108,283 Gb/s – co dowodzi niemal teoretycznego wykorzystania przepustowości infrastruktury sieciowej:

Wnioski

Budowa systemu sieciowego o wysokiej wydajności wymaga zrozumienia zarówno architektury sprzętowej, jak i zasad projektowania systemów. Chociaż biblioteki takie jak NCCL zapewniają doskonałe rozwiązania dla typowych wzorców, czasami w przypadku specyficznych wymagań konieczne są rozwiązania niestandardowe.

Nasza podróż pokazuje, że osiągnięcie niemal teoretycznej wydajności sieci jest możliwe przy starannym zwróceniu uwagi na architekturę systemu, topologię sprzętową i różne techniki optymalizacji. Kluczem jest nie tylko zrozumienie poszczególnych komponentów, ale także tego, jak współdziałają one, tworząc kompletny system.

Pełna analiza techniczna tej podróży, w tym szczegóły implementacji i przykłady kodu, jest dostępna w naszym repozytorium open source.

Cytat

@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}
}