Gradient Accumulation

Chcesz wyliczyć jedną poprawkę modelu na dużej grupie przykładów, ale cała grupa nie mieści się w pamięci. Możesz przetworzyć kilka małych porcji i połączyć ich wskazówki przed zmianą wag.

Gradient Accumulation zbiera gradienty, czyli informacje o wpływie wag na błąd, z kilku przebiegów. Wagi pozostają niezmienione do końca zbierania, a optymalizator — algorytm aktualizacji — wykonuje dopiero jeden wspólny krok.

Cztery równe porcje po osiem przykładów dają efektywną grupę 32. Żeby odpowiadać średniej dużej grupy, trzeba właściwie uśrednić wkłady. Przy nierównych porcjach znaczenie ma liczba ocenianych przykładów lub tokenów, czyli kawałków tekstu.

To nie jest to samo co cztery kolejne aktualizacje wag. Nie każda operacja modelu zachowuje się identycznie po podziale grupy, dlatego równoważność z jedną wielką porcją wymaga sprawdzenia założeń.

Mechanizm i szczegóły

Przy równych porcjach po $b$ przykładów i $k$ przebiegach jeden krok obejmuje $bk$ przykładów. Z $p$ replikami Data Parallelism jest to $bkp$, jeśli każda dostaje odrębne dane. Dokumentacja PyTorch 2.11, Gradient accumulation pokazuje dodawanie gradientów i aktualizację dopiero po całej porcji.

Średnia musi mieć właściwy mianownik

Jeżeli $g_j$ jest gradientem średniej straty z porcji liczącej $n_j$ przykładów, gradient średniej po wszystkich przykładach wynosi:

$$g=\sum_j \frac{n_j}{\sum_i n_i}g_j.$$

Dzielenie każdej straty przez liczbę porcji wystarcza przy równych rozmiarach. Przy nierównych trzeba uwzględnić liczebności. Dla Cross-entropy uśrednianej po tokenach liczą się oceniane tokeny, również z uwzględnieniem Loss masking. To wynika z liniowości sumy i różniczkowania.

Własny przykład poniżej używa czterech celów $[0,0,0,8]$, jednej wagi $w=1$ i straty $\ell=\frac12(w-y)^2$. Gradienty to $[1,1,1,-7]$, więc średnia wynosi $-1$. Jeden krok SGD z learning rate 0,25 daje $w=1{,}25$, niezależnie od podziału danych. Zobacz, co zepsuje podział 3+1 i zwykła średnia dwóch średnich.

Zgodność z dużym batchem zakłada tę samą funkcję straty i te same parametry podczas wszystkich przebiegów. Operacje zależne od całego batcha, losowość i zaokrąglenia mogą ją naruszyć. Wykonanie kroku optymalizatora po każdej porcji zmienia już sam algorytm.

Przy DDP można odroczyć synchronizację przez no_sync(). Kontekst musi obejmować również forward; pierwszy przebieg poza nim synchronizuje zgromadzone gradienty. Opisuje to dokumentacja DistributedDataParallel.