A Cohere dinamikus spekulatív dekódolással gyorsítaná az LLM-eket
A Cohere olyan dinamikus spekulatív dekódolási módszert mutatott be, amely a hardver terheléséhez igazítja, hány tokent készítsen elő egy kisebb modell. A vállalat szerint a megoldás a rögzített beállítású módszernél magas kötegméretnél is jobb teljesítményt adhat.
- A DSD a hardverterhelés alapján változtatja a vázlattokenek számát.
- Sűrű modelleknél az optimális K a kötegméret növekedésével csökken.
- A Command A modellen a DSD BS 128 és BS 256 mellett körülbelül 23 százalékkal gyorsabb volt a fix K=3 módszernél.
- A Cohere a fejlesztést a vLLM-be is hozzájárult.
A spekulatív dekódolás működése
A nagy nyelvi modellek következtetése során a rendszer jellemzően egyszerre egy tokent generál. A spekulatív dekódolás ezt úgy gyorsítja, hogy egy kisebb, úgynevezett vázlatmodell több tokent javasol, amelyeket a nagy célmodell egyetlen lépésben ellenőriz.
Ha a célmodell elfogadja a javaslatok egy részét, egy lépésben több token készülhet el. A módszer a grafikus processzor számítási kapacitása és memória-sávszélessége közötti különbséget használja ki. Kis kötegméretnél a modell súlyainak memóriából való betöltése jelenti a fő korlátot, ezért a számítási egységek jelentős része kihasználatlan maradhat. A spekulatív dekódolás ilyenkor több tokent küld ellenőrzésre a célmodellhez.
Nagyobb kötegméretnél azonban a mátrixszorzások válnak meghatározóvá, vagyis a rendszer számítási korlátba ütközik. Ebben a tartományban a spekulatív dekódolás akár lassabb is lehet a hagyományos következtetésnél.
A fix tokenmennyiség problémája
A hagyományos spekulatív dekódolásban a vázlatmodell által előállított tokenek száma, a K értéke rögzített. Ez termelési rendszerekben korlátozza a módszer használhatóságát, mert a kötegméret folyamatosan változhat, és ritkán marad alacsony.
A Cohere szerint a probléma a megerősítéses tanulásnál is jelentős. A tanítási folyamat gördítési szakasza, amelyben a rendszer a tanításhoz szükséges adatokat és jutalomjeleket állítja elő, akár az erőforrások 85 százalékát is felhasználhatja. A gondolkodó modellek hosszú, elnyújtott generálásokat produkálhatnak, amelyekben egyetlen kérés sokáig fut, és feltartja a köteg többi elemét. A spekulatív dekódolás segíthetne az ilyen generálásoknál, de nagy kötegméretnél romolhat az áteresztőképessége.
A DSD a hardverhez igazítja a K értékét
A Hardware-Aware Dynamic Speculative Decoding, röviden DSD, futás közben változtatja a vázlatmodell által előállított tokenek számát. A Cohere leírása szerint memóriasávszélesség által korlátozott működésnél növeli K értékét, számítási korlátnál pedig csökkenti.
Sűrű modelleknél ez általában azt jelenti, hogy alacsony kötegméretnél nagyobb, magas kötegméretnél kisebb K az optimális érték. A kötegméret növekedésével az optimális K monoton csökken. A szakértői keverék, vagyis MoE modelleknél más a helyzet. Alacsony kötegméretnél a megerősítés további szakértői komponensek betöltését okozhatja, ezért K alacsony értékről indul. Közepes kötegméretnél növekedhet, mivel ekkor már szinte minden szakértő betöltődött, magas kötegméretnél pedig ismét csökken, amikor a rendszer számítási korlátba ütközik.
Az optimális értéket a Cohere az úgynevezett jó teljesítmény, vagyis goodput alapján választja ki. Ennek egyszerűsített képlete az elfogadási hossz, az AL, osztva a tokenek közötti késleltetéssel, az ITL-lel. A vállalat offline profilozással méri ezeket, majd a kiválasztott K értékeket keresőtáblában tárolja. A táblázat később a futásidejű AL- és ITL-statisztikákkal is kiegészíthető, így változó munkaterheléshez alkalmazkodhat.
Mérési eredmények és vLLM-integráció
A Cohere az MT-Bench adathalmazon hasonlította össze a hagyományos, spekulatív dekódolás nélküli beállítást, a fix K=3 értékkel működő spekulatív dekódolást és a dinamikus K-t használó DSD-t. A minták számát 20 × BS értékre növelték, ami 20, egyenként BS méretű hullámot jelentett.
A Command A sűrű modellnél a DSD alacsony kötegméretnél elérte a K=3 beállítású spekulatív dekódolás gyorsulását. BS 64 és 128 mellett gyorsabb volt a másik két beállításnál, BS 256 mellett pedig a hagyományos modell tokenenkénti, felhasználónkénti teljesítményét érte el. A Cohere számai szerint DSD BS 128 és BS 256 mellett egyaránt körülbelül 23 százalékkal volt gyorsabb a fix K=3 módszernél. A hagyományos beállításhoz képest BS 128 mellett 7,5 százalékos, BS 256 mellett 1,82 százalékos gyorsulást mértek.
A Command A+ MoE modellnél a DSD és a rögzített beállítás hasonló gyorsulást hozott. A DSD a kötegmérettartomány nagy részében K=3 értéket választott, BS 16 és 32 között pedig K=5-re váltott. Ez nem eredményezett további gyorsulást, amit a Cohere azzal magyaráz, hogy az EAGLE fejlesztőfejet egyetlen időlépésre tanították, majd több lépésben használták újra. A vállalat szerint az EAGLE-3 vagy a DFlash típusú újabb módszerek ebben a helyzetben egyértelműbb javulást mutathatnak.
A Cohere az év elején a vLLM-be is hozzájárult a fejlesztéssel. A vállalat szerint az integráció kihívását az jelentette, hogy a DSD-nek együtt kell működnie az olyan optimalizációkkal, mint az aszinkron ütemezés és a teljes CUDA Graph használata. Az aszinkron ütemezés a következő időlépés CPU-s ütemezését átfedi az aktuális GPU-s modellfuttatással. A DSD azonban eltérő számú vázlattokent generálhat és ellenőrizhet az egyes időlépésekben, ezért módosítja a vLLM korábbi feltételezését.
