Gradient Accumulation: effektive Batches ohne Verlust der korrekten Normalisierung

Worum es geht

Bei der Gradient Accumulation werden mehrere Microbatches verarbeitet, bevor die Modellparameter aktualisiert werden. Ihre Gradienten fliessen gemeinsam in einen einzigen Optimizer-Schritt ein. So lässt sich ein grösserer effektiver Batch verwenden, wenn ein entsprechend grosser physischer Batch nicht in den Gerätespeicher passt. In PyTorch addieren Backward-Aufrufe ihre Werte zu den bereits vorhandenen Parametergradienten, bis diese zurückgesetzt werden. [1][2]

Funktionsweise

Bei gleich grossen Microbatches und einem Mittelwert-Loss über gleich gewichtete Beispiele wird jeder Loss durch die Anzahl Microbatches im Akkumulationsfenster geteilt. Die Summe dieser Gradienten ergibt dann den Gradienten des gemeinsamen Mittelwert-Loss. Bei ungleichen Grössen ist nach der tatsächlichen Anzahl beitragender Beispiele zu gewichten. Bei Token-Losses zählen nur gültige, nicht maskierte Tokens; Klassengewichte oder andere Reduktionen können einen anderen Nenner erfordern. Wer die Mittelwerte der Microbatches unbesehen mittelt, verändert die Zielfunktion. Diese Normalisierungsregeln ergeben sich aus der Ableitung des beabsichtigten gewichteten Mittelwerts.

Praktischer Ablauf

Zu Beginn eines Fensters die Gradienten zurücksetzen, für jeden normalisierten Microbatch Forward- und Backward-Pass ausführen und anschliessend genau einen Optimizer-Schritt durchführen. Danach die Gradienten für das nächste Fenster zurücksetzen. Getrennte Forward-Passes erfordern in der Regel nicht, frühere Berechnungsgraphen aufzubewahren. [2] Das letzte, unvollständige Fenster explizit behandeln. Bei Mixed Precision den Gradient-Scale während der gesamten Akkumulation konstant halten, nach dem letzten Backward-Pass einmalig unskalieren und erst dann den kombinierten Gradienten clippen und den Schritt ausführen. Den Scaler an dieser Grenze aktualisieren. [1] Einen Scheduler, der pro Optimizer-Update konfiguriert ist, nur dann weiterschalten, wenn tatsächlich ein Update erfolgt.

Anschauungsbeispiel

Angenommen, ein Gerät verarbeitet vier Microbatches mit je acht Beispielen. Der effektive Batch umfasst dann 32 Beispiele, jeder Microbatch-Mittelwert erhält also das Gewicht 1/4. Ein letztes Fenster mit 8, 8, 8 und 2 Beispielen umfasst 26 Beispiele; die zugehörigen Gewichte betragen 8/26, 8/26, 8/26 und 2/26. Es handelt sich um reine Arithmetik, nicht um ein gemessenes Trainingsergebnis.

Grenzen und Überprüfung

Batch Normalization verwendet die Statistiken des jeweiligen physischen Microbatches; die Akkumulation bildet einen einzelnen grossen Batch deshalb nicht automatisch nach. [3] Auch Zufallsoperationen und die Summationsreihenfolge bei Gleitkommazahlen können die Ergebnisse verändern. Wird jeder Microbatch einzeln geclippt, ändert das in der Regel das kombinierte Update. Das Verfahren verringert den gleichzeitig benötigten Speicher für Aktivierungen; Modellgewichte und Optimizer-Zustände belegen weiterhin Speicher, und die zusätzlichen Durchläufe können die Laufzeit erhöhen.

Vor einem langen Trainingslauf ein akkumuliertes Update mit einem Update über den kombinierten Batch vergleichen, und zwar an einem kleinen deterministischen Modell ohne batchabhängige Schichten. Dabei Loss-Normalisierung, Gradientenabweichungen innerhalb einer geeigneten Toleranz, Anzahl Updates und Verhalten bei unvollständigen Fenstern prüfen. Es handelt sich um empfohlene Prüfungen, nicht um Tests, die für diesen Artikel durchgeführt wurden.

Quellen

[1] PyTorch: Automatic Mixed Precision examples [2] PyTorch: torch.Tensor.backward [3] PyTorch: BatchNorm1d