Distributed Llama 2 en CPUs

Dist. Llama 2 en CPUs

Un ejemplo práctico de inferencia a granel en hardware de uso común utilizando Python, a través de llama.cpp y PySpark.

Imagen de autor vía DALL-E

¿Por qué?

Este ejercicio trata sobre el uso de Llama 2, un LLM (Large Language Model) de Meta AI, para resumir muchos documentos a la vez. La sumarización escalable de texto no estructurado, semi-estructurado y estructurado puede existir como una característica por sí misma, y también formar parte de tuberías de datos que alimentan a modelos de aprendizaje automático.

Específicamente, queremos demostrar la viabilidad simultánea de:

  • Ejecutar Llama 2 en CPUs (es decir, eliminar las limitaciones de capacidad de GPU)
  • Integración fluida de un LLM con Apache Spark (parte clave de los ecosistemas de Big Data)
  • No utilizar puntos finales de terceros (es decir, los modelos deben ejecutarse localmente debido a la infraestructura aislada o los requisitos de confidencialidad)

¿Cómo?

¡Gran parte del trabajo duro ya está hecho por nosotros!

El proyecto llama.cpp permite ejecutar LLM simplificados en CPUs mediante la reducción de la resolución (“cuantización”) de sus pesos numéricos. Estos archivos de modelo listos para usar están fácilmente disponibles.

A continuación, los enlaces llama-cpp-python proporcionan un acceso sencillo para utilizar llama.cpp desde Python.

Por último, el método applyInPandas() de Spark (documentación) permite dividir fuentes de datos gigantes en fragmentos del tamaño de Pandas y procesarlos de forma independiente. Hay que tener en cuenta que este enfoque puede ser un patrón antiestético si las funciones vectorizadas de Spark pueden lograr el mismo resultado, pero en nuestro caso, básicamente estamos utilizando Spark como un simple orquestador para escalar nuestro uso de llama.cpp. Es probable que haya formas más eficientes de utilizar llama.cpp en el procesamiento por lotes, pero este enfoque resulta atractivo debido a la simplicidad y los beneficios automáticos de la tolerancia a fallos y la escalabilidad de Spark.

Plan

Licencia Creative Commons (CC BY-SA 3.0)

Como una prueba divertida, utilizaremos Llama 2 para resumir “Guerra y Paz” de León Tolstói, una novela de más de 1200 páginas con más de 360 capítulos. Trataremos cada capítulo como un documento. Hay que tener en cuenta que Llama 2 ya “conoce” la novela; preguntarle sobre un personaje clave genera esta salida (utilizando llama-2–7b-chat.ggmlv3.q8_0.bin):

Promoción: Describe brevemente al personaje Anna Pavlovna en 'Guerra y Paz'Respuesta: Anna Pavlovna es un personaje importante en la novela "Guerra y Paz" de León Tolstói. Es una mujer rica, inteligente y encantadora conocida por su belleza y perspicacia. A pesar de su belleza y encanto, también se la representa como manipuladora y astuta, utilizando a menudo su inteligencia y perspicacia para conseguir lo que quiere. A lo largo de la novela, se muestra a Anna Pavlovna como un personaje complejo y multifacético, con cualidades admirables y defectos. Es particularmente cercana al personaje principal, Pierre Bezukhov, y desempeña un papel importante en su desarrollo personal y emocional.

Pasos:

  1. Instalar el modelo de chat cuantizado 7B y llama-cpp-python.
  2. Descargar la novela, dividirla por capítulos y crear un DataFrame de Spark.
  3. Dividir por capítulos y generar resúmenes.

Instalación

Configurar un clúster de Spark está fuera de nuestro alcance; asumiré que tiene Spark funcionando localmente, a través de un servicio gestionado (como Synapse o Elastic Map Reduce), o una implementación personalizada como Kubernetes.

Existen dos artefactos que deben instalarse en todos los nodos de trabajo, ya sean máquinas físicas, máquinas virtuales o pods en un grupo sin servidor:

  • El modelo LLama 2 en formato GGML (ubicado en /models)
  • El módulo llama-cpp-python (instalado mediante pip)

Estamos utilizando la versión 7B chat “Q8” de Llama 2, que se encuentra aquí. Los enlaces de descarga pueden cambiar, pero una configuración de un solo nodo, “bare metal” es similar a la siguiente:

Asegúrese de poder utilizar el modelo a través de python3 y este ejemplo. Para resumir, cada contexto de Spark debe poder leer el modelo desde /models y acceder al módulo llama-cpp-python.

Procesando el Texto de la Novela

Los comandos Bash a continuación descargan la novela e imprimen las cuentas de palabras.

A continuación, leemos el archivo de texto en Python, eliminando el encabezado y el pie de página de Project Gutenberg. Dividiremos en la expresión regular CHAPTER .+ para crear una lista de cadenas de capítulos y crearemos un DataFrame de Spark a partir de ellos (este código asume una SparkSession llamada spark).

El código debería producir la siguiente salida:

número de capítulos = 365máximo de palabras por capítulo = 3636+------------------------------------------------------------+-------+|                                                        text|chapter|+------------------------------------------------------------+-------+|\n\n"Bueno, príncipe, entonces Génova y Lucca son ahora simplemente una familia...|      1||\n\nLa sala de dibujo de Anna Pávlovna se estaba llenando gradualmente. T...|      2||\n\nLa recepción de Anna Pávlovna estaba en pleno apogeo. Las señoritas...|      3||\n\nJusto en ese momento, otro visitante entró en la sala de dibujo: P...|      4||\n\n"Y qué opinas de esta última comedia, la cor...|      5||\n\nDespués de agradecer a Anna Pávlovna por su encantadora soiree,...|      6||\n\nSe escuchó el ruido de un vestido de mujer en la habitación de al lado...|      7||\n\nLos amigos guardaron silencio. A ninguno le importaba comenzar a hablar...|      8||\n\nEran más de la una cuando Pierre dejó a su amigo....|      9||\n\nEl príncipe Vasíli cumplió la promesa que le había hecho al príncipe...|     10|+------------------------------------------------------------+-------+

¡Genial! Ahora tenemos un DataFrame con 365 filas, cada una conteniendo el texto completo y el número del capítulo. El último paso es crear un nuevo DataFrame con resúmenes de cada capítulo.

Procesamiento con Spark

A continuación se muestra el código Python para generar un resumen de un solo capítulo (observe la llamada a limit(1) para devolver una sola fila). La explicación aparece debajo del fragmento de código:

La función llama2_summarize() es el código que se aplica por grupo en Spark. Dado que estamos agrupando por la columna chapter, la función se llama en cada fila de capítulo; el argumento df es simplemente un DataFrame de Pandas con una sola fila. Tenga en cuenta que estamos leyendo el modelo en cada llamada de llama2_summarize(); esto es un atajo que estamos tomando por simplicidad, pero no muy eficiente.

Por último, utilizando Spark, hacemos el groupby() y llamamos a applyInPandas(), estableciendo el esquema para incluir el resumen del capítulo y el número.

La salida (reformateada para mayor legibilidad) se ve así:

resumenEl capítulo trata sobre una conversación entre el príncipe Vasíli Kurágin y Anna Pávlovna Schérer, una socialité conocida y favorita de la emperatriz Márya Fëdorovna. Están discutiendo varios asuntos políticos, incluida la posibilidad de guerra con Francia y el papel de Austria en el conflicto. El príncipe Vasíli espera asegurar un puesto para su hijo a través de la emperatriz viuda, mientras que Anna Pávlovna está entusiasmada con el potencial de Rusia para salvar a Europa de la tiranía de Napoleón. La conversación también toca asuntos personales, como la insatisfacción del príncipe Vasíli con su hijo menor y la sugerencia de Anna Pávlovna de casar a su hijo pródigo Anatole con una heredera adinerada.capítulo1

(Ten en cuenta el uso de Napoleon a pesar de que no aparece en el capítulo. Una vez más, este es un ejercicio divertido en lugar de un ejemplo realista que utiliza documentos verdaderamente inéditos.)

El tiempo de ejecución para esta prueba de un solo capítulo es de aproximadamente 2 minutos en una máquina virtual de 64 núcleos. Hay muchas opciones que hemos pasado por alto y que afectan al tiempo de ejecución, como el tamaño/quantización del modelo y los parámetros del modelo. El resultado clave es que, escalando adecuadamente nuestro clúster de Spark, podemos resumir todos los capítulos en unos pocos minutos. Por lo tanto, es posible procesar cientos de miles (¡o incluso millones!) de documentos diariamente utilizando grandes clústeres de Spark compuestos por máquinas virtuales económicas.

Resumen

Ni siquiera hemos mencionado ajustar los parámetros estándar de LLM como temperature y top_p, que controlan la “creatividad” y la aleatoriedad de los resultados, o la ingeniería de la consigna, que prácticamente es una disciplina propia. También elegimos el modelo Llama 2 7B sin justificación; podría haber modelos más pequeños y más eficientes o familias de modelos más adecuados para nuestro caso de uso particular.

En cambio, hemos mostrado cómo distribuir fácilmente cargas de trabajo de LLM (quantizadas) utilizando Spark con un esfuerzo bastante mínimo. Los siguientes pasos podrían incluir:

  • Carga/caché más eficiente de modelos
  • Optimización de parámetros para diferentes casos de uso
  • Consignas personalizadas