AWS opisało architekturę uczenia ze wzmocnieniem dla modeli mixture-of-experts (MoE), łączącą Amazon EKS, Elastic Fabric Adapter (EFA) i Amazon S3. Firma podała również wynik pomiaru: konfiguracja z DeepEP działającym przez EFA osiągnęła w określonym porównaniu o 40% większą łączną przepustowość rolloutów. Jest to wynik przedstawiony przez AWS dla opisanego układu, dlatego jego zakres wyznaczają szczegóły testowanych konfiguracji.

Podział pracy w klastrze

W opisanym klastrze EKS poszczególne grupy instancji mają różne zadania. Instancje z GPU generują rollouty, wykonują inferencję modelu nagrody i trenują politykę. Instancje z CPU obsługują środowiska oraz wstępne przetwarzanie. Instancje zoptymalizowane pod kątem pamięci służą jako bufory doświadczeń i pamięci podręczne punktów kontrolnych. Ten podział pokazuje, gdzie doświadczenia powstają, gdzie są tymczasowo przechowywane i jak trafiają do treningu.

Doświadczenia z rolloutów wpływają do bufora na instancjach zoptymalizowanych pod kątem pamięci. Trening polityki pobiera z niego partie danych i publikuje zaktualizowane punkty kontrolne. Punkty kontrolne oraz ukończone artefakty są zapisywane w Amazon S3. Według AWS EKS odpowiada za planowanie, skalowanie, odzyskiwanie po awariach i koordynację grup pracowników. Te funkcje spinają generowanie doświadczeń z kolejnymi etapami treningu.

Komunikacja między ekspertami

W obrębie instancji z GPU architektura używa NVLink i NVSwitch, a do komunikacji między instancjami używa EFA. DeepEP ma jądro dispatch, które kieruje tokeny do zdalnych ekspertów, oraz jądro combine, które zbiera przetworzone tokeny. Opisana ścieżka między węzłami korzysta z libfabric przez EFA. AWS podaje, że przeniosło mechanizmy komunikacyjne DeepEP do libfabric i zapewniło natywną obsługę EFA w DeepEP v2.

Każdy z tych elementów pełni inną rolę w opisanym przepływie pracy. EKS koordynuje grupy pracowników. Bufor przechowuje doświadczenia na potrzeby treningu, a S3 punkty kontrolne i gotowe artefakty. NVLink i NVSwitch obsługują komunikację wewnątrz instancji GPU, natomiast EFA łączy instancje. Operacje dispatch i combine w DeepEP dotyczą przesyłania tokenów do ekspertów i zbierania wyników. Taki opis pozwala odczytać wynik pomiaru w kontekście całej konfiguracji.

Co obejmuje wynik 40%

Porównanie przedstawione przez AWS wykorzystało 48 instancji P5en uruchamiających bardzo rzadki model MoE. Szesnaście instancji przeznaczono na trening, a 32 na inferencję. Konfiguracją wyjściową był zestaw oparty na Slime, bez równoległości ekspertów przyspieszanej przez EFA. Ulepszony zestaw wykorzystywał DeepEP przez EFA. AWS podaje dla tych dwóch konfiguracji wzrost łącznej przepustowości rolloutów uczenia ze wzmocnieniem o 40%.

Podana wartość dotyczy łącznej przepustowości rolloutów w tym konkretnym porównaniu. Dostępne informacje nie potwierdzają takiego samego wyniku dla innych modeli, typów instancji ani wdrożeń. Nie pozwalają też wydzielić części wzrostu wynikającej wyłącznie z DeepEP przez EFA, ponieważ porównywane zestawy oprogramowania różniły się również innymi komponentami. Wniosek dotyczy więc zaobserwowanej różnicy między opisanymi konfiguracjami.

Przerwane zadania i wymagania

AWS opisuje ponadto obsługę przerwanych zadań rolloutów na Spot Instances: zadanie może wrócić do kolejki i zostać przydzielone ponownie, podczas gdy inni pracownicy nadal generują doświadczenia. To wyjaśnia sposób ponownego rozdzielenia pracy po przerwaniu. Źródło nie podaje zmierzonych oszczędności kosztów związanych ze Spot Instances, więc z tego mechanizmu nie wynika konkretna kwota oszczędności.

Wśród wymagań instrukcja wymienia Amazon EKS 1.31 lub nowszy, instalator EFA 1.49 z wtyczką AWS OFI NCCL, DeepEP 2.0.0, NCCL 2.31.2, SGLang 0.5.17 oraz PyTorch 2.12.1 z CUDA 13.0. Wersje te opisują kontekst programowy przedstawionego rozwiązania. Architekturę, konfiguracje i wynik pomiaru AWS przedstawiło w swoim artykule technicznym.