NVIDIA bemutatta a TensorRT 11.0‑t, amely új, natív több‑eszközös (multi‑device) inferencia támogatást hoz a TensorRT futtatókörnyezetbe. A cél a generatív AI‑munka terhelések skálázása több GPU‑ra anélkül, hogy elvesznének a TensorRT által nyújtott kritikus optimalizációk, mint a kernel fúziók, memória tervezés és kvantizálás, amelyek fontosak a produkciós telepítésekhez.
A multi‑device inferencia funkció összekapcsolva a Torch‑TensorRT‑rel lehetővé teszi, hogy fejlesztők nagy PyTorch modelleket konvertáljanak és telepítsenek a keretrendszeren kívül, áttörve az egyetlen GPU memóriájából és számítási kapacitásából fakadó korlátokat. A TensorRT 11.0 multi‑device képessége a NVIDIA Developer Portal‑ról letölthető.
NCCL mint a terjesztett inferencia szállító rétege
A NVIDIA Collective Communications Library (NCCL) magas teljesítményű kollektív műveleteket biztosít több‑GPU és több‑node környezetekhez, és automatikusan kiválasztja a topológiához optimális adatszállítást (NVLink, NVSwitch, PCIe, InfiniBand). A TensorRT a NCCL‑lel közvetlen integráción keresztül örökli ezeket a szállítási optimalizációkat, amikor multi‑device inferenciát futtat.
Az új multi‑device funkció a teljes NVIDIA NCCL kollektív művelet‑készletet támogatja: AllReduce, Broadcast, Reduce, AllGather, ReduceScatter, AlltoAll, Gather és Scatter.
Párhuzamosítási stratégiák elosztott inferenciához
Az elosztott inferencia többféle párhuzamosítási stratégiával írható le; mindegyik más kompromisszumot hoz a memória megtakarítás, a számítási skálázódás és a kommunikációs overhead között. A leggyakoribb stratégiák a tensor párhuzamosítás és a kontextus párhuzamosítás.
Tensor párhuzamosítás
Tensor párhuzamosításnál egy réteg súlyait darabolják szét GPU‑k között; minden GPU egy részét számolja ki a réteg mátrixszorzatának, majd kollektív művelettel egyesítik a részleges eredményeket a teljes kimenet előállításához. Ez csökkenti az egyes eszközökre jutó súlyméretet, és gyakran az egyetlen lehetőség, ha egy egyedi réteg súlyai meghaladják egy GPU memóriáját, függetlenül a bemeneti szekvencia hosszától vagy a batch mérettől.
Transformerekben az oszlop‑ és sorparalel vetületek (például QKV oszlop‑, valamint az attention kimenet és MLP down‑projekció sor‑paralelizálása) úgy szervezhetők, hogy minden blokkhoz csak egy AllReduce szükséges legyen, így a kommunikációs költség korlátok között marad.
Kontextus párhuzamosítás
A kontextus párhuzamosítás a bemeneti szekvenciát darabolja a szekvencia dimenzió mentén GPU‑kra. Minden GPU a szekvencia egy szeletét dolgozza fel; kollektív műveletek teszik a teljes szekvenciát elérhetővé, amikor szükséges (például az attention során). Ez különösen hatékony hosszú szekvenciás munkafolyamatoknál, ahol az attention számítás és memóriaigénye kvadratikusan nő a szekvencia hosszával.
Ez a megközelítés természetes választás difúziós és DiT modellekhez is, amelyek bidirekcionális attention‑t használnak, és így elkerülik a kauzális maszkból adódó terheléseltolódási problémákat.
NVIDIA TensorRT 11.0 támogatja az IDistCollectiveLayer primitíveket, amelyeket a különböző párhuzamosítási stratégiák igényelnek. Az alábbiakban részletesebben a kontextus párhuzamosításra fókuszálunk, mivel az ez a domináns költségtényező a modern generatív média pipeline‑okban: a hosszú szekvenciás attention.
Kontextus párhuzamosítás a generatív média pipeline‑oknál
Difúziós alapú kép‑ és videógenerálásnál a számítás és memória nagy része az attention blokkokban történik, amelyek hosszú token‑sorozatokon futnak. Egy nagyfelbontású kép latent vagy több képkockás videó klip blokkonként tízezres token‑sorozatokat eredményezhet, és az attention költsége kvadratikusan nő a szekvencia hosszával.
A kontextus párhuzamosítás több megvalósítási módot is kínál, amelyek különböző kompromisszumokkal rendelkeznek:
AllGather KV
Az AllGather KV megközelítésnél a rankek megosztják egymással a key (K) és value (V) részeiket egy AllGather kollektív művelettel, majd lokálisan kiszámítják az attentiont, így minden rank képes a teljes szekvenciára attentiont végezni. Ez egy plusz kollektívet igényel per attention blokkért, de a Q × Kᵀ lokális szorzat mérete arányosan csökken a rankek számával. Videó- és nagyfelbontású képdiffúzióknál ez a költség minden denoising lépésnél kedvezően oszlik el: a kommunikáció per lépés az AllGather‑re korlátozódik, míg a számítási és memória megtakarítás minden attention rétegben és lépésben érvényesül.
Ring Attention
A Ring Attention a kommunikáció és számítás átfedésére épít: a K és V folyamatosan „átfolynak” a GPU‑kon egy körtopológiában, miközben minden GPU a saját lokális Q‑ját dolgozza. Ez csökkenti a memóriaigényt is, mert online softmax alkalmazása esetén a teljes K és V tenzorok nem neednek teljesen materializálódni bármelyik GPU‑n.
DeepSpeed Ulysses
Nagyon hosszú kontextusoknál (tízezres token‑tartományok) a DeepSpeed Ulysses alternatív implementációt kínál: először a mintákat a szekvencia dimenzió mentén partícionálja a résztvevő GPU‑k között. Az attention előtt all‑to‑all kollektív műveletet hajt végre a partícionált Q, K és V tenzorokon, így minden GPU a teljes szekvencia hosszához jut, de csak egy nem-átfedő attention head részhalmazához. Ez párhuzamos attention számítást tesz lehetővé, majd egy második all‑to‑all művelet gyűjti össze az eredményeket a headek mentén, és újrapartícionálja azokat a szekvencia dimenzió szerint.
Benchmarkok: média generálás kontextus párhuzamosítással C++‑ban
A publikált benchmarkok multi‑device TensorRT inferenciát mértek média generálási munkafolyamatokra, C++ produkciós telepítési célra. Két reprezentatív pipeline szerepelt a tesztekben: egy videó‑generációs pipeline az NVIDIA Cosmos 3‑ra építve, és egy képgenerálási pipeline FLUX.1‑re alapozva.
Ezeket a modelleket PyTorch‑ban fejlesztették, majd a Torch‑TensorRT segítségével kimenet‑független, TensorRT engine‑ekké konvertálták, amelyeket C++ inferencia alkalmazásokban telepítettek. A benchmarkok összehasonlították az end‑to‑end késleltetést különböző kontextus párhuzamosítási stratégiákkal: AllGather KV, Ring Attention és Ulysses. Minden mérés egyetlen node‑on, 8 GPU‑val történt.
Video generálás – NVIDIA Cosmos 3
A NVIDIA Cosmos modellplatform és a Cosmos3‑Nano modell multimodális bemenetek (szöveg, kép, videó) alapján képeket, videót és hangot képes generálni. A méréseknél alkalmazott példafájl alapján a DeepSpeed Ulysses stratégia bizonyult egyértelműen legjobbnak, amikor a modell rendkívül hosszú kontextussal dolgozik (tízezres token‑tartományok).
Kép generálás – FLUX.1
A FLUX.1‑dev modell (Black Forest Labs) szöveges leírásokból képeket generál. A „a beautiful photograph of Mt. Fuji during cherry blossom” promptot használva a benchmarkok azt mutatták, hogy Ulysses itt is nyerő stratégia; ugyanakkor a Ring Attention jól skálázódott akár 4 GPU‑ig is.
Hogyan kezdjünk el a TensorRT multi‑device funkciójával
A TensorRT multi‑device inferenciát támogat, így egy hálózat több GPU‑n futhat integrált terjesztett kommunikációs primitívekkel. A fő munkafolyamat hasonló a single‑device TensorRT‑hez, az eltérés az, hogy a hálózat most tartalmazhat disztribúciós kommunikációs rétegeket (IDistCollectiveLayer).
Főbb előfeltételek és lépések összefoglalva:
- Töltsük le a TensorRT 11‑et a NVIDIA Developer Portal‑ról és telepítsük a leírt utasítások szerint.
- Szerezzünk egy single‑node, több‑GPU gépet.
- Telepítsük az OpenMPI‑t a fejlesztési környezetben (bare metal vagy konténer).
- A hálózat szintjén a INetworkDefinition::addDistCollective metódussal adhatók hozzá a kollektív műveletek (például ALL_REDUCE vagy ALL_GATHER), majd állítsuk be a világméretet (setNbRanks).
A cikk forrásában bemutatott példa kódrészletek illusztrálják az INetworkDefinition létrehozását, egy kollektív réteg hozzáadását, a builder/config használatát engine építéshez, execution context létrehozását, a bemeneti/kimeneti tensorok GPU címzésére vonatkozó bindelést, a kommuniktor beállítását és az enqueueV3 meghívást. Fontos megjegyezni, hogy a NCCL kommuniktornak érvényesnek kell maradnia a hozzá tartozó execution context teljes élettartama alatt.
A mintaprogram futtatásának példája 8 GPU‑val OpenMPI alatt (a környezeti változók beállításával meg kell adni a ranket és a world size‑ot, illetve egy fájlt az NCCL ID‑hoz):
mpirun -np 8 bash -lc 'export TRT_MY_RANK=$OMPI_COMM_WORLD_RANK; export TRT_WORLD_SIZE=$OMPI_COMM_WORLD_SIZE; export TRT_NCCL_ID_FILE=/tmp/nccl_id.txt; ./sample_dist_collective --op all_reduce'
További anyagok
A cikk több hivatkozást is javasol további elmélyülésre, köztük a NVIDIA Collective Communications Library (NCCL), a Context Parallelism for Scalable Million‑Token Inference, Ring Attention és DeepSpeed Ulysses témacikkeket, valamint a NVIDIA TensorRT és Torch‑TensorRT dokumentációkat és GitHub repozitóriumokat.
Összefoglalásként a TensorRT 11.0 multi‑device támogatása és a NCCL‑integráció lehetővé teszi, hogy nagy generatív modellek produkciós környezetben több GPU‑t használjanak megőrizve a TensorRT optimalizációit; a különböző kontextus párhuzamosítási stratégiák közül a DeepSpeed Ulysses bizonyult a leghatékonyabbnak extrém hosszú kontextus esetén a bemutatott benchmarkok alapján.
