Activation Checkpointing

Activation Checkpointing zmniejsza pamięć potrzebną do treningu: zachowuje wybrane aktywacje, a pozostałe odtwarza, gdy obliczanie gradientu ich potrzebuje. Płaci za to dodatkowymi obliczeniami. Aktywacje są pośrednimi wynikami przejścia przez model. Ten checkpoint nie jest kopią wag na dysku służącą do wznowienia treningu.

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.

Wykorzystuję treści generowane przez AI jako część mojego codziennego procesu nauki.