Efektywna i przenośna komunikacja Mixture-of-Experts

Prezentujemy wysokowydajną, przenośną bibliotekę open-source do komunikacji Mixture-of-Experts (MoE), która osiąga 10-krotnie większą wydajność w porównaniu ze standardowymi prymitywami komunikacyjnymi All-to-All.

Nasza implementacja zawiera szereg kluczowych innowacji technicznych, które zapewniają doskonałą wydajność komunikacji MoE:

Komunikacja inicjowana przez GPU (IBGDA): Obsługuje bezpośrednią komunikację GPU-to-NIC, znacząco redukując opóźnienia poprzez pominięcie zaangażowania procesora CPU.

Nakładanie się komunikacji i obliczeń: Architektura dzielonego jądra z oddzielnymi etapami wysyłania i odbierania umożliwia prowadzenie obliczeń jednocześnie z transferami sieciowymi.

Najwyższa wydajność w pojedynczym węźle: 2,5-krotnie niższe opóźnienia niż w przypadku dotychczas najszybszej implementacji w konfiguracjach z pojedynczym węzłem.

Wydajna i przenośna wydajność wielowęzłowa: Nasza implementacja osiąga prędkości do 10 razy większe niż standardowa komunikacja all-to-all. Mimo że jest około 2 razy wolniejsza od wysoce wyspecjalizowanych implementacji, nasze podejście oferuje lepszą przenośność między wersjami NVSHMEM a środowiskami sieciowymi (NVLink, CX-7 i EFA).

Biblioteka jest w pełni open-source i dostępna pod adresem https://github.com/ppl-ai/pplx-kernels.

W tym artykule zgłębiamy wyzwania związane z równoległością ekspertów w wielkoskalowych modelach MoE i opisujemy nasze podejście do efektywnego przesyłania i łączenia tokenów w rozproszonych środowiskach GPU.

Wstęp

Modele Mixture-of-Experts (MoE), takie jak DeepSeek R1 i Mixtral 8x7B, ulepszają modele gęste poprzez ograniczenie liczby wag aktywowanych dla każdego tokena. Zamiast osiągać wysoką liczbę parametrów poprzez zwiększanie liczby warstw lub rozmiaru warstw liniowych w wielowarstwowym perceptronie (MLP) każdej warstwy dekodera, modele MoE zastępują tradycyjny MLP wieloma ekspertami i routerem. Na przykład, z 671 mld parametrów modelu DeepSeek R1, tylko 37 mld jest mnożonych przez dany token podczas wnioskowania. Znacznie zmniejsza to ilość obliczeń wymaganych do zdekodowania tokena i poprawia opóźnienia w porównaniu z modelem gęstym.

Modele MoE stanowią dodatkowe wyzwanie dla wnioskowania w porównaniu z modelami gęstymi. Podczas gdy sami eksperci są małymi warstwami MLP, każda warstwa dekodera zawiera router, który decyduje, do których ekspertów wysyłany jest token, przy czym każdy token jest wysyłany do wielu ekspertów. Router jest zazwyczaj małą warstwą liniową wytwarzającą rozkład prawdopodobieństwa. Zazwyczaj wybierani są eksperci z wynikami top-K, a ostateczna aktywacja jest obliczana jako średnia ważona, sumująca wyniki ekspertów pomnożone przez wagę pochodzącą z rozkładu prawdopodobieństwa routera.

Aby zminimalizować opóźnienia w rozproszonych systemach MoE, można wykorzystać równoległość między wieloma urządzeniami, ale wprowadza to wyzwania komunikacyjne. Modele takie jak Mixtral 8x7B lub Llama-70B mieszczą się na 8 urządzeniach w obrębie jednego węzła, korzystając z szybkich i nisko-opóźnieniowych połączeń NVLink (do 900 Gb/s). Jednak większe modele wymagają shardingu ekspertów pomiędzy wieloma węzłami przy użyciu InfiniBand (z szczytową przepustowością 400 Gb/s), co wprowadza dodatkowe wyzwania związane z opóźnieniami.

Powiększ

W tym artykule badamy problem równoległości ekspertów, opisując naszą implementację jąder GPU, które skutecznie wysyłają tokeny do urządzeń obsługujących przypisanych ekspertów i zbierają je z powrotem, obliczając ważoną sumę wyników ekspertów. Choć istnieją agresywnie zoptymalizowane implementacje ukierunkowane na konkretny sprzęt sieciowy, prezentujemy wysoce przenośną implementację, która opiera się na minimalnym zestawie prymitywów NVSHMEM, a mimo to zapewnia wydajność 10-krotnie większą niż standardowa komunikacja all-to-all. Nasza implementacja osiąga najnowocześniejszą wydajność w konfiguracjach z pojedynczym węzłem, zachowując przy tym doskonałą przenośność w różnych środowiskach sieciowych.

Sharding i równoległość

Dla efektywnego wnioskowania wagi modelu muszą być przechowywane w pamięci urządzenia, przy jednoczesnym pozostawieniu wystarczającej przestrzeni na aktywacje, pamięci podręczne KV i inne bufory wymagane przez przebieg w przód przez model. Najpotężniejsze modele przekraczają pojemność nawet najwydajniejszych procesorów graficznych, dlatego wnioskowanie musi być rozproszone na wiele urządzeń, które mogą łącznie przechowywać wagi modelu. W oparciu o schematy shardingu wag należy stosować różne schematy komunikacji i obliczeń, aby odciążyć obliczenia i zsynchronizować urządzenia.

Powiększ

Równoległość ekspertów (EP), zilustrowana na pierwszym rysunku, równolegli jedynie obliczenia ekspertów. Różni eksperci są przypisywani do różnych urządzeń, które przechowują ich wagi. Po routingu tokeny są wysyłane do odpowiedniego urządzenia, a wyniki są następnie gromadzone i sumowane. Złożoność routingu zależy od stopnia równoległości w innych częściach modelu: replikacja innych warstw mogłaby całkowicie wyeliminować potrzebę routingu, ponieważ każda ranga może wybierać tokeny z lokalnie zreplikowanej tabeli routingu. Jeśli jednak tylko jedna z rang przeprowadza routing, wymagane jest rozgłaszanie (broadcast) w celu wysłania tokenów, indeksów i wag do odpowiednich ekspertów. Wreszcie, all-gather lub all-to-all broadcast synchronizuje wyjściowe tokeny z dowolną rangą, która kontynuuje wykonywanie modelu. Taka implementacja jest stosunkowo prosta, ponieważ torch udostępnia już wymagane prymitywy, chociaż pewne korzyści można by uzyskać poprzez połączenie akumulacji z operacją gather zbierającą tokeny od ekspertów.

Powiększ

Równoległość tylko ekspertów nie skaluje się idealnie, ponieważ węzły w klastrze mogą być bezczynne, gdy model wykonuje warstwy niebędące ekspertami, takie jak uwaga (attention), norma i próbkowanie. Jednak obliczenia tych warstw, przede wszystkim uwagi, również mogą czerpać korzyści z równoległości tensorów (TP). Większość modeli opiera się na wielogłowicowej uwadze, co oznacza, że głowice uwagi i odpowiadające im projekcje Q, K i V mogą być również shardowane między urządzeniami, replikując lub zbierając wycinki między różnymi warstwami. Jeśli uwaga jest rozproszona na wszystkie urządzenia, all-gather może zsynchronizować aktywacje, umożliwiając replikację routingu, co wymaga prymitywów synchronizacji podobnych do przypadku równoległości ekspertów dla wydajnej implementacji. Istnieją jednak ograniczenia równoległości na tym poziomie, ponieważ zmniejszenie liczby głowic uwagi poniżej określonego progu przyniesie malejące zyski.

Powiększ

Aby jak najlepiej wykorzystać wszystkie urządzenia i wspierać bardzo wysoki stopień równoległości ekspertów, do 128 lub 256 GPU, wymagana jest równoległość danych (DP). W ramach tego schematu urządzenia, na które podzielony jest model, są grupowane w celu jednoczesnego obsługiwania żądań, obliczania uwagi i utrzymywania pamięci podręcznych KV shardowanych w ramach ich lokalnej grupy. Wiele instancji tych grup równoległych współpracuje przy ocenie ekspertów, z których każda hostuje inny podzbiór wag ekspertów. W oparciu o liczbę głowic uwagi grupa może zazwyczaj obejmować wielkość całego węzła, ponieważ komunikacja wewnątrz węzła jest szybsza. Na przykład, na rysunku powyżej, jedna ranga DP niezależnie obsługuje dwa żądania, A i B, zajmując się uwagą, normą i wszelkimi innymi operacjami księgowymi dla żądań. Druga ranga DP przetwarza odrębne żądanie C. Jednak pierwszy węzeł hostuje połowę ekspertów, podczas gdy drugi węzeł drugą połowę, więc po routingu tokeny z A i B mogą zostać wysłane do drugiego węzła i odwrotnie. Prowadzi to do rzadkiego problemu komunikacyjnego: każde urządzenie może wysyłać różną liczbę tokenów do dowolnej innej rangi docelowej. Istniejące prymitywy z torcha, głównie all_to_all, nie są szczególnie dobrze przystosowane, ponieważ mogą wymagać pewnej formy dopełnienia (padding) lub synchronizacji GPU-to-CPU i rozgłaszania metadanych. Aby skutecznie wdrożyć komunikację, wymagane są niestandardowe jądra do przetwarzania informacji o routingu i inicjowania komunikacji z urządzeń bez interwencji procesora CPU. Po routingu jądro wysyłające musi wysłać tokeny do rang, do których zostały skierowane, podczas gdy po stronie łączenia aktywacje należące do żądań w bieżącej grupie DP muszą zostać zebrane. Ponadto praca musi być zbalansowana w obrębie grupy DP, która może mieć wiele urządzeń, w celu zminimalizowania rozgłaszania i poprawnej replikacji tokenów we wszystkich rangach dla kolejnych warstw.

NVSHMEM

NVSHMEM to specyficzna dla NVIDIA implementacja OpenSHMEM, zapewniająca przenośne funkcje komunikacji między urządzeniami, które abstrahują od złożoności sprzętu. Interfejs API może wyrażać odczyty i zapisy między urządzeniami, które są mapowane na prymitywy poszczególnych warstw transportowych. Obsługiwane są zarówno NVLink, jak i RDMA, co daje pewien stopień przenośności i pozwala jądrom NVSHMEM działać wewnątrz węzła lub w klastrze wielu węzłów. W naszych jądrach używamy zdalnych zapisów do pamięci i operacji sygnalizacji w celu przesyłania danych i synchronizacji uczestniczących urządzeń.

Operacje NVSHMEM opierają się na koncepcji pamięci symetrycznej: działają na buforach, które zostały przydzielone na wszystkich urządzeniach uczestniczących w komunikacji między urządzeniami. Lokalne urządzenie może użyć adresu pochodzącego z własnego lokalnego bufora wraz z indeksem rangi docelowej, aby określić miejsce docelowe operacji. Poniższy rysunek ilustruje tę koncepcję: oba procesory graficzne przydzielają symetryczne bufory tej samej wielkości, zachowując wskaźniki src do nich. Pierwsza ranga chce wysłać 3 liczby całkowite do drugiej, umieszczając je na początku bufora: nvshmem_int_put_nbi wyprowadza adres początkowy z lokalnego bufora, określając urządzenie docelowe. Druga ranga wyprowadza przesunięcie ze swojego własnego bufora, wysyłając jeden element do pierwszego urządzenia, przesunięte o jeden. Podczas gdy adresy docelowe muszą zawsze być symetrycznymi buforami przydzielonymi przy użyciu nvshmem_alloc, bufory źródłowe mogą być dowolnymi regionami pamięci urządzenia, pod warunkiem, że są wstępnie zarejestrowane w NVSHMEM.

Podczas gdy NVSHMEM zapewnia szeroką gamę prymitywów, nasze jądra opierają się tylko na 3 funkcjach, budując na nich całą synchronizację i grodzenie.

Powiększ

nvshmemx_putmem_signal_nbi_warp: Przesyła blok danych z jednego urządzenia do drugiego, jednocześnie ustawiając flagę na zdalnym urządzeniu. Operacja albo ustawia (NVSHMEM_SIGNAL_SET), albo zwiększa (NVSHMEM_SIGNAL_ADD) 64-bitową lokalizację. Flaga jest aktualizowana po przesłaniu całego bloku pamięci. Jeśli zdalne urządzenie zaobserwowało zmianę flagi, może bezpiecznie uzyskać dostęp do bufora w swojej własnej pamięci. Ta funkcja jest przydatna do łączenia transferu danych i synchronizacji.

nvshmemx_signal_op działa na pojedynczej lokalizacji pamięci, zazwyczaj 64-bitowej fladze, atomowo ustawiając ją lub zwiększając. Jest przydatna przy przesyłaniu metadanych i synchronizacji urządzeń.

nvshmem_uint64_wait_until jest używana po stronie odbiorczej sygnału, aby odpytywać flagę, dopóki zdalne urządzenie jej nie zaktualizuje.

Zapewnienie jakiejkolwiek formy porządkowania operacji poprzez nvshmem_quiet lub nvshmem_fence jest kosztowne, dlatego nasza implementacja jąder unika używania jakichkolwiek innych prymitywów NVSHMEM i minimalizuje wszelkie zależności poza pojedynczymi parami wysyłania-odbierania. Podczas wysyłania danych zawsze używamy nieblokującej wersji funkcji, bez czekania, aż dane zostaną nawet wysłane z lokalnej rangi. Inna, ukryta synchronizacja w grupie jąder zapewnia, że te bufory nie zostaną destrukcyjnie nadpisane, podczas gdy karty sieciowe potrzebują do nich dostępu.

NVSHMEM używa GPUDirect RDMA, dzięki czemu karta sieciowa może bezpośrednio odczytywać i zapisywać pamięć GPU w celu wymiany danych bez angażowania procesorów CPU. Co więcej, na kartach sieciowych ConnectX, NVSHMEM obsługuje komunikację inicjowaną przez GPU (znaną również jako Infiniband GPUDirect Async lub IBGDA), która pozwala GPU na bezpośrednie przesyłanie operacji RDMA do karty sieciowej i bezpośrednie odpytywanie kolejki zakończeń karty sieciowej. Na platformach, które nie obsługują komunikacji inicjowanej przez GPU, NVSHMEM uruchamia "wątek" proxy na procesorze CPU hosta, aby inicjować komunikację i odpytywać zakończenie w imieniu GPU. Program NVSHMEM jest przenośny niezależnie od tego, czy komunikacja inicjowana przez GPU jest obsługiwana, czy nie. Jednak komunikacja inicjowana przez GPU znacząco skraca opóźnienia, ponieważ całkowicie omija objazd do procesora CPU.

Przenośne jądra

Implementujemy komunikację MoE poprzez parę jąder wysyłających (dispatch) i zbierających (combine). Jądra dispatch są odpowiedzialne za odczytywanie tokenów i informacji o routingu na każdej randze, wysyłając je do odpowiednich ekspertów. Jądra combine zbierają aktywacje wytworzone przez ekspertów i wysyłają je z powrotem do ich rang źródłowych, obliczając jednocześnie średnią ważoną od wybranych ekspertów na podstawie wag obliczonych przez router. Jądra są dalej podzielone na komponent wysyłający i odbierający, aby umożliwić nakładanie się transferów danych z obliczeniami. Jądra wysyłające są w pełni nieblokujące i niesynchronizujące: po prostu wysyłają wszystkie zapisy do zdalnych urządzeń. Z drugiej strony, jądra combine tylko odczytują z pamięci, czekając, aż wszystkie wymagane dane zostaną przesłane. Po wysłaniu pracy do kart sieciowych, podczas gdy dane są przesyłane asynchronicznie przez łącze, procesory graficzne mogą wykonywać inne przydatne zadania lokalnie, takie jak stosowanie wspólnych ekspertów lub obliczanie uwagi.

Każda para jąder ma własne symetryczne bufory pamięci, przy czym każda ranga przydziela prywatną przestrzeń dyskową do odbierania danych od wszystkich innych rang. Maksymalna liczba tokenów, które mogą zostać wysłane z grupy DP, jest ustalona w ramach rang, co również ustawia górną granicę, którą każda ranga może otrzymać dla każdego lokalnego eksperta od każdej rangi DP. Pozwala to rangom wysyłającym na wyprowadzenie unikalnego adresu na randze docelowej, do którego mają zapisać, bez konieczności jakiejkolwiek synchronizacji między nimi. Po odebraniu danych jądra przetasowują tokeny, aby ułożyć je w pamięci w optymalnym formacie wymaganym przez jądra odbierające. Chociaż bufory mają znaczny wymiar, są ponownie wykorzystywane we wszystkich kolejnych warstwach modelu.

Powiększ

Jedyna forma globalnej synchronizacji, jak zilustrowano na rysunku powyżej, jest zaimplementowana w jądrze combine-receive. Gdy dane zostaną wysłane z send, ranga powraca i kontynuuje wykonywanie, gdy otrzyma wszystkie potrzebne tokeny od innych rang, i wchodzi do jąder combine send, gdy tylko zostaną przetworzone. Bariera w combine-receive zapewnia, że żadna ranga nie może wyprzedzić i rozpocząć dispatch-send, podczas gdy jakakolwiek inna ranga wciąż czeka na odebranie danych, co potencjalnie powoduje destrukcyjne nakładanie się. Synchronizacja dla prostoty odbywa się w jądrze combine; w ramach optymalizacji tę samą barierę można zaimplementować między jądrami dispatch receive a combine send, jeszcze bardziej redukując opóźnienia poprzez danie kartom sieciowym więcej czasu na ustabilizowanie transferów, podczas gdy GPU wykonuje inne jądra.

Oba jądra są podzielone na wszystkie dostępne SM urządzeń: podczas gdy jądra dispatch send i combine receive muszą równoleglić się według maksymalnej liczby tokenów na rangę (max_m), jądra dispatch receive i combine send muszą obsługiwać tę liczbę tokenów od każdej rangi i lokalnego eksperta, dla maksimum max_m * (num_experts // EP) * (EP // TP).

Wysyłanie (Dispatch)

Jądro dispatch jest odpowiedzialne za wysyłanie tokenów i czynników skalowania bloków do partnerów, którzy hostują ekspertów, do których tokeny zostały skierowane. Strona wysyłająca opiera się na specjalizacji warp, aby zrównoleglić dwa zadania: agregację informacji o routingu w celu określenia, ile tokenów jest wysyłanych z każdej rangi, oraz pakowanie tokenów w komunikaty, które mają zostać wysłane do zdalnego urządzenia. Strona odbiorcza najpierw czeka na odebranie wszystkich liczników tokenów, a następnie odczytuje i rozpakowuje ładunki z pamięci współdzielonej do tensorów, które zostaną przekazane do jąder Grouped GEMM znajdujących się w dalszej części potoku. Przechowuje również informacje w buforach współdzielonych z jądrami combine, aby wskazać, gdzie każdy token powinien zostać odesłany. Ta informacja jest wymagana, ponieważ jądro odbiorcze tasuje tokeny w ciągłych buforach w niedeterministycznej kolejności.

Powiększ

W części wysyłającej pojedynczy warp 32 wątków w każdym bloku odpowiada za odczytanie informacji o routingu i zliczenie liczby tokenów, które każdy docelowy ekspert ma otrzymać od tej rangi, wysyłając liczbę plus jeden przy użyciu nvshmem_signal_op. Licznik jest zwiększany o jeden, ponieważ przejście z zera na wartość niezerową po stronie zdalnej oznacza otrzymanie liczników od rangi źródłowej. Równolegle pozostałe warpy współpracują, aby skopiować tokeny do pamięci symetrycznej we wszystkich blokach jednocześnie, pakując aktywacje, czynniki skalowania i ich indeks na lokalnej randze w ciągły blok. Indeks jest wymagany przez nadawcę combine w celu określenia adresu, pod który zostanie zapisany token. Następnie, po upewnieniu się, że wszystkie dane zostały skopiowane przez barierę w grupach warp, warpy ponownie działają niezależnie, każdy wysyłając ten sam bufor do innego eksperta równolegle. Tokeny są wysyłane przy użyciu nvshmemx_putmem_signal_nbi_warp, co również atomowo zwiększa liczbę wysłanych tokenów z lokalnej rangi na zdalnym urządzeniu. W obrębie grupy DP, ponieważ każda ranga posiada replikę tokena do wysłania, wysyłanie jest zbalansowane równomiernie, przy czym każde urządzenie wysyła podzbiór tokenów.

Po stronie odbiorczej wszystkie jądra najpierw synchronizują się z rangami wysyłającymi, czekając na przesłanie całkowitej liczby tokenów. Następnie wszystkie czekają, aż atomowo zwiększone liczby wysłanych tokenów ustabilizują się do całkowitych liczb, co oznacza, że wszystkie ładunki z rang źródłowych zostały również wysłane, dzięki semantyce wywołania putmem. Jądra odpytują liczniki przy użyciu nvshmem_uint64_wait_until, zrównoleglając operację we wszystkich blokach i wątkach. Następnie bariera międzyblokowa zapewnia, że żaden blok nie odczytuje buforów, chyba że wszystkie dane zostały poprawnie odebrane. Rozproszone między blokami i zsynchronizowane za pomocą atomowo zwiększanego licznika, tokeny są kopiowane z pamięci symetrycznej do regularnych ciągłych buforów, które stają się danymi wejściowymi dla jąder downstream implementujących ekspertów. Ranga źródłowa, indeks eksperta i indeks tokena są przechowywane oddzielnie, dokładnie określając lokalizację, w której jądro combine musi wysłać aktywacje. Mimo że tokeny z grupy DP są wysyłane z różnych urządzeń, wszystkie są pogrupowane razem, aby zostać przekazane do odpowiedniego eksperta.

Łączenie (Combine)

Jądra combine są ponownie podzielone na połowy wysyłające i odbierające: nadawcy kopiują niekwantyzowane 16-bitowe aktywacje do zdalnych urządzeń, a jądra odbierające czekają na przesłanie danych i lokalnie obliczają średnią ważoną wszystkich wkładów ekspertów. Ponadto działają również jako bariera, aby zsynchronizować sekwencję dispatch-combine: każda ranga ustawia flagę u każdego partnera przy wejściu do jąder wysyłających, przy czym jądra odbierające nie mogą powrócić, chyba że zaobserwują ustawienie flagi. Opóźnienie synchronizacji jest minimalne, ponieważ nakłada się na rzeczywistą komunikację i obliczenia.

Powiększ

Po stronie wysyłającej jądra przeglądają listę tokenów przypisanych do wszystkich lokalnych ekspertów równolegle, zapisując je do bufora na randze docelowej. Ranga docelowa, indeks eksperta i indeks tokena są odczytywane z buforów dla poszczególnych tokenów zapełnionych przez jądra scater po otrzymaniu tokenów. Każdy nadawca ma swój własny prywatny region pamięci na eksperta, do którego może zapisać, jak wskazano na rysunku powyżej, unikając konieczności synchronizacji. Podobnie jak w przypadku dispatch, combine atomowo zwiększa liczniki na token na randze docelowej, aby wskazać otrzymanie danych: gdy licznik odpowiada liczbie ekspertów, do których token został wysłany, zawartość tokena może zostać odczytana.

W jądrze odbierającym lista tokenów jest przeglądana równolegle w wielu blokach, czekając na nadejście ich zawartości poprzez odpytywanie flagi ustawionej przez operację sygnalizacji. Po nadejściu, ładunki są odczytywane z prywatnych buforów, przy czym tabela routingu wskazuje, z którego bufora czytać i jaką wagę przypisać każdemu ekspertowi. Wyniki są następnie zapisywane do zewnętrznie przydzielonych tensorów, a jądro kończy wykonywanie, gdy wszystkie urządzenia przejdą przez barierę. W obrębie grupy DP wszystkie rangi otrzymują kopię każdej aktywacji eksperta, aby obliczyć własne repliki.

Testy porównawcze (Benchmarks)

Oceniamy nasze jądra w klastrze 128 procesorów graficznych rozmieszczonych w 16 węzłach połączonych przez InfiniBand i karty sieciowe CX-7. Próbujemy zarówno GPUDirect Async (IBGDA), jak i Reliable Connection (RC) z proxy procesora CPU. Porównujemy je z wysoce zoptymalizowaną implementacją DeepSeek, a także z gęstymi prymitywami dostarczonymi przez PyTorch (przez NCCL) lub NVSHMEM.

Powiększ

Chociaż w gęstej operacji NVSHMEM all-to-all wydajność IBRC i IBGDA jest podobna (6378 µs vs 6180 µs), IBGDA jest znacząco szybsza przy rzadkich jądrach. Sumując Dispatch i Combine, IBGDA wykorzystuje 902 µs, podczas gdy IBRC zajmuje 3223 µs - 3,6-krotna poprawa opóźnienia. Podczas gdy all-to-all jest ograniczone przepustowością, rzadkie jądra rozgłaszają o rzędy wielkości mniej danych, będąc ograniczonymi przez opóźnienia. Poprzez inicjowanie transferów sieciowych bezpośrednio z GPU, bez wymagania proxy procesora CPU do koordynacji GPU i karty sieciowej, opóźnienie end-to-end jest znacząco zredukowane.

Podczas gdy nasze przenośne jądra są około 2 razy wolniejsze od wysoce zoptymalizowanych jąder DeepSeek, poprawiają opóźnienia o 10x w porównaniu do gęstych jąder (902 µs vs 9944 µs).

Dodatkowo, podział na komponenty nadawca-odbiorca pozwala również ukryć część opóźnień, w przeciwieństwie do prymitywów bibliotecznych.

Powiększ

W pojedynczym węźle (EP8) NVSHMEM wykorzystuje NVLINK do transportu, zapewniając niższe opóźnienia i wyższą przepustowość niż sieci międzywęzłowe. Nasze przenośne jądra są około 2,5-krotnie szybsze od DeepEP w pojedynczym węźle (186µs vs 481 µs).

Dalsze ulepszenia

Opisane tutaj jądra przewyższają wbudowane prymitywy frameworków ML i oferują przyzwoitą wydajność bez nadmiernej specjalizacji pod konkretne transporty między urządzeniami, takie jak NVLink czy InfiniBand. Oprócz już wspomnianych możliwości, dalsze zyski wydajnościowe można osiągnąć poprzez zastąpienie prymitywów komunikacyjnych bardziej wyspecjalizowanymi wersjami. Na przykład, w NVLink, użycie pamięci symetrycznej można zastąpić buforami współdzielonymi między urządzeniami, eliminując część kopiowania i pozwalając na bardziej precyzyjne schematy synchronizacji, między poszczególnymi tokenami zamiast partii tokenów. W InfiniBand, implementacja mogłaby uzyskiwać dostęp bezpośrednio do bazowych par kolejek, eliminując potrzebę przesyłania zarówno oczekiwanej, jak i wysłanej liczby tokenów oraz opróżniania tokenów z kolejek do buforów wyjściowych w miarę ich przychodzenia. Jednak takie implementacje wiążą się z kosztem przenośności, podczas gdy nasze implementacje pozwalają na ocenę klastrów o określonych konfiguracjach przed podjęciem wysiłku rozwojowego w celu optymalizacji pod konkretny sprzęt.

Wnioski

Zaprezentowaliśmy wysokowydajną, przenośną bibliotekę do komunikacji MoE, która osiąga 10-krotnie szybszą wydajność w porównaniu ze standardową komunikacją all-to-all, zachowując jednocześnie kompatybilność z różnymi konfiguracjami sprzętowymi. W wdrożeniach z pojedynczym węzłem z NVLink, nasze rozwiązanie wykazuje 2,5-krotnie niższe opóźnienia niż poprzednie implementacje.

Nasze podejście równoważy wydajność z przenośnością dzięki kluczowym innowacjom, w tym wsparciu komunikacji inicjowanej przez GPU, architekturze dzielonego jądra umożliwiającej nakładanie się obliczeń i komunikacji oraz efektywnemu przesyłaniu tokenów przy użyciu minimalnych prymitywów NVSHMEM. Mimo że jest około 2 razy wolniejsza od wysoce wyspecjalizowanych implementacji w konfiguracjach wielowęzłowych, nasza biblioteka oferuje lepszą elastyczność w różnych środowiskach sieciowych (NVLink, CX-7 i EFA).

W miarę jak modele MoE będą się dalej skalować, skuteczne strategie komunikacyjne, takie jak nasza, staną się coraz ważniejsze dla praktycznego wdrożenia. Nasza w pełni open-source implementacja jest dostępna pod adresem https://github.com/ppl-ai/pplx-kernels.

Bibliografia

DeepEP

Raport techniczny DeepSeek-V3

NVIDIA NVSHMEM

Cytowanie

@online{perplexity2025EfficientAndPortableMixtureOfExpertsCommunication,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Phi Hoang},
  title = {{Efficient and portable mixture-of-experts communication}},
  date = {2025-04-02},
  url = {https://www.perplexity.ai/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}