Bienvenue dans le monde de la pré-tokenisation à très haute vitesse. Si vous entraînez ou affinez des modèles de langage, vous savez que la préparation des données et la tokenisation peuvent vite devenir un point de blocage frustrant. C'est précisément pour résoudre ce problème que Marcel Rød a développé Gigatoken, une bibliothèque open source écrite en Rust avec des liaisons Python, capable de tokeniser du texte à une vitesse mesurée en gigaoctets par seconde.
La plupart des tokeniseurs actuels s'appuient sur des moteurs d'expressions régulières pour découper le texte en sous-mots avant d'appliquer l'algorithme Byte-Pair Encoding. Bien que des outils comme HuggingFace Tokenizers ou Tiktoken soient déjà optimisés et exécutés en multithread via Rust, ils conservent un surcoût non négligeable. Lorsque vous devez traiter des terabytes ou des petabytes de texte brut issus du Web, ce temps de calcul s'accumule et ralentit considérablement les pipelines de données.
Gigatoken réinvente ce processus en appliquant des optimisations bas niveau poussées. Au lieu de s'en remettre à un moteur Regex classique pour la pré-tokenisation, la bibliothèque utilise des instructions vectorielles SIMD adapées aux architectures x86 modernes et ARM.
L'outil gère le cache et les threads avec brio. Un système de hiérarchie de cache conserve les correspondances des mots déjà croisés, évitant de recalculer l'encodage des termes fréquents. Les threads travaillent de manière quasi indépendante avec un minimum de synchronisation et la lecture des fichiers peut s'effectuer directement en Rust pour contourner les lenteurs de l'interpréteur Python.
Les résultats mesurés sur différents processeurs sont impressionnants. Sur une machine équipée d'AMD EPYC à 144 cœurs, Gigatoken atteint des débits d'encodage dépassant 24 Go/s sur des modèles comme GPT-2 ou Llama 3, soit une accélération de près de 1000 fois par rapport aux implémentations standards. Même sur une puce Apple M4 Max, le débit dépasse régulièrement 8 Go/s. À ce rythme, la totalité du corpus Common Crawl, estimé à environ 130 billions de tokens, pourrait être tokenisée en moins de six heures et demie.
Adopter Gigatoken ne nécessite pas de réécrire l'ensemble de votre code. L'outil s'installe d'une simple commande pip install gigatoken et propose un mode de compatibilité ascendante.
Pour une transition sans effort, vous pouvez encapsuler votre tokeniseur HuggingFace ou Tiktoken existant via les méthodes de conversion intégrées. Si vous recherchez les performances maximales, l'API native permet de lire directement les fichiers texte depuis le disque en Rust, minimisant ainsi tout passage de structures de données lourdes entre Python et le moteur sous-jacent.
Gigatoken montre qu'il reste une marge de progression considérable dans l'optimisation des outils de traitement du langage naturel. En tirant pleinement parti du matériel et des architectures de processeurs actuels, ce projet transforme une tâche autrefois chronophage en une étape quasi instantanée.