Włączanie modeli o bilionie parametrów w AWS EFA

W Perplexity wykorzystujemy najlepsze modele dla naszego produktu, naszych interfejsów API oraz naszych zespołów badawczych. Duże modele typu Mixture-of-Experts typu open-source, takie jak Kimi-K2, stanowią szczególne wyzwanie, ponieważ największe węzły wnioskowania z 8 procesorami graficznymi NVIDIA H200 nie mogą ich efektywnie obsłużyć, co wymaga wdrożeń wielowęzłowych. Prezentujemy zestaw kerneli dla równoległości ekspertów, które osiągają najnowocześniejsze opóźnienia na ConnectX-7, przewyższając wydajność DeepEP. Te same kernele są również pierwszymi, które osiągają opłacalne opóźnienia w AWS Elastic Fabric Adapter (EFA), umożliwiając wdrażanie modeli o bilionie parametrów.

Wypróbuj nasze kernele na GitHub i przeczytaj pełny artykuł badawczy na arXiv.

Wprowadzenie

Architektura Mixture-of-Experts (MoE) stała się standardem skalowania do setek miliardów lub bilionów parametrów, oferując jednocześnie rozsądne opóźnienia w czasie wnioskowania. Modele MoE zastępują gęstą warstwę transformatora zestawem ekspertów oraz małą warstwą liniowego routingu, która decyduje, z którym ekspertem pomnożyć token. Dzięki zgrupowaniu tokenów powiązanych z tym samym ekspertem problem staje się rzadki i łatwo równoległy, ponieważ różni eksperci mogą być teraz hostowani na różnych procesorach graficznych rozproszonych między wieloma węzłami.

W przeciwieństwie do innych form równoległości, takich jak Tensor Parallelism (TP) i Data Parallelism (DP), które można łatwo zaimplementować przy użyciu bibliotek komunikacji zbiorowej, takich jak torch.distributed i NCCL, routing MoE wiąże się z rzadką komunikacją peer-to-peer. Kernel dyspozytorski musi podzielić zestaw przychodzących tokenów i rozesłać je do rang hostujących ekspertów, podczas gdy jego problem dualny, kernel kombinowany, musi otrzymać ich przetworzone odpowiedniki i obliczyć ich ważoną średnią. Te transfery nie mapują się na operacje zbiorowe i wymagają dostosowanych kerneli w celu uzyskania niskiego opóźnienia.

W Perplexity rozwijamy zarówno przenośne kernele międzywęzłowe, jak i wyspecjalizowane kernele wewnątrz węzła, prezentując naszą pracę bardziej szczegółowo w poprzednim artykule na temat komunikacji MoE. Pokazaliśmy również w innym wpisie na blogu, że wdrożenie wielowęzłowe może przynieść niższe opóźnienia i wyższą przepustowość dla dużych modeli MoE.

W tym artykule dogłębnie analizujemy implementację naszych nowych kerneli międzywęzłowych, które znacznie poprawiają wydajność naszych wcześniejszych prac opartych na NVSHMEM. Oprócz zaprezentowania najnowocześniejszych wyników na ConnectX-7, demonstrujemy również, że w ramach naszego środowiska uruchomieniowego wnioskowania kernele mogą być skutecznie wykorzystywane do obsługi największych modeli open-source.

Wnioskowanie przez InfiniBand

Równoległość tensorowa i równoległość ekspertów mogą efektywnie skalować się w ramach jednego węzła, ponieważ NVLink oferuje wysoką przepustowość (900 GB/s) przy mikrosekundowych opóźnieniach. Chociaż systemy NVIDIA GB200 i GB300 są wyposażone w wystarczającą pamięć, aby pomieścić wagi i pamięci podręczne KV modeli o bilionie parametrów, wdrożenia na instancji AWS p5en z maksymalnie 8 procesorami graficznymi H200 są ograniczone do 1120 GB HBM, co musi zostać podzielone między wagi modelu a pamięci podręczne KV. W konsekwencji niektóre modele wymagają wdrożeń międzywęzłowych przez InfiniBand.

Chociaż InfiniBand jest nadal niesamowicie szybki, oferuje tylko 400 Gb/s (50 GB/s) przepustowości, co dodaje od dziesiątek do setek mikrosekund opóźnienia dla routingu MoE. Na szczęście dodane opóźnienia można ukryć poprzez zastosowanie współdzielonych ekspertów, mikrobatchingu i nakładania się obliczeń z komunikacją, aby zapewnić opłacalną przepustowość przy konkurencyjnym koszcie. DeepSeek zapoczątkował te techniki poprzez DeepSeek-R1, wdrożony na klastrach procesorów graficznych H800 i adapterach NVIDIA ConnectX-7 InfiniBand.

W klastrach AWS sytuacja jest inna ze względu na zastosowanie ich niestandardowych kart Elastic Fabric Adapters (EFA). Mimo że zapewniają one szczytową przepustowość 400 Gb/s przy operacjach zbiorowych typowo używanych w zadaniach treningowych, nieco ustępują ConnectX-7 pod względem rozmiarów komunikatów wymienianych podczas wysyłania i łączenia MoE. Dodatkowo EFA nie obsługuje GPUDirect Async, co wymaga wątku pośredniczącego CPU łączącego procesor graficzny z kartą sieciową (NIC) w celu zainicjowania transferów. Poza zwiększoną złożonością, dodatkowe transakcje PCIe dodają mikrosekundy narzutu do wszystkich transferów. Frameworki takie jak NVSHMEM pozwalają twórcom kerneli na dostarczanie implementacji niezależnych od urządzenia, jednak ogólne interfejsy API generują narzuty, które czynią wdrożenia nieefektywnymi pod względem kosztów i opóźnień.

Pplx-kernels oparte na NVSHMEM

Nasze poprzednie kernele abstrahowały podstawową implementację transportu poprzez NVSHMEM. Polegały one na nvshmemx_putmem_signal_nbi_warp do przesyłania tokenów jeden po drugim, używając liczników atomowych do synchronizacji peerów. Chociaż oferowało to rozsądną wydajność na kartach sieciowych ConnectX-7 z IBGDA, implementacja IBRC ConnectX-7 oparta na proxy była znacznie wolniejsza, podczas gdy implementacja EFA nie była w stanie zapewnić opóźnień poniżej milisekundy.

Wykorzystaliśmy kernele jako środowisko testowe dla bazowego sprzętu, jednocześnie ustalając najbardziej odpowiedni interfejs do udostępnienia naszym modelom. Postawiliśmy hipotezę, że częste transfery i interakcje między kernelami a kartami sieciowymi lub wątkami proxy pogarszały wydajność. Ponadto nasza implementacja nie opierała się na żadnych gwarancjach ze strony sprzętu poza niezawodnością, co sugerowało, że moglibyśmy przyspieszyć działanie poprzez zmniejszenie narzutu prymitywów synchronizacji.

W ślad za tymi obserwacjami opracowaliśmy nasze kernele oparte na EFA, udostępniające ten sam interfejs, optymalizując interakcje między kernelem urządzenia a wątkiem proxy hosta pod kątem najwyższej wydajności.

MoE all-to-all przez EFA

Powiększ

Nowe kernele wykorzystują hybrydową architekturę CPU-GPU, w której kernele GPU wchodzą w interakcję z modelem na urządzeniu, a wątek proxy działający na procesorze CPU hosta obsługuje interakcje z kartą sieciową. Nowe kernele ponownie wykorzystują ten sam TransferEngine, który stworzyliśmy dla transferów pamięci podręcznej KV.

Routing MoE jest wykonywany przez kernel dyspozytorski, który otrzymuje tokeny i ekspertów per-token w celu skierowania ich z modelu, wysyłając je przez sieć do rang hostujących ekspertów. Po tym, jak zgrupowane kernele GEMM pomnożą tokeny z wagami ekspertów, kernel kombinowany przenosi tokeny do ich rang źródłowych i oblicza ich średnią ważoną przed zwróceniem ich do modelu.

Zarówno kernele dyspozytorskie, jak i kombinowane są podzielone na część wysyłającą i odbiorczą, aby ułatwić nakładanie się operacji i mikrobatching. Nadawca tasuje tokeny do buforów wysyłania, dzięki czemu można je zapisać do peerów za pomocą pojedynczej operacji zapisu RDMA. Po wykonaniu nadawcy procesor graficzny może wykonywać inną pracę (taką jak mnożenie współdzielonych ekspertów), podczas gdy transfery przez sieć oczekują w tle. Kernele odbiorcze blokują się do momentu odebrania wszystkich transferów. W tle wątek roboczy po stronie hosta obsługuje transfery: jest sygnalizowany przez GPU, gdy bufory wysyłania zostały wypełnione, po czym wysyła zapisy RDMA. Gdy karta sieciowa zasygnalizuje ukończenie, odblokowuje odbiornik, który zwraca tokeny do modelu.

Kernele są ściśle sprzężone z hostem, komunikując się ze sobą poprzez zunifikowaną pamięć i GDRCopy. Zunifikowana pamięć CUDA pozwala na mapowanie pamięci urządzenia lub pamięci hosta w przestrzenie adresowe obu, transparentnie obsługując komunikację przez PCIe.

Podczas gdy zunifikowana pamięć jest odpowiednia dla transferów masowych, polling o niskim opóźnieniu odbywa się poprzez GDRCopy. GDRCopy wykorzystuje GPUDirect RDMA do odczytu i zapisu w pamięci GPU przy mikrosekundowym opóźnieniu.

Oprócz RDMA, w ramach węzła kernele wykorzystują również NVLink do wymiany tokenów. Może to odciążyć znaczną część ruchu, do jednej czwartej dla EP32 i do jednej ósmej dla EP64, osiągając ogólnie niższe opóźnienia.

Bufory wysyłania i odbierania

Projekt buforów wysyłania dla kernela dyspozytorskiego ma kluczowe znaczenie dla osiągnięcia opłacalnych opóźnień z udziałem proxy hosta. Jeśli istnieje N rang, z których każda wysyła T tokenów, każdy skierowany do R ekspertów z E, ranga może otrzymać do N * T * max(R, E / N) tokenów w najgorszym możliwym przypadku, jeśli wszystkie są do niej kierowane. Ponadto to minimalne ograniczenie można osiągnąć tylko wtedy, gdy nadawcy układają swoje zapisy w sposób ciągły, ponieważ prywatne bufory odbierania zużyłyby zbyt dużo pamięci. Aby nadawcy mogli układać swoje zapisy kolejno bez konfliktów, pełna wiedza o informacjach o routingu musi zostać rozesłana do każdej rangi. W konsekwencji nasze kernele dyspozytorskie najpierw wymieniają liczby tokenów na eksperta, po czym każdy z nich może określić, dokąd trafiają jego zapisy i z jakiego przesunięcia otrzymują tokeny od każdej rangi.

Sama wymiana informacji o routingu skutkowałaby niepotrzebnie wysokimi opóźnieniami, ponieważ każda ranga musiałaby czekać na nadejście zapisów przed wysłaniem własnych. Zamiast tego rezerwujemy niewielką ilość prywatnej przestrzeni na nadawcę w każdej randze odbierającej i wysyłamy część tokenów wraz z zapisami informacji o routingu. Zapewnia to pełne wykorzystanie przepustowości przez większość procesu wysyłania. Po wymianie tras pozostałe tokeny są przesyłane za pomocą pojedynczego zapisu na zdalnego peera. Wymiana informacji o routingu i bufory o stałym rozmiarze, wraz z transferami je wypełniającymi, zostały zilustrowane na poniższym rysunku.

Powiększ

Kernele kombinowane są znacznie prostsze, ponieważ informacje o routingu są już dla nich dostępne z etapu wysyłania. Mogą one kopiować tokeny w pojedynczym zapisie na peera, po czym odczytują z bufora odbierania, aby obliczyć średnią ważoną tokenów do tensorów wyjściowych.

W przypadku NVLink bufory wysyłania i odbierania w ramach węzła są mapowane do peerów, aby mogli oni bezpośrednio wypychać lub pobierać tokeny z sąsiednich rang.

Wysyłanie i łączenie

Poniższe rysunki szczegółowo opisują czas i interakcje między GPU, CPU i kartą sieciową, pokazując jedną iterację wysyłania i łączenia przeplataną z kernelami modelu.

Praca rozpoczyna się od wysyłania dyspozytorskiego, gdzie kernel agreguje informacje o routingu i oblicza liczby tokenów na eksperta poprzez akumulację atomową w pamięci współdzielonej w ramach jednego bloku. Gdy liczby tokenów są dostępne, są one kopiowane do pamięci hosta, a proxy jest natychmiast powiadamiane o rozpoczęciu wysyłania informacji o routingu do wszystkich innych rang peerów. Równolegle dyspozycja kontynuuje obliczanie sumy skumulowanej liczb tokenów, identyfikując przesunięcia wymagane do spakowania tokenów do ciągłych buforów wysyłania per-peer. Gdy przesunięcia są dostępne, wszystkie tokeny do wysłania przez EFA są pakowane do buforów wysyłania. Proxy jest powiadamiane o umieszczeniu zapisów w buforach wysyłania międzywęzłowych peerów. Wewnątrz węzła kernele również natychmiast wypychają tokeny do prywatnych buforów w tym samym węźle poprzez NVLink.

Powiększ

Powiększ

Gdy ranga otrzyma informacje o routingu od wszystkich peerów, centralizuje je i oblicza różne przesunięcia niezbędne do identyfikacji tokenów do wysyłania i łączenia. Ta praca jest wykonywana na procesorze CPU, podczas gdy transfery już oczekują, więc nie jest to krytyczne dla opóźnień. Następnie, gdy wszystkie przesunięcia są znane, zapisy dla pozostałych tokenów, które przekraczają rozmiar bufora prywatnego, są wysyłane, a faza wysyłania dyspozycji jest zakończona.

Po stronie odbiorczej kernel najpierw pobiera tokeny z tego samego węzła, które wykraczają poza pojemność buforów prywatnych poprzez NVLink. Następnie blokuje się, aż proxy hosta zasygnalizuje zakończenie wszystkich transferów. Gdy sygnał zostanie odebrany, tokeny są tasowane z buforów odbierania do tensora wyjściowego przekazanego do grupy GEMM. Ponieważ bufory są ponownie używane między wysyłaniem a łączeniem, bariery NVLink i RDMA zapewniają, że bufory odbiornika zostały odczytane przed zezwoleniem następnemu kernelowi na ich nadpisanie.

Kernele kombinowane postępują według mniej więcej podobnego wzoru: tasowanie tokenów, sygnalizowanie proxy i czekanie na zakończenie. Jednak metadane do wysyłania ich zapisów są łatwo dostępne, a wszystkie zdalne przesunięcia są znane, więc wszystkie tokeny przesłane wewnątrz węzła są wypychane przez NVLink do odbiorców, a tokeny międzywęzłowe są przesyłane za pomocą pojedynczego zapisu. Przed zwróceniem wyników kernel akumuluje tokeny z buforów odbiornika, zapisując je w tensorach wyjściowych.

TransferEngine dla routingu MoE

Aby przygotować prototyp naszej implementacji, początkowo ponownie wykorzystaliśmy Transfer Engine zbudowany do zdezagregowanego prefill i dekodowania w celu wydawania i śledzenia zapisów RDMA. Jednak w przeciwieństwie do zdezagregowanego dekodowania, gdzie każdy węzeł wysyła dużą liczbę zapisów do tego samego peera, dyspozycja MoE wysyła małe transfery do wszystkich peerów w jednej operacji. Od razu wydajność nie była optymalna, a nasza implementacja pozostawała w tyle za DeepEP na ConnectX-7 o około 20 us. W konsekwencji wyspecjalizowaliśmy TransferEngine, aby obsługiwał dwie operacje: scatter i barrier, w celu optymalizacji interakcji z dużą liczbą peerów.

Zarówno operacje scatter, jak i barrier udostępniają grupy peerów, które można wstępnie zarejestrować, aby wstępnie wypełnić struktury danych używane przez podstawową implementację transportu (szablonowanie Work Request (WR)). Scatter kopiuje różne wycinki z tego samego bufora źródłowego do różnych przesunięć we wszystkich członkach grupy peerów. Barrier sygnalizuje wszystkich peerów za pomocą operacji immediate. Kolejkowanie i obsługa tych operacji zbiorczo oszczędza cenne mikrosekundy na ścieżce krytycznej. Na EFA, gdzie dwie karty sieciowe 200 Gb/s agregują się, zapewniając przepustowość 400 Gb/s, shardujemy grupy peerów między dwie karty sieciowe zamiast shardowania bajtów każdego transferu.

Chociaż początkowo rozwijaliśmy TransferEngine wyłącznie na EFA przy użyciu libfabric, kernele MoE działają na ConnectX-7 bez zmiany ani jednej linii kodu po dodaniu wsparcia dla ConnectX-7 w TransferEngine przy użyciu libibverbs. W porównaniu do EFA opartej na Scalable Reliable Datagram (SRD), ConnectX-7 wymaga dodatkowej pracy przy konfiguracji połączenia i zarządzaniu peerami. Używamy pary kolejek Unreliable Datagram (UD), aby uruchomić Reliable Connection (RC). Tworzymy dwie pary kolejek dla każdego zdalnego peera, jedną dla dwustronnego SEND/RECV i jedną dla jednostronnego WRITE_IMM, ponieważ zarówno RECV, jak i WRITE_IMM zużywają operację RECV w kolejności, podczas gdy obie operacje wymagają różnych rozmiarów bufora. Ponadto rozluźniamy kolejność transakcji PCIe między ConnectX-7 a pamięcią GPU. Podobnie jak w przypadku EFA, implementujemy optymalizację szablonowania WR. Dodatkowo, podczas wysyłania wielu małych komunikatów do tego samego peera, używamy łańcuchowania WR, aby poprawić szybkość komunikatów.

Ponieważ praca po stronie GPU jest identyczna w obu adapterach, przesunęliśmy podstawową implementację transportu do jej granic, aby przekroczyć wydajność DeepEP, podkreślając, że różnica między routingiem MoE EFA a ConnectX-7 jest określona przez możliwości podstawowego sprzętu.

Ewaluacja

Opóźnienie wysyłania/łączenia

Oceniliśmy opóźnienia wysyłania i łączenia w systemach H200 wyposażonych w adaptery ConnectX-7 i EFA w standardowej konfiguracji DeepSeek-V3 i Kimi-K2 z ukrytym wymiarem 7168, blokowo skalowanym wysyłaniem fp8 i łączeniem bf16. Rozważamy rozmiar batcha 128 dla wysyłania i 4096 dla łączenia. Na ConnectX-7 porównujemy się z DeepEP i naszymi poprzednimi kernelami pplx opartymi na NVSHMEM, używając zarówno GPUDirect Async (IBGDA), jak i proxy hosta (IBRC). Osiągamy połączone opóźnienia 459 us, 582 us i 692 us odpowiednio na EP16, EP32 i EP64. Przewyższamy inne kernele oparte na EFA polegające na ogólnych proxy, takich jak UCCL-EP, które osiągają 519 us, 966 us i 1159 us.

Powiększ

Podczas dekodowania, mimo że nasze kernele są o kilka mikrosekund wolniejsze od DeepEP, osiągamy ogólnie niższe opóźnienia, ponieważ nasze kernele kombinowane są znacznie szybsze. Chociaż używamy proxy hosta, ze względu na nasze transfery zbiorcze przekraczamy wydajność wszystkich implementacji bez proxy opartych na IBGDA. Pokazujemy również, że implementacja proxy ma znaczenie, ponieważ nasze wysoce wyspecjalizowane proxy jest rząd wielkości szybsze niż ogólna implementacja NVSHMEM. Pomimo faktu, że EFA jest dwa razy wolniejsza niż ConnectX-7 podczas przesyłania pakietów 256 KB, które są bliskie typowemu rozmiarowi pakietu wymienianego przez routing MoE, różnica między dwoma adapterami jest ogólnie mniej istotna.

Powiększ

Mimo że nasze kernele zoptymalizowane pod kątem dekodowania mogą skalować się do rozmiarów batcha prefill bez żadnych dodatkowych poprawek, wydajność pozostaje w tyle za DeepEP. Wynika to głównie z faktu, że akumulujemy tokeny tylko po stronie odbiorczej łączenia, podczas gdy DeepEP akumuluje w nadawcy, aby zmniejszyć ilość przesyłanych danych. Jednak ta optymalizacja może mieć wpływ na dokładność. Ponadto, ponieważ kernele przesyłają dane w jednym zapisie, wykorzystujemy więcej pamięci na bufory wysyłania i odbierania.

Przepustowość Kimi-K2 / DeepSeek-V3

Aby udowodnić opłacalność naszych kerneli, benchmarkujemy kompleksową wydajność naszego wewnętrznego silnika wnioskowania obsługującego DeepSeek-V3 (671 mld parametrów) i Kimi-K2 (1 bln parametrów) na wielu instancjach p5en wyposażonych w Elastic Fabric Adapters (EFA) i procesory graficzne H200. Uruchamiamy model w trybie weryfikacji szkicu i raportujemy opóźnienie przy założeniu wskaźnika akceptacji 80% przy użyciu jednej warstwy MTP.

W przypadku DeepSeek-V3 porównujemy się z bazą intranode EP=8 z wydajnymi kernelami wysyłania i łączenia tylko przez NVLink. Wdrożenia międzywęzłowe ogólnie dorównują lub przekraczają wydajność wdrożenia jednowęzłowego, szczególnie w przypadku średniej wielkości batchy, co pokazuje, że skalowanie między węzłami przynosi poprawę ogólnej przepustowości.

W przypadku Kimi-K2 kernele umożliwiają obsługę. Ze względu na swój rozmiar model nie mieści się w tym samym węźle przy użyciu Data Parallelism (DP) na węźle H200, więc do obsługi modelu wymagane jest wdrożenie międzywęzłowe. Ponieważ model ma mniej głów atencji (attention heads), osiągamy lepsze opóźnienia niż DeepSeek-V3 i możemy obsługiwać model na EFA przy opłacalnym opóźnieniu.

Powiększ

Powiększ

Rozważaliśmy również mikrobatching, chociaż poprawił on opóźnienia tylko przy bardzo dużych rozmiarach batchy lub przy użyciu więcej niż jednego tokena szkicu. Aby skorzystać z dzielenia problemów zgrupowanych GEMM na dwa, rozmiary batchy tokenów muszą być wystarczająco gęste. Większość czasu nadal jest spędzana na wysyłaniu i łączeniu MoE, co wskazuje, że wszelkie dalsze ulepszenia kerneli w konsekwencji poprawią opóźnienia i ułatwią mikrobatching.

Przyszłe prace

Ściśle współpracujemy z inżynierami AWS, aby poprawić wydajność na EFA. Ze swojej strony będziemy śledzić aktualizacje libfabric, które zmniejszą narzut płaszczyzny danych. Ponadto poeksperymentujemy z efa-direct, aby jeszcze bardziej zmniejszyć narzuty przestrzeni użytkownika i bezpośrednio komunikować się ze sprzętem.

Odniesienia

Raport techniczny DeepSeek R1

Raport techniczny Kimi K2

DeepEP

UCCL-EP

Poprawa wydajności sieci systemów HPC przy użyciu NVIDIA Magnum IO NVSHMEM i GPUDirect Async

Cytat

@online{perplexity2025EnablingTrillionParameterModelsOnAwsEfa,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{Enabling trillion-parameter models on AWS EFA}},
  date = {2025-11-04},
  url = {https://www.perplexity.ai/hub/blog/enabling-trillion-parameter-models-on-aws-efa}
}