A PyTorch FBTriton kerneljei gyorsítják a táblás embeddingeket

Egy GPU-műveletben több embeddingtábla

A Table-Batched Embedding, röviden TBE, sok táblán végez embeddingkeresést és összevonást egyetlen GPU-indításban. A megközelítés csökkenti az indítási többletet és javítja a memóriahasználatot. A PyTorch blogbejegyzése szerint ezek az alapvető műveletek több ezer, felosztott GPU-n futó ajánlórendszerekben kapnak szerepet.

Az FBTriton előre- és visszaterjesztési kerneleket használ. Az előreirányú művelet összegyűjti az indexelt sorokat, szükség esetén mintánkénti súlyokkal szorozza őket, FP32-ben halmozza az eredményt, majd egy D szélességű összevont kimenetet ír ki. FP32 súlyok esetén a felhalmozás FP64-ben történik a nagy D értékek melletti pontosság megőrzésére.

Kétféle előreirányú megvalósítás

A fejlesztők általános gather, vagyis adatbegyűjtő útvonalat, valamint egy kis táblákhoz készült gyorsított útvonalat építettek. Az általános kernel a feldolgozást programokra osztja, amelyek a T jellemzőkön belül hurkolnak, ahelyett hogy minden B és T kombinációhoz külön rácsot indítanának. A hangolt, két baget kezelő útvonal programonként nyolc baget dolgoz fel bizonyos nagy, nem VBE és nem FP32 munkaterheléseknél.

A speciális, hisztogramot és tenzormagokat használó útvonal akkor választható, ha egy jellemző E értéke legfeljebb 64, a D értéke 64 és 128 közötti, az L értéke legalább 64, FP16 súlyokat és FP32 kimenetet használ a konfiguráció, továbbá nincs mintánkénti súlyozás és VBE. Ez az útvonal 16 baget kezel programonként, a fennmaradó indexeket skaláris feldolgozás kezeli.

A PyTorch a határellenőrzés gyorsítását is bemutatja. B200 GPU-n a frissített CUDA-validáció a vállalat mérése szerint legfeljebb 1,24-szeres gyorsulást ér el ezen a komponensen. A bekapcsolható opció alapértelmezés szerint ki van kapcsolva, és TorchRecen keresztül érhető el.

A visszaterjesztés a terheléselosztást célozza

A TBE visszaterjesztése minden érintett, egyedi tábla- és sorpárhoz összegzi a batch azon pozícióiból érkező gradienseket, amelyek az adott sort használták, majd egyetlen optimalizálófrissítést alkalmaz. A műveletben nincs tenzormag, viszont jelentős adatmozgatást és redukciót igényel, miközben a munkaterhelés kiegyensúlyozása nehéz.

A transpose_embedding_input a batch bemenetét olyan futásokra fordítja meg, amelyek egy egyedi sort és az azt használó mintákat párosítják. Ezt a műveletet az előreirányú szakaszba emelték, így kikerül a visszaterjesztés kritikus útvonaláról. A futás hossza alapján három kernel dolgozik: a rövid futásokat kezelő short_run 256-nál kisebb értékeknél, a legalább 256 elemű futásokhoz használt grad_accum + apply, valamint Blackwellen, nagyon nagy batcheknél a fused változat.

A hosszú futásokat 256 lekérdezéses részekre bontják. Így egy kétmillió lekérdezésből álló futás nagyjából nyolcezer alprogramra osztható, ami a PyTorch szerint jobban ki tudja tölteni a gépet. A TLX Blackwell CLC-megoldása hardveres munkalopással segíti a terheléselosztást, a TLX által biztosított fence pedig lehetővé teszi, hogy az utolsó részprogram alkalmazza a frissítést.

Mért eredmények és felhasználói jelentőség

A PyTorch egy nagy B200-konfiguráció példáján azt írja, hogy az indextranszponálást, rendezést és futamhossz-kódolást az előreirányú szakaszba helyező megoldásnál az előreirányú idő 22,844 ms-ről 33,252 ms-re nőtt, a visszaterjesztésé 56,693 ms-ről 32,931 ms-re csökkent. A kombinált késleltetés 79,537 ms-ről 66,183 ms-re mérséklődött, ami 16,8 százalékos csökkenés.

A változtatások azokat a TorchRec-alapú ajánlórendszeres munkaterheléseket célozzák, ahol sok embeddingtábla, változó futáshossz és jelentős adatmozgatás találkozik. Egyes konfigurációk, köztük a súlyozott, változó batchméretű, AMD-s és transpose-hoist esetek, továbbra is standard validációra esnek vissza. Az FBTriton jelenlegi tervezése a PyTorch szerint további optimalizálási lehetőségeket is hagy.