Az Aleph Alpha 512 B200 GPU-n skálázta 30 milliárd paraméteres modelljét

Kisebb kísérletekkel szűkítették a keresési teret

Az Aleph Alpha szeptember 30-án közzétett kutatási bejegyzésében azt mutatja be, hogyan készítették fel a 30B-A3B MoE modellt 512 NVIDIA B200 GPU-n történő előtanításra. A cég saját, optimalizált torchtitan-változatát használta. A célrendszer 64, InfiniBanddel összekötött csomópontból állt, csomópontonként 8, NVLinken kapcsolódó GPU-val.

A megközelítés lényege, hogy nem rögtön a teljes GPU-kapacitáson próbálták ki a lehetséges beállításokat. Először kisebb léptékű kísérletekkel szűkítették a keresési teret, és csak ezután növelték a GPU-k számát. Az Aleph Alpha szerint egy nagyméretű, naiv rácskeresés túl költséges lenne, mivel számos, a hatékonyságot befolyásoló paramétert kellene egyszerre vizsgálni.

A sebességet és a kihasználást külön mérik

A bejegyzés három mérőszámot használ. A TPS, vagyis a GPU-nkénti másodpercenként feldolgozott tokenek száma, azt mutatja meg, milyen gyorsan zajlik a tanítás. Ezt a mutatót a feldolgozott tokenek száma, a falióra szerinti idő és a GPU-k száma alapján számítják ki. Az Aleph Alpha ezt tekinti elsődleges optimalizálási céljának, mert közvetlenebb képet ad a tanítási sebességről.

Az MFU, azaz a modell FLOPS-kihasználása, ugyanezt a teljesítményt a hardver elméleti számítási plafonjához viszonyítja. A vállalat BF16 csúcsteljesítményt használ nevezőként, mivel a tanítás BF16 formátumban zajlik, az FP32 gradienseket leszámítva. A harmadik mutató a memóriahasználat, amely a tanítás során lefoglalt csúcsmemória és az eszköz teljes memóriájának aránya. A számításhoz a PyTorch lefoglalt memóriáját veszik alapul, mert ez határozza meg a memóriahiány miatti hibák kockázatát.

Három beállítást hangoltak a számítási terheléshez

Az Aleph Alpha célja az volt, hogy a GPU-k lehetőleg számítási korlátba ütközzenek. Ilyenkor a teljesítményt elsősorban a matematikai műveletek sebessége határozza meg, és nem az, hogy milyen gyorsan lehet az adatokat a GPU-k között mozgatni. A kommunikáció ugyanis önmagában nem viszi előre a modell számítását, miközben az adatokat váró GPU-erőforrások kihasználatlanul maradhatnak.

A kutatásban három szabadsági fokot vizsgáltak: a párhuzamosítási sémát, az aktivációk ellenőrzőpontozását és a helyi kötegméretet. A párhuzamosítás meghatározza, hogyan oszlanak meg a GPU-k között a modell súlyai, gradiensei, az optimalizátor állapotai és az aktivációk. A vizsgálat az adatok párhuzamosítására, valamint a teljesen felosztott adathalmaz-párhuzamosításra, vagyis az FSDP-re, más néven ZeRO-3-ra is kitér. FSDP esetén a súlyokat, a gradienseket és az optimalizátor állapotait több GPU között osztják fel, amelyek együtt alkotnak egy teljes modellpéldányt.

Mit jelent ez a felhasználóknak és a fejlesztőknek?

Az Aleph Alpha eredménye szerint 512 GPU-n is közel lineárisan lehet növelni a teljesítményt, ha a skálázás előtt megfelelően szűkítik és hangolják a tanítási konfigurációt. A 35,3 százalékos MFU a vállalat mérésében 6 százalékkal maradt el a tökéletes lineáris skálázástól.

A módszer jelentőségét a cég nem kizárólag a több száz GPU-val rendelkező laboratóriumok számára látja. A bejegyzés szerint számos használt eszköz és ötlet méretfüggetlen, így kisebb számítási környezetben is alkalmazható. A hatékonyabb előtanítás csökkentheti a tanításhoz szükséges időt, miközben a kutatás középpontjában elsősorban a hierarchikus skálázási módszer állt. A rendszer- és kernelszintű optimalizálások nagy része, amelyek lehetővé tették a 35,3 százalékos MFU-t, nem tartozott a bejegyzés fő témájához.