Ce projet est encore en développement. Certaines fonctionnalités peuvent ne pas être encore implémentées et la documentation peut être incomplète.
Un système d'inférence LLM petit mais puissant, conçu pour la recherche.
Performances équivalentes à vLLM avec seulement 2 000 lignes de code (2 % de vLLM).
Il existe de nombreux frameworks open source pour le service LLM, notamment HuggingFace Transformers, vLLM, LightLLM, DistServe et DeepSpeed-MII. Pourquoi SwiftLLM ?
La raison en est que ces cadres sont conçus pour la production , plutôt que pour la recherche . Ils sont équipés de nombreuses fonctionnalités, telles que plus de 100 supports de modèles, divers supports matériels, LoRA, quantification, multimodal, mise en cache de préfixes, recherche de faisceaux, etc. Bien qu'il s'agisse d'une solution tout-en-un pour la production, leur base de code est trop volumineuse et complexe à comprendre et à modifier (par exemple, vLLM contient plus de 100 000 lignes de code), ce qui rend difficile son utilisation à des fins de recherche. Leur fardeau historique constitue également un problème.
SwiftLLM est conçu pour être un système d'inférence LLM petit mais puissant, conçu à des fins de recherche . « Minuscule » signifie qu'il ne conserve que les fonctionnalités essentielles à la recherche, « puissant » signifie qu'il ne fait aucun compromis sur les performances et enfin « rapide » signifie qu'il est facile à comprendre et à modifier. Tout en prenant en charge les fonctionnalités de base (voir la liste ci-dessous) et en étant capable d'atteindre des performances équivalentes à celles de vLLM, la base de code de SwiftLLM contient moins de 2 000 lignes de code (~ 2 % de vLLM), écrites en Python et OpenAI Triton (un DSL pour l'écriture Noyaux CUDA), ce qui le rend facile à lire, modifier, déboguer, tester, étendre et peut être facilement intégré à vos idées de recherche nouvelles et brillantes.
Actuellement, SwiftLLM prend en charge les fonctionnalités suivantes :
Et nous prévoyons d'ajouter la prise en charge des fonctionnalités suivantes à l'avenir :
Pour garder la base de code petite, nous ne prendrons pas en charge les fonctionnalités suivantes. Si vous souhaitez les utiliser dans votre projet de recherche, vous devrez peut-être les mettre en œuvre vous-même :
N'oubliez pas que SwiftLLM n'est PAS une solution tout-en-un pour la production. Il est conseillé de le considérer comme une « base » pour votre projet de recherche, et vous devrez peut-être implémenter certaines fonctionnalités par vous-même. Nous vous encourageons, mon cher chercheur, à lire le code, à le comprendre, à le modifier et à l'étendre pour l'adapter à vos besoins de recherche.
L'architecture de SwiftLLM peut être divisée en deux parties principales : le plan de contrôle et le plan de données .
En bref, le plan de contrôle décide « quoi calculer » ou « comment planifier », tandis que le plan de données décide « comment calculer » ou « comment implémenter » et effectue le calcul concret. Ils fonctionnent de manière maître-travailleur : le plan de contrôle agit comme un maître, qui effectue la planification et la coordination de haut niveau et envoie les tâches au plan de données, qui agit comme un travailleur, qui effectue le calcul de bas niveau.
Le code du plan de contrôle réside dans le répertoire swiftllm/server
, y compris des composants tels que Engine
, Scheduler
, le serveur API et TokenizationEngine
. Le code du plan de données réside dans le répertoire swiftllm/worker
, y compris les descriptions du graphe de calcul (dans swiftllm/worker/model.py
), l'implémentation des couches dans le modèle (dans swiftllm/layers
) et les noyaux OpenAI Triton ( vous pouvez imaginer des "noyaux" comme des fonctions exécutées sur le GPU) (dans swiftllm/kernels
).
Prenons comme exemple le serveur API jouet (situé dans swiftllm/server/api_server.py
) :
EngineConfig
pour créer un Engine
.Engine.initialize
, où il crée le Scheduler
, le TokenizationEngine
et un ensemble de travailleurs (actuellement un seul puisque le parallélisme Tensor n'est pas pris en charge). Ensuite, il ordonne au travailleur d'exécuter profile_num_blocks
pour calculer le nombre de blocs GPU, après quoi le moteur ordonne à tous les travailleurs d'allouer leur cache KV et leur échange KV.Engine.start_all_event_loops
. À chaque étape de la boucle, le moteur interroge le planificateur pour connaître le prochain lot de requêtes à calculer, ordonne au travailleur d'effectuer un échange d'entrée/sortie, puis envoie le lot au travailleur pour qu'il calcule. Actuellement, le plan de contrôle ( Engine
) et le plan de données ( LlamaModel
) résident sur le même nœud. Une fois le parallélisme tenseur/parallélisme pipeline mis en œuvre, le plan de données peut être distribué sur plusieurs nœuds.
Nous proposons deux manières d'utiliser SwiftLLM : en utilisant à la fois le plan de contrôle et le plan de données, ou en utilisant uniquement le plan de données.
Si votre idée est suffisamment simple ou élégante pour pouvoir être intégrée de manière transparente au plan de contrôle existant, vous pouvez utiliser à la fois le plan de contrôle et le plan de données. Dans un autre cas, où vous souhaitez implémenter une idée splendide, vous pouvez uniquement exploiter le plan de données et implémenter vous-même un nouveau plan de contrôle.
Commençons par configurer l'environnement :
packaging
via pip install packaging
Et puis vient l'installation :
git clone https://github.com/interestingLSY/swiftLLM.git
cd
dans le dépôt ( cd swiftLLM
) et installez d'autres dépendances via pip install -r requirements.txt
.pip install -e .
pour installer SwiftLLM dans votre environnement.pip install -e csrc
Voici quelques exemples :
.bin
et .safetensors
sont pris en charge. Supposons que le poids de votre modèle soit stocké dans /data/to/weight/
.python3 examples/offline.py --model-path /data/to/weight
. Cet exemple utilise uniquement le plan de données. Si vous envisagez d'utiliser SwiftLLM sans le plan de contrôle, c'est un bon point de départ.Engine
, vous pouvez essayer python3 examples/online.py --model-path /data/to/weight
. C'est un excellent exemple si vous envisagez d'utiliser à la fois le plan de contrôle et le plan de données.swiftllm/server/api_server.py
. Il lance un serveur API et fournit une interface de type vLLM pour le service en ligne. Bien qu'il soit petit (les plus petits peuvent aussi être adorables !), SwiftLLM ne fait aucun compromis sur les performances. Nous avons évalué SwiftLLM sur plusieurs scénarios et démontrons que SwiftLLM peut atteindre des performances équivalentes, voire meilleures, par rapport à vLLM.
Le premier scénario est "une seule opération avant", dans laquelle nous alimentons le modèle avec un lot d'entrées et le laissons générer un jeton de sortie (équivalent à une opération "avant"). Il s’agit de l’opération de base de l’inférence LLM (à la fois en ligne et hors ligne), ses performances sont donc cruciales.
Ici, nous utilisons le modèle LLaMA-3 7B avec le GPU NVIDIA A100 80G PCIE / RTX 4090 sous précision FP16. Les résultats sont présentés ci-dessous (plus le niveau est bas, mieux c'est) :
On peut voir que SwiftLLM peut atteindre des performances équivalentes (voire même surpasser) celles de vLLM avec les mêmes paramètres.
Le deuxième scénario est le « service en ligne », dans lequel nous démarrons un serveur API, échantillonnons des invites à partir d'un ensemble de données réel et laissons le modèle générer des complétions. C'est le scénario dans lequel LLM est utilisé dans des applications du monde réel telles que les chatbots ou la complétion de code.
Ici, nous utilisons l'ensemble de données ShareGPT pour échantillonner les invites et utilisons un processus poisson avec différents lambdas pour simuler différents taux d'arrivée de demandes. Les résultats sont présentés ci-dessous (plus le niveau est bas, mieux c'est) :
On peut voir que sur A100 80G PCIE, SwiftLLM peut atteindre des performances équivalentes à celles de vLLM, tandis que sur RTX 4090, SwiftLLM surpasse considérablement vLLM (principalement à cause de cela, notre plan de contrôle a une surcharge inférieure).