Introducción Detallada al Proyecto TensorRT-LLM

Resumen del Proyecto

TensorRT-LLM es una biblioteca de código abierto desarrollada por NVIDIA, diseñada específicamente para optimizar el rendimiento de inferencia de los Modelos de Lenguaje Grandes (LLM) en GPUs de NVIDIA. Ofrece una API de Python fácil de usar para definir LLM y soporta las técnicas de optimización más avanzadas para ejecutar inferencias de manera eficiente en GPUs de NVIDIA.

Características Principales

1. Técnicas de Optimización Avanzadas

TensorRT-LLM proporciona diversas funciones de optimización avanzadas, que incluyen:

• Kernels de atención personalizados: Implementaciones de mecanismos de atención especialmente optimizados.
• Procesamiento por lotes dinámico (Inflight Batching): Procesamiento en tiempo real de secuencias de entrada de diferentes longitudes.
• Caché KV paginado: Gestión eficiente de la caché de claves y valores.
• Decodificación especulativa: Acelera la generación prediciendo múltiples tokens.
• Soporte para múltiples cuantificaciones: FP8, FP4, INT4 AWQ, INT8 SmoothQuant, entre otros.

2. Detalles de las Técnicas de Cuantificación

TensorRT-LLM ofrece un conjunto de herramientas de cuantificación unificado líder en la industria, que acelera significativamente el despliegue de Deep Learning/IA Generativa en hardware de NVIDIA, manteniendo la precisión del modelo.

Principales métodos de cuantificación:

• FP8: Generalmente ofrece el mejor rendimiento y precisión en escenarios de inferencia con lotes grandes, adecuado para tamaños de lote ≥16.
• INT8 SmoothQuant: Suavizado de pesos y cuantificación de canales INT8, con escalado a nivel de tensor para el rango de activación.
• INT4 AWQ: Reescalado de pesos y cuantificación a nivel de bloque a INT4, recomendado para escenarios de inferencia con lotes pequeños (tamaño de lote ≤4).
• W4A8 AWQ: Cuantificación de pesos a INT4 y de activaciones a INT8.

Mejora del rendimiento:

Según las pruebas de rendimiento, las técnicas de cuantificación pueden ofrecer mejoras significativas:

• Cuantificación FP8: En comparación con la base FP16, el modelo Llama 3 8B puede obtener una aceleración de 1.45x, y el modelo 70B una aceleración de 1.81x.
• INT4 AWQ: En escenarios con tamaño de lote de 1, el modelo 70B puede lograr una mejora de rendimiento de hasta 2.66x.
• Optimización de memoria: Todas las versiones cuantificadas del modelo Llama 3 70B pueden ejecutarse en una sola GPU NVIDIA H100, mientras que la precisión FP16 requiere al menos dos GPUs.

3. Soporte Multi-GPU y Multi-Nodo

TensorRT-LLM incluye pasos de preprocesamiento y postprocesamiento, así como primitivas de comunicación multi-GPU y multi-nodo, logrando un rendimiento de inferencia de LLM innovador a través de una API de definición de modelos de código abierto sencilla.

4. Amplio Soporte de Hardware

TensorRT-LLM es compatible con GPUs basadas en las arquitecturas NVIDIA Hopper, NVIDIA Ada Lovelace y NVIDIA Ampere. En particular:

• GPU H100: Soporta la conversión automática y los kernels optimizados para el formato FP8.
• GPU H200: Puede alcanzar casi 12,000 tokens/segundo en Llama2-13B.
• Serie RTX: Soporta la inferencia de modelos grandes en GPUs de consumo.

Instalación y Uso

Instalación con Docker (Recomendado)

# Ejecutar el contenedor Docker preconstruido
docker run --ipc host --gpus all -it nvcr.io/nvidia/tensorrt-llm/release

Ejemplo de Uso de la API de 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
    
    # Soporta nombres de modelos de HuggingFace, rutas de modelos HF locales o puntos de control de cuantificación del optimizador de modelos TensorRT
    llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", 
              build_config=build_config)
    
    # Ejemplos de prompts
    prompts = [
        "Hello, my name is",
        "The capital of France is",
        "The future of AI is",
    ]
    
    # Crear parámetros de muestreo
    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}")

Despliegue de Servicio en Línea

# Iniciar un servidor compatible con OpenAI
trtllm-serve --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --port 8000

Flujo de Trabajo de Cuantificación

Comandos Básicos de Cuantificación

# Cuantificación FP8
python quantize.py --model_dir $MODEL_PATH --qformat fp8 --kv_cache_dtype fp8 --output_dir $OUTPUT_PATH

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

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

# Cuantificación automática (combinación de varios métodos)
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

Modelos Compatibles

TensorRT-LLM soporta una gran cantidad de arquitecturas de LLM populares, incluyendo, entre otras:

• Serie Llama: Llama 2, Llama 3, Llama 3.1, Llama 3.3
• Serie Falcon: Incluyendo Falcon-180B
• Serie GPT: Arquitecturas relacionadas con ChatGPT
• Serie Gemma: Modelos de código abierto de Google
• Serie Mixtral: Modelos de expertos mixtos
• Serie DeepSeek: Incluyendo DeepSeek R1
• CodeLlama: Modelo especializado en generación de código

Integración del Ecosistema

Ecosistema NVIDIA

• NVIDIA NeMo: Marco de trabajo de extremo a extremo para construir, personalizar y desplegar aplicaciones de IA generativa.
• Triton Inference Server: Servidor de inferencia de grado de producción.
• NVIDIA Dynamo: Marco de servicio de inferencia distribuida a escala de centro de datos.

Integración de Terceros

• HuggingFace Hub: Proporciona modelos precuantificados.
• LlamaIndex: Desarrollo de aplicaciones RAG.
• SageMaker LMI: Inferencia gestionada por AWS.

Puntos de Referencia de Rendimiento

Ejemplos de mejora del rendimiento:

• En comparación con plataformas CPU: Velocidad de inferencia aumentada hasta 36 veces.
• En comparación con RTX sin optimización: Velocidad de LLM aumentada hasta 4 veces en plataformas Windows RTX.
• Falcon-180B: Inferencia en una sola GPU H200 utilizando INT4 AWQ.
• Llama-70B: Aumento de velocidad de 6.7 veces en comparación con A100.

Recomendaciones de Mejores Prácticas

Selección del Método de Cuantificación

Elija el método de cuantificación adecuado según los diferentes escenarios:

1. Inferencia con lotes pequeños (tamaño de lote ≤4):

   • Se recomienda utilizar métodos de cuantificación de pesos (como INT4 AWQ).
   • La principal consideración son las limitaciones de ancho de banda de memoria.

2. Inferencia con lotes grandes (tamaño de lote ≥16):

   • Priorice la cuantificación FP8, que generalmente ofrece el mejor rendimiento y precisión.
   • Si los resultados no son satisfactorios, pruebe INT8 SmoothQuant, y luego AWQ y/o GPTQ.

3. Aplicaciones de dominio específico:

   • Para aplicaciones altamente especializadas como el autocompletado de código, se recomienda utilizar conjuntos de datos relevantes para el dominio para la calibración.

Ventajas Técnicas

1. Facilidad de uso: Proporciona una API de Python de alto nivel que simplifica el proceso de definición y optimización de LLM.
2. Rendimiento: Incluye todas las técnicas de optimización principales, como fusión de kernels, cuantificación, optimización en tiempo de ejecución, etc.
3. Escalabilidad: Soporta varios escenarios de despliegue, desde una sola GPU hasta múltiples nodos.
4. Compatibilidad: Integración profunda con PyTorch y soporte para los principales ecosistemas de inferencia.
5. Código abierto: Completamente de código abierto, con desarrollo continuo impulsado por la comunidad.

Desarrollo Futuro

TensorRT-LLM mejora la facilidad de uso y la escalabilidad a través de una API de definición de modelos modular de código abierto para definir, optimizar y ejecutar nuevas arquitecturas y mejoras, lo que permite una fácil personalización a medida que evolucionan los LLM.

Las direcciones de desarrollo continuo del proyecto incluyen:

• Mayor soporte para arquitecturas de modelos.
• Técnicas de cuantificación más avanzadas.
• Mejor capacidad de escalado multi-nodo.
• Integración más estrecha con el ecosistema.