Az AWS új MoE-optimalizációkkal gyorsítja a SageMaker AI-t

A fused MoE a számítási folyamat több lépését egyesíti

A Mixture-of-Experts, röviden MoE, ritkán aktivált architektúra, amely nagy nyelvi modellek kapacitását képes növelni úgy, hogy az egyes tokenek feldolgozásához szükséges számítási mennyiség alacsony marad. Az AWS szerint az MoE-műveletek a teljes számítási idő jelentős részét adhatják, ezért a hatékonyságuk közvetlenül befolyásolja a válaszok késleltetését és a költségeket.

A fused MoE egyik fontos optimalizációja a tokenek útválasztását, az adatok átrendezését és a szakértői hálózatok számításait egyetlen, kötegelt GPU-műveletbe vonja össze. Ez kiváltja az egymást követő, különálló ciklusműveleteket, ugyanakkor összetettebb kernelimplementációkat és hibakeresést tesz szükségessé. Az AWS a bejegyzésben egy reprezentatív MoE-modellen keresztül mutatja be a fejlesztéseit.

A mátrixok felcserélése jobban kihasználja a Hopper GPU-kat

Az első optimalizáció az NVIDIA Hopper architektúra WGMMA, vagyis Warp Group Matrix Multiply Accumulate utasításainak működéséhez igazítja a mátrixszorzást. A WGMMA egy 128 szálból álló együttműködő warp-csoporton hajtja végre a műveletet, és az első operandushoz tartozó M dimenziót 64-es rögzített mérethez igazítja. Ha ez a dimenzió 64-nél kisebb, a számítási kapacitás egy része kihasználatlan marad.

Ez különösen a dekódolási szakaszban jelent gondot, amikor a rendszer felhasználói sorozatonként egyszerre csak egy tokent dolgoz fel. Ilyenkor az M dimenzió gyakran 1. A hagyományos elrendezésben az egyik mátrix keskeny tokenvektorrá válik, miközben a szakértői súlyok nagy, sűrű blokkot alkotnak. A fused MoE esetében az N dimenzió ráadásul az aktív szakértők számától és a szakértői köztes mérettől függően több tízezer vagy akár több százezer paramétert is jelenthet.

Az AWS ezért a mátrixszorzás algebrai tulajdonságát kihasználva felcseréli az operandusok szerepét, és transzponált formában hajtja végre a műveletet. Így a nagy szakértői dimenzió kerül az első operandus helyére, ahol könnyebben kitölti a Hopper hardverének 64 soros elrendezését. A megoldás az aszinkron előbetöltést is segíti, mivel a nagyméretű súlyblokkok folyamatosabban továbbíthatók a megosztott memóriából a regiszterekbe.

Az AWS szerint ez a módosítás 2 és 4 százalék közötti, teljes folyamatra vetített kimeneti teljesítménynövekedést eredményezett nagyméretű MoE-modelleken H200 GPU-val. Az optimalizáció bekerült a vLLM projektbe, és a 0.24.0 vagy újabb verziókban érhető el.

FP8 esetén a Triton backend bizonyult gyorsabbnak

A második fejlesztés az MoE-számításokhoz használt backend kiválasztását módosítja. A vLLM MoE-orákuluma különböző kvantálási eljárásokat a Triton, a FlashInfer és CUTLASS kombinációja, illetve más backende­k felé irányíthat. Az AWS profilozása szerint Hopper GPU-kon az FP8 tenzorkvantálást használó modelleknél a Triton backend következetesen gyorsabb volt az alapértelmezett FLASHINFER_CUTLASS útvonalnál.

Az AWS magyarázata szerint a Triton egyetlen hardveres menetben egyesíti az FP8-skálázást, a mátrixszorzást és a szakértők eredményeinek összesítését. Az adatokat egyszer olvassa be a nagy sávszélességű memóriából, majd a GPU helyi regisztereiben és számítási egységeiben dolgozza fel, mielőtt továbbítja a végső mátrixhoz. Ez elkerüli azokat a globális memóriába és memóriából történő köztes olvasási és írási műveleteket, amelyek hátrányt jelenthetnek a CUTLASS számára.

A Triton dinamikus ütemezési elrendezéseket is használ a változó tokencsoportokhoz. Ez a dekódolásnál jellemző, nagyon kis kötegméreteknél is hatékonyabban osztja el a töredezett tokenadatokat az aktív szakértői súlyok között. Az AWS szerint a backend útválasztási szabályának módosítása, kernelkód átírása nélkül, 10 százalékos teljes folyamatra vetített teljesítménynövekedést hozott FP8-kvantált, fused MoE modelleken H200 GPU-val.

Mit jelent ez a SageMaker AI felhasználóinak?

Az AWS bejegyzése szerint a bemutatott eredmények két különböző beavatkozásból származnak. Az egyik közvetlenül a kernel matematikai elrendezését igazítja a Hopper hardveréhez, a másik pedig az FP8-kvantált MoE-modellekhez megfelelő backend kiválasztását végzi. A vállalat célja, hogy ezekkel a módosításokkal javítsa a fused MoE modellek GPU-kihasználtságát, válaszidejét és teljesítményét az Amazon SageMaker AI környezetében.

A mátrixelrendezéshez kapcsolódó fejlesztés a vLLM 0.24.0 vagy újabb verziójában érhető el. A backendváltás konfigurációs módosítás, így az AWS szerint ehhez nem kellett kernelkódot változtatni.