Trenowanie z mieszaną precyzją za pomocą AMP

Trenowanie w trybie mixed precision wykonuje większość operacji na 16-bitowych liczbach zmiennoprzecinkowych (FP16) zamiast 32-bitowych (FP32). FP16 zajmuje o połowę mniej pamięci i działa szybciej na współczesnych rdzeniach tensorowych GPU, podczas gdy kilka wrażliwych operacji pozostaje w FP32 dla zachowania stabilności.

Mixed precision zapewnia:

• około 2 razy mniejsze zużycie pamięci, co umożliwia używanie większych batchy lub modeli
• Szybsze mnożenie macierzy na rdzeniach tensorowych
• Brak lub niewielki spadek końcowej dokładności modelu, jeśli rozwiązanie zostanie poprawnie zastosowane

FP16 ma wąski zakres wartości. Małe wartości gradientów mogą ulec underflow i stać się zerem, niepostrzeżenie zatrzymując uczenie. To główne wyzwanie, które musi rozwiązać mixed precision, i powód, dla którego nie można po prostu przekonwertować wszystkiego do FP16.

PyTorch Automatic Mixed Precision (AMP) obsługuje te szczegóły za pomocą dwóch narzędzi: autocast dobiera precyzję dla poszczególnych operacji, a GradScaler zapobiega underflow gradientów.

Przebieg w przód należy umieścić w autocast(). W jego obrębie PyTorch automatycznie uruchamia każdą operację z użyciem najbezpieczniejszej precyzji: mnożenie macierzy w FP16, a redukcje, takie jak softmax i obliczanie funkcji straty, w FP32.

GradScaler przeciwdziała underflow, mnożąc funkcję straty przez duży współczynnik skalowania przed propagacją wsteczną. Przesuwa to małe gradienty do zakresu reprezentowalnego w FP16; współczynnik jest usuwany przed krokiem optymalizatora.

scaler.scale(loss).backward() skaluje funkcję straty w górę, a następnie uruchamia propagację wsteczną. Wynikowe gradienty również są przeskalowane, dzięki czemu małe wartości nie zanikają.

scaler.step(optimizer) najpierw przywraca gradientom ich rzeczywistą skalę, a następnie wywołuje optimizer.step(), ale tylko wtedy, gdy nie pojawiły się wartości inf ani NaN. Jeśli gradienty uległy przepełnieniu, krok zostaje pominięty.

scaler.update() dostosowuje współczynnik skalowania na następną iterację: zwiększa go, gdy kroki kończą się powodzeniem, i zmniejsza po wystąpieniu przepełnienia. To adaptacyjne skalowanie automatycznie zapewnia stabilność trenowania.

Po połączeniu wszystkich elementów otrzymują Państwo kompletny krok trenowania w trybie mixed precision.

Warto pamiętać o następujących zasadach:

• Wywołują Państwo backward() tylko dla przeskalowanej funkcji straty
• autocast obejmuje wyłącznie przebieg w przód, a nie propagację wsteczną ani krok optymalizatora
• AMP najlepiej sprawdza się na GPU z rdzeniami tensorowymi; na starszym sprzęcie korzyści są mniejsze

Proszę sprawdzić swoją wiedzę o AMP.

Poznali Państwo trenowanie w trybie mixed precision z użyciem AMP:

• autocast() dobiera FP16 lub FP32 dla każdej operacji podczas przebiegu w przód
• GradScaler skaluje funkcję straty, aby zapobiec underflow gradientów
• Cykl obejmuje scaler.scale(loss).backward(), scaler.step(optimizer) oraz scaler.update()
• Rezultat: około 2 razy mniejsze zużycie pamięci i szybsze trenowanie