Jak trenować agentów opartych na LLM za pomocą RL na własnych GPU

Trenowanie modeli językowych za pomocą uczenia ze wzmocnieniem od dawna kojarzyło się z ogromnymi klastrami i zamkniętymi laboratoriami. Podczas gdy SFT (Supervised Fine-Tuning) stało się rutyną dla inżynierów ML, niezawodne uruchamianie PPO lub GRPO dla agentów wielokrokowych zawsze stanowiło spore wyzwanie. Albo musiałeś pisać własny kod opakowujący wokół vLLM i DeepSpeed, albo godzić się z ograniczeniami gotowych frameworków.

Naukowcy z Berkeley Sky Computing Lab postanowili uprościć ten proces i udostępnili projekt SkyRL jako open source. To pełnoprawny framework obejmujący cały cykl fine-tuningu RL: od generowania trajektorii w izolacji po aktualizację wag modelu w czasie rzeczywistym.

Z czego składa się framework

Deweloperzy nie budowali monolitycznego systemu. Zamiast tego podzielili projekt na kilka niezależnych warstw, które można łatwo łączyć pod kątem konkretnych zadań.

Centralna warstwa skyrl odpowiada za orkiestrację treningu na Twoim sprzęcie. Łączy silnik treningowy z rozproszoną inferencją. Główną cechą tej warstwy jest wsparcie dla otwartego standardu API Tinker. Możesz napisać skrypt treningowy raz, a następnie uruchomić go na serwerze lokalnym z GPU NVIDIA lub AMD, a także w chmurze.

Warstwa skyrl-gym dostarcza interfejs do tworzenia środowisk treningowych. W przeciwieństwie do klasycznego Gymnasium dla robotyki, tutaj środowiska są dostosowane do interakcji tekstowych i kodowych. Po wyjęciu z pudełka otrzymujesz wsparcie dla baz danych SQL, interpreterów Pythona, wyszukiwarek i terminala Linux.

Warstwa skyrl-agent koncentruje się na zadaniach długoterminowych. Pomaga trenować agentów, którzy wykonują dziesiątki sekwencyjnych kroków, wywołują zewnętrzne narzędzia i dostosowują swoje działania na podstawie odpowiedzi środowiska.

Wyniki w praktyce

Prawdziwy framework ocenia się nie po liczbie abstrakcji, lecz po realnych metrykach. Zespół projektu zaprezentował możliwości biblioteki na zadaniu Text-to-SQL.

Autorzy wzięli standardowy model z 7 miliardami parametrów i przeprowadzili trening RL przez potok SkyRL- SQL. Dataset zawierał tylko 653 przykłady treningowe. Model nie tylko zapamiętywał poprawne zapytania SQL — nauczył się je wykonywać w prawdziwym DBMS, otrzymywać błędy wykonania i poprawiać składnię w następnym kroku. W rezultacie ten niewielki model przewyższył komercyjne usługi GPT-4o i o4-mini pod względem dokładności generowania na wyspecjalizowanych benchmarkach.

Innym interesującym przykładem jest integracja ze środowiskiem Harbor do treningu agentów wiersza poleceń. Model uczy się pracować w prawdziwym terminalu: edytować pliki, uruchamiać testy, znajdować błędy w repozytoriach i wdrażać środowiska bez udziału człowieka.

Szczegóły implementacji technicznej

W tradycyjnym treningu RL modeli inferencja i trening często blokują się nawzajem. Model generuje odpowiedzi, proces się zatrzymuje, obliczane są gradienty, aktualizowane są wagi i dopiero wtedy rozpoczyna się nowa iteracja. SkyRL implementuje w pełni asynchroniczny tryb z tzw. in-flight weight updates.

Podczas gdy aktor generuje nowe tokeny i wchodzi w interakcję ze środowiskiem, proces treningowy jednocześnie aktualizuje parametry modelu. Nowe wagi są przesyłane do usługi inferencji na bieżąco, co dramatycznie redukuje przestoje kosztownych akceleratorów.

Biblioteka nie próbuje wymyślać koła na nowo. Pod maską wykorzystuje sprawdzone rozwiązania ze społeczności open-source: pomysły z veRL, OpenRLHF i Search-R1.

Gdzie projekt znalazł już zastosowanie

Mimo młodego wieku, framework stał się już fundamentem dla kilku zewnętrznych projektów badawczych.

Zespół ze Stanford wykorzystał SkyRL do stworzenia Biomni-R0. To agent zdolny do rozwiązywania złożonych zadań biomedycznych i budowania łańcuchów logicznych na poziomie eksperckim. Projekt CodeScout zastosował framework do precyzyjnej lokalizacji błędów w kodzie w benchmarku SWE- Bench. Projekty takie jak OpenThoughts-Agent i Endless Terminals również rozwijane są na bazie tej biblioteki, gdzie RL jest wykorzystywane do generowania zadań proceduralnych w terminalach Linux bez ręcznego etykietowania danych.

Od czego zacząć eksperymenty

Będziesz potrzebować serwera Linux z wieloma GPU i zainstalowanym PyTorch. Proces instalacji jest standardowy dla projektów Python:

git clone https://github.com/NovaSky-AI/SkyRL.git
cd SkyRL
pip install -e .

Następnie możesz uruchomić jeden z gotowych przepisów treningowych. Na przykład, aby uruchomić najprostszy cykl RL na zadaniach matematycznych lub SQL, wystarczy wywołać odpowiedni skrypt z folderu examples:

python -m skyrl.train --config-name=math_ppo

Jeśli chcesz zintegrować własne środowisko, po prostu odziedzicz po klasie API Gymnasium i zdefiniuj metody step() i reset(). Model otrzyma obserwacje jako tekst lub odpowiedzi systemowe i zwróci wygenerowany tekst jako akcję.

Czy warto wdrożyć to w swoich projektach

SkyRL będzie przydatny dla zespołów R&D i inżynierów, którzy wyrośli poza standardowy SFT i chcą wycisnąć maksymalną wydajność z modeli średniej wielkości. Uczenie LLM korzystania z narzędzi w określonej domenie poprzez RL to realny sposób na pokonanie ogólnych modeli proprialnych pod względem dokładności i kosztów zapytań.

Główna uwaga: projekt jest w aktywnym rozwoju. Repozytorium ma około 400 otwartych zgłoszeń, a architektura poszczególnych modułów nadal ewoluuje. Jeśli potrzebujesz gotowego rozwiązania typu „just works