Współczesne uczenie przez wzmacnianie (RL) w kontekście dużych modeli językowych wymaga ciągłego balansowania między wydajnością a stabilnością procesu treningowego. Głównym wyzwaniem jest asynchroniczna natura systemów RL, w których procesy próbkowania (samplery) generują dane, podczas gdy trener aktualizuje politykę modelu. Powstaje wówczas zjawisko rozbieżności polityki, wynikające z używania nieaktualnych danych oraz błędów kwantyzacji, co w skrajnych przypadkach prowadzi do degradacji wyników modelu.
Wyzwania kwantyzacji w RL
Kwantyzacja, czyli redukcja precyzji obliczeń, jest kluczowa dla zwiększenia przepustowości sprzętowej. Choć formaty takie jak NVFP4 pozwalają na znaczące przyspieszenie operacji na nowoczesnych procesorach graficznych (np. NVIDIA Rubin), ich zastosowanie w RL jest znacznie trudniejsze niż w standardowym uczeniu wstępnym (pretraining). W uczeniu wstępnym sygnał gradientu jest gęsty i uśredniany, natomiast w RL gradient jest z natury zaszumiony, co sprawia, że błędy kwantyzacji mogą łatwo zdestabilizować proces optymalizacji.
Aby rozwiązać ten problem, opracowano nową recepturę niskoprecyzyjnego uczenia przez wzmacnianie, która zachowuje dynamikę treningu typową dla wyższych precyzji. Kluczowe interwencje obejmują:
Stabilizacja gradientów i wydajność
Tradycyjne podejście, w którym faza wsteczna odbywa się w pełnej precyzji BF16, często prowadzi do rozbieżności z fazą forward, jeśli ta jest skwantyzowana. Zastosowanie dekwantyzacji wag w fazie wstecznej pozwala na użycie operacji BF16 przy jednoczesnym uwzględnieniu ograniczeń wynikających z kwantyzacji NVFP4. Dzięki temu rozwiązaniu udało się wyeliminować nagłe skoki normy gradientu, które występowały w mniej zaawansowanych konfiguracjach.
Wdrożenie techniki 4/6, opartej na adaptacyjnym skalowaniu bloków, pozwoliło na dalszą poprawę precyzji bez zauważalnego wpływu na czas generowania próbek. Optymalizacja ta została zintegrowana z bibliotekami TransformerEngine oraz FlashInfer, co zapewnia bitową zgodność między trenerem a samplerem. Warto zauważyć, że optymalizacje te redukują szczytowe zużycie pamięci podczas treningu o około 70% (dla przykładowych obciążeń redukcja z ok. 600 zł do ok. 180 zł w przeliczeniu na koszt zasobów pamięciowych przy założeniu standardowych stawek za jednostki obliczeniowe).
Selektywna precyzja i praktyczne wdrożenie
Nie wszystkie warstwy modelu wymagają tej samej precyzji. Zgodnie z przyjętą strategią, kluczowe elementy architektury, takie jak „współdzieleni eksperci” w modelach MoE (Mixture-of-Experts) oraz ostatnie 15% warstw modelu, utrzymywane są w wyższej precyzji (BF16). Takie podejście pozwala na zachowanie stabilności przy jednoczesnym czerpaniu korzyści z wydajności NVFP4 w pozostałych 97% parametrów modelu.
Praktycznym efektem tych prac jest możliwość bezpośredniego serwowania modeli skwantyzowanych do NVFP4 bez konieczności przeprowadzania kosztownej kalibracji czy dodatkowych procesów konwersji. Wystarczy użyć odpowiedniej flagi w środowisku SGLang, aby uzyskać wydajność zbliżoną do modeli pełnoprecyzyjnych przy znacznie mniejszym zapotrzebowaniu na zasoby sprzętowe. Opracowane rozwiązania są dostępne jako otwarte oprogramowanie, co umożliwia społeczności badawczej łatwiejsze eksperymentowanie z niskoprecyzyjnym uczeniem przez wzmacnianie.
