Az NVIDIA Transformer Engine 10-szer gyorsabbá teszi a MoE tréninget JAX-ban
Az NVIDIA Transformer Engine és a JAX 10.4-szeres sebességnövekedést ért el a Mixture of Experts (MoE) tréningben NVIDIA GB200-on, a DeepSeek-V3 modell teljesítménye 103-ról 1068 TFLOPS/GPU-ra emelkedett.
Technikai optimalizálások
A szakértői párhuzamosságot (expert parallelism) az NCCL EP gyorsítja, amely egyesíti az elküldési és összesítő szakaszokat, valamint deduplikálja a tokeneket a hálózati forgalom csökkentése érdekében. További optimalizálások, mint a JAX host offloading és az XLA multistreaming, csökkentik a memória-szűk keresztmetszeteket és átfedik a kommunikációt. (Nvidia Developer)
A Dropless MoE minden tokent feldolgoz, megőrzi a modell minőségét, szemben a kapacitásalapú MoE-val, amely tokeneket dob vagy kitölt. A Transformer Engine új, csoport-aware MXFP8 kvantálást és optimalizált EP műveleteket kínál az elküldéshez és összesítéshez.
Skálázhatóság és elérhetőség
A DeepSeek-V3 671B modell tréningje során az optimalizált rendszer 1024 GPU-n 97%-os skálázási hatékonyságot tart fenn NVIDIA GB300 NVL72 hardveren.
Az optimalizált JAX MoE útvonal reprodukálásához próbálja ki az NVIDIA NGC MaxText konténert Transformer Engine-nel. A részletes beállításokért tekintse meg a MaxText MoE Configuration útmutatót, a könyvtár képességeinek megértéséhez pedig a Transformer Engine dokumentációját.