A TRL (PR #7017, TRL v1.14) AsyncGRPOTrainer mostantól képes LoRA adapterrel dolgozni és csak az adaptert szinkronizálni vLLM felé. Egy rank‑1 LoRA adapter mérete néhány megabájt, így adaptereket Storage Bucketen keresztül lehet megosztani Hugging Face Jobs között a nagy súlyú (GB-os) súlyszinkronizálás helyett. A megvalósításban a trainer és a vLLM replikák külön HF Jobokon futnak, egy közös bucketet mountolva, és egy kis proxy végzi az authentikációt, a rolloutok affin alapú irányítását és az adapter betöltések broadcastelését.
Miért működik jól a LoRA RL-ben?
A LoRA (alacsony rangú adapterek) különösen hatékony a policy‑gradient RL esetén: Thinking Machines „LoRA Without Regret” eredményei szerint rank‑1 adapter képes megközelíteni a teljes finomhangolás teljesítményét. Az indoklás információelméleti: egy epizód előnye (advantage) csak ~O(1) bit információt ad, így kis kapacitású adapterek is elegendők a hasznos minták befogadására.
Rendszerarchitektúra: HF Jobs, Storage Bucket és proxy
A hagyományos megközelítésben trainer és generációs munkák ugyanazon a fürtön, megosztott fájlrendszeren vagy NCCL csoportban futnak. Hugging Face Jobs esetén azonban egy Job egyetlen VM-en fut, nincs lokális megosztott lemez és nincs közvetlen node‑közi NCCL. Megoldás: Storage Buckets mountolása minden Jobba POSIX fájlrendszerként (hf-mount), így az adapterek fájlként kerülnek át a trainerből a vLLM szerverekbe.
A felállás röviden:
- egy trainer Job futtatja az AsyncGRPOTrainer-t LoRA-val (FSDP használata opcionális),
- két vLLM Job, mindegyik a bázismodellt és az aktuálisan publikált adaptereket szolgálja ki,
- egy Storage Bucket mountolva mindhárom Jobban ugyanarra az abszolút elérési útra,
- egy kis proxy a trainer Jobon, amely hozzáadja az Authorization fejléceket, route‑olja a kéréseket a megfelelő replikához (KV-prefix alapú affinity) és broadcasteli az adapter betöltéseket.
A trainer adaptert fájlként menti: <output_dir>/.vllm_lora/trl-policy-v{N}, atomikus átnevezéssel publikálja, majd POST-olja a vLLM /v1/load_lora_adapter végpontjára a path-ot. A vLLM a fájlrendszerről tölti be az adaptert, ezért mindkét Jobnak ugyanazt az útvonalat kell látni a bucket mounton keresztül.
vLLM beállítások és adapter‑slotok
A vLLM replikákhoz az alábbi kulcsfontosságú pontok:
- vLLM verzióként v0.27.1-et használtak (runtime LoRA flagok és load endpoint erre a verzióra vonatkoznak),
- a max_staleness=4 konfiguráció miatt minden replikának egyszerre szolgálnia kell az aktuális policy‑t plusz a négy korábbi verziót; ezért --max-loras 6 kellett (4+2 csereidő miatt). Ha kevesebb slot van, vLLM evictálhatna még futó rolloutot használó policy-ket.
Példa parancs (részlet):
- VLLM_ALLOW_RUNTIME_LORA_UPDATING=1
- VLLM_SERVER_DEV_MODE=1
- --enable-lora --max-lora-rank 1 --max-loras 6
A döntés mellettük: verziózott adapter‑nevek használata megakadályozza, hogy a KV prefix cache összekeverje különböző súlyverziók előtöltéseit.
Dátumok, modellek és dataset
- TRL v1.14 (PR #7017) hozta az AsyncGRPO Trainer LoRA támogatását.
- vLLM image: vllm/vllm-openai:v0.27.1.
- Modell: Qwen/Qwen2.5-Math-1.5B.
- Dataset: sail/Sanity-Test-R1D-1.5B (Qi et al., „Defeating the Training-Inference Mismatch via FP16”, 2025), amely 1,460 kérdést tartalmaz a kiválasztott sávon belüli nehézséggel.
- LoRA: rank r=1, alpha=2, a recept a sail-sg/Precision-RL és oat/scripts/lora alapján.
A proxy szerepe és a routing logikája
A proxy két okból szükséges:
- HF Job által kitett portokhoz Authorization: Bearer <HF token> fejléc kell — a proxy teszi hozzá ezt a trénertől függetlenül.
- Több külön gépen futó replikát szeretnénk használni; a proxy broadcastolja a state‑változtató kéréseket (adapter betöltés, pause/resume/unload) minden replikára, és ugyanakkor egyedi lekéréseket irányít a legvalószínűbb cache‑tartó replikához.
Routing: vLLM a prefix KV cache-t 16 tokenes blokkokban tárolja. A router ugyanígy blokkokra vágja a promptot, blokkonként chained hash-t számol (blake2b, 16-token blokkok, adapter névvel seedelve). A router megjegyzi, hogy melyik replikának mely hash blokkok a tulajdonai (successor set‑tel, maximum két folytatással). A döntés szabályai:
- affinity hit: egy replikának specifikus (nem common) vezető blokkjai vannak és nincs túlságosan túlterhelve (<= 8 requests előnye a legkevésbé terhelthez képest), ekkor oda megy a kérés;
- spill: ha specifikus blokkjai vannak de túlságosan előrébb, akkor a legkevésbé terheltre küldjük;
- unmatched: ha senkinek nincs specifikus blokkja, a legkevésbé terheltre (round-robin tie esetén).
Common blokkok: olyan blokkok, amelyeket minden replikának szolgáltak vagy több utóda van; ezek kerülendők, mert egy közös chat template minden új promptnál hamis találatot eredményezne.
A routing célja, hogy a G=8 rollouthoz tartozó kérések ugyanahhoz a replikához jussanak, így a prefill egyszer történik és a többi kérés újrahasználja a cache‑t.
Adapter broadcast és transzkakciós viselkedés
Az adapter betöltést all‑or‑nothing módon kezelik: a proxy egyszerre küldi a /v1/load_lora_adapter kérést minden replikának, figyelve arra, ha valamelyik replikán a bucket mount még nem látja a könyvtárat ("No adapter found for <path>"). Ilyenkor újrapróbálja csak azt a replikát; ha más hibába ütközik, visszagörgeti a sikeres betöltéseket úgy, hogy sehol se létezzen részlegesen a policy név.
A proxy /health csak akkor ad 200‑at, ha minden replikáé egészséges.
Kísérleti eredmények: 5 futás és bottleneck vadászat
A kísérletek 500 lépésen (steps) futottak, 128 completions per step és 8 rollouts per prompt, egy h200x2 Job a trainernek és két h200 Job a két vLLM replikának. Költség: ~20 USD/óra mindhárom Job együtt.
Főbb mérőszámok és tapasztalatok:
- Weight sync (126 sync total): korai konfigurációban a teljes sync p50 ~30.8 s volt, később optimalizálva ~8.5 s (min 6.6, max 9.2). A szüneteltetés a replikákon ~0.3 s, adapter save/bucket ~0.6–1.1 s, míg a replikák elfogadása dominálta a latenciát (~7–29 s).
- Adapter betöltések: összesen 252 adapter load (126 sync × 2 replikák). Mind sikeresek, többségük harmadik próbálkozásra, hat darab másodikra.
- Routing végén (64,728 rollouts): routed [31928, 32800], affinity 54712, spilled 820, unmatched 9196. Ez 14.2% unmatched, 84.5% affinity, 1.3% spills; az elméleti minimum unmatched 12.5% (mivel 8 rolloutsból legalább egy cold).
Részletes futások (összefoglalva):
- run 1 (r1-dp2) — reference: trainer a bottleneck; p50 step 22.9 s, forward+backward 21.9 s, rollout queue ~476/512, trainer MFU 3.9%, 500 steps 3 h 27 min.
- run 2 (r1-dp2-tb16k) — token-budget packing (token_budget=16384, grad_accum=6): microbatches csökkennek 64→6, samples_per_row ~1→~12.7, perf step 23 s → 5.9 s, MFU 3.9% → 19%, generation throughput nő (4.6k→25k tok/s). Bottleneck megszűnik, most a generáció válik szűk keresztmetszetté.
- run 3 (r1-dp2-tb16k-nockpt) — kikapcsolt gradient_checkpointing: forward+backward csökken 5.6 s → 4.6 s, MFU 23%, a bottleneck átkerül generációra, rollout wait nő ~0.02→0.6 s.
- run 4 (r1-dp3-tb16k-nockpt) — három replikával próbálták javítani a generációt; weight sync javult (7.6 s→5.8 s) de generált token/s csak kismértékben nőtt (25k→26k). Ok: a rollout/inflight limit (128) megosztódott 3 GPU között, ami gyakorlatilag ugyanazt a per‑GPU terhelést eredményezte (44+43+41). Tehát a kliens-korlát lett a problémás rész.
- run 5 (r1-dp3-inflight384) — max_inflight_tasks=384 és queue_maxsize=768: mindhárom replikára 128 request jut; a queue ~690/768, perf step p50 22.9 s → 4.8 s, forward+backward 21.9→4.6 s, weight_sync p50 8.5→6.2 s, staleness mean 1.5→2.0. 500 step p50 ideje 3 h 27 min → 53 min (3.9× gyorsulás). Reward görbék nagyon hasonlóak: run 1 reward 0.145→0.438, run 5 reward 0.145→0.416.
Végső táblázat (500 steps, run1 vs run5 — kiemelt számok):
- wall clock: 3 h 27 min → 53 min
- perf/step_s p50: 22.9 s → 4.8 s
- perf/fwd_bwd_s p50: 21.9 s → 4.6 s
- perf/mfu_fwd_bwd: 3.9% → 23.5%
- samples trained: 64,000 → 84,078
- perf/weight_sync_s p50: 8.5 s → 6.2 s
- sample/staleness_mean: 1.5 → 2.0
A tanulság: sorvezető metrikák (AsyncGRPO logged metrics) és kis rendszerparaméter‑módosítások (token packing, checkpointing kikapcsolása, in-flight limit emelése, proxy retry csökkentése) jelentősen felgyorsították a folyamot anélkül, hogy eltávolodtak volna a tanulási minőségtől vagy a policy‑konzisztenciától (ratio ≈ 1.000 végig).
Megvalósítható lépések, reprodukció
A szerzők reprodukciós repositoryt biztosítanak. Röviden:
- git clone https://github.com/AmineDiro/hfjobs-lora-buckets
- hf auth login
- futtatás: run_all.sh script különböző környezeti változókkal (MAX_STEPS, TOKEN_BUDGET, GRAD_ACCUM, GRADIENT_CHECKPOINTING, PROXY_LORA_RETRY_S, MAX_INFLIGHT, QUEUE_MAXSIZE). Egy rövid smoke run ~15 perc, a teljes 500-step run korábban ~3.5 óra, végső optimalizált ~55 perc.
Következtetés
A TRL AsyncGRPOTrainer LoRA támogatása lehetővé teszi, hogy a tanítás és az inferencia fizikailag külön VM‑eken fusson Hugging Face Jobs környezetben anélkül, hogy minden frissítésnél gigabájtos modelleket mozgatnánk. A Storage Bucket mount és egy egyszerű proxy a gyakorlatban elegendő infrastrukturális eszköz a hatékony adapter‑csere és a generációs routing megvalósításához. A részletes metrikák segítették megtalálni és megszüntetni a bottleneckeket: token‑packing, checkpointing kikapcsolása és nagyobb párhuzamos in‑flight korlátok együttesen 3.9× gyorsulást eredményeztek 500 step alatt, miközben a politika konzisztenciája megmaradt.
Hivatkozások és források
- TRL AsyncGRPOTrainer PR #7017 és TRL v1.14 (LoRA adapter és adapter-only vLLM sync)
- vllm/vllm-openai:v0.27.1
- Thinking Machines, "LoRA Without Regret" (September 2025)
- Qi et al., "Defeating the Training-Inference Mismatch via FP16" (2025) — forrása a Sanity datasetnek
- repository: aminediro/hfjobs-lora-buckets (reprodukciós szkriptek és run_all.sh)
A cikk a szerzők méréseire és a futtatott trackio dashboard adataira épül; minden szereplő, verzió és mérés a fent felsorolt forrásokból származik.
