Activation Checkpointing

Podczas treningu model musi pamiętać wiele wyników pośrednich, aby później wyliczyć poprawki wag. Te wyniki zajmują tyle miejsca, że cała praca nie mieści się na karcie. Możesz zachować mniej danych, a brakujące później policzyć ponownie.

Activation Checkpointing przechowuje wybrane aktywacje — wyniki pośrednich obliczeń — i odtwarza inne, gdy są potrzebne do liczenia gradientu, czyli wpływu wag na błąd.

Jeżeli zachowasz wejście pewnego fragmentu sieci, możesz ponownie przejść jego warstwy przy obliczeniach wstecz. Zamiast trzymać wszystkie kroki naraz, płacisz dodatkowymi obliczeniami za rekonstrukcję.

Ten checkpoint nie jest kopią modelu na dysku do wznowienia treningu. Oszczędność pamięci zależy od miejsc zachowania wyników. Powtórzone obliczenia muszą być zgodne, co wymaga uwagi przy losowych operacjach i stanie zmienianym podczas przejścia.

Mechanizm i szczegóły

Chen et al., Training Deep Nets with Sublinear Memory Cost, §4.1–4.3 opisują podział grafu obliczeń na fragmenty: zachowane wejścia pozwalają ponownie wykonać fragment podczas backward. Dla odpowiednio podzielonego łańcucha warstw autorzy uzyskują koszt pamięci aktywacji rzędu $O(\sqrt n)$ zamiast $O(n)$. To wynik przy założeniach pracy, nie obietnica takiego spadku całej pamięci GPU.

Usunięcie wyniku nie musi zmieniać gradientu

We własnym przykładzie liczymy:

$a=x^2,\qquad b=a+1,\qquad z=b^2.$

Dla $x=2$ otrzymujemy $a=4$, $b=5$ i $z=25$. Do obliczenia pochodnej ostatniego kwadratu potrzebujemy $b$. Jeśli go usuniemy, trzeba najpierw odtworzyć $a$ i $b$ z zachowanego $x$. Potem reguła łańcuchowa daje ten sam wynik:

$\frac{dz}{dx}=2b\cdot1\cdot2x=40.$

Przejdź backward krok po kroku. Porównaj zachowanie $b$ z jego usunięciem i zmień wejście także na zero lub liczbę ujemną.

Licznik pokazuje wartości zachowane po forward na potrzeby backward: tutaj $x,b$ albo tylko $x$. Pomija wspólny wynik $z$, gradienty i chwilowe bufory. Nie jest pomiarem maksymalnego zużycia pamięci; podczas odtwarzania znowu powstają wartości pośrednie. Demo wykonuje tylko niezbędne ponowne operacje, nie cały forward.

Odtworzenie musi być wierne

Gdy odtwarzany fragment używa losowości, np. dropout, trzeba zadbać o zgodny stan generatora. Zmiana zachowania funkcji między forward a backward może dać błędny gradient. Dokumentacja PyTorch, torch.utils.checkpoint opisuje te ograniczenia oraz wariant zatrzymujący odtwarzanie po uzyskaniu potrzebnych aktywacji.

Metodę można stosować do bloków Transformer i łączyć z Pipeline Parallelism. Ogranicza pamięć aktywacji; nie usuwa wag ani stanów optymalizatora. Zysk pamięci i koszt czasu zależą od miejsca wybranych checkpointów.