vLLM PagedAttention para una inferencia LLM 24 veces más rápida

vLLM PagedAttention for 24x faster LLM inference.

Una forma más eficiente de calcular la atención del Transformer durante la inferencia

PagedAttention para una consulta “el gato está durmiendo en la cocina y el perro está”. Los pares de tensores clave-valor para el cálculo de la atención se almacenan en bloques virtualmente contiguos mapeados a bloques no contiguos en la memoria GPU. — Imagen del autor

Casi todos los modelos de lenguaje grandes (LLM) dependen de la arquitectura neuronal Transformer. Aunque esta arquitectura es elogiada por su eficiencia, tiene algunos cuellos de botella computacionales bien conocidos.

Durante la decodificación, uno de estos cuellos de botella es el cálculo de la atención con pares de tensores clave-valor para cada token de la entrada. Todos estos tensores deben almacenarse en la memoria.

Nota: no explicaré en este artículo cuál es el papel de estos pares clave-valor. Es uno de los aspectos más complicados e interesantes de la arquitectura Transformer. Si no lo sabes, te recomiendo encarecidamente que leas “The Illustrated Transformer” de Jay Alammar.

A medida que los LLM aceptan entradas cada vez más largas, por ejemplo, el LLM Claude acepta entradas de hasta 100.000 tokens, la memoria consumida por estos tensores puede volverse muy grande.

Almacenar todos estos tensores en la memoria de manera ingenua conduce a una sobre-reservación y fragmentación de la memoria. Esta fragmentación puede hacer que el acceso a la memoria sea muy ineficiente, especialmente para secuencias largas de tokens. En cuanto a la sobre-reservación, el sistema lo hace para asegurarse de haber asignado suficiente memoria para los tensores, incluso si no la consume toda.

Para aliviar estos problemas, UC Berkeley propone PagedAttention.

PagedAttention se implementa en vLLM (licencia Apache 2.0), que es implementado por LMSYS, una organización para investigación abierta fundada por estudiantes y profesores de UC Berkeley con la ayuda de UCSD y CMU.

En este artículo, explicaré qué es PagedAttention y por qué acelera significativamente la decodificación. Muestro al final del artículo cómo empezar con vLLM para aprovechar PagedAttention para la inferencia y el servicio de LLM en tu ordenador.

PagedAttention para Transformer

Kwon et al. (2023) proponen PagedAttention.

El objetivo es almacenar tensores clave-valor de manera más eficiente en los espacios no contiguos de la VRAM de la GPU.

En resumen, la idea detrás de PagedAttention es crear bloques virtuales contiguos mapeados a bloques físicos en la memoria GPU.

Cada bloque está diseñado para almacenar tensores de pares clave-valor para un número predefinido de tokens. Todos los bloques son virtualmente contiguos y mapeados a bloques físicos no contiguos, asignados a demanda durante la inferencia, en la memoria de la GPU fragmentada. También se crea una simple tabla de índices en la memoria para asociar lo virtual con lo físico.

El kernel de PagedAttention recupera estos bloques según sea necesario. Esto es eficiente porque el sistema recupera números más pequeños de tensores clave-valor debido al tamaño limitado de los bloques.

Tomemos la siguiente consulta como ilustración:

el gato está durmiendo en la cocina y el perro está

Tenemos tensores clave-valor para cada token. Con PageAttention, podemos (arbitrariamente) establecer el tamaño del bloque en 4. Cada bloque contiene 4 tensores clave-valor, excepto el último que contiene solo 3 tensores clave-valor. Los bloques son virtualmente contiguos pero no necesariamente contiguos en la memoria de la GPU, como se ilustra en la figura de la introducción de este artículo.

Para el cálculo de la atención, para cada token de consulta, el sistema recupera el bloque uno por uno, como se ilustra a continuación.

Ilustración de bloques virtuales que contienen tensores de clave-valor para hasta 4 tokens — Imagen del autor

Al recuperar los tensores clave-valor por bloques, en lugar de toda la secuencia de tensores, el cálculo de la atención es mucho más rápido.

Muestreo paralelo para inferencia

Otra ventaja de PagedAttention es que los bloques virtuales pueden ser compartidos cuando se realiza el muestreo durante la inferencia. Todas las secuencias generadas en paralelo a través del muestreo o la búsqueda de haz pueden usar los mismos bloques virtuales, evitando duplicados.

En sus experimentos, LMSYS observó una reducción del 55% en el uso de memoria para la decodificación de la búsqueda de haz.

Rendimiento de PagedAttention informado por LMSYS

Antes de probarlo por nosotros mismos, echemos un vistazo al rendimiento informado por los autores (UC Berkeley/LMSYS) al utilizar PagedAttention implementado en vLLM en comparación con la biblioteca de inferencia de generación de texto desarrollada por Hugging Face.

Rendimiento de los modelos de LLaMa para tareas de completado de salida para la biblioteca original de Hugging Face (HF), biblioteca de inferencia de generación de texto (TGI) y vLLM con PagedAttention (vLLM) - Gráficos de UC Berkeley y LMSYS

vLLM parece ser mucho más rápido según estos resultados, especialmente en el caso de múltiples completados de salida. La diferencia entre TGI y vLLM aumenta con modelos más grandes. Esto es esperado ya que los modelos más grandes requieren más memoria y, por lo tanto, son más impactados por la fragmentación de memoria.

En general, vLLM es hasta 24 veces más rápido que la biblioteca Hugging Face Transformers.

Nota: De hecho, también estoy impresionado por la mejora de HF a TGI. Todavía no he cubierto TGI en mi blog, pero probablemente escribiré una guía al respecto. TGI se utiliza en producción en Hugging Face. Si bien parece mucho más lento que vLLM, TGI tiene otras ventajas, como el soporte para muchos más modelos y características.

Cómo configurar vLLM en su computadora

Nota: vLLM aún no es compatible con CUDA 12. Use una versión inferior, como la 11.8.

En esta sección, solo explicaré los conceptos básicos de cómo configurar y ejecutar vLLM en su computadora. Para un uso más avanzado, puede consultar la documentación de vLLM.

Mientras escribo este artículo, vLLM solo admite algunos tipos de modelos:

  • GPT-2
  • GPT-NeoX y basado en Pythia
  • Basado en LLaMa
  • Basado en OPT

Puede agregar el soporte de otros modelos siguiendo estas instrucciones.

En el código a continuación, uso Dolly V2 (licencia MIT). Es un modelo de chat basado en Pythia y entrenado por DataBricks.

Elegí la versión más pequeña con 3 mil millones de parámetros. Puede ejecutar una GPU de consumo con 24 GB de VRAM, por ejemplo, una nVidia RTX 3080/3090.

La forma más sencilla de instalar vLLM es con pip:

pip install vllm

Nota: Esto debería tardar hasta 10 minutos.

Pero en mi caso, tanto en mi computadora como en Google Colab, pip no pudo instalar la biblioteca vllm. Los autores de vLLM confirman que hay un problema con algunas versiones de nvcc y entornos. Sin embargo, para la mayoría de las configuraciones, pip debería instalar vLLM sin ningún problema.

Si se encuentra en la misma situación que yo, la solución alternativa es simplemente utilizar una imagen de Docker. Esta funcionó para mí:

docker run --gpus all -it --rm --shm-size=8g nvcr.io/nvidia/pytorch:22.12-py3

Nota: Una vez en el docker, los autores recomiendan eliminar Pytorch antes de instalar vLLM: pip uninstall torch. Luego, “pip install vllm” debería funcionar.

Luego, podemos comenzar a escribir en Python.

Primero necesitamos importar vllm, y luego cargamos el modelo con vllm. La inferencia se activa con llm.generate().

from vllm import LLM
prompts = ["Háblame de la gravedad"] #Puedes poner varias indicaciones en esta listallm = LLM(model="databricks/dolly-v2-3b")  # Cargar el modelooutputs = llm.generate(prompts)  # Desencadenar la inferencia

También puedes usar vLLM para servir LLMs. Funciona de manera similar a TGI. También es mucho más simple que ejecutar el servidor de inferencia NVIDIA Triton que describí en un artículo anterior.

Primero necesitas iniciar el servidor:

 python -m vllm.entrypoints.openai.api_server --model databricks/dolly-v2-3b

Nota: El servidor escuchará en el puerto 8000. Asegúrate de que esté disponible o cámbialo en el archivo de configuración de vLLM.

Luego, puedes consultar el servidor con indicaciones de la siguiente manera:

curl http://localhost:8000/v1/completions \    -H "Content-Type: application/json" \    -d '{        "model": "databricks/dolly-v2-3b",        "prompt": "Háblame de la gravedad",        "max_tokens": 200    }'

¡Y eso es todo! Tienes un servidor LLM muy eficiente funcionando en tu computadora.

Conclusión

PagedAttention acelera significativamente la inferencia. Es otro paso hacia la IA más asequible con LLM.

En experimentos posteriores, confirmé que vLLM es especialmente eficiente con lotes de indicaciones. Para aprovechar al máximo vLLM, considera optimizar tu estrategia de lotificación para la inferencia.

Mientras que la búsqueda de haz con grandes haces podría haber sido prohibitiva con el cómputo de atención estándar, la búsqueda de haz con PagedAttention es más rápida y eficiente en memoria.

Uno de mis próximos experimentos será combinar PagedAttention con QLoRa para reducir el uso de memoria. Debería ser sencillo. Hacer funcionar LLMs en hardware para consumidores sería aún más eficiente.

Si te gusta este artículo y estás interesado en leer los siguientes, la mejor manera de apoyar mi trabajo es convertirte en miembro de Zepes usando este enlace:

Como miembro de Zepes, una parte de tu cuota de membresía va a los escritores que lees, y tienes acceso completo a todas las historias…

Zepes.com

Si ya eres miembro y quieres apoyar este trabajo, solo sígueme en Zepes.

We will continue to update Zepes; if you have any questions or suggestions, please contact us!

Share:

Was this article helpful?

93 out of 132 found this helpful

Discover more

Inteligencia Artificial

IA Generativa en la Salud

Introducción La inteligencia artificial generativa ha ganado repentinamente tracción en los últimos años. No es sorpr...

Inteligencia Artificial

Descodificación del habla

Trabajando para traducir pensamientos a habla o acciones a través de una interfaz cerebro-computadora.

Inteligencia Artificial

De Experimentos 🧪 a Despliegue 🚀 MLflow 101 | Parte 01

Imagínate esto tienes una nueva idea de negocio y los datos que necesitas están a tu alcance. Estás emocionado/a por ...

Ciencia de Datos

Consejos de Matplotlib para mejorar instantáneamente tus visualizaciones de datos - Según Storytelling with Data

Cómo mejorar una visualización de datos de Matplotlib y seaborn basada en las lecciones aprendidas en Storytelling wi...

Ciencia de Datos

AI, Gemelos Digitales para Desatar la Próxima Ola de Innovación en la Investigación del Clima

La inteligencia artificial y la computación acelerada ayudarán a los investigadores del clima a lograr los milagros q...

Inteligencia Artificial

Cómo elegir la estrategia adecuada de visualización de datos para tu proyecto

Nuevas herramientas y paquetes van y vienen, pero la gramática básica de la visualización de datos sigue siendo incre...