Minimálja a használatlan gyorsítók használatát: A helyi RL-munkák időszámítással történő egymásba kapcsolása az LLM-ben

A megerősítéses tanulás (RL) matematikája a nagy nyelvi modellek (LLM) utáni képzés során hírhedten kemény. Ahogy a határvonalas AI laborok a megerősítéses tanuláson alapuló algoritmusokat, például a Group Relative Policy Optimization-t (GRPO) használva a következtetések és kódolási modellek határait feszítik, rendszeresen találkoznak komoly architekturális és infrastrukturális korlátokkal. Miközben az ipar fókusza továbbra is a nyers gyorsító kapacitás megszerzésére összpontosít, az infrastrukturális hatékonyság ugyanolyan kritikus a szükséges magas sebesség eléréséhez, amely lehetővé teszi több RL feladat futtatását, és a modellek intelligencia szintjének növelését. Nagy léptékben a disztribúciós RL komoly erőforrás-beszorulásokkal küzd, mivel a szinkron mintavételezés és képzés szigorúan sorrendben történik, így a tréner és a mintavevő erőforrások váltakozva tétlenül ülnek. Eközben az aszinkron architektúrák megpróbálják átfedni ezeket a fázisokat, de a trénerek még mindig gyakran tapasztalnak tétlenséget, míg a konkrét trajektória-batch-ek befejezésére várnak, mielőtt elkezdenék a következő ciklust. Ma bemutatunk egy megoldást erre a strukturális pazarlásra: együttműködő időrészezés az llm-d projekt keretein belül. Azáltal, hogy a diszkrét RL lépéseket — mint például a mintavételezési gördülések és a gradienstanulás — dinamikus, ütemezhető entitásokként kezeljük, képesek vagyunk összefonni a független RL feladatokat megosztott fizikai hardverre. Kezdeti mérőszámaink azt mutatják, hogy ez a platformszintű multiplexelés növeli az aggregált gyorsítók kihasználtságát egy ~40%-os alapvonalról 70%-ra anélkül, hogy hatással lenne a modell konvergenciájára vagy pontosságára. Ez javítja az ár-érték arányt és jelentősen csökkenti az összesített költséget az idővel felhalmozódó számítási pazarlás megszüntetésével. A szinkron beállításokban a platform összefonja a mintavevőket és a trénereket, hogy minimalizálja a váltakozó tétlenségi időablakokat, míg az aszinkron munkaterhelések időrészezést használnak fel a fragmentált tétlenségi rések dinamikus visszanyerésére és kihasználására az RL-tréner iterációk között. A blog egészében ismertetjük az időrészezés megoldását, részletezve a technikai folyamatokat, a jelenlegi kiadást és a jövőbeli ütemtervet. llm-d az RL infrastrukturális hatékonyság érdekében (a nagyobb kép) Kezdetben vártuk a nagy léptékű RL utáni képzés súlyos infrastrukturális szűk keresztmetszeteit, és befektettünk az RL munkaterhelések infrastrukturális hatékonyságának javításába. Az llm-d-t egy rendkívül kompozíciós infrastrukturális stack-ként építettük ki az inference, az ügynöki és az RL munkaterhelések számára, amely az gyorsítók tétlen idejének megszüntetésére összpontosít. Az llm-d stack az RL számára a következőket tartalmazza: Átvitelen alapuló inference (llm-d-router): Érett, termelési tesztelt motor, amelyet az RL munkaterhelések között telepítettek, és amelynek célja a gördülési generálás átvitelének maximalizálása a pipeline folyamatos telítettsége érdekében. Nagy sebességű Ügynöki Homokozó (recept): Tesztelve a méret és sűrűség tekintetében, és segítve a biztonságos, másodperc alatti eszközhasználat és izolált kódvégrehajtás biztosításában a gördülési generálás és az értékelések során. Az Ügynöki Homokozó szolgál a jutalomjelzés generálásának nagy sebességű beömlő csatornájaként, biztosítva, hogy a Homokozó soha ne legyen a késleltetési szűk keresztmetszet, amely éhezteti a time-sliced NVIDIA GPU-kat. Alap pipeline primitívek: A megbízhatóság és a sebesség küzdelme érdekében a Súlyterjesztési Interfészt (WPI) építjük, valamint a RL-re vonatkozó általános megfigyelhetőség és megbízhatóság javítására összpontosítunk. Az RL ciklusok hatékonysági problémája A disztribúciós RL utáni képzés darabos, folyamatos ciklusban működik, amely váltakozik a generálás (mintavételezési gördülések) és az optimalizálás (gradiens frissítések) között. Mivel a hagyományos felhőinfrastruktúra folyamatos, állandó munkaterhelésekhez van tervezve, a standard Kubernetes klaszterek nem tudják alkalmazkodni ehhez a váltakozó ritmushoz. Nagy léptékben ez a strukturális ritmus két hatalmas rendszerszintű hatékonysági problémát vezet be: Tétlen gyorsítók: Mivel ezek a fázisok egymás után következnek, a GPU klaszterek az életciklusuk 40%-60%-ában teljesen tétlenek (0% kihasználtság). A trénerek tétlenül várakoznak a mintavételezési gördülések befejezésére; a mintavevők tétlenül ülnek a gradienstanulás és a súlyok elosztása alatt. Ez évente milliós tőkepazarlást jelenthet. Zárolt kontextus: Az RL képzés és a mintavevők az egész futamidejük alatt fenntartják a gyorsító allokációkat, még a tétlen fázisok alatt is, mert a NVIDIA CUDA kontextusnak és minden eszköz memóriájának állandóan bent kell maradnia. A standard ütemezők ezeket a podokat statikus, elszigetelt allokációként kezelik, ahelyett, hogy azokat a folyamatban lévő RL hurok váltakozó, fázis-szintű állapotaihoz igazítanák, így értékes hardver marad lezárva még az inaktív fázisok alatt is. Fontos, hogy ez nem csupán egy szinkron RL probléma. Az aszinkron változatok átfedik a generálást és a képzést, de nem mérséklik teljesen a tétlen időt. A generálás továbbra is a RL hurok inherens szűk keresztmetszete, ami azt jelenti, hogy a trénerek gyorsítói éheznek, míg a gördülési adatok felhalmozódására várnak. Minél közelebb fut egy aszinkron feladat az on-policy-hoz, annál nagyobbak ezek a tétlenségi ablakok — a határozott avultság korlátozza, hogy a generálás és a képzés mennyire válhat el egymástól, leállítva a pipelinet, amikor a friss gördülések nincsenek készen. Hogyan segít az együttműködő időrészezés (RL feladatok összefonása) Az RL feladatok során a tétlen gyorsítók megszüntetése érdekében az llm-d projekt keretein belül az együttműködő időrészezés lehetővé teszi az infrastruktúra számára, hogy dinamikusan összefonja a független RL feladatokat megosztott hardverblokkokra, ahelyett, hogy a hardvert arra kényszerítené, hogy várjon az upstream fázisokra. Ez segít növelni az aggregált gyorsítók kihasználtságát anélkül, hogy megváltoztatná az alapvető modell konvergenciáját vagy pontosságát. Amikor az A feladat tétlen marad egy fázishatárnál a szinkron RL-ben (vagy megáll a friss gördülési adatoknál az aszinkron RL-ben), az infrastruktúra időrészezi a fizikai gyorsítókat, átkapcsolva az B feladat aktív mintavételezési vagy képzési fázisára. A háttérben az átkapcsolás egy ellenőrző pont/helyreállítás: Az A feladat teljes eszközállapotát ellenőrzik az gyorsító memóriájából a gazda DRAM-ba, majd az B feladat korábban mentett állapotát állítják vissza a helyére. Mivel egyszerre csak egy feladat állapota foglalja el a gyorsítót, a lépések biztonságosan váltakoznak keretrendszerszintű zavarok vagy memóriahiányos (OOM) hibák nélkül. Időrészezés: Magas szintű architektúra Az időrészező rendszer architektúrája három rétegre oszlik: munkaterhelés-szintű (alkalmazási logika), klaszter-szintű (koordináció), és csomópont-szintű (hardverkezelés). Munkaterhelés-szintű réteg (alkalmazási futás) Itt fut a felhasználó kódja — tanulási ciklusok, inference szerverek és RL keretrendszerek. Az új kiegészítés az időrést kezelő klienskönyvtár, amely két gRPC API-t tesz elérhetővé az időrészezőn: acquire() az exkluzív gyorsító hozzáférés kérése, és yield() a visszaadására. A felhasználó minden gyorsítóval érintkező fázist ezekkel a hívásokkal keretezi, hogy jelezze a fázishatárokat az ütemezőnek. Minden más — az ML keretrendszer (PyTorch FSDP, vLLM, stb.), a CUDA kontextus, a gyorsító memória allokációk — változatlanul fut. Klaszter-szintű réteg (vezérlés és ütemezés) Ez a réteg dönt arról, hogy melyik feladat kap gyorsító hozzáférést és mikor. Azok a feladatok, amelyek ugyanazokat a fizikai gyorsítókat osztják meg — például két RL feladat, amelyek összefonódnak ugyanazon GPU csomópontok készletén — egy csoportba kerülnek. Minden csoport számára az időrészező fenntart egy zárolási sort — egy rendezett listát a gyorsítók exkluzív hozzáférésére váró feladatokról. Csak a sor elején álló feladat birtokolja a zárat és fut a hardveren; az összes többi feladat várakozik, leszorítva az acquire() hívásukra. Amikor a futó feladat yield() hívást végez, az ütemező átadja a zárat a következő feladatnak a sorban, és koordinált kontextusváltást indít el az összes csomóponton a csoportban. A jövőben egy munkaterhelés elhelyező optimalizáló képes lesz profilozni a munkaterhelési fázis mintákat, és automatikusan párosítani a feladatokat az egymásnak megfelelő tétlen fázisokkal, eltávolítva az igényt arra, hogy a felhasználó kifejezetten jelezze a feladatcsoportosítást. Csomópont-szintű réteg (hardver és adat sík elválasztás) Ez a réteg végzi el a checkpoint/restore váltást minden gyorsítócsomópontban. A snapshot ügynök, egy privilégiumos DaemonSet, irányelveket kap az ütemezőtől, és lefordítja azokat hardver-szintű műveletekre — gyorsító folyamatok szüneteltetése, eszközállapot sorozatának rögzítése a gazda DRAM-ba, és helyreállítása, amikor a feladat újra hozzáférést nyer. Az ügynök egy plugin által vezérelt backend interfész köré épül, az első implementáció a cuda-checkpoint (továbbiak várhatók). A jövőbeni back-endek gyorsabb snapshot mechanizmusokat és szelektívebb megközelítéseket vezetnek be, például a LoRA adapterek specifikus memória címek elhagyását a teljes eszközállapot helyett. Az ügynök maga úgy van tervezve, hogy önállóan fusson Kubernetesen kívül a bare metal és Slurm környezetekben. Az áramlás: Hogyan működik mindez együtt Amikor egy munkaterhelés befejezi a jelenlegi gyorsítói fázisát, az időrészező klienskönyvtára yield() hívással értesíti az időrészezőt a hozzáférés visszaadásáról. Az ütemező megkezdi a kontextusváltást, az irányelveket elküldve a snapshot ügynöknek minden csomópontban a csoportban. Az ügynök megfagyasztja a visszaadó munkaterhelés folyamatait és átmozdítja eszközállapotát az gyorsító memóriából a gazda DRAM-ba. A gyorsítók elhagyásával az ütemező megadja a csoport zárát a következő munkaterhelésnek, amely a sorban várakozik. Utasítja a Snapshot Ügynököket ezeken a csomópontokon, hogy állítsák vissza azt a munkaterhelést, amelynek korábban mentett állapotát a gazda DRAM-ból vissza az gyorsító memóriába, majd feloldja a munkaterhelés függő acquire() hívását. A munkaterhelés pontosan ott folytatja a végrehajtást, ahol abbahagyta — nincs tartály újraindítás, nincs keretrendszer újraírás, nincs modell újratöltés a tárolóból. A visszaadó munkaterhelés melegen marad a gazda DRAM-ban. Amikor az ütemező ismét megadja neki a zárat, a Snapshot Ügynökök végrehajtják ugyanazt a váltást visszafelé. Fejlesztői élmény (ügyfél oldalon) A kutatók a fő modellezési logikára szeretnének összpontosítani, nem pedig az alacsony szintű CUDA kontextus váltással vagy egyedi ütemezési ciklusokkal küzdeni. Ha Ray-t vagy hasonló platformot használsz az RL feladatod ütemezésére, az időrészezésnek minimális hatása lesz az ügyfél oldalán. Valójában lehet, hogy egyáltalán nincs hatás az ügyfél oldalán, ha külön sorba állítod a képzési és mintavételezési feladatokat a platform szintjén. Jelenlegi kiadás és jövőbeli kilátások Ma kiadjuk a teljes időrészező stack-et: a Snapshot Ügynököt, a Gyorsító Ütemezőt és a Python klienskönyvtárakat, mindegyikhez felhasználói útmutató készült az időrészezés integrálásához az RL munkaterhelésekbe. Kulcsfontosságú ütemtervi kiemelések közé tartozik: Késleltetés és állapotoptimalizálás: A Snapshot Ügynök bővítése gyorsabb ellenőrző pont/helyreállító backendekkel a kontextusváltási többleterő csökkentése érdekében, a kiválasztott memória régiók snapshottingjához alkalmazás-érzékeny backendekkel (pl. LoRA adapterek cseréje a teljes modell súlyok helyett). Automatizált ütemezés és beilleszkedés: Egy automatizált ütemező bevezetése, amely profilozza a futó folyamatokat, azonosítja az időrészezhető struktúrákat, és dinamikusan kezeli a feladatok elhelyezését. Hardverek közötti kompatibilitás: Az adat sík támogatásának kiterjesztése a GPU-kon túl TPUs-ra és egyedi gyorsító architektúrákra. Kezdés Érdekes, rendkívül optimalizált RL infrastruktúra létrehozása szoros együttműködést igényel a munkaterheléseket nagy léptékben működtető mérnökökkel és kutatókkal. Ha jelenleg alacsony GPU kihasználással, szinkronizációs késésekkel vagy bonyolult ütemezési logikával küzdesz az utáni képzési folyamataidban, az időrészezés segíthet. A támogatott erőforrásokkal indíthatod a következő műveleteket, és ne felejtsd el hátrahagyni a visszajelzéseidet! Kezdj el időrészezést használni az RL futásod során azonnal ezekkel a felhasználói útmutatókkal. Próbáld ki az llm-d-router-t (kubernetes natív) vagy az RL Scheduler-t (python könyvtár) a mintavételezési átvitel javítása érdekében az RL generálási fázis során. Fedezd fel a Weight Propagation Interface repót. Csatlakozz a #sig-rl csatornához az llm-d Slacken. Járulj hozzá azzal, hogy megosztod a referenciaimseratcióidat, benchmarkjaidat és szélsőséges eseteidet, hogy segíts nekünk tovább finomítani ezt az utat. Köszönet Dolev Ish Am-nak és Bogdan Berce-nek a blogbejegyzéshez való hozzájárulásukkal.