LLM Inference Optimization - AI Engineering

Bienvenue à tous. Nous avons abordé le neuvième chapitre du livre d'ingénierie IA de Chip, axé sur l'optimisation de l'inférence. Ce chapitre marque un changement de perspective par rapport aux chapitres précédents, qui se concentraient sur la qualité des générations du modèle. Désormais, l'objectif est de rendre le modèle rapide et économique pour les utilisateurs et les clients. Il existe deux perspectives principales : celle de l'utilisateur, qui privilégie la rapidité, et celle du fournisseur, qui cherche à réduire les coûts tout en maintenant une vitesse raisonnable. L'inférence est distincte de l'entraînement. L'entraînement, impliquant la rétropropagation, les optimiseurs comme Adam, et la gestion de la mémoire pour les activations et les moments, est beaucoup plus exigeant. Ce chapitre se concentre uniquement sur la phase d'inférence, où le modèle est figé. Si l'on utilise une API, la plupart des préoccupations liées à l'optimisation sont gérées par le fournisseur. Cependant, si l'on héberge son propre modèle, que ce soit pour des raisons de coût, d'exécution locale ou de personnalisation, il devient crucial de comprendre les facteurs d'optimisation. Deux goulots d'étranglement principaux limitent la vitesse d'inférence : le calcul et la bande passante mémoire. Le calcul implique un grand nombre d'opérations en virgule flottante, comme les multiplications de matrices. Un modèle avec des milliards de paramètres nécessite des téraoctets de poids à charger pour chaque jeton généré. La bande passante mémoire devient alors un facteur limitant. L'intensité arithmétique, définie comme le rapport entre les opérations en virgule flottante et les octets chargés, est un concept clé. Une faible intensité arithmétique, typique du décodage où des opérations vectorielles sont effectuées, rend le système limité par la bande passante mémoire. À l'inverse, lors du pré-remplissage, où de grandes matrices sont multipliées, l'intensité arithmétique est plus élevée et le calcul devient le goulot d'étranglement. Le graphique "roofline" illustre cette relation : à faible intensité arithmétique, la performance est limitée par la bande passante mémoire ; au-delà d'un certain seuil, elle est limitée par la capacité de calcul. Comprendre son propre profil d'intensité arithmétique permet d'identifier le goulot d'étranglement. Contrairement aux modèles plus anciens qui produisaient leur sortie en un seul passage, les LLM génèrent un jeton à la fois, ce qui allonge considérablement le temps d'inférence pour des sorties plus longues. Le processus d'inférence se décompose en deux phases : le pré-remplissage (traitement du prompt initial et mise en cache des clés-valeurs) et le décodage (génération de jetons un par un). Le pré-remplissage, impliquant des multiplications matricielles importantes, est limité par le calcul. Le décodage, impliquant des multiplications de matrices par des vecteurs, est limité par la bande passante mémoire. Cette différence explique la tarification distincte des jetons d'entrée et de sortie par les fournisseurs d'API. Concernant les métriques, celles axées sur l'utilisateur incluent la latence, le temps jusqu'au premier jeton (TTFT) et le temps par jeton de sortie (TPOT). Le TTFT représente le temps de traitement du prompt, tandis que le TPOT mesure le temps entre chaque jeton généré. La latence totale est la somme du TTFT et de la somme des TPOT. Le débit mesure la quantité de travail traitée par unité de temps, mais le "goodput" est une métrique plus pertinente car elle prend en compte la qualité de la réponse, en éliminant les réponses lentes ou insatisfaisantes. D'un point de vue technique, les fournisseurs utilisent souvent les métriques MFU (Model Flops Utilization) et MBU (Model Bandwidth Utilization) pour évaluer l'utilisation du calcul et de la bande passante mémoire. Il est important de noter que les chiffres de performance annoncés par les fabricants de matériel sont souvent optimisés pour des charges de travail spécifiques et difficiles à atteindre en pratique. Un MFU ou MBU de 50% est souvent considéré comme excellent. Des accélérateurs d'IA tels que les GPU et les TPU sont essentiels pour l'inférence. Les GPU, avec leur mémoire à large bande passante et leur SRAM proche des cœurs de calcul, gèrent plusieurs threads et utilisent des caches pour optimiser la vitesse. Les TPU, quant à eux, privilégient une allocation de l'espace sur la puce pour le calcul et s'appuient sur des compilateurs intelligents pour ordonnancer les opérations de manière prédéterminée. Les architectures plus récentes, comme les TPU v8, se rapprochent des GPU en offrant plus de flexibilité. La mémoire est hiérarchisée : la SRAM sur la puce est la plus rapide, suivie de la mémoire à large bande passante des GPU, puis de la DRAM des CPU. Décharger des données de la mémoire GPU vers la mémoire principale ralentit considérablement le modèle. Des techniques comme la quantification permettent de réduire la taille des modèles, par exemple en passant de 16 bits à 4 bits, avec une perte de performance souvent minime grâce à des méthodes intelligentes qui ajustent la précision par couche. Le cache KV (clé-valeur) est crucial pour l'inférence des LLM. Sa taille augmente linéairement avec la longueur du contexte et la taille du lot. Des techniques comme l'attention à fenêtre glissante, l'attention linéaire (Mamba, RWKV) et le partage de cache KV visent à réduire la consommation mémoire et à accélérer le calcul. L'attention paginée est une innovation majeure qui virtualise la mémoire GPU, résolvant les problèmes de fragmentation et d'allocation inefficace de la mémoire. Le traitement par lots est également optimisé. Le traitement par lots continu, où les étapes de décodage sont considérées comme des unités de travail distinctes, permet une meilleure utilisation des ressources. Le "ragged batching" évite l'ajout de jetons de remplissage inutiles en masquant intelligemment les données, réduisant ainsi les calculs superflus. Les modèles Mixture of Experts (MoE) sont conçus pour accélérer les couches MLP en répartissant les experts sur plusieurs GPU, bien que cela puisse entraîner une sous-utilisation si le nombre d'utilisateurs est faible. Le parallélisme de données, le parallélisme de pipeline et le parallélisme d'experts sont combinés pour servir de très grands modèles sur des centaines de GPU. Le décodage spéculatif permet de traiter plusieurs jetons en une seule passe, accélérant ainsi l'inférence. Les "custom kernels" et les compilateurs intelligents permettent d'optimiser des opérations spécifiques qui ne sont pas nativement prises en charge par le matériel. Enfin, il est important de noter que l'exécution du même calcul sur du matériel différent peut produire des résultats légèrement différents en raison de la précision finie des nombres à virgule flottante. Cela souligne l'importance des évaluations rigoureuses après toute modification ou optimisation du modèle.