Artykuł
Niższe opóźnienia i wyższa przepustowość dzięki wielowęzłowej implementacji DeepSeek
W większości systemów opóźnienie i przepustowość to często cele konfliktowe, które wymagają kompromisów podczas projektowania i wdrażania. Na przykład w gęstych dużych modelach językowych zwiększenie rozmiaru partii może poprawić przepustowość, ale także zwiększa opóźnienie; zwiększenie równoległości tensorów w ramach jednej maszyny może zmniejszyć opóźnienie, ale zmniejsza liczbę replik, co prowadzi do niższej przepustowości.
Modele mieszane ekspertyz (MoE), takie jak DeepSeek-V3/R1, niedawno wykazały doskonałe możliwości modelowe i wydajność operacyjną. Na przykład model DeepSeek-V3/R1 ma łącznie 671 miliardów parametrów, ale każdy token używa tylko 37 miliardów parametrów podczas wnioskowania. Ta architektura modelu przedstawia zarówno wyzwania, jak i możliwości dla systemów wnioskowania.
Ten artykuł pokazuje, że w przeciwieństwie do konwencjonalnych systemów, modele MoE, takie jak DeepSeek-V3/R1 mogą jednocześnie osiągać wyższą przepustowość i niższe opóźnienie, korzystając z większej liczby GPU w wdrożeniach multi-node w większości scenariuszy.
Architektury wdrożeń
Z powodu dużej liczby małych ekspertów w modelu, wdrożenia muszą być rozłożone na wiele urządzeń. Rozważaliśmy zarówno wdrożenia jedno-węzłowe na jednym węźle z 8xH200 GPU, jak i wdrożenia multi-node na 8xH100 GPU.
Obie architektury wdrożeń wykorzystują Parallelizm Danych, zorganizowany przez nasz wewnętrzny harmonogram żądań. Implementacja parallelizmu danych polega na uruchamianiu wielu instancji silnika wnioskowania, z których każda działa niezależnie, aby obsługiwać i zarządzać żądaniami. Harmonogram żądań, który komunikuje się z silnikiem przez GRPC, odpowiada za rozłożenie żądań jak najbardziej równomiernie, jednocześnie ułatwiając ponowne użycie KV, wysyłając żądania z częściowym dopasowanym prefiksem do serwerów zawierających pamięć cache. Instancje silnika nie obejmują wielu węzłów. Mogą opcjonalnie używać równoległości tensorów do dzielenia uwagi na wiele urządzeń. Instancje są połączone przez NVLink w przypadku jednego węzła lub InfiniBand dla przypadku multi-node, dostarczając i zbierając ekspertów.
Konfiguracja wdrożenia na jednym węźle zapewnia wyższe opóźnienia w przypadku małych rozmiarów partii; jednak wydajność gwałtownie spada w warunkach zwiększonego obciążenia.
Aby wdrożyć silnik obsługi, uruchamiamy jeden pod na węzeł, hostując wiele instancji silnika. PyTorch odpowiada za zestawienie komunikacji rozproszonej i negocjowanie inicjalizacji NVSHMEM. Do komunikacji polegamy na niestandardowych jądrach CUDA opisanych w wcześniejszej publikacji na blogu. Implementacja dwóch wdrożeń jest praktycznie identyczna, z modelem wybierającym odpowiednie jądra do użycia w oparciu o tkaninę implementującą równoległość ekspertów.
Techniki paralelizacji
Zanim przejdziemy do naszych porównań wydajności, istotne jest zrozumienie kluczowych strategii paralelizacji, które sprawiają, że wdrożenie masywnych modeli MoE, takich jak DeepSeek-V3/R1, jest możliwe.
Równoległość tensorów
W wnioskowaniu LLM, równoległość tensorów (TP) jest zazwyczaj używana do zmniejszenia zużycia pamięci i obliczeń na GPU, zmniejszając tym samym opóźnienie. Zwykle możemy podzielić projekcje liniowe w warstwach uwagi i MLP wzdłuż wierszy lub kolumn oraz operacje uwagi wzdłuż wymiaru głowy uwagi.
Dzięki TP, architektura Llama-3 nie ma zduplikowanych obliczeń dla operacji liniowych projektowania i uwagi na różnych GPU, co jest idealną metodą zgniatania. Jednak w modelach DeepSeek-V3/R1 TP nie może tego osiągnąć.
Modele DeepSeek-V3/R1 używają Wielokrotnej Latentnej Uwagi (MLA). Warstwa MLA najpierw używa liniowej projekcji kv_a_proj do obliczenia wektora latentnego, a następnie używa innej projekcji liniowej kv_b_proj do przekształcenia go w przestrzeń każdej głowy uwagi. Ponieważ wszystkie głowy uwagi dzielą ten sam wektor latentny, TP nie może podzielić wektora latentnego, więc wszystkie rangi TP muszą replikować parametry i obliczenia kv_a_proj i kv_b_proj. Podobnie, ponieważ MLA przechowuje wektor latentny w pamięci KV, każda ranga TP przechowuje identyczną kopię pamięci KV.
Pomimo pewnych duplikacji w MLA, równoległość tensorów nadal zapewnia częściowe zmniejszenie zapotrzebowania na obliczenia, co czyni ją cenną w scenariuszach wymagających dużych szybkości wyjściowych.
Równoległość ekspertów
Modele DeepSeek-V3/R1 zastępują warstwy MLP warstwami MoE. Warstwa MoE ma 256 kierowanych ekspertów i jednego eksperta współdzielonego. Każdy token jest wysyłany do 8 różnych kierowanych ekspertów do obliczeń, a wyniki są ważone sumowane. Każdy token również oblicza w ekspercie współdzielonym, a wynik jest dodawany do wyniku od kierowanych ekspertów.
Równoległość ekspertów (EP) służy jako typowe podejście do dzielenia warstw MoE, z każdym GPU zarządzającym 256 / EP kierowanymi ekspertami, jednocześnie utrzymując kopię eksperta współdzielonego. W porównaniu z TP, zaletą EP jest to, że może rozłożyć obliczenia na większą liczbę GPU, zmniejszając zużycie obliczeń i pamięci na GPU.
Przed wykonaniem obliczeń ekspertów wszystkie GPU muszą wykonać komunikację AllToAll, aby wysłać tokeny do GPU, gdzie znajdują się odpowiedni eksperci; po obliczeniach ekspertów, potrzebna jest kolejna komunikacja AllToAll, aby zebrać wyniki obliczeń z różnych GPU i wykonać ważenie. Wdrożyliśmy zoptymalizowaną wersję tych dwóch jąderek komunikacyjnych AllToAll, Dispatch i Combine, używając NVSHMEM. W naszym wcześniejszym poście na blogu szczegółowo opisaliśmy implementację, a nasze jądra zostały opublikowane jako open source na GitHubie.
Równoległość danych
Dzięki EP możemy rozłożyć obliczenia MoE na 128 lub nawet więcej GPU. Jednak obliczenia MLA nie mogą być podzielone z EP. W tym momencie możemy wprowadzić równoległość danych (DP). Każda grupa DP posiada pełną kopię warstwy MLA. Każda grupa DP przyjmuje różne dane wejściowe i wykonuje obliczenia warstwy MLA niezależnie.
DP warstwy MLA i TP można połączyć, z jedną grupą DP dzieloną na wiele rang TP. EP warstwy MoE można połączyć z DP/TP warstwy MLA. EP = DP * TP. Na przykład, na 16 maszynach, EP128 DP32 TP4 oznacza rozdzielenie kierowanych ekspertów na 128 GPU, z każdą grupą po 4 GPU tworzącą grupę DP, dla łącznie 32 niezależnych grup DP.
Pojedynczy węzeł a multi-node
Parametry 671B DeepSeek przekraczają pojemność pamięci pojedynczej maszyny 8-GPU H100 (80 GB * 8), ale pojedyncza maszyna 8-GPU H200 może w pełni pomieścić cały model (141 GB * 8). Używając konfiguracji EP8 DP8 TP1, model używa około 100 GB pamięci na GPU, pozostawiając około 40 GB na pamięć KV i inne wyniki pośrednie. Jeden token zajmuje 70 272 bajty pamięci KV. Zakładając, że każde żądanie ma 5 000 tokenów, każde GPU może obsłużyć około 100 żądań.
Chcieliśmy zrozumieć różnice w wydajności między wdrożeniami jedno-węzłowymi a multi-node w różnych konfiguracjach. Użyliśmy jednej maszyny H200 do wdrożenia jedno-węzłowego i do 16 maszyn H100 do wdrożeń multi-node. Dla każdego środowiska wdrożeniowego użyliśmy kombinacji TP 1, 2, 4, 8 i wielkości partii na GPU 1, 2, 4, 8, 16, 32, 64, 128. Założyliśmy, że każde żądanie ma długość pamięci KV 5 000 tokenów. Przyjęliśmy również Multi-Token Prediction (MTP), które przewiduje 1 dodatkowy token (tj. długość zapytania każdego żądania to 2) i konserwatywnie założyliśmy wskaźnik akceptacji 60%. Poniższa figura pokazuje przepustowość i prędkość wyjściową dla różnych konfiguracji.
Oś pozioma przedstawia prędkość wyjściową na żądanie w tokenach/s. Oś pionowa używa skali logarytmicznej do pokazania przepustowości na maszynę w tokenach/s. Oznaczyliśmy Pareto Frontier dla każdej konfiguracji EP różnymi kolorowymi liniami.
W scenariuszach z ekstremalnie wysokimi wymaganiami dotyczącymi prędkości wyjściowej, używanie jednego-węzłowego EP8 DP1 TP8 z rozmiarem partii 1 może osiągnąć prędkość wyjściową przekraczającą 100 tokenów/s, ale przepustowość jest niezwykle niska, co odpowiada prędkości wyjściowej. W tym scenariuszu cała partia ma tylko 2 tokeny, które mogą być dostarczone do maksymalnie 2*8=16 ekspertów, aktywując łącznie maksymalnie 57 miliardów parametrów.
W zakresie prędkości wyjściowej 80-40 tokenów/s, wraz ze wzrostem przepustowości, prędkość wyjściowa znacząco spada. W przeciwieństwie do tego, EP128 ma około 5 razy wyższą przepustowość niż wdrożenie jedno-węzłowe przy tej samej prędkości wyjściowej.
Zjawisko to można wyjaśnić, analizując zachowanie wdrożeń jedno-węzłowych: zwiększenie rozmiaru partii bezpośrednio koreluje ze wzrostem liczby aktywowanych ekspertów. Gdy rozmiar partii wynosi 1, średnia liczba aktywowanych ekspertów na GPU to 2 * 8 / 8 = 2. Gdy partia jest wystarczająco duża, wszyscy eksperci są aktywowani, co oznacza, że każde GPU aktywuje 256 / 8 = 32 ekspertów. Aktywacja większej liczby ekspertów powoduje, że GPU musi wczytać więcej parametrów z pamięci, znacznie zwiększając ciśnienie na szerokość pasma pamięci. Ponieważ faza dekodowania dużych modeli językowych już jest zablokowana przez przepustowość pamięci, a nie przez wydajność obliczeniową, zwiększenie rozmiaru partii w wdrożeniu jedno-węzłowym znacznie zmniejsza prędkość wyjściową.
Porównanie czterech konfiguracji wdrożeń multi-node (EP16, EP32, EP64 i EP128) ujawnia, że wyższe wartości EP przesuwają Pareto Frontier w kierunku jednoczesnych ulepszeń w przepustowości i prędkości wyjściowej.
Używanie wyższej liczby EP oznacza, że każde GPU ma mniej przydzielonych ekspertów. Na przykład, EP128 oznacza, że każde GPU jest odpowiedzialne za 256 / 128 = 2 ekspertów, więc ciśnienie na szerokość pasma pamięci jest znacznie zmniejszone. Innymi słowy, używając większej liczby EP, skutecznie zyskujemy więcej przepustowości pamięci. Kiedy rozmiar partii na GPU jest mniejszy niż 64, zwiększenie rozmiaru partii nie wpływa znacząco na prędkość obliczeń ekspertów, ponieważ zwiększenie liczby danych wejściowych nie znacznie zwiększa ciśnienia na szerokość pasma pamięci. Dlatego obserwujemy, że przy użyciu EP128 zwiększanie rozmiaru partii nie wpływa znacząco na prędkość wyjściową.
Co ciekawe, przy większych rozmiarach partii (64 żądania na GPU) zauważyliśmy nowe zjawisko: przepustowość wdrożenia jedno-węzłowego jest nieznacznie wyższa niż wdrożenia multi-node. Część przyczyny polega na tym, że wewnątrz węzła NVLink ma wyższą przepustowość niż InfiniBand między węzłami. Inna część wynika z ograniczeń naszej implementacji. Przeanalizujemy to zjawisko bardziej szczegółowo później.
Z powodu ograniczeń pojemności pamięci konfiguracja EP8 DP8 TP1 nie może osiągnąć rozmiaru partii 128 na GPU, więc wdrożenie multi-node nadal jest lepszym wyborem w scenariuszach dążących do wyższej przepustowości.
Nakładanie się obliczeń i komunikacji
Jak wspomniano wcześniej w kontekście równoległości ekspertów, GPU są nieaktywne podczas komunikacji warstwy MoE. Aby zmniejszyć marnotrawstwo i obniżyć opóźnienie, musimy znaleźć niezależne od danych zadania obliczeniowe, aby wypełnić ten nieaktywny czas.
Górna część powyższej figury pokazuje przepływ obliczeń jednej warstwy. Obliczenia MoE zależą od Dispatch, a obliczenia następnej warstwy zależą od wyniku Combine.
Umieszczamy współdzielonego eksperta na każdym GPU. W ten sposób obliczenia eksperta współdzielonego nie wymagają komunikacji AllToAll. Dlatego możemy wykonywać obliczenia eksperta współdzielonego natychmiast po wysłaniu Dispatch, a następnie czekać na zakończenie odbioru Dispatch. Nazywamy ten schemat nakładania „Dispatch Overlap”.
Dispatch Overlap oferuje proste wdrożenie i szerokie zastosowanie. Ta technika ukrywa czas obliczeń eksperta współdzielonego we wszystkich rozmiarach EP i rozmiarach partii.
Aby jeszcze bardziej zwiększyć nakładanie się obliczeń i komunikacji, użyliśmy mikropartii wspomnianych w raporcie technicznym DeepSeek, aby przełamać zależność danych. Jak pokazano w dolnej części figury, podzieliliśmy obliczenia jednej warstwy transformacyjnej na 5 etapów:
Etap 1: InputNorm, QKVProj, AppendKV, BMM
Etap 2: BMM, Attn, OProj, PostNorm, Gate
Etap 3: Wysłanie Dispatch, Ekspert współdzielony
Etap 4: Odbiór Dispatch, MoE, Wysłanie Combine
Etap 5: Odbiór Combine
W pierwszych 3 gęstych warstwach transformacyjnych korzystamy z całej partii. W kolejnych 58 warstwach transformacyjnych MoE równomiernie dzielimy partię na dwie mikropartie. Dwa mikroparty są wykonywane naprzemiennie, przesunięte o 3 etapy. Ponieważ między tymi dwoma mikropartiami nie ma zależności danych, możemy przełączyć się do obliczeń innej mikropartii po wysłaniu Dispatch i po wysłaniu Combine.
Rozkład opóźnień
Następnie porównamy efekty nakładania w eksperymencie, a także porównamy różnice w wydajności między wdrożeniem jedno-węzłowym EP8 a wdrożeniem multi-node EP128. Dla łatwiejszego porównania użyliśmy GPU H100 do tego eksperymentu. Użyliśmy TP1, rozmiaru partii 128 na GPU, długości zapytania wynoszącej 2 na żądanie i długości pamięci KV wynoszącej 5000.
Figura powyżej pokazuje całkowity czas spędzony na jednej warstwie transformacyjnej MoE oraz proporcję opóźnienia różnych rodzajów jąderek. Z wyjątkiem Dispatch, Combine i GroupGEMM, czas wykonania innych jąderek powinien być równy w seriach EP8, EP128 NoOverlap i EP128 DispatchOverlap, ponieważ rozmiar partii jest taki sam.
Nakładanie
Najpierw porównajmy efekty trzech metod nakładania. NoOverlap zajęło w sumie 2667 µs, DispatchOverlap zajęło 2651 µs, oszczędzając 16 µs, czyli tylko 0.6%. MicroBatch pokazał bardzo znaczną poprawę, zajmując 1896 µs, co stanowi przyspieszenie o 29%. Czas zarówno Dispatch, jak i Combine zostały znacznie skrócone. Dispatch zmniejszył się z 593 µs do 367 µs, a Combine z 1012 µs do 237 µs.
Należy zauważyć, że dla jąderek obliczeń, podzielenie partii o rozmiarze 128 na dwie partie o rozmiarze 64 zwiększa czas całkowity wykonania. Dlatego mimo że czas spędzony na komunikacji został skrócony o 1001 µs, czas całkowity skrócił się tylko o 771 µs. Wyjaśnimy to za pomocą modelu Roofline w następnej sekcji.
Z tego powodu mikropodział nie zawsze poprawia wydajność.
Figura powyżej pokazuje poprawę wydajności MicroBatch w porównaniu do DispatchOverlap dla rozmiarów partii 4-128. Gdy rozmiar partii jest mniejszy niż 32, MicroBatch zmniejsza wydajność o 5%-40%. Gdy rozmiar partii jest większy lub równy 32, MicroBatch może poprawić wydajność o 10%-35%.
EP8 kontra EP128
Wróćmy do poprzedniej figury i porównajmy EP8 i EP128 MicroBatch. EP8 zajęło 1802 µs w sumie, nieco mniej niż EP128, które zajęło 1896 µs. Oprócz zwiększonego czasu wykonania jądraosh spowodowanego przez MicroBatch, główne różnice dotyczą GroupGEMM używanego do obliczeń MoE i dwóch jąderek komunikacyjnych, Dispatch i Combine.
GroupGEMM EP8 zajęło 555 µs, podczas gdy GroupGEMM EP128 zajęło 270 µs, zmniejszając się o połowę. To jest kluczowa zaleta wdrożeń multi-node.
Niestety, czas spędzony na komunikacji zwiększył się o 213 µs, co znacznie zrównoważyło przewagę GroupGEMM. W oddzielnych testach wydajności naszych jąderek komunikacyjnych odkryliśmy, że mogą one osiągnąć tylko połowę przepustowości Infiniband. Będziemy kontynuować optymalizację naszych jąderek komunikacyjnych.
Inna jądro, która znacznie opóźnia to GEMM. Microbatch zwiększyło GEMM o 95 µs. Przeanalizujemy GEMM bardziej szczegółowo w sekcji Roofline poniżej. Uważamy, że obecna implementacja GEMM nie osiągnęła jeszcze optymalnej wydajności.
Roofline
Model Roofline jest dobrym narzędziem do analizy wydajności jąderek. Jego oś pozioma to Intesity Arithmetyczna, wskaźnik FLOP do bajtów I/O pamięci. Wartość na osi poziomej można obliczyć bezpośrednio z semantyki jądra. Oś pionowa reprezentuje osiągniętą wydajność, obliczaną przez podzielenie FLOP przez opóźnienie benchmarku.
Teoretyczna górna granica wydajności jądra jest bezpośrednio determinowana przez specyfikacje GPU. Maksymalna wydajność FP8 H100 wynosi 1979 TFLOP/s, przedstawiona jako linia pozioma w modelu Roofline. Przepustowość pamięci H100 wynosi 3,35 TB/s, przedstawiona jako nachylenie linii przechodzącej przez punkt początkowy. Dwie linie wyznaczają limity wydajności dla jąderek ograniczonych przez obliczenia i jąderek ograniczonych przez przepustowość pamięci.
Poniżej omówimy wydajność jąderek GroupGEMM i GEMM.
GroupGEMM
Jądro GroupGEMM w MoE wykonuje następujące obliczenia: Istnieje g grup w sumie, i-ta grupa ma m_i tokenów, wykonując mnożenie macierzy [m_i, k] x [k, n] -> [m_i, n]. W testach wydajności zakładamy, że liczba tokenów w każdej grupie jest taka sama, oznacza jako m_i = m. Następnie liczba FLOP dla GroupGEMM wynosi 2 * g * m * k * n, a bajty I/O pamięci wynoszą g * (m * k + n * k + m * n).
W modelu DeepSeek-V3/R1 jest 256 ekspertów, a każdy token jest wysyłany do 8 ekspertów w celu obliczeń. Zakładając rozmiar partii 128, długość zapytania wynoszącą 2, korzystając z konfiguracji EP128 DP128, średnia liczba tokenów otrzymanych przez każdego eksperta (czyli m) wynosi 128 * 2 * 8 * 128 / 256 = 1024. Podobnie, możemy obliczyć m dla innych konfiguracji i rozmiarów partii.
Użyliśmy implementacji GroupGEMM DeepGEMM do testowania wydajności. Punkty testowe pokryły kombinacje konfiguracji EP8, EP16, EP32, EP64, EP128 z TP1 i rozmiarów partii 1-128.
Figura powyżej pokazuje model Roofline dla GroupGEMM przy różnych konfiguracjach EP. Różne EP odpowiadają różnym liczbom grup. Rysunek ilustruje prawie pokrywające się linie wydajności, wskazując, że wydajność GroupGEMM jest przede wszystkim determinuje się przez całkowitą liczbę tokenów (przedstawioną jako g * m).
Gwiazdy oznaczają punkty danych odpowiadające rozmiarowi partii 128 na GPU dla każdej konfiguracji EP. Porównując te gwieździste punkty danych, możemy zauważyć, że w miarę jak EP rośnie (a DP rośnie równocześnie), liczba tokenów na eksperta m również się zwiększa. Przy EP8, m=128, podczas gdy przy EP128, m=2048.
Gdy m rośnie, Intensity Arithmetyczna także rośnie. W większości konfiguracji GroupGEMM jest ograniczone przez przepustowość pamięci, więc zwiększenie m poprawia wydajność.
GEMM
Jądro GEMM odpowiada Liniowym Projekcjom w modelu, takim jak Q/K/V/O Projekcje. Dla mnożenia macierzy [m, k] x [k, n] -> [m, n], liczba FLOP wynosi 2 * m * k * n, a bajty I/O pamięci wynoszą m * k + n * k + m * n. Możemy również przetestować opóźnienie dla rozmiarów partii 1-128.
Figura powyżej pokazuje model Roofline dla GEMM przy różnych konfiguracjach EP. Możemy zauważyć, że wydajność GEMM jest ograniczona przez przepustowość pamięci. W miarę jak rozmiar partii rośnie, Intensity Arithmetyczna również rośnie, co zatem poprawia wydajność.
Mikropartie
Podczas korzystania z mikropartii, dzielimy partię równomiernie na dwie części. Z dwóch powyższych figur możemy zauważyć, że gdy m staje się m/2, wydajność mnożenia macierzy maleje. Dlatego wykonanie dwóch mnożeń macierzy o rozmiarze m/2 trwa dłużej niż wykonanie jednego mnożenia macierzy o rozmiarze m.
Wielokrotna predykcja tokenów
Na całej długości tego artykułu założyliśmy korzystanie z Wielokrotnej Predykcji Tokenów (MTP) dla spekulatywnego dekodowania. MTP zmienia długość zapytania na żądanie z 1 na 2. Dla mnożenia macierzy jest to równoważne ze zmianą m na m * 2, co zatem zwiększa wydajność mnożenia macierzy. Z drugiej strony, jeśli narysujemy model Roofline dla MLA, odkryjemy, że zwiększenie długości zapytania znacznie zwiększa wydajność MLA.
Dlatego korzystanie z MTP odgrywa ważną rolę w efektywności modelu.
Implementacja i optymalizacje
W tej sekcji przedstawimy kilka szczegółów dotyczących implementacji i optymalizacji naszego modelu DeepSeek-V3/R1.
Kwantyzacja
DeepSeek-V3/R1 był natywnie trenowany na FP8 z użyciem schematu kwantyzacji na bloczki, z wagami kwantyzowanymi statycznie i aktywacjami kwantyzowanymi na bieżąco. Zamiast statycznie obliczać współczynniki skali na kanał lub macierz, współczynniki skali są obliczane na wektorach 128-elementowych dla aktywacji i kafelkach 128x128 elementów dla macierzy, ograniczając degradację dokładności z powodu kwantyzacji.
W Perplexity polegamy na połączeniu jąder CUDA i Triton, aby wspierać wnioskowania, z CUDA używanym dla najbardziej wrażliwych na wydajność i rzadko modyfikowanych jądrek (takich jak uwaga i GEMM), a Triton implementującym szeroki zakres aktywacji, normalizacji i jądrek narzędziowych. Triton pozwolił nam szybko dostosować jądra do schematu kwantyzacji na bloczki.
Dla warstw liniowych i MoE mieszamy jądra Deep GEMM z naszymi własnymi jądrokami Triton GEMM, ponieważ zauważyliśmy, że dla pewnych wymiarów macierzy i niskich rozmiarów partii Split-K zapewnia niższe opóźnienie. Jeśli niekwantyzowana warstwa wykonuje mnożenie (M, K) x (K, N), wymaga (M x ceil_div(K, 128)) x (ceil_div(K, 128), ceil_div(N, 128)) współczynnikami skali dla kwantyzacji na bloczki. Dla kwantyzacji na bloczki współczynniki skali dla aktywacji są obliczane na bieżąco, zamiast być wstępnie kalibrowanymi. Ponieważ współczynniki skali aktywacji są agregowane tylko wzdłuż wymiaru K, a nie wzdłuż wymiaru M, jądra wymagają tylko niewielkich zmian, aby obsługiwać ten schemat.
Funkcja aktywacji SiLU używana przez DeepSeek-V3/R1 wymagała istotnych zmian, aby wspierać CUDA graphs, kwantyzację na bloczki i dynamicznie kierowane liczby tokenów. Kwantyzacja na bloczki może być problematyczna, ponieważ wprowadza redukcje poziome, jednak jądro już dzieliło aktywacje wzdłuż wymiaru ukrytego na bloki o rozmiarze 1024 elementów. W jednym bloku tensor do kwantyzacji był dalej dzielony na bloki 128-elementowe, aby obliczyć największą wartość bezwzględną, z Tritonem generującym wydajne maksymalne redukcje między-liniowe, dodając minimalne opóźnienie.
Aby wspierać routing MoE w ramach CUDA graphs, jądra muszą być świadome informacji o routingu, wskazujących liczbę tokenów na eksperta, zamiast planować pracę w oparciu o rozmiar buforów, które były przydzielane do obsługi maksymalnej liczby tokenów. Nie możemy podzielić problemu na podstawie wymiarów tensora wejściowego, dlatego uruchamiamy stałą liczbę trwałych jąderek, które czytają informacje o routingu, aby określić, ile tokenów jest zaludnionych, i dzielą pracę przetwarzania aktywacji między siebie dynamicznie.
Już udostępniliśmy część naszych jąderek do projektu FlashInfer, a w przyszłości będziemy otwierać jeszcze więcej naszego kodu.
Warstwa MLA
Wykorzystujemy FlashInfer do obliczeń MLA. FlashInfer wspiera elastyczne ustawienia Page Table i ekstremalnie wysoką wydajność.
Połączyliśmy q_a_proj i kv_a_proj w jeden qkv_a_proj. Opóźnienie spadło z 15.4 µs + 14.8 µs = 30.2 µs do 16.7 µs.
Antysymetryzowaliśmy kv_b_proj na dwie macierze, k_b_proj i v_b_proj. Napisaliśmy jądro FP8 Block Quantized BMM dla obliczeń związanych z tymi dwoma macierzami.
Rejestr Cuda
Rejestr Cuda może znacznie zmniejszyć opóźnienie uruchamiania jąderek, co jest kluczowe dla wydajności. Tworzymy rejestr Cuda dla każdego rozmiaru partii.
Przed opracowaniem naszego jądra AllToAll, używaliśmy torch.all_to_all_single() do komunikacji AllToAll. Ta operacja wymaga, aby wszystkie GPU używały tego samego rozmiaru partii. Jednak różne grupy DP mogą używać różnych rozmiarów partii.
Aby zapewnić, że all_to_all_single() jest zgodne z różnymi grupami DP używającymi różnych rozmiarów partii, najpierw używaliśmy operacji allreduce() przed każdym uruchomieniem modelu, aby uzyskać maksymalny rozmiar partii wśród wszystkich grup DP. Następnie sprawialiśmy, że wszystkie grupy DP używają tego rozmiaru partii do uruchamiania.
Chociaż podejście to zapewnia, że możemy używać grafu Cuda, ma trzy wady. Po pierwsze, wymaga dodatkowej operacji allreduce(). Po drugie, grupy DP z mniejszymi rozmiarami partii są zmuszone do dopełniania. Po trzecie, sprawia, że nasz kod implementacji jest skomplikowany.
Po zaimplementowaniu naszego własnego jądra AllToAll, nie musimy już wymagać, aby wszystkie GPU używały tego samego rozmiaru partii. Dlatego nie musimy już wykonywać dodatkowych operacji allreduce() ani dopełniać rozmiarów partii.
Router MoE
Router MoE został zaimplementowany w Triton, opierając się na zmodyfikowanym sortowaniu pochodzącym z standardowej biblioteki, które także śledzi indeksy posortowanych elementów. Implementacja dzielona jest między wszystkimi modelami MoE, ponieważ routowanie Mixtral jest specjalnym przypadkiem ścieżek DeepSeek, gdzie grupa Top-K jest taka sama jak grupa wszystkich ekspertów. Sparse kernels bezpośrednio konsumują indeksy Top-K i wyniki, podczas gdy gęste schematy Dispatch/Combine polegające na all-to-all wymagają, aby informacje o routowaniu były agregowane na eksperta, a nie na podstawie tokenów.
Praca przyszłościowa
W przyszłej pracy planujemy dalszą optymalizację wydajności modelu DeepSeek.
Najważniejszą następną optymalizacją jest rozdzielenie wypełniania. Faza wypełniania i faza dekodowania modelu DeepSeek-V3/R1 mają bardzo różne charakterystyki obliczeniowe. Obie mogą korzystać z różnych strategii optymalizacji i schematów wdrażania.
Dla warstwy MLA, w fazie dekodowania używamy Flutter Matrix Absorption, aby zmniejszyć liczbę FLOPów obliczeń MLA. W fazie wypełniania najpierw projektowanie wektora latentnego w przestrzeń K/V, a następnie obliczanie w formie Multi-Head Attention (MHA) działałoby lepiej.
Jeśli wypełnianie i dekodowanie odbywa się na tym samym GPU, aby zmniejszyć wpływ wypełniania na prędkość wyjściową dekodowania, zazwyczaj używamy podzielenia wypełniania na kilka segmentów, aby podzielić zapytanie na kilka segmentów dla wypełniania. Ponieważ pamięć KV przechowuje wektor latentny, trudniej jest konwertować MLA na formę MHA.
Dla warstwy MoE, w fazie dekodowania, używamy największego możliwego EP i DP, aby zwiększyć liczbę tokenów wejściowych na eksperta, co poprawia wydajność GroupGEMM. W fazie wypełniania, ponieważ liczba tokenów jest już wystarczająco duża, GroupGEMM jest już ograniczone przez obliczenia. Dlatego dla wypełniania możemy używać mniejszych EP i DP.
Jeśli wypełnianie i dekodowanie odbywa się na tym samym GPU, tak długo, jak którakolwiek z grup DP wykonuje wypełnianie, opóźnienie warstw MoE na wszystkich GPU wzrośnie, znacząco wpływając na prędkość wyjściową dekodowania.
Oprócz rozdzielania wypełniania, planujemy także optymalizować następujące aspekty:
Wydajność AllToAll: Nasze jądro AllToAll obecnie osiąga tylko 1/3 przepustowości Infiniband. Będziemy kontynuować optymalizację tego jadra.
Spekulatywne dekodowanie w stylu EAGLE: W powyższych danych założyliśmy użycie spekulatywnego dekodowania do przewidywania 1 tokenu. EAGLE może wykorzystać strukturę drzewa do przewidywania wielu tokenów, zwiększając długość akceptacji, co może znacząco zwiększyć prędkość wyjściową.
Jądro GEMM: W Modelu Roofline pokazanym wcześniej możemy zauważyć, że wydajność jądra GEMM jest nadal daleka od teoretycznego limitu. Będziemy kontynuować optymalizację tego jadra.
GB200 NVL72: W najnowszym rozwiązaniu GB200 NVL72 od NVIDIA, 72 Blackwell GPU są połączone za pomocą szybkiego NVLink. Dla modeli architektury MoE jest to bardzo duża szansa i wyzwanie.
Wnioski
Wdrożenie multi-node modeli DeepSeek MoE osiąga to, co zazwyczaj jest niemożliwe przy gęstych LLM: jednocześnie poprawiając zarówno przepustowość, jak i opóźnienie. Poprzez rozdzielenie ekspertów na więcej GPU, zmniejszamy nacisk na przepustowość pamięci na urządzenie, co pozwala na szybsze przetwarzanie i wyższe przepustowości systemu. Nasze eksperymenty pokazują konfiguracje EP128 osiągające do 5x wyższą przepustowość przy równoważnych prędkościach wyjściowych w porównaniu do wdrożeń jedno-węzłowych.
Techniki nakładania obliczeń i komunikacji, takie jak mikropartie, znacznie redukują narzuty komunikacji multi-node, a nasza implementacja pokazuje do 40% przyspieszenia. Nasze niestandardowe jądra komunikacyjne AllToAll i zoptymalizowane implementacje jąderek umożliwiły efektywne wdrażanie modelu o 671 miliardach parametrów.
W miarę jak architektury MoE zyskują na popularności dzięki swojej wydajności, te strategie wdrożeń dostarczają wartościowych wskazówek dla skutecznego skalowania takich modeli.