Jak trenować ogromne modele uczenia ze wzmocnieniem, nie tracąc przy tym zmysłów

Jeśli chodzi o uczenie ze wzmocnieniem (RL) dla dużych modeli językowych, większość programistów od razu myśli o DeepSeek-R1 lub OpenAI o1. Ale gdy tylko próbujesz odtworzyć coś podobnego "w domu" lub na klastrze firmowym, natrafiasz na ścianę. Standardowe biblioteki są albo zbyt wolne, albo zawodzą, gdy próbujesz rozproszyć trening na kilka węzłów.

Niedawno natknąłem się na Prime-RL od zespołu Prime Intellect. Ci ludzie wydali framework specjalnie zaprojektowany do skalowania RL na tysiące GPU. To nie jest tylko kolejne opakowanie PyTorch, ale pełnowartościowe środowisko do trenowania modeli "agentowych", które potrafią rozumować, pisać kod i korzystać z narzędzi.

Pod maską

Krótko mówiąc, Prime-RL to pomost między efektywnym treningiem a szybkim generowaniem odpowiedzi. Główny problem konwencjonalnego RL polega na tym, że model musi ciągle coś generować (wnioskowanie), aby ocenić swoje działania, a następnie aktualizować wagi (trening). W Prime-RL ten proces jest w pełni asynchroniczny. Podczas gdy niektóre GPU obliczają gradienty, inne już generują nowe przykłady.

Deweloperzy zaimplementowali obsługę FSDP2 do dystrybucji wag oraz vLLM do reaktywnego wnioskowania. Dla tych, którzy pracują z ogromnymi modelami Mixture-of-Experts (MoE), dodali Expert Parallelism (EP) i Context Parallelism (CP). To drugie jest kluczowe, jeśli musisz dostarczyć modelowi ogromne logi lub długie łańcuchy rozumowania.

Jak ten projekt pomaga w praktyce

Repozytorium ma mnóstwo gotowych przykładów i to właśnie to czyni je atrakcyjnym. Zamiast zgadywać, jak skonfigurować rzeczy, możesz po prostu wziąć szablon.

Obsługa potentatów

Out of the box framework dobrze współpracuje z rodzinami Qwen 3, GLM-5, a nawet MiniMax. I nie jest to tylko obsługa modeli Hugging Face — to są zoptymalizowane implementacje. Na przykład GLM-5 wykorzystuje wnioskowanie FP8 i rozdzielenie zasobów między generowaniem a treningiem (rozdzielenie P/D). Jeśli masz dostęp do GPU H100 lub B200, pozwoli ci to wycisnąć maksymalną wydajność ze sprzętu.

Środowiska agentowe

Prime-RL jest zintegrowany z hubem środowisk 2. Oznacza to, że możesz trenować modele nie tylko do odpowiadania na pytania, ale także do faktycznego rozwiązywania zadań w środowiskach takich jak SWE-bench (naprawianie błędów w kodzie) czy grania w złożone gry jak Wordle, aby ćwiczyć logikę.

Elastyczność konfiguracji

Zamiast przepisywać kod za każdym razem, gdy zmieniają się hiperparametry, wszystko zostało przeniesione do plików konfiguracyjnych TOML. Jest to wygodne: jeden plik opisuje architekturę modelu, parametry optymizatora i ustawienia serwera wnioskowania.

Jak uruchomić i przetestować

Aby zacząć, wystarczy nawet pojedynczy GPU na poziomie RTX 3090/4090, chociaż projekt jest wyraźnie ukierunkowany na klastry. Instalacja jest dość prosta dzięki 3:

0

Po instalacji możesz przeprowadzić szybki test treningu modelu dla zadania odwracania 4 tekstu. To świetny sposób, aby sprawdzić, czy wszystkie sterowniki i powiązania CUDA są prawidłowo skonfigurowane, zanim wynajmiesz sto H100.

Polecenie uruchamiające prosty trener RL wygląda tak:

1

Kto powinien się temu przyjrzeć

Projekt będzie najbardziej interesujący dla tych, którzy wyrośli ze standardowych skryptów z bibliotek takich jak TRL i chcą czegoś bardziej poważnego do produkcji. Jeśli pracujesz nad tworzeniem własnych modeli "rozumujących" lub automatyzacją pisania kodu, Prime-RL zaoszczędzi tygodnie pisania kodu infrastruktury.

Oczywiście bariera wejścia jest tutaj wyższa niż w przypadku Keras. Musisz rozumieć, jak działają Slurm lub Kubernetes i umieć konfigurować rozproszony trening. Ale dokumentacja w folderze 5 jest dość szczegółowa: znajdziesz tam przewodniki po algorytmach (na przykład o stratze AIPO) i o skalowaniu.

Podsumowując, mamy potężne, choć wymagające pod względem sprzętu narzędzie, które sprawia, że proces trenowania dużych agentów RL staje się przewidywalny i szybki. Jeśli planujesz trenować coś większego niż 7B parametrów przy użyciu uczenia ze wzmocnieniem, Prime-RL zdecydowanie warto dodać do zakładek.