Gradient checkpointing: jak oszczędza pamięć GPU?
Tłumaczenie automatyczne
Ten artykuł został automatycznie przetłumaczony z angielskiego oryginału.
Gradient checkpointing zapisuje mniej aktywacji pośrednich podczas przejścia w przód i oblicza brakujące wartości ponownie podczas obliczeń wstecznych. Używaj go, gdy aktywacje sprawiają, że partia treningowa lub długość sekwencji nie mieści się w pamięci GPU, a dodatkowe obliczenia są akceptowalne.
Zmniejsza pamięć zajmowaną przez zapisane aktywacje. Nie usuwa wag modelu, stanu optymalizatora ani wszystkich buforów tymczasowych.
Co jest obliczane ponownie?
Obliczenia wsteczne potrzebują wyników pośrednich z przejścia w przód, aby obliczyć gradienty. Bez checkpointingu biblioteka zachowuje wymagane wyniki. Z checkpointingiem zapisuje wybrane dane wejściowe lub wyniki pośrednie i wykonuje odpowiednie operacje ponownie, gdy ich wartości są potrzebne.
Na przykład biblioteka może zapisać dane wejściowe grupy warstw transformera i usunąć ich wewnętrzne aktywacje. Podczas obliczeń wstecznych wykonuje tę grupę ponownie, korzystając z zapisanych danych wejściowych. Zapisane dane wejściowe stanowią checkpoint aktywacji, który różni się od checkpointu modelu zapisanego na dysku w celu przywrócenia jego stanu.
Pierwotna analiza Chena i współautorów przedstawia strategię z równymi segmentami, która wymaga O(sqrt(n)) pamięci aktywacji dla sieci o n warstwach i około jednego dodatkowego przejścia w przód na krok treningowy. Rzeczywiste transformery mają aktywacje o różnych rozmiarach, różne implementacje mechanizmu uwagi i zasady działania bibliotek. Zmierz uzyskane zużycie pamięci i czas zamiast traktować oszacowanie asymptotyczne jako prognozę dla wdrożenia.
Włącz integrację przewidzianą dla twojego środowiska
W przypadku modeli obsługiwanych przez Hugging Face Trainer ustawienie gradient_checkpointing=True włącza ogólny mechanizm. Przy FSDP sprawdź zamiast tego integrację checkpointingu aktywacji w bibliotece. Aktualna dokumentacja Transformers Trainer zaleca activation_checkpointing dla FSDP, ponieważ ogólny gradient checkpointing może powodować dodatkową operację all-gather podczas obliczeń wstecznych.
Porównaj tę samą konfigurację treningu z checkpointingiem i bez niego. Zapisz szczytową ilość pamięci przydzielonej na GPU, liczbę przykładów lub tokenów na sekundę oraz to, czy docelowa długość sekwencji mieści się w pamięci. Sprawdź również, czy gradienty pozostają poprawne, gdy model używa niestandardowych operacji lub treningu adapterów.
Jeśli stan optymalizatora zajmuje większość pamięci, rozważ zamiast tego podział stanu między urządzenia lub dostrajanie efektywne pod względem liczby parametrów; ponowne obliczanie aktywacji nie usunie tej alokacji.
Sekcja o gradient checkpointingu w LLM Engineering Guide wyjaśnia obliczenia podzielone na segmenty.