HBM-korlátok csökkentése JAX host-offloaddal nagy nyelvi modellek betanításakor

Nagy nyelvi modellek (LLM) betanításakor a GPU magas sávszélességű memóriája (HBM) gyakran válik korlátozó tényezővé még azelőtt, hogy a számítási kapacitás teljességgel kihasznált volna. A modell súlyai, gradientek, optimalizáló állapota, kommunikációs pufferek és a köztes aktivációk egyszerre versengenek az HBM-ért. Modellméret, szekvenciahossz és batchméret növekedésével az HBM-kapacitás gyakran lesz a skálázás fő szűk keresztmetszete.

Mi a host offloading és hogyan működik JAX-ban?

A host offloading a JAX nyílt forrású Python könyvtárban azt jelenti, hogy kiválasztott aktivációkat a forward pass alatt áthelyeznek rögzített (pinned) host memóriába, és a backward pass során visszaáramoltatják őket a GPU-ra. Ez alternatívája az activation rematerializationnek: ahelyett, hogy egyes aktivációkat újraszámolnának a visszafelé lépésben, azok újratöltődnek a host memóriából.

Miért különösen előnyös ez NVIDIA Grace Blackwell rendszereken?

A host offloading különösen jól működik az NVIDIA Grace Blackwell rendszereken, mert az NVIDIA Grace CPU és az NVIDIA Blackwell GPU NVLink‑C2C kapcsolattal rendelkezik, amely 900 GB/s kétirányú sávszélességet biztosít — így a pinned host memória praktikus tárolóhely lehet kiválasztott aktivációk számára. A Vera CPU és Rubin GPU konfigurációk tovább növelik ezt az átbocsátást 1.8 TB/s‑re. Magas CPU–GPU átviteli sebességre van szükség, de önmagában a sávszélesség még nem elég: az aktivációk átvitelének átfedésben kell történnie a hasznos GPU-munkával ahhoz, hogy teljesítménynyereséget hozzon.

Mérések: MaxText munkaterheletek NVIDIA GB200 NVL72 rendszereken

A kísérletek a MaxText JAX LLM betanító keretrendszert használták, amely az XLA (Accelerated Linear Algebra) fordítóval skáláz. Az összes eredményt NVIDIA GB200 NVL72 rendszereken, 128 GPU-val mérték. Vizsgált munkaterheltek:

• Llama 3.1 405B: sűrű (dense) decoder‑only transformer, címezett QKV aktiváció offload vizsgálatára, fix batch mellett.
• DeepSeek‑V3 671B: sparse mixture‑of‑experts (MoE) modell multihead latent attentionnel (MLA), áteresztőképesség és memória‑kapacitás hatások vizsgálatára.

DeepSeek‑V3 671B offload politika

A DeepSeek‑V3 671B 61 dekóder rétegből áll: az első három réteg sűrű MLP blokkokat használ, a maradék MoE blokkokat. Az offload politika a domináns, ismétlődő MoE dekóder rétegben egyes MLA query és key/value projekciók köztes eredményeit, valamint MoE up projekciók nagy intermediátjait helyezi át host memóriába — ezek az aktivációk elég nagyok ahhoz, hogy meghatározzák, milyen batch‑konfigurációk férnek el.

Teljesítmény‑javulás

A MaxText az áteresztést TFLOPs/s/device‑ben adja meg (modell FLOP‑ok per lépés osztva a mért lépésidővel).

• DeepSeek‑V3 671B: a konfiguráció, amely engedélyezte az offloadingot, a Latency Hiding Schedulert (LHS) és a pipelined átviteleket, 908.2 TFLOPs/s/device‑et ért el. Ez 57% gyorsulást jelent azonos batch mellett az activation rematerializationhoz képest, és 67.7% gyorsulást az offloading nélküli LHS/pipelining nélküli futáshoz képest.

A nagy MoE és MLA aktivációs lábnyom miatt a pipelined átvitelek engedélyezése különösen fontos volt a teljes áteresztés javításához — ellentétben a sűrű Llama munkaterhelttel, ahol önmagában az LHS sokszor elég a késleltetés elrejtéséhez. Az eredmények hátterében a szoros szoftver‑hardver együtttervezés áll: Blackwell rendszereken az XLA egyedi ütemezési opciói együttműködnek a hardverrel, aszinkron adatmozgatást biztosítva dedikált copy streameken keresztül.

Batch‑méret növelése (DeepSeek példája)

A DeepSeek‑V3 671B összehasonlítása több aktiváció‑elhelyezési politikára:

• Konfiguráció 1: Host offload + LHS + pipelined offloading: mikrobatch 8, globális batch 1024, 908.2 TFLOPs/s/device, GPU csúcs 165.2 GiB, host memória 145.1 GiB.
• Konfiguráció 2: Nincs offload, LHS, activation rematerialization: mikrobatch 8, globális batch 1024, 578.3 TFLOPs/s/device, GPU csúcs 151.3 GiB, host memória 0 GiB.
• Konfiguráció 3: Host offload, nincs LHS, nincs pipelined: mikrobatch 8, globális batch 1024, 541.6 TFLOPs/s/device, GPU csúcs 145.6 GiB, host memória 145.1 GiB.
• Konfiguráció 4: Nincs offload, LHS, mentés eszközön (save on device): mikrobatch 2, globális batch 256, 425.3 TFLOPs/s/device, GPU csúcs 113.3 GiB.
• Konfiguráció 5: Nincs offload, LHS, mentés eszközön: mikrobatch 8, globális batch 1024: OOM (nem futott el).

Ebből látható, hogy host offloading tette lehetővé a mikrobatch=8, globális batch=1024 konfigurációt, amely különben OOM hibát eredményezett volna. Az offload politikával nagy intermediátokat (MLA, MoE, MLP) helyeztek át a host memóriába, így több HBM maradt a modell állapotának, kommunikációs puffereknek és futásidejű munkaterületeknek. Fontos megjegyzés: LHS és pipelined átvitelek engedélyezése esetén a GPU csúcs memória 165.2 GiB volt, míg LHS nélkül 145.6 GiB — ez az extra HBM a több másolási puffer és előtöltött aktiváció tárolásából adódik: tudatos trade‑off a kapacitás és az átfedés/átviteli teljesítmény között.

Llama 3.1 405B eredmények

A Llama 3.1 405B kísérlet 10 lépésen futott szintetikus adatokkal: batch size 2, szekvenciahossz 8192, FSDP=128, bfloat16 aktivációk és NVFP4 4‑bites súlykvantálás.

• Alap (no offload, LHS on): 2,669 TFLOPs/s/device, GPU 149.6 GiB, host 0 GiB.
• QKV offload + LHS: 2,746 TFLOPs/s/device (2.9% növekedés), GPU 149.9 GiB, host 70.9 GiB.
• QKV offload, LHS off: 2,569 TFLOPs/s/device, GPU 139.7 GiB, host 70.9 GiB.
• QKV offload + LHS + pipelining: 2,718 TFLOPs/s/device, GPU 151.0 GiB, host 70.9 GiB.

Ebben a sűrű modell, fix batch példában az LHS önmagában elrejtette a legtöbb átviteli késleltetést, így az offloading által hozott nyereség kisebb volt, de mégis számottevő: az offloading lecserélte a backward pass QKV rematerializationt olyan átvitelekre, amelyek átfedhetnek compute‑tal és kommunikációval. A 70.9 GiB host memória a QKV aktivációk összesített mérete a 126 rétegben, nem pedig egyszeri GPU‑memória megtakarítás.

Mikor érdemes host offloadingot használni?

• Hasznos, ha az HBM korlátozza a modellméretet, a szekvenciahosszt vagy a batchméretet, és ha a kiválasztott tenzorok elég nagyok ahhoz, hogy csökkentsék az HBM‑nyomást.
• Különösen hatásos, ha helyettesíti az erőforrás‑igényes aktiváció‑rematerializációt, vagy lehetővé tesz nagyobb batch‑konfigurációt.
• Teljesítményfüggő: a host offloading akkor működik a legjobban, ha van elég párhuzamos számítás, kommunikáció vagy más független munka, ami elrejti az átviteli késleltetést. NVIDIA‑nál az XLA dedikált copy streamjei, LHS és pipelined host offloading segítik ezt az átfedést.

Mikor nem érdemes: ha a tenzorok kicsik, ha kevés független munka érhető el az átfedéshez, vagy ha a munkaterhelés más erőforrásban (nem memóriában) szűk keresztmetszetes. Statikus becslések helyett mindig validálni kell futásidő memóriahasználatot, mert a statikus számítások nem feltétlenül tartalmazzák az NVIDIA Collective Communications Library (NCCL) scratch területet, NVIDIA cuDNN attention workspace‑t és a keretrendszer által kezelt puffereket.

Hogyan kezdjenek hozzá a fejlesztők?

• Induljon egy kis, reprezentatív JAX betanítással. Válasszon nagy aktivációkat a drága forward utakból, engedélyezze az offloadot, és mérje a futásidejű GPU memóriát, host memória használatot és az end‑to‑end lépésidőt.
• A JAX host offloadingról szóló nyílt forrású tutorialokat használja (jax.remat, checkpoint szabályok, memory_kind="pinned_host", illetve paraméter és optimizer állapotok offloadolásához jax.device_put()).

A MaxText kísérletek az NGC JAX konténert használták; a DeepSeek‑V3 671B futtatások a ghcr.io/nvidia/jax:deepseek_v3_maxtext konténert alkalmazták, amely tartalmaz további MaxText integrációt és Transformer Engine MoE permutációs optimalizációkat. Az XLA‑ban használt optimizált flagek például:

• --xla_gpu_enable_latency_hiding_scheduler=true
• --xla_gpu_enable_pipelined_host_offloading=true
• --xla_gpu_experimental_parallel_async_compute_limit=8

Használjon profilozó eszközöket, például NVIDIA Nsight Systems‑t annak ellenőrzésére, hogy a device→host és host→device másolások valóban átfednek‑e a számítással és az NCCL kommunikációval.

Összefoglalás

A host offloading erős eszköz, ha az HBM kapacitás korlátozza a modell méretét, a kontextushosszt vagy a batchméretet, és ha a kiválasztott aktivációk elég nagyok a memória‑nyomás csökkentéséhez. A leghasznosabb, amikor kiváltja az aktivációk rematerializációját, vagy lehetővé teszi nagyobb batch konfigurációk futását. Az optimális eredményhez érdemes LHS‑t és pipelined host offloadot engedélyezni, és mindig valós futtatásokon mérni a memória‑ és időhatékonyságot.

Elismerések

Köszönettel tartozunk Jaroslav Sevciknek, Sevin Varoglu‑nak, Tj Xu‑nak, Haixin Liu‑nak, Abhinav Goel‑nek, Md Fahim Faysal Khan‑nak, Stefano Bosisio‑nak és Jinxin Yang‑nak a technikai hozzájárulásokért.