Turboalimentando el Rendimiento de Meta Llama 3 con NVIDIA TensorRT-LLM y Servidor de Inferencia NVIDIA Triton

por Anjali Shah

Nos complace anunciar la compatibilidad con la familia de modelos Meta Llama 3 en NVIDIA TensorRT-LLM, lo que acelera y optimiza el rendimiento de la inferencia de LLM. Puede probar inmediatamente Llama 3 8B y Llama 3 70B, los primeros modelos de la serie, a través de una interfaz de usuario del navegador. O bien, a través de endpoints de API que se ejecutan en una pila de NVIDIA totalmente acelerada desde el catálogo de API de NVIDIA, donde Llama 3 se empaqueta como un NVIDIA NIM con una API estándar que se puede implementar en cualquier lugar.

Los grandes modelos de lenguaje son computacionalmente intensivos. Su tamaño los hace caros y lentos de ejecutar, especialmente sin las técnicas adecuadas. Hay muchas técnicas de optimización disponibles, como kernel fusion y la cuantificación para optimizaciones en tiempo de ejecución, como las implementaciones de C++, el almacenamiento en caché de KV, el procesamiento por lotes continuo en tránsito y la atención paginada. Los desarrolladores deben decidir qué combinación ayuda a su caso de uso. TensorRT-LLM simplifica este trabajo.

TensorRT-LLM es una biblioteca de código abierto que acelera el rendimiento de inferencia en los LLM más recientes en las GPU NVIDIA. NeMo, un framework integral para crear, personalizar e implementar aplicaciones de IA generativa, utiliza TensorRT-LLM y Servidor de Inferencia NVIDIA Triton para implementaciones de IA generativa.

TensorRT-LLM utiliza el compilador de deep learning NVIDIA TensorRT. Incluye los últimos kernels optimizados para implementaciones de vanguardia de FlashAttention y masked multi-head attention (MHA) para la ejecución de modelos LLM. También consta de pasos previos y posteriores al procesamiento y primitivas de comunicación de varias GPU/nodos en una API de Python simple y de código abierto para un rendimiento de inferencia de LLM innovador en GPU.

Para tener una idea de la biblioteca y cómo usarla, repasemos un ejemplo de cómo usar e implementar Llama 3 8B con TensorRT-LLM y Servidor de Inferencia Triton.

Para obtener una visión más detallada, incluidos diferentes modelos, diferentes optimizaciones y ejecución de varias GPU, consulte la lista completa de ejemplos de TensorRT-LLM.

Introducción a la Instalación

Comenzaremos clonando y compilando la biblioteca TensorRT-LLM siguiendo las instrucciones de instalación específicas del sistema operativo con el comando pip. Esta es una de las formas más fáciles de compilar TensorRT-LLM. Como alternativa, la biblioteca se puede instalar mediante un dockerfile para recuperar dependencias.

Los siguientes comandos extraen la biblioteca de código abierto e instalan las dependencias necesarias para instalar TensorRT-LLM dentro del contenedor.

git clone -b v0.8.0 https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM

Recuperación de los Pesos del Modelo

TensorRT-LLM es una biblioteca para la inferencia de LLM. Para usarlo, debe proporcionar un juego de pesas entrenadas. Se puede extraer un conjunto de pesos de repositorios como Hugging Face Hub o NVIDIA NGC. Otra opción es usar sus propios pesos de modelo entrenados en un framework como NeMo.

Los comandos de esta publicación extraen automáticamente los pesos (y los archivos tokenizadores) para la variante ajustada por instrucciones del modelo Llama 3 de 8 mil millones de parámetros del Hugging Face Hub. También puede descargar los pesos para usarlos sin conexión con el siguiente comando y actualizar las rutas en comandos posteriores para que apunten a este directorio:

git lfs install
git clone https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct

Tenga en cuenta que el uso de este modelo está sujeto a una licencia particular. Acepte los términos y autentíquese con HuggingFace para descargar los archivos necesarios.

Ejecución del Contenedor TensorRT-LLM

Lanzaremos un contenedor docker base e instalaremos las dependencias requeridas por TensorRT-LLM.

# Obtain and start the basic docker image environment.
docker run --rm --runtime=nvidia --gpus all --volume ${PWD}:/TensorRT-LLM --entrypoint /bin/bash -it --workdir /TensorRT-LLM nvidia/cuda:12.1.0-devel-ubuntu22.04

# Install dependencies, TensorRT-LLM requires Python 3.10
apt-get update && apt-get -y install python3.10 python3-pip openmpi-bin libopenmpi-dev

# Install the stable version (corresponding to the cloned branch) of TensorRT-LLM.
pip3 install tensorrt_llm==0.8.0 -U --extra-index-url https://pypi.nvidia.com

Compilación del Modelo

El siguiente paso en el proceso es compilar el modelo en un motor de TensorRT con pesos de modelo y una definición de modelo escrita en la API de Python de TensorRT-LLM.

El repositorio de TensorRT-LLM incluye varias arquitecturas de modelos y usamos la definición del modelo Llama. Para obtener más detalles y complementos y cuantificaciones más robustos disponibles, consulte este ejemplo de Llama y la documentación de precisión.

# Log in to huggingface-cli
# You can get your token from huggingface.co/settings/token
huggingface-cli login --token *****

# Build the Llama 8B model using a single GPU and BF16.
python3 examples/llama/convert_checkpoint.py --model_dir ./Meta-Llama-3-8B-Instruct \
            --output_dir ./tllm_checkpoint_1gpu_bf16 \
            --dtype bfloat16

trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_bf16 \
            --output_dir ./tmp/llama/8B/trt_engines/bf16/1-gpu \
            --gpt_attention_plugin bfloat16 \
            --gemm_plugin bfloat16

Cuando creamos la definición del modelo con la API de TensorRT-LLM, construimos un gráfico de operaciones a partir de primitivas de TensorRT que forman las capas de nuestra red neuronal. Estas operaciones se asignan a kernels específicos que son programas preescritos para la GPU.

El compilador de TensorRT puede barrer el gráfico para elegir el mejor kernel para cada operación y cada GPU disponible. También puede identificar patrones en el gráfico en los que varias operaciones son buenas candidatas para fusionarse en un único kernel fusionado, lo que reduce la cantidad necesaria de movimiento de memoria y la sobrecarga de lanzar varios kernels de GPU.

Además, TensorRT crea el gráfico de operaciones en un NVIDIA CUDA Graph que se puede iniciar al mismo tiempo. Esto reduce aún más la sobrecarga de lanzamiento de kernels.

El compilador de TensorRT es eficiente para fusionar capas y aumentar la velocidad de ejecución, sin embargo, hay algunas fusiones de capas complejas, como FlashAttention, que implican intercalar muchas operaciones juntas y que no se pueden detectar automáticamente. Para esos, podemos reemplazar explícitamente partes del gráfico con complementos en tiempo de compilación. En nuestro ejemplo, incluimos el complemento gpt_attention, que implementa un kernel de atención fusionado similar a FlashAttention, y el complemento gemm, que realiza la multiplicación de matrices con acumulación FP32. También llamamos a nuestra precisión deseada para el modelo completo como FP16, coincidiendo con la precisión predeterminada de los pesos que descargamos de HuggingFace.

Cuando terminemos de ejecutar el script de compilación, deberíamos esperar ver los siguientes tres archivos en la carpeta /tmp/llama/8B/trt_engines/bf16/1-gpu:

  • rank0.engine es la salida principal de nuestro script de compilación, que contiene el gráfico ejecutable de operaciones con los pesos del modelo incrustados.
  • config.json incluye información detallada sobre el modelo, como su estructura general y precisión, así como información sobre qué plug-ins se incorporaron al motor.

Ejecución del Modelo

Entonces, ahora que tenemos nuestro motor modelo, ¿qué podemos hacer con él?

El archivo del motor contiene la información para ejecutar el modelo. TensorRT-LLM incluye un tiempo de ejecución de C++ altamente optimizado para ejecutar archivos de motor y administrar procesos como tokens de muestreo de la salida del modelo, administrar la caché de KV y procesar solicitudes por lotes juntos.

Podemos usar el tiempo de ejecución directamente para ejecutar el modelo localmente, o podemos implementar usando Servidor de Inferencia Triton en el entorno de producción para compartir el modelo con varios usuarios.

Para ejecutar el modelo localmente, podemos ejecutar el siguiente comando:

python3 examples/run.py --engine_dir=./tmp/llama/8B/trt_engines/bf16/1-gpu --max_output_len 100 --tokenizer_dir ./Meta-Llama-3-8B-Instruct --input_text "How do I count to nine in French?"

Implementación con Servidor de Inferencia Triton

Más allá de la ejecución local, también podemos usar Servidor de Inferencia Triton para crear una implementación lista para producción de nuestro LLM. El backend del servidor de inferencia de Triton para TensorRT-LLM utiliza el tiempo de ejecución de C++ de TensorRT-LLM para una ejecución de inferencia de alto rendimiento. Incluye técnicas como el procesamiento por lotes en tránsito y el almacenamiento en caché de KV paginado que proporcionan un alto rendimiento con baja latencia. El backend de TensorRT-LLM se ha incluido con Servidor de Inferencia Triton y está disponible como un contenedor prediseñado en NGC.

En primer lugar, debemos crear un repositorio de modelos para que Servidor de Inferencia Triton pueda leer el modelo y los metadatos asociados.

El repositorio tensorrtllm_backend incluye la configuración de un repositorio de modelos requerido en all_models/inflight_batcher_llm/ que podemos replicar.

En el directorio hay cuatro subcarpetas que contienen artefactos para diferentes partes del proceso de ejecución del modelo. Las carpetas preprocessing/y postprocessing/contienen scripts para el backend de Python de Servidor de Inferencia Triton. Estos scripts sirven para tokenizar las entradas de texto y destokenizar las salidas del modelo para convertir entre cadenas y los identificadores de token en los que opera el modelo.

El folder tensorrt_llm es donde colocaremos el motor de modelo que compilamos anteriormente. Y, por último, el folder ensemble define un conjunto de modelos que vincula los tres componentes anteriores y le indica al Servidor de Inferencia Triton cómo hacer fluir los datos a través de ellos.

Extraiga el repositorio de modelos de ejemplo y copie el modelo que compiló en el paso anterior.

# After exiting the TensorRT-LLM docker container
cd ..
git clone -b v0.8.0 https://github.com/triton-inference-server/tensorrtllm_backend.git
cd tensorrtllm_backend
cp ../TensorRT-LLM/tmp/llama/8B/trt_engines/bf16/1-gpu/* all_models/inflight_batcher_llm/tensorrt_llm/1/

A continuación, debemos modificar los archivos de configuración desde el esqueleto del repositorio con la ubicación del motor del modelo compilado. También debemos actualizar los parámetros de configuración, como el tokenizador, para usar y controlar la asignación de memoria para la memoria caché de KV al procesar por lotes las solicitudes de inferencia.

#Set the tokenizer_dir and engine_dir paths
HF_LLAMA_MODEL=TensorRT-LLM/Meta-Llama-3-8B-Instruct
ENGINE_PATH=tensorrtllm_backend/all_models/inflight_batcher_llm/tensorrt_llm/1

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/preprocessing/config.pbtxt tokenizer_dir:${HF_LLAMA_MODEL},tokenizer_type:auto,triton_max_batch_size:64,preprocessing_instance_count:1

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/postprocessing/config.pbtxt tokenizer_dir:${HF_LLAMA_MODEL},tokenizer_type:auto,triton_max_batch_size:64,postprocessing_instance_count:1

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/tensorrt_llm_bls/config.pbtxt triton_max_batch_size:64,decoupled_mode:False,bls_instance_count:1,accumulate_tokens:False

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/ensemble/config.pbtxt triton_max_batch_size:64

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt triton_max_batch_size:64,decoupled_mode:False,max_beam_width:1,engine_dir:${ENGINE_PATH},max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:True,enable_kv_cache_reuse:False,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:0

Ahora, podemos poner en marcha nuestro contenedor docker y lanzar el servidor Triton. Debemos especificar el tamaño del mundo, el número de GPU para las que se creó el modelo, y señalar nuestro model_repo que acabamos de configurar.

#Change to base working directory
cd..
docker run -it --rm --gpus all --network host --shm-size=1g \
-v $(pwd):/workspace \
--workdir /workspace \
nvcr.io/nvidia/tritonserver:24.03-trtllm-python-py3

# Log in to huggingface-cli to get tokenizer
huggingface-cli login --token *****

# Install python dependencies
pip install sentencepiece protobuf

# Launch Server

python3 tensorrtllm_backend/scripts/launch_triton_server.py --model_repo tensorrtllm_backend/all_models/inflight_batcher_llm --world_size 1

Envío de Solicitudes

Para enviar solicitudes de inferencia y recibir finalizaciones del servidor en ejecución, puede utilizar una de las bibliotecas cliente del Servidor de Inferencia Triton o enviar solicitudes HTTP al endpoint generado.

El siguiente comando curl muestra una prueba rápida para solicitar finalizaciones por parte del servidor en ejecución y el script de cliente con más funciones se puede revisar para comunicarse con el servidor.

curl -X POST localhost:8000/v2/models/ensemble/generate -d \
'{
"text_input": "How do I count to nine in French?",
"parameters": {
"max_tokens": 100,
"bad_words":[""],
"stop_words":[""]
}
}'

Conclusión

TensorRT-LLM proporciona herramientas para optimizar y ejecutar de manera eficiente grandes modelos de lenguaje en GPU NVIDIA. Servidor de Inferencia Triton es ideal para implementar y servir de forma eficiente grandes modelos de lenguaje, como Llama 3.

Con esta guía de introducción y comience su viaje aprovechando las herramientas de código abierto para usar Llama 3 y muchos otros grandes modelos de lenguaje.

NVIDIA AI Enterprise, una plataforma de software de IA de extremo a extremo que incluye TensorRT, pronto incluirá TensorRT-LLM, para la inferencia de IA de misión crítica con seguridad, estabilidad, capacidad de administración y soporte de nivel empresarial.