A tokenizers könyvtár következő nagy verziója (v1) alapvetően a teljesítményre összpontosít: a cél az volt, hogy a tokenizálás könnyű és a munkafolyamatokhoz skálázódó legyen, és hogy a GPU-k sose várakozzanak a CPU által végzett tokenizálás miatt. A kiadási jelöltet a fejlesztők részletes mérésekkel hasonlították össze a korábbi v0.23 verzióval és más népszerű alternatívákkal.
Miért számít a tokenizálás teljesítménye
Bár hagyományosan a tokenizálás számításigénye elenyésző volt a modellezéshez képest, a modellek felgyorsulása és a munkaterhelések növekedése megváltoztatta az egyensúlyt. Nagy adatkészleteken történő tréning, sok egyidejű kérés kiszolgálása vagy hosszú bemenetek ismételt feldolgozása esetén a tokenizálás lehet a modell táplálását akadályozó tényező.
Milyen eredmények születtek
A fejlesztők mérési eredményei szerint a tokenizers v1 encode útvonala tíz modellcsaládon át vizsgálva egy Apple M4 Max gépen egy szálon 3–30-szor gyorsabb, mint a v0.23. A legalacsonyabb javulás a t5-base, a legnagyobb a gpt2 esetén volt. Több szál használatakor a skálázódás körülbelül 76%-os volt lineáris ideálhoz képest nyolc munkáson.
A v1 változtatások célja továbbá az volt, hogy az output tokenazonosítók, az API, a szókincs és a merge-rankek változatlanok maradjanak: a v1 ugyanazokat az ID-ket állítja elő, mint a v0.23.
Melyik részeken dolgoztak
A tokenizálás négy szakaszban történik: normalizálás, pre-tokenizálás (split), modell (merge/átalakítás tokenokra és ID-kra) és post-processzálás. A v1 fejlesztései mind a négy szakaszra kiterjedtek, de különösen a modell lépésnél hoztak jelentős optimalizációkat. A főbb változtatások:
- Workspace split: a korábbi egyetlen crate több darabra oszlott (tk-encode, tk-serialize, tk-convert, tk-train), így a futtatáshoz csak az szükséges részek kerülnek linkelésre.
- No-alloc model: a merge munkamennyiség egy hívó által birtokolt scratch puffert használ, így a ciklus nem érinti az allokátort.
- Bitcannon: a szétválasztási mintázat Boolean műveletekké alakult bitfolyamokon, SIMD utasításokkal, a regex motor helyett.
- Merge-loop átírás: a merge során keletkező darabok intruzív kétszintű láncolt listaként egy előre foglalt pufferen belül élnek; így egy merge két index frissítésével jár, ahelyett hogy adatot mozgatna.
- Word cache: egy thread-local gyorsítótár térképezi a pre-token bájtokat kész ID-kra, így egy ismétlődő szó feldolgozása egyszeri merge-dzsel megúszható.
- Native parallelism: egy megosztott tokenizer több szálról egyszerre kódol; minden szál a saját al-pooljából kap scratch puffert és word cache-et, megszüntetve a globális lock alatti sorban állást.
A split: bitfolyamok a regex helyett
A BPE-alapú modellek a bemenetet pre-tokenökre bontják, és korábban ezt egy fix reguláris kifejezés végezte, ami minden encode hívásnál a regex-motort foglalkoztatta. Mivel a minta modell-specifikus és futásidőben változatlan, kézzel írt implementációval helyettesíthető. A bitcannon megközelítése a bemenet bájtjait párhuzamos bitfolyamokként kezeli és 64 bájtonként SIMD regiszter-műveletekkel találja meg a határokat, így a határok boolean műveletekből adódnak, nem karakterenkénti átjárásból. Ez nagy gyorsulást eredményez, ha a tokenizer mintája a támogatott grammák közé tartozik; ha nem, a regex-út marad és a gyorsulás kiesik.
A word cache
A valós szövegben sok szó ismétlődik. Mivel a BPE egy adott pre-tokenre mindig ugyanazt az ID-sort adja, v1 eltárolja egy thread-local cache-ben a már feldolgozott pre-token bájtokhoz tartozó ID-kat. Későbbi előfordulások így kihagyhatják a költséges merge folyamatot. A cache előnyei a bemenet ismétlődési mintázatától függenek: sok ismétlődésnél jelentős megtakarítás érhető el; kevés ismétlődésnél a lookup költségei ellensúlyozhatják az előnyt.
A merge-ciklus átírása
A korábbi implementáció minden pre-tokenre új memóriát allokált és új prioritási sort épített. V1-ben a modell újrahasználja a hívó által biztosított scratch puffert, a szimbólumokat sík tömbben tárolja és a szomszédos szimbólumokat a tömbindexekkel kapcsolja össze. A jelöltek párok 64 bites értékbe vannak csomagolva, amelynek felső bitjeiben a merge-rank található; így két jelölt összehasonlítása egyszerű egész-szám összehasonlítássá válik, mellőzve a feltételes elágazásokat.
Mérések metodikája
A fejlesztők részletes szabályokat alkalmaztak, hogy a különböző motorok összehasonlítása konzisztens legyen: egyetlen időzítő ciklus minden motornál, a szókincs betöltése külön van időzítve, az output ID-k FNV-1a hash-einek egyezését ellenőrzik, mediánok csak a közösen lefutott sejtekre vonatkoznak, minden ismétlés új folyamatban fut, munkavállalókat fizikai magokra tűzték rá, és külön munkák mérik a host-host variációt. Különbség van az ismételt egy dokumentum kódolása és a különböző dokumentumok sorozatos kódolása között: az előbbi a cache teljes kihasználtságát mérheti, az utóbbi az új szavak megjelenését is.
Összegzés: mi adta a nyereséget
A 3–30× gyorsulás és a többközpontú skálázás több különböző optimalizáció együttműködésének eredménye: a kézzel írt splitter a regex helyett, a thread-local word cache, az allocator-mentes merge-ciklus, és az elő-tokenek kötegelt feldolgozása egyetlen modellhívásban. Ezek mindegyike a tokenizáció különböző pontjain csökkenti a végzett munkát.
Mi következik
A következő prioritás több modellcsalád támogatása: további modelleket terveznek átmozgatni az új merge-ciklusra még 1.0.0 előtt. A cél, hogy a transformers és az ökoszisztéma többi része is profitáljon ezekből a fejlesztésekből. Később, a 1.0.0 után további lehetőségeket vizsgálnak, például GPU-alapú kódolást és kötegelt dekódolást (tok-devices), amely opcionális komponensként nagy kötegekhez lenne hasznos.
Hogyan próbálható ki
A v1 release candidate elérhető a crates.io-on (cargo add tokenizers --pre). Az API változatlan, tehát a meglévő hívások ugyanazt az interfészt használják. A trainekhez tartozó C++ függőség alapértelmezettként be van kapcsolva; ha csak encode-olásra van szükség, ki lehet kapcsolni a training részt (–no-default-features --features http). Az encode és encode_batch hívások maradtak, az encode_batch skálázódása szolgáltatja a többmagos méréseket.
Köszönetnyilvánítás
A munkát az aktív nyílt forrású tokenizálási projektek (például gigatoken, tiktoken, kitoken, tokie, fastokens, wordchipper, ai-tokenizer és sok más) eredményeire építve végezték. IBM, NVIDIA és az ExecuTorch csapat hozzájárulásait és hardvertesztjeit külön megköszönték, mert ezek szélesebb platformtámogatást tettek lehetővé.
A poszt és az abban szereplő ábrák a tokbench mérési eredményeiből származnak, és a támogatás bővülésével frissítik őket.
