Az NVIDIA bemutatja, hogyan gyorsítható az Mixture-of-Experts (MoE) architektúrák gyakorlati tréningje biológiai alapmodellek esetén a Transformer Engine (TE) optimalizált primitívjeivel és a BioNeMo MoE receptekkel. A bemutatott megoldások — mint a GroupedLinear, a MXFP8 alacsony precizitású formátum és a GroupedMLP kernel — csökkentik a kernel indítási költségeket, mérséklik a memóriaigényt és növelik a GPU-k kihasználtságát.
Miért van különbség a dense és az MoE modellek között
A hagyományos, sűrű (dense) transformerben minden token átmegy az összes rétegen, így a modell képességeinek növelése egyenesen arányosan növeli a számítási igényt tréning és inferencia alatt. Az MoE architektúrák ezzel szemben sok alhálózatot ("expert") tartalmaznak, de minden tokenhez csak egy kis részhalmazt aktiválnak, ami hatékonyabb kapacitásnövelést tesz lehetővé. Az előnyök gyakorlati megvalósítása azonban függ az implementációtól: a szétszórt expert-számítás csökkentheti a GPU-k kihasználtságát, a routing kommunikációs overhead-et okozhat, és a nagyobb paraméterméret memória- illetve elosztott tréning kihívásokat hoz.
A Transformer Engine szerepe
A Transformer Engine (TE) célja ezen szűk keresztmetszetek kezelése: optimalizált primitíveket ad grouped expert-számításhoz, kernel-fúzióhoz és alacsony precizitású tréninghez. Ezek a lehetőségek különösen fontosak a biológiai alapmodellek számára, amelyek nővekvő paraméterszámot és hosszú szekvenciákat használnak, mert javítják a GPU-hatékonyságot miközben lehetővé teszik a modellszélesség növelését.
Főbb technikai megoldások a receptben
1) Fragmentált expert-kermek problémája
Egy egyszerű, naiv MoE-implementáció minden szakértőt külön Python-ciklusban hív meg, ami sok különálló kernel indítást eredményez. A TE GroupedLinear primitívje ehelyett egyetlen hívásban végrehajt több lineáris transzformációt úgy, hogy összegyűjti az expert-súlyokat és a bemeneti tokeneket. Mivel az egyes expertek különböző számú tokent kaphatnak, a GroupedLinear per-expert token-számokat (split_sizes) is elfogad. Ez a grouped GEMM útvonalat használja ahelyett, hogy egy PyTorch Linear műveletet indítana minden egyes expertre, ezzel csökkentve az indítási és ütemezési overhead-et. Példa: a GroupedLinear létrehozása és használata úgy, hogy minden expert megőrzi saját súlytenzorát, és a hívás megkapja a per-expert token-számokat.
2) Modellméret és aktivációs memória
Az MoE növeli a teljes paraméterkapacitást, ráadásul a genomikai feladatok gyakran hosszú szekvenciákat használnak, ami jelentős aktivációs memóriaigényt eredményez a tréning során. A BioNeMo recept támogatja az FP8 és MXFP8 formátumokat a Transformer Engine-en keresztül, amelyek 8 bites ábrázolást használnak a BF16 (16 bites) helyett. A MXFP8 fontos különbsége a skálázás szemcsézettsége: blokkonként (32 egymást követő érték) ad skálafaktort, ami segít megőrizni a numerikus tartományt és pontosságot. NVIDIA Blackwell GPU-kon a MXFP8 hardveresen gyorsított, így a MXFP8 GEMM-ek speciális Tensor Core utasításokat használhatnak.
3) Kvantálás overhead az alacsony precizitású tréningben
Bár a legtöbb számítás 8 bites pontossággal történik, a modell „master” súlyai 16 bites formátumban maradnak, ezért a tréning-keretrendszernek kvantálási és dekvantálási lépéseket kell végrehajtania a formátumok közti konverzióhoz. Egy naiv útvonal ezeket külön műveletekként hajtja végre, ezért a TE a fuzionált MLP utat kínálja: a GroupedLinear, a ScaledSwiGLU és a lecsatoló GroupedLinear összevonásával csökkenthető az extra memória-mozgatás és a köztes eredmények materializálása. Példakód a TE Sequential API használatára, amely felismeri a GroupedLinear → ScaledSwiGLU → GroupedLinear mintát és egy fuzionált ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 kernelre cseréli azt.
Gyakorlati eredmények
Az NVIDIA által közölt benchmarkban a BioNeMo recept több optimalizálása együttvéve jelentős sebességnövekedést eredményezett. Egy nyolc GPU-s méréssorozaton (nyolc NVIDIA B200 Tensor Core GPU) a recept a Hugging Face alapértelmezett implementációhoz képest legfeljebb 2.21× átviteli teljesítményt (throughput) ért el a Mixtral-8x7B tréning esetén.
Hogyan futtatható a recept
A dokumentált indítási javaslatok szerint érdemes először a két-GPU-s L0_sanity konfigurációval ellenőrizni az expert-parallelizmust és a környezetet:
torchrun --nproc_per_node=2 train_fsdp2_ep.py --config-name L0_sanity
Ha ez sikeres, skálázzunk az Mixtral-8x7B konfigurációra EP=8 és MXFP8 precizitással nyolc GPU-n:
torchrun --nproc_per_node=8 train_fsdp2_ep.py --config-name L1_8x7B_ep checkpoint.ckpt_dir=/path/to/ckpt
Válasszuk a BF16 vagy MXFP8 módot a GPU-k és memóriakövetelmények alapján. A data-parallel és expert-parallel méretek szorzatának egyenlőnek kell lennie a rendelkezésre álló GPU-k számával. A recept README-je tartalmazza a részletes indítási, checkpoint és benchmark parancsokat.
Feltételek és követelmények
- Ismeret Pythonból, PyTorch-ból és elosztott tréning koncepciókból
- NVIDIA CUDA-képes környezet (a recepthez tartozó Dockerfile is használható vagy a követelmények telepíthetők)
- Legalább két GPU expert-parallelizmushoz; NVIDIA Blackwell GPU-k szükségesek a fuzionált MXFP8 GroupedMLP kernel használatához
Köszönetnyilvánítás
A recept és a bemutatott munkák mögött szerepelnek: Sudhakar Singh US, Varun Thumbe US, Santosh Santosh US, Timur Rvachov US és Chris Hoge US.
