TwELL: új CUDA-kernel gyorsítja a LLM-ek betanítását 21,9%-kal, inferenciáját 20,5%-kal
A TwELL nevű megoldás a feedforward rétegek számításait teszi olcsóbbá a strukturálatlan ritkaság kihasználásával, anélkül, hogy a modell architektúráját megváltoztatná.
A nagy nyelvi modellek (LLM) skálázása rendkívül költséges, mivel a tokenek feldolgozásának és a gradiens-számításoknak több mint kétharmada a feedforward rétegeken keresztül zajlik. Ezek a rétegek a nagyobb modellek teljes FLOP-számának több mint 80%-át teszik ki — írja a MarkTechPost.
A Sakana AI és a NVIDIA kutatói egy új módszert dolgoztak ki, amely közvetlenül ezt a szűk keresztmetszetet célozza. A TwELL névre keresztelt rendszer nem az architektúrát módosítja, hanem a feedforward rétegek számításait teszi lényegesen olcsóbbá a strukturálatlan ritkaság kihasználásával.
A ritkaság kihasználása
A transzformer modellek feedforward blokkjában az input tokenek csak kis része aktiválódik, a többi neuron nullát ad vissza az aktivációs függvény után. Ezt aktivációs ritkaságnak nevezik, de a korábbi GPU-optimalizációk nem tudták hatékonyan kihasználni, mivel a sűrű mátrixszorzásokra voltak kihegyezve.
Ritka formátumok hatékonysága
A TwELL egy új, csempe-alapú ritka formátum, amelyet kifejezetten a GPU-kernelek működéséhez igazítottak. Ez a megoldás a nem nulla értékeket helyben, csempéken belül csomagolja, így a konstrukció az existing gate projection kernel epilógusába integrálható, extra kernel indítása, memóriahozzáférés vagy szinkronizációs overhead nélkül.
A teljesítményteszt eredmények szerint a 2 milliárd paraméteres modelleknél 20,5%-os inferencia és 21,9%-os betanítási sebességnövekedést értek el, miközben az energiafogyasztás 17%-kal csökkent. A modell március végén kerül nyilvánosságra a NVIDIA által.