Transfer wag do RL po szkoleniu w poniżej 2 sekundy

Ostatnio osiągnęliśmy 1,3-sekundową aktualizację parametrów między maszynami dla modelu Kimi-K2 (1T parametrów), przesyłając wagi z 256 procesorów graficznych (GPU) szkoleniowych (BF16) do 128 procesorów graficznych (GPU) wnioskujących (FP8).

W asynchronicznym dostrajaniu wzmacniającym (reinforcement learning fine-tuning) szkolenie i wnioskowanie odbywają się na oddzielnych procesorach graficznych (GPU). Po każdym kroku szkoleniowym nowe wagi muszą zostać przesłane do węzłów wnioskowania. Wiele istniejących frameworków zajmuje kilka sekund, a nawet minut, w przypadku modeli z bilionami parametrów.

Dzięki wykorzystaniu komunikacji punkt-punkt RDMA jesteśmy w stanie sprawić, że transfer wag będzie niezwykle szybki, bez zmiany silnika wnioskowania, a także sprawić, że kod będzie łatwiejszy do napisania i utrzymania.

Nasze rozwiązanie opiera się na RDMA WRITE, jednostronnym prymitywie, w którym źródło zapisuje dane bezpośrednio do pamięci GPU docelowego.

python

def rdma_write(src_ptr, dst_ptr, size, src_mr, dst_mr):
    # Write from local [src_ptr, src_ptr+size) to remote [dst_ptr, dst_ptr+size).
    # src_mr and dst_mr contains Memory Region metadata of both sides.
    ...

Strona docelowa nie otrzymuje nawet powiadomienia o transferze. Zapewnia to niskie opóźnienia, wysoką przepustowość oraz transfery typu zero-copy inicjowane przez węzły szkoleniowe, bez konieczności stosowania logiki sterującej w węzłach wnioskowania.

Gromadzenie metadanych – Kontroler zbiera metadane parametrów ze wszystkich szkoleniowych i wnioskujących procesorów graficznych (GPU).

Obliczanie harmonogramu – Kontroler oblicza statyczny harmonogram transferu wag, mapując, który procesor graficzny (GPU) szkoleniowy wysyła dany parametr do którego procesora graficznego (GPU) wnioskującego oraz w jakiej kolejności.

Dystrybucja harmonogramu – Kontroler wysyła harmonogram do wszystkich szkoleniowych procesorów graficznych (GPU).

Wykonanie – Po każdym kroku szkoleniowym kontroler wysyła sygnał do szkoleniowych procesorów graficznych (GPU) o rozpoczęciu transferu.

Po zdefiniowaniu wysokopoziomowego przepływu pracy, kluczowym wyzwaniem jest efektywne wykonywanie transferów wag na skalę bilionów parametrów. Tutaj opisujemy szczegóły ścieżki wykonania.

Parametry w szkoleniu są dystrybuowane zgodnie z rozmieszczeniem FSDP. Używając full_tensor(), wszystkie procesory graficzne (GPU) w DeviceMesh mogą zrekonstruować pełny parametr, dzięki czemu wszystkie mogą służyć jako źródło transferu wag.

Wiele rozłącznych DeviceMesh tworzy grupę mesh. Ponieważ DeviceMesh w tej samej grupie są rozłączne, ich transfery nie kolidują i mogą przebiegać w pełni równolegle. Pomiędzy grupami mesh wprowadzamy globalną barierę w celu wymuszenia kolejności.

Potok przenoszenia wag nakłada się na kopiowanie host-to-device, składanie tensorów, fuzję i kwantyzację oraz transfery RDMA między MeshGroups.

Transfer każdego tensora parametrów traktujemy jako zadanie. Proces transferu wag wykorzystuje wiele typów zasobów sprzętowych, dlatego dzielimy zadanie transferu wag na różne etapy potoku, które nakładają się w czasie:

memcpy z hosta do urządzenia — jeśli FSDP przenosi wagi do procesora (CPU)

Przygotowanie parametrów — Rekonstrukcja pełnej wagi za pomocą full_tensor(), zastosowanie fuzji projekcji, kwantyzacja w razie potrzeby.

Transfer RDMA — Zapis zero-copy do pamięci zdalnego procesora graficznego (GPU) wnioskowania

Bariera globalna — Po zakończeniu wszystkich wywołań full_tensor() synchronizacja między grupami mesh za pomocą GLOO przez Ethernet.

W implementacji utrzymujemy kolejkę FIFO zadań dla każdego etapu potoku. Gdy zadanie na początku kolejki kończy etap, jest przenoszone na koniec kolejki następnego etapu.

full_tensor() i inne operacje GPU wprowadzają dodatkowe zużycie pamięci GPU. Aby uniknąć błędu braku pamięci (OOM), rozpoczynamy wykonanie zadania tylko wtedy, gdy bieżące zadania w toku zajmują mniej tymczasowej pamięci GPU niż konfigurowalny limit.

Kilka decyzji projektowych sprawia, że nasz system działa znacznie szybciej i jest łatwiejszy w utrzymaniu niż powszechne rozwiązania typu open-source.

Transfer wag oparty na Rank0 zbiera wagi, zanim rozdystrybuuje je do węzła inferencyjnego.

Komunikacja punkt-punkt przenosi wagi z obu węzłów treningowych do interfejsów sieciowych węzła inferencyjnego.

Powszechnym schematem jest kierowanie wszystkich parametrów przez procesory graficzne (GPU) rank-0: zbieranie na szkoleniowym rank-0, wysyłanie do wnioskującego rank-0, a następnie ponowne rozpraszanie. Szybko staje się to wąskim gardłem, ograniczonym przepustowością PCIe pojedynczego GPU i karty sieciowej (NIC) (np. 400 Gb/s ≈ 50 GB/s).

W przeciwieństwie do tego, nasza konfiguracja punkt-punkt pozwala każdemu szkoleniowemu procesorowi graficznemu (GPU) wysyłać dane bezpośrednio do każdego wnioskującego procesora graficznego (GPU), nasycając całą strukturę sieci, a nie tylko pojedyncze łącze.

Transfer oparty na RPC używa bufora po stronie inferencji i wywołań procesów roboczych, natomiast transfer jednostronny zapisuje wagi bezpośrednio, bez udziału procesu roboczego inferencji.

Niektóre systemy polegają na wywoływaniu metody update_weight() silnika wnioskowania dla każdego tensora. Oznacza to inwazyjne zmiany w kodzie wnioskowania oraz narzut wynikający z wywołań RPC, serializacji i koordynacji płaszczyzny sterowania.

Dzięki prymitywowi RDMA WRITE aktualizujemy wagi w sposób niewidoczny w pamięci GPU wnioskowania, bez dodatkowych kopii. Nie jest wymagany żaden komunikat płaszczyzny sterowania ani logika sterująca procesora (CPU). Nie jest wymagana modyfikacja silnika wnioskowania.

Proces transferu wag może wykorzystywać cztery rodzaje zasobów sprzętowych: (1) Przesyłanie danych host-urządzenie, (2) obliczenia GPU dla fuzji projekcji i kwantyzacji, (3) sieć RDMA dla płaszczyzny danych, (4) Ethernet dla płaszczyzny sterowania.

Nasz projekt dzieli zadania transferu wag na etapy potoku, umożliwiając łatwe nakładanie się na siebie różnych zasobów sprzętowych.

Niektóre implementacje ponownie obliczają harmonogram transferu przy każdym kroku szkoleniowym, wielokrotnie gromadząc metadane i rozsyłając instrukcje. Dodaje to niepotrzebne opóźnienia płaszczyzny sterowania.

Nasz harmonogram jest obliczany raz przy inicjalizacji. Każda iteracja szkoleniowa po prostu powtarza plan: kontroler wysyła sygnał „start”, a procesory graficzne (GPU) podążają wyznaczonymi trasami. Wykonanie jest przewidywalne i lekkie.

Kuszące jest powiązanie całego procesu aktualizacji wag w jedną monolityczną funkcję: gromadzenie metadanych, dopasowywanie nazw, zbieranie wewnątrz węzła, fuzja projekcji, kwantyzacja, dzielenie świata komunikacji, transfer sieciowy między węzłami. Trudno to poprawnie zaprogramować, a jeszcze trudniej zoptymalizować.

W naszych pracach inżynieryjnych rozdzieliliśmy te kroki jako osobne komponenty. Każdy komponent może być testowany jednostkowo, analizowany i optymalizowany w izolacji.

Szybki, niezawodny transfer wag jest kluczowym elementem składowym dostrajania RL na dużą skalę. Łącząc prymityw RDMA WRITE, statyczny harmonogram transferu oraz potokowe wykonanie, skróciliśmy aktualizacje parametrów w modelu z bilionem parametrów do zaledwie 1,3 sekundy w Kimi-K2. Podejście to jest łatwe do zrozumienia, utrzymania i pozwala uniknąć wąskich gardeł tradycyjnych projektów.

BibTeX

@online{perplexity2025WeightTransferForRlPostTrainingInUnder2Seconds,
  author = {{Perplexity} and Lequn Chen and Erin McKnight and Vivian Wang},
  title = {{Weight transfer for RL post-training in under 2 seconds}},
  date = {2025-09-24},
  url = {https://www.perplexity.ai/hub/blog/weight-transfer-for-rl-post-training-in-under-2-seconds}
}

RDMA WRITE: transfery jednostronne

Przepływ pracy na wysokim poziomie

Wykonanie transferu wag

DeviceMesh i grupy mesh

Potok zadań

Kontrola użycia pamięci GPU

Dlaczego jest to szybkie i proste

Komunikacja punkt-punkt

Jednostronny transfer danych

Potokowanie

Statyczny harmonogram

Czysta separacja

Podsumowanie

Cytowanie