NVIDIA prezentuje Molt: nowe ramy dla agentowego uczenia ze wzmocnieniem w PyTorch

Zespół NeMo firmy NVIDIA zaprezentował Molt – nowy framework do agentowego uczenia ze wzmocnieniem (RL), natywny dla PyTorch. Jego głównym celem jest zredukowanie kosztów związanych z ciągłymi modyfikacjami algorytmów w badaniach RL, które w tradycyjnych frameworkach często wymagają przeprojektowania wielu warstw kodu.

Kompaktowa konstrukcja i dostępność

Molt wyróżnia się nietypowym podejściem do projektowania, stawiając na kompaktowość. Cała baza kodu ma być na tyle zwięzła, aby badacz mógł ją w pełni zrozumieć, a asystent kodowania AI mógł ją analizować w całości. Według pomiarów, Molt zawiera około 8,6 tys. linii kodu RL, co stanowi znaczną redukcję w porównaniu do innych frameworków, takich jak verl (około 62 tys. linii) czy slime (25 tys. linii). OpenRLHF jest nieco mniejszy, z 7,2 tys. linii.

Framework jest dostępny na licencji Apache 2.0 i zawiera skrypty Slurm oraz prekonfigurowany kontener. Chociaż Molt jest pozycjonowany jako infrastruktura badawcza, a nie usługa produkcyjna, jego wymagania sprzętowe są znaczące. Dostępne przepisy zakładają wykorzystanie dwóch węzłów z ośmioma procesorami graficznymi H100 każdy, z podziałem na osiem GPU do treningu i osiem do generowania danych (rollout).

Takie wymagania sprawiają, że Molt jest przeznaczony dla:

• zaawansowanych laboratoriów badawczych,
• dobrze finansowanych startupów AI zajmujących się post-treningiem,
• korporacyjnych grup badawczych AI w sektorach finansów, opieki zdrowotnej i robotyki, które trenują agentów w środowiskach własnościowych,
• laboratoriów akademickich z dostępem do wielowęzłowych systemów H100/H200.

Potencjalne zastosowania Molt obejmują:

• agentów do wieloetapowego użycia narzędzi,
• agentów do wykonywania kodu,
• środowiska wizualno-językowe (np. przepis geo3k),
• pętle nagród z LLM działającym jako sędzia,
• destylację polityk na mniejsze modele studenckie.

Architektura i mechanizmy działania

Molt integruje się z istniejącymi narzędziami, takimi jak Ray do zarządzania zasobami i asynchronicznych kolejek, vLLM do generowania danych oraz NVIDIA AutoModel z FSDP2 do treningu. Ważne jest, że żadne z tych komponentów nie zostało zmodyfikowane (forked), co oznacza, że ulepszenia w ich oryginalnych wersjach są łatwo integrowane poprzez aktualizację kontenera.

Środowisko uruchomieniowe Molt składa się z puli agentów, zestawu silników vLLM za routerem zapytań oraz pojedynczego aktora polityki podlegającego treningowi. Pula strumieniowa utrzymuje grupy promptów w trakcie przetwarzania, aby silniki nie były bezczynne podczas treningu aktora. Częściowe generowanie danych (partial rollout) wstrzymuje silniki, rozsyła fragmenty aktora za pomocą NCCL bezpośrednio do każdego silnika i wznawia zatrzymane żądania zamiast je odrzucać.

Proces uczenia ze wzmocnieniem w Molt wymaga jedynie wskazania modułu Pythona eksportującego AgentRunner. Cała reszta kodu, w tym definicja nagrody, jest standardowym kodem Pythona.

Framework obsługuje dwa tryby pracy:

• Z Env, framework zarządza pętlą LLM w zgodności z metodą step() z Gymnasium.
• Z ChatAgent, użytkownik kontroluje pętlę za pomocą standardowego SDK OpenAI lub Anthropic. Molt uruchamia serwer loopback, który obsługuje oba protokoły, a każde żądanie jest dekodowane po stronie serwera do dokładnej akumulacji tokenów. Gdy agent o długim horyzoncie kompaktuje swój kontekst i przepisuje prefiks, serwer automatycznie zamyka bieżący segment i otwiera nowy.

Inwarianty poprawności i spójności

Projekt Molt opiera się na trzech inwariantach poprawności:

1. Tożsamość tokenów: Trajektoria jest definiowana przez identyfikatory próbkowanych tokenów, a nie przez ponownie tokenizowany transkrypt.
2. Semantyka wersji polityki: Treningowe tokeny zachowują swoje log-prawdopodobieństwa polityki behawioralnej, a asynchroniczne użycie jest korygowane dla każdego tokena za bramą na poziomie sekwencji.
3. Spójność w przód: Generowanie danych (rollout) i aktor muszą być zgodne co do semantyki modelu.

Ostatni inwariant jest szczególnie istotny dla polityk typu mixture-of-experts (MoE). W takich systemach routery rolloutu i treningu niezależnie wybierają ekspertów, a niewielkie różnice numeryczne mogą zmienić wybory top-k. Molt rozwiązuje ten problem poprzez replay routingu rolloutu, gdzie vLLM zwraca identyfikatory ekspertów dla każdego tokena, a trening odtwarza je w przód.

Znaczenie dla badań nad AI

Wprowadzenie Molt przez NVIDIA stanowi istotny krok w kierunku usprawnienia badań nad agentowym uczeniem ze wzmocnieniem. Dzięki swojej kompaktowej budowie i integracji z zaawansowanymi narzędziami, framework ma potencjał, aby znacząco przyspieszyć cykl eksperymentalny dla badaczy. Umożliwia to szybsze testowanie nowych algorytmów i iteracji, co jest kluczowe w dynamicznie rozwijającej się dziedzinie AI, szczególnie w kontekście złożonych agentów i dużych modeli językowych. Ograniczenie złożoności kodu może również obniżyć barierę wejścia dla nowych badaczy, jednocześnie zwiększając efektywność pracy doświadczonych zespołów.

Źródło: marktechpost.com