Az NVIDIA szerint akár 2,21-szer gyorsabb lehet az MoE-modellek tanítása

Az MoE-modellek skálázásának problémái

A sűrű transzformermodellekben minden token minden rétegen áthalad, ezért a modell képességeinek bővítésével a tanítás és a következtetés számítási igénye is növekszik. A kevert szakértői, vagyis Mixture-of-Experts, MoE-architektúrák ezzel szemben több alhálózatot, úgynevezett szakértőt használnak, de tokenenként csak ezek egy részét aktiválják.

Ez hatékonyabban növelheti a modell kapacitását, az előnyök azonban az implementációtól függenek. A széttagolt szakértői számítások ronthatják a GPU-k kihasználtságát, az útválasztás kommunikációs többletet okozhat, a nagyobb paraméterszám pedig memória- és elosztott tanítási nehézségeket hoz magával. Az NVIDIA Transformer Engine ezekre optimalizált primitívekkel próbál választ adni.

Egyetlen műveletbe szervezett szakértői számítás

Az NVIDIA fejlesztői bejegyzése szerint a Hugging Face alapimplementációja Python-ciklusban dolgozza fel a szakértőket. Ez azt jelenti, hogy minden szakértő külön lineáris műveletet és külön kernelindítást kap. A Transformer Engine GroupedLinear ehelyett több szakértő mátrixműveletét egyetlen csoportosított műveletként indítja el.

A megoldás megtartja az egyes szakértők saját súlytömbjeit, miközben a bemeneti tokeneket és a szakértőnként eltérő tokenszámokat együtt kezeli. Az NVIDIA szerint ez csökkenti a kernelindítás és az ütemezés többletterhét. A cég megjegyzi, hogy a Hugging Face Transformers is kínál grouped_mm megoldást, a Transformer Engine azonban az összevont műveletet további optimalizálásokkal is össze tudja kapcsolni.

MXFP8 és összeolvasztott MLP-kernel

A biológiai alapmodellek növekvő paraméterszáma és a genomikai feladatok hosszú szekvenciái az aktivációs memória használatát is megterhelik. A BioNeMo-recept FP8 és MXFP8 tanítást támogat, amelyek 16 helyett 8 biten ábrázolják a súlyokat és az aktivációkat. Az MXFP8 minden 32 egymást követő értékhez külön skálázási tényezőt rendel. Az NVIDIA Blackwell GPU-kon ezt a formátumot hardveresen támogatott Tensor Core utasítások gyorsítják.

Az alacsonyabb pontosságú tanítás közben a keretrendszernek továbbra is kezelnie kell a 16 bites fő súlyokat, ezért kvantálási és visszakvantálási lépésekre van szükség. A Transformer Engine Sequential API a GroupedLinear, a ScaledSwiGLU és a routing súlyozásának műveleteit összevonhatja. A megfelelő műveletsorozat felismerésekor a rendszer a ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 kernelt használja a forward, valamint a hozzá tartozó összevont műveletet a backward számításhoz. Ez csökkenti a keretrendszer többletterhét, és elkerüli egyes köztes eredmények külön létrehozását.

Akár 2,21-szeres áteresztés a benchmarkban

Az NVIDIA nyolc NVIDIA B200 Tensor Core GPU-n végzett tanítási benchmarkjában a BioNeMo-recept akár 2,21-szeres áteresztést ért el a Hugging Face alapjához viszonyítva. A mérés a Mixtral-8x7B tanítási áteresztését vizsgálta.

A bemutatott recept kipróbálásához az NVIDIA legalább két GPU-t ír elő a szakértői párhuzamossághoz. A két GPU-s L0_sanity konfigurációval ellenőrizhető a környezet és a szakértői párhuzamosság működése, ezt követően a Mixtral-8x7B konfiguráció nyolc GPU-n, EP=8 beállítással és MXFP8 pontossággal futtatható. Az összevont MXFP8 GroupedMLP-kernel használatához NVIDIA Blackwell GPU-k szükségesek.