Etapy DeepSpeed ZeRO: który stan treningu jest dzielony?
Tłumaczenie automatyczne
Ten artykuł został automatycznie przetłumaczony z angielskiego oryginału.
ZeRO-1 dzieli stan optymalizatora między GPUs pracujące równolegle na danych, ZeRO-2 dodatkowo dzieli gradienty, a ZeRO-3 również parametry modelu. Przetestuj najniższy etap, przy którym zadanie mieści się w pamięci i osiąga akceptowalną przepustowość. Wyższy etap oszczędza więcej pamięci stanu modelu, ale może wymagać przesyłania większej ilości danych.
Te etapy zmieniają miejsce przechowywania stanu treningu. Nie zmniejszają automatycznie pamięci zajmowanej przez aktywacje.
Przykład zużycia pamięci z jawnymi założeniami
Załóżmy, że trening używa wag i gradientów w FP16 oraz głównych kopii wag Adama w FP32 i dwóch momentów optymalizatora. Na każdy parametr przypadają wtedy 2 bajty na wagi, 2 na gradienty i 12 na stan optymalizatora: łącznie 16 bajtów.
Dla 7.5 miliarda parametrów na ośmiu GPUs obliczenie pamięci stanu modelu daje:
| Etap | Stan dzielony między GPUs | GB stanu modelu na GPU |
|---|---|---|
| Zwykła równoległość danych | Żaden | 120 |
| ZeRO-1 | Stan optymalizatora | 41.25 |
| ZeRO-2 | Stan optymalizatora i gradienty | 28.125 |
| ZeRO-3 | Wszystkie trzy składniki | 15 |
To dziesiętne GB. Nie obejmują aktywacji, tymczasowo zebranych parametrów, narzutu alokatora ani buforów środowiska wykonawczego. Stos treningowy, który przechowuje dane z inną precyzją, wymaga innych obliczeń.
Publikacja o ZeRO wyjaśnia podział stanu i analizuje komunikację. ZeRO-3 odtwarza potrzebne parametry przez operacje all-gather; analiza pełnego podziału w publikacji daje objętość komunikacji około 1.5 razy większą niż w standardowej równoległości danych. Ten stosunek wynika z analizy przy jej założeniach i nie jest gwarancją dla każdej sieci i konfiguracji.
Wybierz na podstawie profilu pamięci
Jeśli parametry mieszczą się w pamięci, a błąd powoduje stan optymalizatora, przetestuj etap 1 lub 2. Jeśli w pełni replikowane parametry się nie mieszczą, przetestuj etap 3. Jeśli dominują aktywacje, połącz wybrany etap z checkpointingiem aktywacji lub mniejszą partią treningową.
DeepSpeed obsługuje też przenoszenie stanu do CPU i NVMe. Sprawdź aktualną konfigurację ZeRO i zmierz koszt transferu na rzeczywistym hoście i nośniku danych. Dodatkowa pojemność uzyskana przez przeniesienie stanu może pomóc, ale powolne transfery mogą obniżyć przepustowość.
Sekcja o ZeRO w LLM Engineering Guide zawiera wykonywalne obliczenie dla innych liczb parametrów i grup GPU.