Présentation détaillée du projet TensorRT-LLM

Aperçu du projet

TensorRT-LLM est une bibliothèque open-source développée par NVIDIA, spécialement conçue pour optimiser les performances d'inférence des grands modèles de langage (LLM) sur les GPU NVIDIA. Elle offre une API Python facile à utiliser pour définir les grands modèles de langage et prend en charge les techniques d'optimisation les plus avancées pour exécuter efficacement l'inférence sur les GPU NVIDIA.

Fonctionnalités clés

1. Techniques d'optimisation avancées

TensorRT-LLM offre une variété de fonctionnalités d'optimisation avancées, notamment :

• Noyaux d'attention personnalisés : Implémentations optimisées spécifiquement pour les mécanismes d'attention.
• Traitement par lots dynamique (Inflight Batching) : Traitement en temps réel de séquences d'entrée de différentes longueurs.
• Mise en cache KV paginée : Gestion efficace du cache clé-valeur.
• Décodage spéculatif (Speculative Decoding) : Accélération de la génération en prédisant plusieurs tokens.
• Prise en charge de diverses quantifications : FP8, FP4, INT4 AWQ, INT8 SmoothQuant, etc.

2. Détails des techniques de quantification

TensorRT-LLM fournit une boîte à outils de quantification unifiée, leader de l'industrie, qui accélère considérablement le déploiement du deep learning/IA générative sur le matériel NVIDIA, tout en maintenant la précision du modèle.

Principales méthodes de quantification :

• FP8 : Offre généralement les meilleures performances et précision dans les scénarios d'inférence à grands lots, adaptée aux tailles de lot ≥ 16.
• INT8 SmoothQuant : Lissage des poids et quantification INT8 par canal, étalonnage des plages d'activation au niveau du tenseur.
• INT4 AWQ : Recalibrage des poids et quantification par bloc en INT4, recommandé pour les scénarios d'inférence à petits lots (taille de lot ≤ 4).
• W4A8 AWQ : Quantification des poids en INT4, quantification des activations en INT8.

Améliorations des performances :

Selon les benchmarks, les techniques de quantification peuvent apporter des améliorations significatives des performances :

• Quantification FP8 : Accélération de 1,45x pour le modèle Llama 3 8B et de 1,81x pour le modèle 70B, par rapport à la référence FP16.
• INT4 AWQ : Jusqu'à 2,66x d'amélioration des performances pour le modèle 70B dans les scénarios avec une taille de lot de 1.
• Optimisation de la mémoire : Toutes les versions quantifiées du modèle Llama 3 70B peuvent fonctionner sur un seul GPU NVIDIA H100, tandis que la précision FP16 nécessite au moins deux GPU.

3. Prise en charge multi-GPU et multi-nœuds

TensorRT-LLM inclut des étapes de pré-traitement et post-traitement ainsi que des primitives de communication multi-GPU multi-nœuds, permettant des performances d'inférence LLM révolutionnaires grâce à une API de définition de modèle open-source simple.

4. Large prise en charge matérielle

TensorRT-LLM prend en charge les GPU basés sur les architectures NVIDIA Hopper, NVIDIA Ada Lovelace et NVIDIA Ampere. En particulier :

• GPU H100 : Prend en charge la conversion automatique et les noyaux optimisés pour le format FP8.
• GPU H200 : Peut atteindre près de 12 000 tokens/seconde sur Llama2-13B.
• Série RTX : Prend en charge l'inférence de grands modèles sur les GPU grand public.

Installation et utilisation

Installation Docker (recommandée)

# Exécuter le conteneur Docker pré-construit
docker run --ipc host --gpus all -it nvcr.io/nvidia/tensorrt-llm/release

Exemple d'utilisation de l'API LLM

from tensorrt_llm import BuildConfig, SamplingParams
from tensorrt_llm._tensorrt_engine import LLM

def main():
    build_config = BuildConfig()
    build_config.max_batch_size = 256
    build_config.max_num_tokens = 1024
    
    # Prend en charge les noms de modèles HuggingFace, les chemins de modèles HF locaux ou les checkpoints de quantification de l'optimiseur de modèles TensorRT
    llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", 
              build_config=build_config)
    
    # Exemple de prompts
    prompts = [
        "Hello, my name is",
        "The capital of France is",
        "The future of AI is",
    ]
    
    # Créer des paramètres d'échantillonnage
    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
    
    for output in llm.generate(prompts, sampling_params):
        print(f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}")

Déploiement de service en ligne

# Lancer un serveur compatible OpenAI
trtllm-serve --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --port 8000

Flux de travail de quantification

Commandes de quantification de base

# Quantification FP8
python quantize.py --model_dir $MODEL_PATH --qformat fp8 --kv_cache_dtype fp8 --output_dir $OUTPUT_PATH

# Quantification INT4 AWQ
python quantize.py --model_dir $MODEL_PATH --qformat int4_awq --awq_block_size 64 --tp_size 4 --output_dir $OUTPUT_PATH

# Quantification INT8 SmoothQuant
python quantize.py --model_dir $MODEL_PATH --qformat int8_sq --kv_cache_dtype int8 --output_dir $OUTPUT_PATH

# Quantification automatique (combinaison de plusieurs méthodes)
python quantize.py --model_dir $MODEL_PATH --autoq_format fp8,int4_awq,w4a8_awq --output_dir $OUTPUT_PATH --auto_quantize_bits 5 --tp_size 2

Modèles pris en charge

TensorRT-LLM prend en charge un grand nombre d'architectures LLM populaires, y compris, mais sans s'y limiter :

• Série Llama : Llama 2, Llama 3, Llama 3.1, Llama 3.3
• Série Falcon : Y compris Falcon-180B
• Série GPT : Architectures liées à ChatGPT
• Série Gemma : Modèles open-source de Google
• Série Mixtral : Modèles à experts mixtes
• Série DeepSeek : Y compris DeepSeek R1
• CodeLlama : Modèle dédié à la génération de code

Intégration de l'écosystème

Écosystème NVIDIA

• NVIDIA NeMo : Cadre de bout en bout pour la construction, la personnalisation et le déploiement d'applications d'IA générative.
• Triton Inference Server : Serveur d'inférence de qualité production.
• NVIDIA Dynamo : Cadre de service d'inférence distribuée à l'échelle du centre de données.

Intégrations tierces

• HuggingFace Hub : Fournit des modèles pré-quantifiés.
• LlamaIndex : Développement d'applications RAG.
• SageMaker LMI : Inférence gérée par AWS.

Benchmarks de performance

Exemples d'améliorations de performance :

• Par rapport aux plateformes CPU : Vitesse d'inférence multipliée par 36.
• Par rapport aux RTX sans optimisation : Vitesse des grands modèles de langage multipliée par 4 sur les plateformes Windows RTX.
• Falcon-180B : Inférence sur un seul GPU H200 en utilisant INT4 AWQ.
• Llama-70B : 6,7 fois plus rapide que l'A100.

Conseils de bonnes pratiques

Choix de la méthode de quantification

Choisir la méthode de quantification appropriée en fonction des différents scénarios :

1. Inférence à petits lots (taille de lot ≤ 4) :

   • Recommandation : Utiliser des méthodes de quantification des poids (comme INT4 AWQ).
   • Considération principale : Limitations de la bande passante mémoire.

2. Inférence à grands lots (taille de lot ≥ 16) :

   • Priorité : Choisir la quantification FP8, qui offre généralement les meilleures performances et précision.
   • Si les résultats ne sont pas satisfaisants : Essayer INT8 SmoothQuant, puis AWQ et/ou GPTQ.

3. Applications spécifiques à un domaine :

   • Pour les applications hautement spécialisées comme la complétion de code, il est recommandé d'utiliser des jeux de données spécifiques au domaine pour l'étalonnage.

Avantages techniques

1. Facilité d'utilisation : Fournit une API Python de haut niveau, simplifiant la définition et l'optimisation des LLM.
2. Performance : Inclut toutes les techniques d'optimisation courantes, telles que la fusion de noyaux, la quantification, l'optimisation d'exécution, etc.
3. Évolutivité : Prend en charge tous les scénarios de déploiement, du GPU unique aux configurations multi-nœuds.
4. Compatibilité : Intégration profonde avec PyTorch, prend en charge les principaux écosystèmes d'inférence.
5. Open-source : Entièrement open-source, développement continu axé sur la communauté.

Développement futur

TensorRT-LLM améliore la facilité d'utilisation et l'évolutivité grâce à une API de définition de modèle modulaire open-source, permettant de définir, d'optimiser et d'exécuter de nouvelles architectures et fonctionnalités, et de s'adapter facilement à l'évolution des LLM.

Les directions de développement continu du projet incluent :

• Prise en charge de plus d'architectures de modèles.
• Techniques de quantification plus avancées.
• Meilleure évolutivité multi-nœuds.
• Intégration plus étroite de l'écosystème.