Implementación de LLM en producción utilizando TensorRT LLM | de Trivedi | febrero de 2024

Tutorial práctico de Python

Hay dos pasos para implementar un modelo usando TensorRT-LLM:

1. compilar el modelo
2. Implementar el modelo compilado como un punto final de API REST

Paso 1: compilar el modelo

Para este tutorial, trabajaremos con Mistral 7B Instrucción v0.2. Como se mencionó anteriormente, la fase de compilación requiere una GPU. Descubrí que la forma más sencilla de compilar un modelo es en una computadora portátil de Google Colab.

TensorRT LLM se admite principalmente en GPU Nvidia de gama alta. Ejecuté Google Colab en una GPU A100 de 40 GB y también usaré la misma GPU para la implementación.

!git clone https://github.com/NVIDIA/TensorRT-LLM.git
%cd TensorRT-LLM/examples/llama

• Clonar el Repositorio git TensorRT-LLM. Este repositorio contiene todos los módulos y scripts que necesitamos para compilar el modelo.

!pip install tensorrt_llm -U --pre --extra-index-url https://pypi.nvidia.com
!pip install huggingface_hub pynvml mpi4py
!pip install -r requirements.txt

• Instale las dependencias necesarias de Python.

from huggingface_hub import snapshot_download
from google.colab import userdata
snapshot_download(
"mistralai/Mistral-7B-Instruct-v0.2",
local_dir="tmp/hf_models/mistral-7b-instruct-v0.2",
max_workers=4
)

• Descargue los pesos del modelo Mistral 7B Instruct v0.2 desde Hugging Face y guárdelos en un directorio local en tmp/hf_models/mistral-7b-instruct-v0.2
• Si miras dentro del tmp/hf_models directorio en Colab, debería ver los pesos del modelo allí.

!python convert_checkpoint.py --model_dir ./tmp/hf_models/mistral-7b-instruct-v0.2 \
--output_dir ./tmp/trt_engines/1-gpu/ \
--dtype float16

• Los pesos del modelo sin procesar no se pueden compilar. En su lugar, deben convertirse a un formato tensorRT LLM específico.
• El convert_checkpoint.py El script toma los pesos Mistral sin procesar y los convierte a un formato compatible.
• El --model_dir es el camino hacia los pesos del modelo sin procesar.
• El --output_dir es el camino hacia los pesos convertidos.

!trtllm-build --checkpoint_dir ./tmp/trt_engines/1-gpu/ \
--output_dir ./tmp/trt_engines/compiled-model/ \
--gpt_attention_plugin float16 \
--gemm_plugin float16 \
--max_input_len 32256

• El trtllm-build El comando compila el modelo. En esta etapa, también puede pasar varios indicadores de optimización. Para simplificar las cosas, no he utilizado ninguna optimización adicional.
• El --checkpoint_dir es el camino hacia el convertido pesos del modelo.
• El --output_dir es donde se guarda el modelo compilado.
• Mistral 7B Instruct v0.2 admite una longitud de contexto de 32 K. He establecido esa longitud de contexto usando el--max_input_length bandera.

Nota: compilar el modelo puede tardar entre 15 y 30 minutos.

Una vez compilado el modelo, puede cargar su modelo compilado en abrazando el centro de la cara. Para poder cargar archivos en Hugging Face Hub, necesitará un token de acceso válido que tenga ESCRIBIR acceso.

import os
from huggingface_hub import HfApi
for root, dirs, files in os.walk(f"tmp/trt_engines/compiled-model", topdown=False):
for name in files:
filepath = os.path.join(root, name)
filename = "/".join(filepath.split("/")[-2:])
print("uploading file: ", filename)
api = HfApi(token=userdata.get('HF_WRITE_TOKEN'))
api.upload_file(
path_or_fileobj=filepath,
path_in_repo=filename,
repo_id="<your-repo-id>/mistral-7b-v0.2-trtllm"
)

• Este código carga el modelo compilado, el .motor archivo, para abrazar la cara bajo su identificación de usuario.
• Reemplace la en el código con su repositorio de cara de abrazo, que suele ser su identificación de usuario de cara de abrazo.

¡Impresionante! Eso termina la parte de compilación del modelo. En el paso de implementación.

Paso 2: implementar el modelo compilado

Hay muchas formas de implementar este modelo compilado. Puedes utilizar una herramienta sencilla como API rápida o algo más complejo como el servidor de inferencia tritón.

Cuando se utiliza una herramienta como FastAPI, el desarrollador debe configurar el servidor API, escribir el Dockerfile y configurar CUDA correctamente. Administrar estas cosas puede ser una verdadera molestia y arruina la experiencia general del desarrollador.

Para evitar estos problemas, he decidido utilizar una sencilla herramienta de código abierto llamada Braguero. Truss permite a los desarrolladores empaquetar fácilmente sus modelos con soporte para GPU y ejecutarlos en cualquier entorno de nube. Tiene un montón de características excelentes que hacen que los modelos de contenedores sean muy sencillos:

• Soporte de GPU listo para usar. No es necesario tratar con CUDA.
• Creación automática de Dockerfile. No es necesario que lo escriba usted mismo.
• Servidor API listo para producción
• Interfaz sencilla de Python

El principal beneficio de usar Truss es que puede contener fácilmente un modelo con soporte de GPU e implementarlo en cualquier entorno de nube.

Creando la armadura

Cree o abra un entorno virtual de Python con una versión de Python ≥ 3.8 e instale la siguiente dependencia:

pip install --upgrade truss

(Opcional) Si desea crear su proyecto Truss desde cero, puede ejecutar el comando:

truss init mistral-7b-tensort-llm

Se le pedirá que le dé un nombre a su modelo. Cualquier nombre como Mistral 7B Tensort LLM servirá. La ejecución del comando anterior genera automáticamente los archivos necesarios para implementar un Truss.

Para acelerar un poco el proceso, tengo un repositorio de Github que contiene los archivos necesarios. Clone el repositorio de Github a continuación:

Así es como debería verse la estructura del directorio mistral-7b-tensorrt-llm-truss :

├── mistral-7b-tensorrt-llm-truss
│   ├── config.yaml
│   ├── model
│   │   ├── __init__.py
│   │   └── model.py
|   |   └── utils.py
|   ├── requirements.txt

Aquí hay un desglose rápido de para qué se utilizan los archivos anteriores:

1. El config.yaml se utiliza para establecer varias configuraciones para su modelo, incluidos sus recursos, dependencias, variables ambientales y más. Aquí es donde podemos especificar el nombre del modelo, qué dependencias de Python instalar y qué paquetes del sistema instalar.

2. El model/model.py es el corazón de Truss. Contiene el código Python que se ejecutará en el servidor Truss. En el model.py Hay dos métodos principales: load() y predict() .

• El load El método es donde descargaremos el modelo compilado de Hugging Face e inicializaremos el motor TensorRT LLM.
• El predict El método recibe solicitudes HTTP y llama al modelo.

3. El model/utils.py contiene algunas funciones auxiliares para el model.py archivo. yo no escribí el utils.py archivo yo mismo, lo tomé directamente del Repositorio TensorRT LLM.

4. El requirements.txt contiene las dependencias de Python necesarias para ejecutar nuestro modelo compilado.

Explicación más profunda del código:

El model.py contiene el código principal que se ejecuta, así que profundicemos un poco más en ese archivo. Primero echemos un vistazo a la load función.

import subprocess
subprocess.run(["pip", "install", "tensorrt_llm", "-U", "--pre", "--extra-index-url", "https://pypi.nvidia.com"])
import torch
from model.utils import (DEFAULT_HF_MODEL_DIRS, DEFAULT_PROMPT_TEMPLATES,
load_tokenizer, read_model_name, throttle_generator)
import tensorrt_llm
import tensorrt_llm.profiler
from tensorrt_llm.runtime import ModelRunnerCpp, ModelRunner
from huggingface_hub import snapshot_download
STOP_WORDS_LIST = None
BAD_WORDS_LIST = None
PROMPT_TEMPLATE = None
class Model:
def __init__(self, **kwargs):
self.model = None
self.tokenizer = None
self.pad_id = None
self.end_id = None
self.runtime_rank = None
self._data_dir = kwargs["data_dir"]
def load(self):
snapshot_download(
"htrivedi99/mistral-7b-v0.2-trtllm",
local_dir=self._data_dir,
max_workers=4,
)
self.runtime_rank = tensorrt_llm.mpi_rank()
model_name, model_version = read_model_name(f"{self._data_dir}/compiled-model")
tokenizer_dir = "mistralai/Mistral-7B-Instruct-v0.2"
self.tokenizer, self.pad_id, self.end_id = load_tokenizer(
tokenizer_dir=tokenizer_dir,
vocab_file=None,
model_name=model_name,
model_version=model_version,
tokenizer_type="llama",
)
runner_cls = ModelRunner
runner_kwargs = dict(engine_dir=f"{self._data_dir}/compiled-model",
lora_dir=None,
rank=self.runtime_rank,
debug_mode=False,
lora_ckpt_source="hf",
)
self.model = runner_cls.from_dir(**runner_kwargs)

Que esta pasando aqui:

• En la parte superior del archivo importamos los módulos necesarios, específicamente tensorrt_llm
• A continuación, dentro de la función de carga, descargamos el modelo compilado usando el snapshot_download función. Mi modelo compilado está en el siguiente ID de repositorio: htrivedi99/mistral-7b-v0.2-trtllm . Si cargó su modelo compilado en otro lugar, actualice este valor en consecuencia.
• Luego, descargamos el tokenizador para el modelo usando el load_tokenizer función que viene con model/utils.py .
• Finalmente, usamos TensorRT LLM para cargar nuestro modelo compilado usando el ModelRunner clase.

Genial, echemos un vistazo a predict funcionar también.

def predict(self, request: dict):
prompt = request.pop("prompt")
max_new_tokens = request.pop("max_new_tokens", 2048)
temperature = request.pop("temperature", 0.9)
top_k = request.pop("top_k",1)
top_p = request.pop("top_p", 0)
streaming = request.pop("streaming", False)
streaming_interval = request.pop("streaming_interval", 3)
batch_input_ids = self.parse_input(tokenizer=self.tokenizer,
input_text=[prompt],
prompt_template=None,
input_file=None,
add_special_tokens=None,
max_input_length=1028,
pad_id=self.pad_id,
)
input_lengths = [x.size(0) for x in batch_input_ids]
outputs = self.model.generate(
batch_input_ids,
max_new_tokens=max_new_tokens,
max_attention_window_size=None,
sink_token_length=None,
end_id=self.end_id,
pad_id=self.pad_id,
temperature=temperature,
top_k=top_k,
top_p=top_p,
num_beams=1,
length_penalty=1,
repetition_penalty=1,
presence_penalty=0,
frequency_penalty=0,
stop_words_list=STOP_WORDS_LIST,
bad_words_list=BAD_WORDS_LIST,
lora_uids=None,
streaming=streaming,
output_sequence_lengths=True,
return_dict=True)
if streaming:
streamer = throttle_generator(outputs, streaming_interval)
def generator():
total_output = ""
for curr_outputs in streamer:
if self.runtime_rank == 0:
output_ids = curr_outputs['output_ids']
sequence_lengths = curr_outputs['sequence_lengths']
batch_size, num_beams, _ = output_ids.size()
for batch_idx in range(batch_size):
for beam in range(num_beams):
output_begin = input_lengths[batch_idx]
output_end = sequence_lengths[batch_idx][beam]
outputs = output_ids[batch_idx][beam][
output_begin:output_end].tolist()
output_text = self.tokenizer.decode(outputs)
current_length = len(total_output)
total_output = output_text
yield total_output[current_length:]
return generator()
else:
if self.runtime_rank == 0:
output_ids = outputs['output_ids']
sequence_lengths = outputs['sequence_lengths']
batch_size, num_beams, _ = output_ids.size()
for batch_idx in range(batch_size):
for beam in range(num_beams):
output_begin = input_lengths[batch_idx]
output_end = sequence_lengths[batch_idx][beam]
outputs = output_ids[batch_idx][beam][
output_begin:output_end].tolist()
output_text = self.tokenizer.decode(outputs)
return {"output": output_text}

Que esta pasando aqui:

• El predict La función acepta algunas entradas del modelo, como la prompt , max_new_tokens , temperature etc. Extraemos todos estos valores en la parte superior de la función usando el request.pop método.
• A continuación, formateamos el mensaje en el formato requerido para TensorRT LLM usando el self.parse_input función auxiliar.
• Luego, llamamos a nuestro modelo LLM para generar los resultados utilizando el self.model.generate función. La función generar acepta una variedad de argumentos que ayudan a controlar la salida del LLM.
• También agregué algo de código para habilitar la transmisión produciendo un generator objeto. Si la transmisión está deshabilitada, el tokenizador simplemente decodifica la salida del LLM y la devuelve como un objeto JSON.

¡Impresionante! Eso cubre la parte de codificación. Vamos a contenerlo.

Contenerizar el modelo:

Para ejecutar nuestro modelo en la nube necesitamos contenerlo. Truss se encargará de crear el Dockerfile y empaquetar todo por nosotros, por lo que no tenemos que hacer mucho.

Fuera del mistral-7b-tensorrt-llm-truss directorio crear un archivo llamado main.py . Pegue el siguiente código dentro de él:

import truss
from pathlib import Path
tr = truss.load("./mistral-7b-tensorrt-llm-truss")
command = tr.docker_build_setup(build_dir=Path("./mistral-7b-tensorrt-llm-truss"))
print(command)

Ejecute el main.py archivar y mirar dentro del mistral-7b-tensorrt-llm-truss directorio. Deberías ver que se generan automáticamente un montón de archivos. No necesitamos preocuparnos por lo que significan estos archivos, es simplemente Truss haciendo su magia.

A continuación, construyamos nuestro contenedor usando Docker. Ejecute los siguientes comandos secuencialmente:

docker build mistral-7b-tensorrt-llm-truss -t mistral-7b-tensorrt-llm-truss:latest
docker tag mistral-7b-tensorrt-llm-truss <docker_user_id>/mistral-7b-tensorrt-llm-truss
docker push <docker_user_id>/mistral-7b-tensorrt-llm-truss

¡Dulce! ¡Estamos listos para implementar el modelo en la nube!

Implementando el modelo en GKE

Para esta sección, implementaremos el modelo en Google Kubernetes Engine. Si recuerdas, durante el paso de compilación del modelo ejecutamos Google Colab en una GPU A100 de 40 GB. Para que TensorRT LLM funcione, debemos implementar el modelo exactamente en la misma GPU para realizar inferencias.

No profundizaré en cómo configurar un clúster de GKE ya que no está dentro del alcance de este artículo. Pero proporcionaré las especificaciones que utilicé para el clúster. Aquí están las especificaciones:

• 1 nodo, clúster de Kubernetes estándar (sin piloto automático)
• 1.28.5 versión gke kubernetes
• 1 GPU Nvidia A100 de 40 GB
• Máquina a2-highgpu-1g (12 vCPU, 85 GB de memoria)
• Instalación del controlador GPU administrada por Google (de lo contrario, debemos instalar el controlador Cuda manualmente)
• Todo esto se ejecutará en una instancia puntual.

Una vez configurado el clúster, podemos iniciarlo y conectarnos a él. Una vez que el clúster esté activo y se haya conectado correctamente a él, cree la siguiente implementación de Kubernetes:

apiVersion: apps/v1
kind: Deployment
metadata:
name: mistral-7b-v2-trt
namespace: default
spec:
replicas: 1
selector:
matchLabels:
component: mistral-7b-v2-trt-layer
template:
metadata:
labels:
component: mistral-7b-v2-trt-layer
spec:
containers:
- name: mistral-container
image: htrivedi05/mistral-7b-v0.2-trt:latest
ports:
- containerPort: 8080
resources:
limits:
nvidia.com/gpu: 1
nodeSelector:
cloud.google.com/gke-accelerator: nvidia-tesla-a100
---
apiVersion: v1
kind: Service
metadata:
name: mistral-7b-v2-trt-service
namespace: default
spec:
type: ClusterIP
selector:
component: mistral-7b-v2-trt-layer
ports:
- port: 8080
protocol: TCP
targetPort: 8080

Esta es una implementación estándar de Kubernetes que ejecuta un contenedor con la imagen. htrivedi05/mistral-7b-v0.2-trt:latest . Si creaste tu propia imagen en la sección anterior, continúa y úsala. Siéntete libre de usar el mío de otra manera.

Puede crear la implementación ejecutando el comando:

kubectl create -f mistral-deployment.yaml

Se necesitan unos minutos para que se aprovisione el pod de Kubernetes. Una vez que el pod se esté ejecutando, se ejecutará la función de carga que escribimos anteriormente. Puede verificar los registros del pod ejecutando el comando:

kubectl logs <pod-name>

Una vez cargado el modelo, verá algo como Completed model.load() execution in 449234 ms en los registros del pod. Para enviar una solicitud al modelo a través de HTTP, necesitamos reenviar el servicio. Puede utilizar el siguiente comando para hacer eso:

kubectl port-forward svc/mistral-7b-v2-trt-service 8080

¡Excelente! ¡Por fin podemos empezar a enviar solicitudes al modelo! Abra cualquier script de Python y ejecute el siguiente código:

import requests
data = {"prompt": "What is a mistral?"}
res = requests.post("http://127.0.0.1:8080/v1/models/model:predict", json=data)
res = res.json()
print(res)

Verá un resultado como el siguiente:

{"output": "A Mistral is a strong, cold wind that originates in the Rhone Valley in France. It is named after the Mistral wind system, which is associated with the northern Mediterranean region. The Mistral is known for its consistency and strength, often blowing steadily for days at a time. It can reach speeds of up to 130 kilometers per hour (80 miles per hour), making it one of the strongest winds in Europe. The Mistral is also known for its clear, dry air and its role in shaping the landscape and climate of the Rhone Valley."}

El rendimiento de TensorRT LLM se puede notar visiblemente cuando se transmiten los tokens. A continuación se muestra un ejemplo de cómo hacerlo:

data = {"prompt": "What is mistral wind?", "streaming": True, "streaming_interval": 3}
res = requests.post("http://127.0.0.1:8080/v1/models/model:predict", json=data, stream=True)
for content in res.iter_content():
print(content.decode("utf-8"), end="", flush=True)

Este modelo de mistral tiene una ventana de contexto bastante grande, así que siéntete libre de probarlo con diferentes indicaciones.

Puntos de referencia de rendimiento

Con solo mirar los tokens que se transmiten, probablemente pueda darse cuenta de que TensorRT LLM es realmente rápido. Sin embargo, quería obtener números reales para capturar las ganancias de rendimiento al usar TensorRT LLM. Ejecuté algunos puntos de referencia personalizados y obtuve los siguientes resultados:

Pequeño aviso:

Mensaje medio:

Mensaje grande:

Conclusión

En esta publicación de blog, mi objetivo era demostrar cómo se puede lograr una inferencia de última generación utilizando TensorRT LLM. Cubrimos todo, desde la compilación de un LLM hasta la implementación del modelo en producción.

Si bien TensorRT LLM es más complejo que otros optimizadores de inferencia, el rendimiento habla por sí solo. Esta herramienta proporciona optimizaciones LLM de última generación, al mismo tiempo que es completamente de código abierto y está diseñada para uso comercial. Este marco se encuentra todavía en las primeras etapas y se encuentra en desarrollo activo. El desempeño que vemos hoy solo mejorará en los próximos años.

Espero que hayas encontrado algo valioso en este artículo. ¡Gracias por leer!

¿Disfrutaste esta historia?

Considerar suscribiéndose gratis.

Imágenes

A menos que se indique lo contrario, todas las imágenes son creadas por el autor.