TensorRT-LLM: Una guía completa para optimizar la inferencia de modelos de lenguaje grande para un rendimiento máximo

Como la demanda de modelos de lenguaje grande (LLM) sigue aumentando, garantizar una inferencia rápida, eficiente y escalable se ha vuelto más crucial que nunca. NVIDIA’s (NVDA ) TensorRT-LLM se presenta para abordar este desafío al proporcionar un conjunto de herramientas y optimizaciones poderosas diseñadas específicamente para la inferencia de LLM. TensorRT-LLM ofrece una impresionante variedad de mejoras de rendimiento, como cuantización, fusión de kernels, procesamiento por lotes en vuelo y soporte multi-GPU. Estos avances permiten lograr velocidades de inferencia hasta 8 veces más rápidas que los métodos tradicionales basados en CPU, transformando la forma en que desplegamos LLM en producción.

Esta guía completa explorará todos los aspectos de TensorRT-LLM, desde su arquitectura y características clave hasta ejemplos prácticos para desplegar modelos. Ya sea que sea un ingeniero de AI, desarrollador de software o investigador, esta guía le brindará el conocimiento para aprovechar TensorRT-LLM para optimizar la inferencia de LLM en GPUs de NVIDIA.

Acelerando la inferencia de LLM con TensorRT-LLM

TensorRT-LLM entrega mejoras dramáticas en el rendimiento de la inferencia de LLM. Según las pruebas de NVIDIA, las aplicaciones basadas en TensorRT muestran velocidades de inferencia hasta 8 veces más rápidas en comparación con plataformas solo CPU. Esto es un avance crucial en aplicaciones en tiempo real como chatbots, sistemas de recomendación y sistemas autónomos que requieren respuestas rápidas.

Cómo funciona

TensorRT-LLM acelera la inferencia al optimizar las redes neuronales durante la implementación utilizando técnicas como:

• Cuantización: Reduce la precisión de los pesos y las activaciones, reduciendo el tamaño del modelo y mejorando la velocidad de inferencia.
• Fusión de capas y tensores: Combina operaciones como funciones de activación y multiplicaciones de matrices en una sola operación.
• Ajuste de kernel: Selecciona kernels CUDA óptimos para el cálculo de GPU, reduciendo el tiempo de ejecución.

Estas optimizaciones garantizan que sus modelos LLM se ejecuten de manera eficiente en una amplia gama de plataformas de implementación, desde centros de datos de hipercala hasta sistemas integrados.

Optimizando el rendimiento de la inferencia con TensorRT

Construido sobre el modelo de programación paralela CUDA de NVIDIA, TensorRT proporciona optimizaciones altamente especializadas para la inferencia en GPUs de NVIDIA. Al simplificar procesos como la cuantización, el ajuste de kernel y la fusión de operaciones de tensores, TensorRT garantiza que los LLM puedan ejecutarse con latencia mínima.

Algunas de las técnicas más efectivas incluyen:

• Cuantización: Esto reduce la precisión numérica de los parámetros del modelo mientras mantiene una alta precisión, lo que efectivamente acelera la inferencia.
• Fusión de tensores: Al fusionar múltiples operaciones en un solo kernel CUDA, TensorRT minimiza la sobrecarga de memoria y aumenta el rendimiento.
• Ajuste automático de kernel: TensorRT selecciona automáticamente el mejor kernel para cada operación, optimizando la inferencia para una GPU determinada.

Estas técnicas permiten que TensorRT-LLM optimice el rendimiento de la inferencia para tareas de aprendizaje automático como el procesamiento de lenguaje natural, motores de recomendación y análisis de video en tiempo real.

Acelerando las cargas de trabajo de AI con TensorRT

TensorRT acelera las cargas de trabajo de aprendizaje automático al incorporar optimizaciones de precisión como INT8 y FP16. Estos formatos de precisión reducida permiten una inferencia significativamente más rápida mientras se mantiene la precisión. Esto es particularmente valioso en aplicaciones en tiempo real donde la latencia baja es un requisito crítico.

Las optimizaciones INT8 y FP16 son particularmente efectivas en:

• Transmisión de video: Las tareas de procesamiento de video basadas en IA, como la detección de objetos, se benefician de estas optimizaciones al reducir el tiempo necesario para procesar los fotogramas.
• Sistemas de recomendación: Al acelerar la inferencia para modelos que procesan grandes cantidades de datos de usuarios, TensorRT permite la personalización en tiempo real a gran escala.
• Procesamiento de lenguaje natural (NLP): TensorRT mejora la velocidad de las tareas de NLP como la generación de texto, la traducción y el resumen, lo que las hace adecuadas para aplicaciones en tiempo real.

Desplegar, ejecutar y escalar con NVIDIA Triton

Una vez que su modelo haya sido optimizado con TensorRT-LLM, puede desplegarlo, ejecutarlo y escalarlo fácilmente utilizando NVIDIA Triton Inference Server. Triton es un software de código abierto que admite la ejecución de modelos concurrente, conjuntos de modelos y alto rendimiento. Proporciona un entorno flexible para administrar modelos de AI a gran escala.

Algunas de las características clave incluyen:

• Ejecución de modelos concurrente: Ejecuta múltiples modelos simultáneamente, maximizando la utilización de la GPU.
• Procesamiento por lotes dinámico: Combina múltiples solicitudes de inferencia en un solo lote, reduciendo la latencia y aumentando el rendimiento.
• Entradas de audio/video en streaming: Admite flujos de entrada en aplicaciones en tiempo real, como análisis de video en vivo o servicios de voz a texto.

Esto hace que Triton sea una herramienta valiosa para desplegar modelos optimizados con TensorRT-LLM en entornos de producción, garantizando una alta escalabilidad y eficiencia.

Características principales de TensorRT-LLM para la inferencia de LLM

API de Python de código abierto

TensorRT-LLM proporciona una API de Python altamente modular y de código abierto, simplificando el proceso de definir, optimizar y ejecutar LLM. La API permite a los desarrolladores crear LLM personalizados o modificar los preestablecidos para adaptarlos a sus necesidades, sin requerir un conocimiento profundo de CUDA o marcos de aprendizaje automático.

Procesamiento por lotes en vuelo y atención paginada

Una de las características destacadas de TensorRT-LLM es el procesamiento por lotes en vuelo, que optimiza la generación de texto al procesar múltiples solicitudes de forma concurrente. Esta característica minimiza el tiempo de espera y mejora la utilización de la GPU al procesar por lotes secuencias de forma dinámica.

Además, la atención paginada garantiza que el uso de memoria se mantenga bajo incluso cuando se procesan secuencias de entrada largas. En lugar de asignar memoria contigua para todos los tokens en una secuencia (lo que puede provocar fragmentación de memoria), la atención paginada divide los datos de caché de clave-valor en “páginas” de memoria que se pueden reutilizar dinámicamente, evitando la fragmentación de memoria y mejorando la eficiencia.

Inferencia multi-GPU y multi-nodo

Para modelos más grandes o cargas de trabajo más complejas, TensorRT-LLM admite inferencia multi-GPU y multi-nodo. Esta capacidad permite distribuir los cálculos del modelo en varias GPUs o nodos, mejorando el rendimiento y reduciendo el tiempo total de inferencia.

Soporte para FP8

Con la llegada de FP8 (punto flotante de 8 bits), TensorRT-LLM aprovecha las innovaciones de hardware de NVIDIA en la arquitectura H100 Hopper. FP8 reduce la huella de memoria de los LLM al almacenar pesos y activaciones en un formato de punto flotante de 8 bits, lo que resulta en cálculos más rápidos sin sacrificar mucha precisión. TensorRT-LLM compila automáticamente los modelos para utilizar kernels FP8 optimizados, acelerando aún más los tiempos de inferencia.

Esto hace que TensorRT-LLM sea una opción ideal para despliegues a gran escala que requieren un rendimiento y eficiencia energética de primer nivel.

Arquitectura y componentes de TensorRT-LLM

Comprender la arquitectura de TensorRT-LLM le ayudará a aprovechar al máximo sus capacidades para la inferencia de LLM. Veamos los componentes clave:

Definición del modelo

TensorRT-LLM le permite definir LLM utilizando una API de Python simple. La API construye una representación gráfica del modelo, lo que facilita la gestión de las capas complejas involucradas en las arquitecturas de LLM como GPT o BERT.

Vinculación de pesos

Antes de compilar el modelo, los pesos (o parámetros) deben vincularse a la red. Este paso garantiza que los pesos estén incrustados en el motor de TensorRT, lo que permite una inferencia rápida y eficiente. TensorRT-LLM también permite actualizaciones de pesos después de la compilación, agregando flexibilidad para modelos que necesitan actualizaciones frecuentes.

Coincidencia de patrones y fusión

La fusión de operaciones es otra característica poderosa de TensorRT-LLM. Al fusionar múltiples operaciones (por ejemplo, multiplicaciones de matrices con funciones de activación) en un solo kernel CUDA, TensorRT minimiza la sobrecarga asociada con múltiples lanzamientos de kernel. Esto reduce las transferencias de memoria y acelera la inferencia.

Complementos

Para ampliar las capacidades de TensorRT, los desarrolladores pueden escribir complementos, kernels personalizados que permiten optimizaciones o operaciones específicas no cubiertas por la biblioteca de TensorRT estándar.

Por ejemplo, el complemento Flash-Attention es un kernel personalizado bien conocido que optimiza las capas de atención multi-cabeza en modelos basados en Transformadores. Al utilizar este complemento, los desarrolladores pueden lograr aceleraciones sustanciales en el cálculo de la atención, uno de los componentes más intensivos en recursos de los LLM.

Benchmark: Ganancias de rendimiento de TensorRT-LLM

TensorRT-LLM demuestra ganancias de rendimiento significativas para la inferencia de LLM en varias GPUs de NVIDIA. A continuación, se muestra una comparación de la velocidad de inferencia (medida en tokens por segundo) utilizando TensorRT-LLM en diferentes GPUs de NVIDIA:

Estos benchmarks muestran que TensorRT-LLM ofrece mejoras sustanciales en el rendimiento, particularmente para secuencias más largas.

Práctica: Instalación y compilación de TensorRT-LLM

Paso 1: Crear un entorno de contenedor

Para facilitar el uso, TensorRT-LLM proporciona imágenes de Docker para crear un entorno controlado para compilar y ejecutar modelos.

docker build --pull \
--target devel \
--file docker/Dockerfile.multi \
--tag tensorrt_llm/devel:latest .

docker run --rm -it \
--ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --gpus=all \
--volume ${PWD}:/code/tensorrt_llm \
--workdir /code/tensorrt_llm \
tensorrt_llm/devel:latest

python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt
pip install ./build/tensorrt_llm*.whl

include_directories(${TENSORRT_LLM_PATH}/cpp/include)
link_directories(${TENSORRT_LLM_PATH}/cpp/build/tensorrt_llm)
target_link_libraries(your_project tensorrt_llm)

mkdir -p model_repository/gpt2/1
cp ./trt_engine/gpt2_fp16.engine model_repository/gpt2/1/

name: "gpt2"
platform: "tensorrt_llm"
max_batch_size: 8

<p>input [
{
name: &quot;input_ids&quot;
data_type: TYPE_INT32
dims: [-1]
}
]</p>

<p>output [
{
name: &quot;logits&quot;
data_type: TYPE_FP32
dims: [-1, -1]
}
]</p>

docker run --rm --gpus all \
-v $(pwd)/model_repository:/models \
nvcr.io/nvidia/tritonserver:23.05-py3 \
tritonserver --model-repository=/models

curl -X POST http://localhost:8000/v2/models/gpt2/infer -d '{
"inputs": [
{"name": "input_ids", "shape": [1, 128], "datatype": "INT32", "data": [[101, 234, 1243]]}
]
}'