Skalowanie LLMs za pomocą równoległości wielo-GPU i wielowęzłowej
Tłumaczenie automatyczne Ten artykuł został automatycznie przetłumaczony z angielskiego oryginału.
Obciążenia związane z dużymi modelami wychodzą poza pojedyncze GPU z różnych powodów. Zadanie treningowe może wyczerpać pamięć na stan optymalizatora. Inne może wyczerpać ją na aktywacje dla długich sekwencji. Model, który mieści się w pamięci, nadal może nie osiągać docelowej przepustowości. Każdy z tych problemów wymaga innego podziału i wzorca komunikacji.
To praktyczny przegląd głównych strategii równoległości oraz stojących za nimi ograniczeń, opracowany na podstawie Ultra-Scale Playbook firmy Hugging Face. Celem jest pokazanie, co daje każdy rodzaj podziału, jakie dane są komunikowane i kiedy konieczne stają się kombinacje strategii.
TL;DR. Replikowana równoległość danych zwiększa przepustowość treningu, gdy pojedyncza replika mieści się w pamięci. W pełni shardowana równoległość danych dzieli stan modelu, ale dodaje operacje all-gather parametrów oraz reduce-scatter gradientów. Równoległość tensorowa, potokowa, kontekstowa i ekspercka dzieli odpowiednio obliczenia warstw, głębokość modelu, sekwencję oraz warstwy mixture-of-experts (MoE). Łącz je dopiero po zidentyfikowaniu kluczowego ograniczenia pamięci lub komunikacji.
Ten przewodnik zakłada znajomość propagacji wstecznej, warstw Transformer oraz standardowej pętli treningowej w PyTorch.
Zacznij od dwóch budżetów pamięci
Trening i inferencja mają inny profil zużycia pamięci.
training peak ≈ parameters
+ gradients
+ optimizer state
+ saved activations
+ temporary buffers
+ communication buffers
+ allocator headroom
inference peak ≈ resident weights
+ key/value (KV) cache
+ runtime workspace
+ communication buffers
+ allocator headroom
Model z 70 miliardami parametrów ma dolną granicę wynoszącą 140 GB w zapisie dziesiętnym — i dotyczy ona wyłącznie wag BF16. Ta liczba niewiele mówi o treningu, podczas którego dominować mogą gradienty, stan optymalizatora, wagi master i aktywacje. Nie określa też wymagań związanych z servingiem, gdzie znaczenie mają polityka cache, długość sekwencji, współbieżność batchy i kwantyzacja.
Profiluj dokładną architekturę, precyzję, długość sekwencji, micro-batch, optymalizator, politykę checkpointingu i runtime. Rejestruj szczytową ilość zaalokowanej i zarezerwowanej pamięci, tokeny na sekundę, czas spędzony w kernelach oraz czas widoczny w operacjach collective.
Każdy wymiar równoległości wymaga kompromisu
Dla każdej strategii zadaj pytania: który wymiar tensora jest dzielony, który stan jest replikowany i która operacja collective trafia na ścieżkę krytyczną.
| Strategia | Dzieli | Główna korzyść | Wprowadzana komunikacja |
|---|---|---|---|
| Replikowana równoległość danych | batch | przepustowość treningu | gradient all-reduce |
| W pełni shardowana równoległość danych | parametry, gradienty i stan optymalizatora w ramach grupy data-parallel (DP) | pamięć zajmowana przez stan modelu | parameter all-gather, gradient reduce-scatter |
| Równoległość tensorowa | wymiary macierzy lub attention wewnątrz warstw | wagi warstw i aktywacje | operacje collective wewnątrz bloków transformera |
| Równoległość potokowa | grupy warstw | głębokość modelu i stan poszczególnych etapów | aktywacje przesyłane point-to-point oraz puste miejsca w harmonogramie |
| Równoległość kontekstowa | wymiar sekwencji | pamięć aktywacji dla długich sekwencji | wymiana key/value lub attention między elementami grupy sekwencji |
| Równoległość ekspercka | eksperci MoE i routowane tokeny | pojemność ekspertów na rank | dispatch i łączenie tokenów, zwykle all-to-all |
Redukcja pamięci nie jest stałym mnożnikiem. Zależy od stopnia shardowania, tego, co pozostaje replikowane, tymczasowego stanu bez shardowania, polityki aktywacji, paddingu, niezbalansowania oraz buforów.
Replikowana równoległość danych: przepustowość bez zwiększania pojemności
Replikowana równoległość danych, zwykle nazywana distributed data parallel, utrzymuje pełną replikę treningową na każdym ranku. Dokumentacja DistributedDataParallel w PyTorch opisuje ten model replikacji i synchronizacji gradientów. Każdy rank przetwarza inny micro-batch, a gradienty są synchronizowane przed krokiem optymalizatora.
Używaj tej strategii, gdy pełny stan treningu mieści się w pamięci z bezpiecznym zapasem, a globalny batch można zwiększyć lub dostosować akumulację gradientów. Jej główne zalety to prosta semantyka i dojrzała implementacja, która nakłada obliczenia propagacji wstecznej na redukcję gradientów w bucketach.
Dodawanie ranków może zaszkodzić, gdy lokalny batch stanie się zbyt mały. Może też zaszkodzić, gdy sieć nie jest w stanie ukryć kosztu all-reduce lub występują opóźnienia w dostarczaniu danych wejściowych. Pożądany batch optymalizacyjny może się nie skalować.
W pełni shardowana równoległość danych: pamięć stanu w zamian za operacje collective
W pełni shardowana równoległość danych przechowuje fragmenty parametrów, gradientów i stanu optymalizatora w całej grupie. Artykuł opisujący ZeRO przedstawia ten wzorzec shardowania stanu, a API FSDP2 w PyTorch go implementuje. Parametry warstwy są zbierane operacją all-gather na potrzeby obliczeń, a następnie mogą zostać ponownie rozproszone. Gradienty są zwracane do właścicieli za pomocą reduce-scatter.
Aktualna dokumentacja PyTorch rozróżnia API fully_shard w wersji fully sharded data parallel 2 (FSDP2) od starszego wrappera FullyShardedDataParallel. FSDP2 grupuje komunikację według modułów, do których zastosowano fully_shard, i zaleca stosowanie od dołu do góry, aby grupy warstw mogły nakładać komunikację na obliczenia.
from torch.distributed.fsdp import fully_shard
from torch.optim import AdamW
# Apply bottom-up: each block becomes a communication group.
for block in model.transformer.blocks:
fully_shard(block)
# Shard remaining root parameters such as embeddings and output projection.
fully_shard(model)
# Construct the optimizer after parameters have become sharded distributed tensors (DTensors).
optimizer = AdamW(model.parameters(), lr=learning_rate)
To schemat strukturalny, a nie kompletny launcher. Mesh urządzeń, mixed precision, checkpointing, inicjalizacja, stan optymalizatora i rozproszone checkpointy muszą być zgodne ze stosem treningowym.
Sharding jest atrakcyjny, gdy stan modelu stanowi kluczowe ograniczenie, a obliczenia warstw mogą ukryć wystarczającą część ruchu operacji collective. Może być nieopłacalny w przypadku małych modeli, wolnych połączeń, niewielkich warstw lub układów, w których grupa shardowania przekracza nieodpowiednią granicę topologii.
Równoległość tensorowa: podział obliczeń warstwy
Równoległość tensorowa dzieli operacje algebry liniowej wewnątrz warstwy. Przykładami są projekcje column-parallel i row-parallel. Przewodnik NVIDIA po strategiach równoległości opisuje ten podział na poziomie warstwy. Częściowe wyniki wymagają operacji collective wewnątrz bloków transformera, dlatego opóźnienia i przepustowość mają wielokrotny wpływ podczas propagacji w przód i wstecz.
Używaj tej strategii, gdy warstwa lub jej aktywacje nie mieszczą się w pamięci albo gdy macierze są wystarczająco duże, aby kernele partycjonowane były szybsze niż wykonanie na jednym ranku. Mapuj grupę tensor-parallel do najszybszej dostępnej domeny komunikacyjnej, a następnie wykonuj pomiary. Duży stopień tensor parallelism może zmniejszyć lokalne macierze do rozmiaru, przy którym spada efektywność kerneli, a rośnie narzut operacji collective.
Sequence parallelism jest często łączony z tensor parallelism, aby uniknąć replikowania części obliczeń aktywacji. Różni się od context parallelism działającym na pełnej sekwencji wejściowej modelu.
Równoległość potokowa: podział głębokości i czasu harmonogramu
Równoległość potokowa umieszcza różne grupy warstw na różnych etapach i przesyła między nimi aktywacje. Micro-batche pozwalają etapom pracować współbieżnie. Artykuł GPipe wykorzystuje ten harmonogram dla ogromnych sieci neuronowych.
Zmniejsza ona stan modelu przypadający na etap i może ograniczyć ilość komunikacji przekraczającej wolniejszą granicę w porównaniu z operacjami tensor collective wykonywanymi dla każdej warstwy. Kosztami są puste miejsca w harmonogramie, transfery aktywacji, niezbalansowanie etapów, bardziej złożone harmonogramowanie oraz trudniejsze odzyskiwanie i checkpointing.
Dla prostego, zbalansowanego harmonogramu w stylu GPipe z p etapami i m micro-batchami wyidealizowany udział pustego miejsca w propagacji w przód wynosi w przybliżeniu:
(p - 1) / (m + p - 1)
Rzeczywiste harmonogramy mogą używać wariantów one-forward/one-backward, interleaving lub zero-bubble, a nierówny koszt warstw może zdominować wzór. Granice etapów wyznaczaj na podstawie zmierzonego czasu i zużycia pamięci, a nie równej liczby warstw.
Równoległość kontekstowa: podział aktywacji dla długich sekwencji
Równoległość kontekstowa rozdziela wymiar sekwencji. Dokumentacja NVIDIA dotycząca context parallelism opisuje podział sekwencji oraz wymianę key/value potrzebną do obsługi attention. Każdy rank posiada fragment sekwencji, a attention wymienia informacje niezbędne do zachowania semantyki pełnego kontekstu. Implementacje mogą korzystać z pierścieni point-to-point, all-gather, all-to-all lub kombinacji hierarchicznych.
Zmniejsza to zużycie pamięci aktywacji podczas treningu z długim kontekstem, ale replikuje wagi w grupie kontekstowej i wprowadza komunikację związaną z attention. Korzyść zależy od typu attention, maskowania przyczynowego, długości sekwencji, recomputation oraz sposobu łączenia grup kontekstowych z grupami tensor i data parallel.
Nie wybieraj tej strategii na podstawie uniwersalnego progu 8K, 32K lub 100K. Profiluj zużycie pamięci aktywacji i komunikację attention dla rzeczywistej architektury.
Równoległość ekspercka: tylko dla architektury MoE
Równoległość ekspercka rozdziela ekspertów w warstwach mixture-of-experts. Przewodnik NVIDIA po strategiach równoległości opisuje rozmieszczenie ekspertów i łączenie tej strategii z innymi wymiarami równoległości. Router wysyła reprezentacje tokenów do wybranych ekspertów, a następnie łączy ich wyniki. Dla każdego tokena obliczenia wykonują tylko wybrani eksperci, ale pełne wagi ekspertów nadal wymagają przechowywania i odpowiedniego rozmieszczenia podczas servingu.
Równoległość ekspercka nie jest przełącznikiem optymalizacyjnym dla modelu gęstego. Stanowi część architektury MoE. Jej wyzwania obejmują balansowanie obciążenia, limity pojemności, all-to-all tokenów, tokeny odrzucane lub uzupełniane paddingiem, straty pomocnicze oraz nierównomierne obciążenie wynikające z awarii. Śledź liczbę tokenów na eksperta, entropię routingu, przepełnienie pojemności, czas komunikacji oraz jakość według trasy.
Zbuduj układ na podstawie topologii
Systemy treningowe dużych modeli gęstych, które nie używają grupy expert-parallel, zwykle wykorzystują iloczyn rozmiarów grup data-parallel (DP), tensor-parallel (TP), pipeline-parallel (PP) i context-parallel (CP):
world size = DP × TP × PP × CP
Gdy expert parallelism (EP) jest niezależną grupą, przewodnik NVIDIA po strategiach równoległości oblicza całość następująco:
total GPUs = TP × PP × CP × EP × DP
Używaj mesha obsługiwanego przez framework, zamiast mnożyć nieobsługiwane konfiguracje.
Buduj układ w następującej kolejności:
- Narysuj domeny komunikacyjne: połączenia GPU-GPU, przełączniki, granice non-uniform memory access (NUMA), sieć węzłów, oversubscription i ścieżkę do storage.
- Umieść częste operacje collective wrażliwe na opóźnienia — zwykle TP — w najszybszej odpowiedniej domenie.
- Wybierz grupy FSDP lub replikowanego DP na podstawie pozostałej pojemności i przepustowości.
- Dodaj PP, gdy rozmieszczenie głębokości lub ruch między domenami na tym zyskuje, równoważąc zmierzony czas i pamięć etapów.
- Dodaj CP wyłącznie z powodu ograniczenia sekwencji, a expert parallelism (EP) wyłącznie na potrzeby topologii ekspertów modelu.
- Potwierdź podzielność liczby głów, wymiarów ukrytych, warstw, ekspertów, batcha i sekwencji w kandydackim meshu.
- Zbenchmarkuj kilka poprawnych meshów. Heurystyki uwzględniające topologię wybierają kandydatów, ale nie zwycięzców.
Dwa klastry z taką samą liczbą GPU mogą preferować różne układy, ponieważ różnią się przepustowością połączeń, hierarchią przełączników, podłączeniem CPU i rywalizacją o zasoby sieciowe.
Trening i serving wymagają odrębnych decyzji
Inferencja zwykle nie obejmuje gradientów ani stanu optymalizatora, więc układy treningowe w stylu FSDP nie przenoszą się automatycznie.
W przypadku servingu zadaj pytania:
- Czy jedna replika mieści wagi, KV cache, workspace i docelową współbieżność?
- Czy przepustowość bardziej zyska na większej liczbie niezależnych replik, czy na shardowaniu jednej repliki?
- Czy TP zmniejsza presję na wagi i cache przypadające na rank na tyle, aby zrekompensować komunikację dla każdej warstwy?
- Czy PP jest efektywnie obsługiwany przez model i scheduler żądań?
- Jak prefill i decode obciążają obliczenia, przepustowość pamięci i interconnect?
- Co dzieje się z opóźnieniem ogona, gdy żądania mają różne długości promptów i wyników?
Benchmarkuj kompletny serwer wraz z schedulerem, kwantyzacją, rozkładem kontekstu, polityką batchowania i kształtem ruchu. Liczba tokenów na sekundę podczas treningu nie pozwala przewidzieć czasu do pierwszego tokena ani opóźnienia między tokenami podczas servingu.
Uczciwie mierz skalowanie układu
Dla każdego kandydata rejestruj:
- model, kod, runtime, kernele i tożsamość topologii
- globalny i lokalny batch, rozkład sekwencji oraz liczbę tokenów
- szczytowe zużycie pamięci według kategorii, jeśli jest dostępne
- użyteczne tokeny na sekundę oraz wykorzystanie operacji zmiennoprzecinkowych modelu (FLOP), gdy jest obliczane w spójny sposób
- czas widoczny w operacjach all-reduce, all-gather, reduce-scatter, all-to-all i point-to-point
- przestoje wejścia, czas checkpointingu, zachowanie podczas restartu oraz rozkład opóźnień stragglerów
- loss treningu lub zgodność wyników servingu z baseline’em
Porównuj weak scaling i strong scaling w sposób zamierzony. Strong scaling utrzymuje stałą całkowitą ilość pracy przy zwiększaniu liczby ranków. Weak scaling utrzymuje stałą ilość pracy na rank, więc całkowita ilość pracy rośnie wraz z liczbą ranków. Procent opisany jako „efektywność skalowania” nie ma znaczenia bez określenia mianownika i baseline’u.
Podsumowanie
Równoległość jest mapowaniem zmierzonego wąskiego gardła na wymiar tensora i wzorzec komunikacji. Replikacja, sharding, podział warstw, etapowanie, podział sekwencji i routing ekspertów łagodzą różne ograniczenia oraz tworzą różne tryby awarii.
Zainwentaryzuj obciążenie, narysuj topologię, wygeneruj poprawne meshe i wykonaj profilowanie. Zwycięski układ to taki, który mieści się w pamięci z zapasem i minimalizuje widoczną komunikację dla zadania uruchamianego w praktyce.
Materiały referencyjne
- Dokumentacja PyTorch FSDP2
fully_shard— shardowanie parametrów na poziomie modułów, all-gather, reduce-scatter i zachowanie DTensor. - Dokumentacja PyTorch DistributedDataParallel — trening replikowanego modelu z synchronizowanymi gradientami.
- Strategie równoległości w NVIDIA Megatron Core — wymiary strategii data, tensor, pipeline, context, expert i fully sharded.
- Ultra-Scale Playbook Hugging Face — interaktywny przegląd równoległości dużych modeli na dużą skalę.
- Artykuł DeepSpeed ZeRO — shardowanie stanów optymalizatora, gradientów i parametrów w celu zmniejszenia zużycia pamięci.
- Artykuł GPipe — trening potokowy z micro-batchami i wykonywaniem etapowym.
- Dokumentacja równoległości NERSC — definicje strong scaling i weak scaling oparte odpowiednio na stałej całkowitej pracy lub stałej pracy na rank.