A JAX host offloadinggal enyhíthető a GPU-memória szűkössége

A host memória tehermentesítheti a GPU-t

Az NVIDIA Developer július 10-én ismertetett bejegyzése szerint a nagy nyelvi modellek képzésekor a GPU memóriáját a modell súlyai, a gradiensek, az optimalizáló állapotai, a kommunikációs pufferek és a köztes aktivációk is használják. A modellméret, a sorozathossz és a kötegméret növekedésével a nagy sávszélességű memória, vagyis a HBM kapacitása válhat a skálázás fő korlátjává.

A JAX-ben elérhető host offloading során a rendszer a forward pass közben kiválasztott aktivációkat rögzített host memóriába helyezi, majd a backward pass során visszatölti őket. Ez alternatívát jelent az aktivációk újraszámításával szemben, amelynél a rendszer a szükséges értékeket ismételt számítással állítja elő.

A Blackwell összeköttetése fontos szerepet kap

A technika az NVIDIA szerint különösen előnyös a Grace Blackwell rendszereken. A Grace CPU és a Blackwell GPU között NVLink-C2C kapcsolat biztosít 900 GB/s kétirányú sávszélességet, így a rögzített host memória használható köztes tárhelyként a kiválasztott aktivációkhoz.

A forrás szerint a Vera CPU és a Rubin GPU ezt a kétirányú, koherens sávszélességet 1,8 TB/s-ra növeli. A nagy sávszélesség önmagában azonban nem elegendő. Az aktivációk átvitelét úgy kell ütemezni, hogy az párhuzamosan történjen a számítási és kommunikációs feladatokkal.

A DeepSeek-V3 671B nagyobb kötegmérettel futhatott

Az NVIDIA a MaxText JAX-alapú képzési keretrendszerrel végzett méréseket NVIDIA GB200 NVL72 rendszereken, 128 GPU használatával. A DeepSeek-V3 671B modellnél a host offloading, a Latency Hiding Scheduler, vagyis az LHS, és a csővezetékes átvitel együttes bekapcsolása 908,2 TFLOPs/s eszközönkénti teljesítményt eredményezett.

Ez azonos kötegkonfiguráció mellett 57 százalékkal haladta meg az aktivációk újraszámítását. A host offloading LHS és csővezetékes átvitel nélkül 541,6 TFLOPs/s eszközönkénti értéket ért el, míg az aktivációk újraszámítása 578,3 TFLOPs/s eszközönkénti teljesítményt adott.

A módszer a memória szempontjából is nagyobb konfigurációt tett lehetővé. A host offloadinggal a mikroköteg mérete 8, a globális kötegé 1024 lehetett. Az aktivációk GPU-n történő tárolása mellett a futtatás mikrokötegméret 2 és globális kötegméret 256 mellett fért el. A 8-as mikroköteg és 1024-es globális köteg használatakor offloading nélkül memóriahiány, OOM hiba jelentkezett.

Az optimalizált futtatás 165,2 GiB GPU-memóriát és 145,1 GiB host memóriát használt. Az NVIDIA magyarázata szerint a nagyobb GPU-memóriahasználatot részben az átvitel és az előbetöltés átfedéséhez szükséges pufferek okozzák.

A Llama-mérés kisebb, de mérhető gyorsulást mutatott

A Llama 3.1 405B esetében a vizsgálat QKV, vagyis lekérdezési, kulcs- és értékaktivációk offloadingjára irányult. A teszt batch size 2, 8192-es sorozathossz, 128-as FSDP-beállítás, bfloat16 aktivációk és NVFP4 4 bites súlykvantálás mellett futott.

Az LHS használatával a QKV aktivációk offloadingja 2669-ről 2746 TFLOPs/s eszközönkénti értékre növelte a teljesítményt, ami 2,9 százalékos javulás. Az LHS kikapcsolásakor ez 2569 TFLOPs/s eszközönkénti értékre csökkent. Az eredmények az NVIDIA szerint azt mutatják, hogy a host offloading teljesítményéhez az átvitel és a számítás hatékony átfedése is szükséges.

A fejlesztők a JAX host offloading oktatóanyagában tanulhatják meg az aktivációk kezelését, a jax.remat használatát és a memory_kind="pinned_host" beállítást. Az NVIDIA a profilozáshoz az Nsight Systems használatát javasolja, hogy ellenőrizhető legyen az eszköz és a host közötti másolások átfedése a számítással és az NCCL-kommunikációval.