Az NVIDIA bemutatja a Reshard kollektív kommunikációs hívást, amely a megerősítéses tanulás (RL) során szükséges súlyfrissítést (refit) DTensor‑to‑DTensor átméretezésként (reshard) kezeli. A megközelítés NCCL‑alapú referenciamegvalósítása, az NCCL M2N Reshard, nyílt forráskódként elérhető az NVIDIA/nccl-extensions GitHub‑tárban, és a NeMo RL keretrendszerben is demonstrálták. A Reshard minimalizálja a ténylegesen átküldendő bájtok mennyiségét azzal, hogy közvetlen forrás–cél részeket küld és NVLink‑doméneken belüli kiegyensúlyozott belső cserét használ, ami a benchmarkok szerint 2.7×–5.7× gyorsulást eredményez a hagyományos gather‑and‑broadcast refithez képest.
Miért probléma a súlyrefit az RL post‑training során?
A megerősítéses tanulás utótréningje során a policy modell inferenciája és a tanítás gyakran külön GPU‑csoportokon párhuzamosan fut, hogy hatékony legyen a GPU‑erőforrások kihasználása. A tanító GPU‑kon végzett súlyfrissítéseket időről időre szinkronizálni kell az inferencia GPU‑kkal — ez a művelet a refit. A hagyományos megoldás a gather‑and‑broadcast: a tanító GPU‑k all‑gather művelettel összerakják a teljes paramétert, majd broadcastolnak az inferencia GPU‑kra, amelyek megtartják a saját lokális shard‑jukat.
Ez a megoldás egyszerű, de redundáns kommunikációt eredményez: az inferencia GPU gyakran megkapja a teljes paramétert akkor is, ha csak egy részére (shardjára) van szüksége. A kommunikációs payload aránya a modellmérettel nő — például Kimi K3 körülbelül 2.8 billió paramétert tartalmaz (BF16 ≈ 5.6 TB), a Qwen 3.8 2.4T A95B körülbelül 2.4 billió paramétert (BF16 ≈ 4.8 TB). A kétfázisú (all‑gather + broadcast) logikai payload így hozzávetőlegesen kétszeresse a modell súlyának (pl. ~11.2 TB, illetve ~9.6 TB), ami minden refit során átküldendő adatként jelenik meg.
A megközelítés: refit mint DTensor‑to‑DTensor reshard
A cél az, hogy csak a szükséges minimum adatmennyiséget mozgassuk. A nehézség abból adódik, hogy a tanító és az inferencia példányok eltérő párhuzamosítási konfigurációkat és GPU‑számot használhatnak, így a shard‑határok különböznek. Egy cél‑shard több forrás‑shard részéből épülhet fel, illetve replikáció esetén ugyanaz a súly több forrásnál is elérhető lehet.
Az NVIDIA csapata ezért a PyTorch distributed tensor (DTensor) absztrakciót használja: egy DTensor leírja a tensor globális alakját, egy device mesh‑t és a placements‑okat (hogyan van sharded vagy replikált). A refit tehát átalakítható úgy, hogy a forrás DTensor‑elrendezésből a cél DTensor‑elrendezésbe reshardolunk, függetlenül a konkrét párhuzamosítási konfigurációktól. Például egy 8×8 mátrixot sorok szerint sharded két tanító GPU‑n és oszlopok szerint sharded négy inferencia GPU‑n a megfelelő részek (4×2 blokkok) közvetlenül átküldhetők shard‑ról shard‑ra.
A Reshard kollektív kommunikáció
A Reshard új kollektív kommunikációs hívás, amely tetszőleges DTensor‑to‑DTensor átvitelekre szolgál. Egy Reshard hívás egyszerűen a forrás és cél DTensor leírásokkal adható meg; a könyvtár belső algoritmusai végrehajtják a tényleges adatmozgatást. Az első referenciamegvalósítás az NCCL M2N Reshard, amely az NCCL‑re épül és M forrás GPU‑ról N cél GPU‑ra mozgatja át a tensor értékeit.
NCCL M2N Reshard először meghatározza a pontos metszeteket a forrás és cél shard‑ok között, majd a fizikai GPU hálózati topológiára leképezi a transzfereket. Ahelyett, hogy minden forrás hozzájárulást külön megismételne minden cél replikának, az implementáció NVLink‑domének szerint csoportosítja a replikákat, és a hálózaton minden hozzájárulást csak egyszer továbbít. A helyi NVLink csere segítségével a replikák összeszerelik a teljes shard‑jukat, anélkül, hogy a hálózati forgalmat megdupláznák. Ha replikák több NVLink‑domént is átfednek, a hozzájárulások egy másolata továbbítódik a domének között, és minden doménen belül kiegyensúlyozott lokális csere történik. A hálózati átvitel és az NVLink replikáció pipelinizált, hogy kurtítsák a késleltetést.
Integráció és hibrid megoldás a NeMo RL‑ben
A Reshard implementációját demonstrálták az NVIDIA NeMo RL keretrendszerben. A DTensor leírások biztosítása minden paraméterre a két példány (tanítás és inferencia) számára egyszerű a DTensor céljára tervezett formátum miatt, de a gyakorlatban a rétegtípusok sokfélesége és a különböző shardolási stratégiák megnehezítik az összes eset lefedését.
Ennek alapján a NeMo RL megvalósítása hibrid megközelítést alkalmaz: a transformer feed‑forward projekciós súlyai — amelyek a modern LLM‑ekben a paramétertömeg mintegy 95%-át teszik ki — a Reshard útvonalat követik, míg a kisebb, egyedi paraméterek (például QKV projekciók vagy egyéb speciális súlyok) a hagyományos gather‑and‑broadcast útvonalon frissülnek. Így a nagy tömegű payload hatékonyan reshardolódik, miközben megmarad a széles körű funkcionalitás.
A GPU‑tagok futás közbeni változása esetén a keretrendszerek vagy újraépítik a kommunikációs csoportot, vagy az NCCL shrink/grow API‑kat használják a meglévő kommunikátor módosítására. A NeMo RL jelenleg az egyszerűbb, csoport újra létrehozásán alapuló megközelítést alkalmazza; az NCCL shrink/grow alapú resiliency támogatás későbbi frissítésként szerepel.
Mérések: mennyivel gyorsabb a Reshard?
A NeMo RL‑ben végzett benchmarkok több modellt és platformot fednek le (NVIDIA DGX H100 NVL8 és NVIDIA DGX GB200 NVL72). A táblázat néhány mérési eredményt foglalja össze: a reshard refit latency 2.7×–5.7× gyorsulást ért el a gather‑and‑broadcast megoldáshoz képest.
- NVIDIA DGX H100, Qwen 3 235B, Training → inference: TP4 x PP8 x EP16 → TP8 x DP16: gather‑and‑broadcast 6.4 s → reshard 1.9 s
- NVIDIA DGX H100, DeepSeek v3 671B, PP16 x EP16 → TP32 x DP8: 14.9 s → 2.6 s
- NVIDIA DGX GB200, DeepSeek v3 671B, PP16 x EP16 → TP32 x DP8: 6.1 s → 2.3 s
- NVIDIA DGX GB200, Nemotron Ultra v3, TP8 x EP32 x PP2 → TP8 x DP8: 8.2 s → 2.3 s
Egy refit‑only mikromérésben a Reshard a Qwen3.8 2.4T BF16 modellnél 4.3 másodperc alatt végzett, ami alátámasztja a megközelítés hatékonyságát trilliónyi paraméteres méretekben.
Gyakorlati hozzáférés és beállítás
Az NCCL M2N Reshard referenciakódja nyílt forráskódú, és elérhető az NVIDIA/nccl-extensions GitHub‑tárban. A NeMo RL‑ben a reshard útvonal engedélyezéséhez a konfigurációban a refit_transport: nccl_reshard értéket kell megadni, például:
policy:
generation:
colocated:
enabled: false
refit_transport: nccl_reshard
A NeMo RL refit útmutatója és az NCCL Reshard konfigurációs követelményei részletezik a támogatott backendeket, elrendezéseket és precizitási módokat.
Köszönetnyilvánítás
Az NCCL M2N Reshard és a NeMo RL‑ben megvalósított nccl‑reshard súlyrefit fejlesztéséhez Ke Wen, Ching‑Hsiang Chu, Guyue Huang, Seonjin Na és Nitin hozzájárulását is elismerik.
Miért számít ez?
A Reshard kollektív és az NCCL M2N Reshard referenciaimplementációja egy lépés a nagy, többbilliós‑billió paraméteres modellek skálázható és hatékony refiteljének megvalósítása felé, különösen olyan környezetekben, ahol inferencia és tanítás el van választva. A konkrét kommunikációs optimalizációk — DTensor alapú ábrázolás, NVLink‑domének szerinti csoportosítás és lokális cserék — jelentősen csökkentik a hálózati forgalmat és a refit időt minden tréninglépésnél, ami közvetlenül gyorsabb RL post‑training ciklusokat tesz lehetővé.
