Inferencia rápida y portable de Llama2 en edge heterogéneo

(secondstate.io)

2 puntos por GN⁺ 2023-11-14 | 1 comentarios | Compartir por WhatsApp

Para reducir las grandes dependencias y la carga de despliegue del stack de inferencia centrado en Python, se empaqueta una app Llama2 en Rust+Wasm como un binario de 2 MB para ejecutarla en varios dispositivos
La implementación es una adaptación de llama.cpp para Wasm, y aprovecha la aceleración de hardware local con archivos de modelo GGUF, el plugin GGML de WasmEdge y la API WASI NN
El ejemplo usa un modelo Llama-2-7B-Chat GGUF cuantizado a 5 bits, con un rendimiento de aproximadamente 25 tokens/s en una MacBook M2 de bajo costo y alrededor de 50 tokens/s en una Nvidia A10G
El runtime WasmEdge puede usarse junto con herramientas de contenedores, lo que facilita desplegar el mismo binario Wasm en dispositivos edge, on-premise y en la nube
Las herramientas GGML ya están disponibles, pero siguen en una etapa temprana; aún falta soporte para más plugins de hardware y sistemas operativos, configuraciones de llama.cpp y la API WASI NN en varios lenguajes compatibles con Wasm

Ejecutar inferencia de Llama2 con Rust+Wasm

El stack Rust+Wasm puede usarse como alternativa a Python en inferencia de IA
- Una app Rust+Wasm puede ser 1/100 del tamaño de una en Python y hasta 100 veces más rápida, y puede aprovechar aceleración de hardware en distintos entornos sin cambiar el mismo código binario
Second State creó un programa en Rust simple que ejecuta inferencia del modelo Llama2 a velocidad nativa
- La app binaria compilada a Wasm ocupa 2 MB
- Esta app Wasm es portable entre dispositivos con aceleradores de hardware heterogéneos
- El runtime WasmEdge ofrece un entorno de ejecución seguro en ambientes de nube
- WasmEdge puede usarse con herramientas de contenedores para orquestar y ejecutar apps portables en varios dispositivos

Basado en llama.cpp y GGUF

Este trabajo se basa en llama.cpp, creado por Georgi Gerganov
El programa original en C++ se ajustó para ejecutarse en Wasm
Los archivos de modelo usan el formato GGUF

Procedimiento de ejecución

Instalar WasmEdge y el plugin GGML en un dispositivo Linux o Mac

curl -sSf https://raw.githubusercontent.com/WasmEdge/WasmEdge/… | bash -s -- --plugins wasmedge_rustls wasi_nn-ggml

Descargar la app Wasm precompilada

curl -LO https://github.com/LlamaEdge/LlamaEdge/…

El ejemplo usa un archivo GGUF del modelo Llama2 7B ajustado para chat, cuantizado con pesos de 5 bits

curl -LO https://huggingface.co/second-state/Llama-2-7B-Chat-GGUF/…

Al ejecutar la app de inferencia Wasm con WasmEdge y pasarle el modelo GGUF, se pueden ingresar preguntas de forma interactiva

wasmedge --dir .:. --nn-preload default:GGML:AUTO:Llama-2-7b-chat-hf-Q5_K_M.gguf llama-chat.wasm

Configuración del modelo y ejemplos de rendimiento

llama-chat.wasm configura la forma de interacción con el modelo mediante opciones de línea de comandos
- --ctx-size: tamaño del contexto del prompt, valor predeterminado 512
- --n-predict: cantidad de tokens a predecir, valor predeterminado 1024
- --n-gpu-layers: cantidad de capas a ejecutar en la GPU, valor predeterminado 100
- --batch-size: tamaño de lote para procesar prompts, valor predeterminado 512
- --temp: temperatura de muestreo, valor predeterminado 0.8
- --repeat-penalty: penalización por tokens repetidos, valor predeterminado 1.1
- --prompt-template: soporta llama-2-chat, codellama-instruct, mistral-instruct-v0.1, chatml, deepseek-chat, deepseek-coder, entre otros
- --log-stat: imprime estadísticas
Un ejemplo de ejecución con longitud de contexto 2048, respuesta máxima de 512 tokens e impresión de estadísticas es el siguiente

wasmedge --dir .:. --nn-preload default:GGML:AUTO:Llama-2-7b-chat-hf-Q5_K_M.gguf llama-chat.wasm -c 2048 -n 512 --log-stat

En una MacBook M2 de bajo costo, las respuestas del LLM se muestran por streaming de forma predeterminada y se generan a aproximadamente 25 tokens/s
- El ejemplo de log registró 25.64 tokens/s en 82 ejecuciones según eval time
- kv self size aparece como 1024.00 MB y compute buffer total size como 630.14 MB
En un equipo Nvidia A10G, se ejecuta a aproximadamente 50 tokens/s
- Usa aceleración CUDA en GPU y descarga 35/35 capas a la GPU
- El uso de VRAM aparece como 4474.93 MB para el modelo, 1648.02 MB para el contexto y 6122.95 MB en total
- Registró 50.55 tokens/s en 83 ejecuciones según eval time

Servidor API compatible con OpenAI y apps LLM

También se ofrece un servidor API compatible con OpenAI hecho con Rust y WasmEdge
Con este servidor se pueden crear agentes y apps LLM usando herramientas de desarrollo compatibles con OpenAI, como flows.network
La explicación relacionada está en Build a super lightweight AI agent

Límites del stack de inferencia en Python

Los LLM como Llama2 suelen entrenarse con frameworks basados en Python, como PyTorch, Tensorflow y JAX
Sin embargo, se considera que Python está en desventaja en las aplicaciones de inferencia, que representan alrededor del 95% del cómputo de IA
Los paquetes de Python pueden volverse difíciles de configurar y usar por sus dependencias complejas
Las imágenes Docker de Python o PyTorch suelen crecer a varios GB o decenas de GB, lo que aumenta la carga para la inferencia de IA en servidores o dispositivos edge
Python es mucho más lento que lenguajes compilados como C, C++ y Rust, y puede ser hasta 35,000 veces más lento
La mayoría de las cargas de trabajo reales se delegan a bibliotecas compartidas nativas debajo de wrappers de Python
- Esta estructura es buena para demos, pero se considera difícil modificar sus componentes internos para necesidades específicas de negocio
La alta dependencia de bibliotecas nativas y la gestión compleja de dependencias reducen la portabilidad de los programas de IA en Python que aprovechan capacidades de hardware específicas de cada dispositivo
También hay casos en los que paquetes de Python usados comúnmente en toolchains de LLM entran en conflicto entre sí por requisitos de versión de pydantic
- llama-cpp-python requiere pydantic 2.0.1 y no funciona con <=2.0
- fastapi y chromadb requieren pydantic 1.9.2 y no funcionan con >=2.0

Ventajas de Rust+Wasm

El stack Rust+Wasm puede usarse como una infraestructura unificada de computación en la nube que abarca dispositivos, edge cloud, servidores on-premise y nubes públicas
En aplicaciones de inferencia de IA, puede ser una alternativa sólida al stack de Python
Ultraligero
- La aplicación de inferencia ocupa 2 MB incluyendo todas las dependencias
- Es menos del 1% del tamaño de un contenedor típico de PyTorch
Ejecución rápida
- Puede alcanzar velocidad nativa de C/Rust en todo el flujo: preprocesamiento, cálculos de tensores y posprocesamiento
Portabilidad
- La misma aplicación en bytecode Wasm se ejecuta en las principales plataformas de cómputo
- También soporta aceleración de hardware heterogénea
Facilidad de configuración, desarrollo y despliegue
- Se reducen las dependencias complejas
- Se puede compilar un único archivo Wasm en una laptop con herramientas estándar y desplegarlo en varios entornos
Seguridad y preparación para la nube
- El runtime Wasm está diseñado para aislar código de usuario no confiable
- Puede gestionarse con herramientas de contenedores y desplegarse en plataformas cloud native

Estructura del programa de inferencia en Rust

El programa de inferencia de la demo está escrito en Rust y compilado a Wasm
El código fuente Rust principal tiene unas 40 líneas
El programa Rust se encarga de lo siguiente
- Gestionar la entrada del usuario
- Dar seguimiento al historial de conversación
- Convertir el texto al chat template de Llama2
- Ejecutar la inferencia con la API WASI NN
Para compilarlo directamente, instalar el compilador de Rust y el target de compilación wasm32-wasi

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
rustup target add wasm32-wasi

Luego descargar el proyecto fuente y compilar el archivo Wasm con cargo

git clone https://github.com/second-state/llama-utils
cd llama-utils/chat/
cargo build --target wasm32-wasi --release
cp target/wasm32-wasi/release/llama-chat.wasm .

Despliegue en la nube y en edge

Una vez que se tiene el archivo de bytecode Wasm, puede desplegarse en cualquier dispositivo que soporte el runtime WasmEdge
Actualmente, el plugin GGML soporta Linux genérico y Ubuntu Linux
- CPU x86 y ARM
- GPU Nvidia
- Apple M1/M2/M3
El plugin GGML de WasmEdge usa automáticamente la aceleración de hardware del dispositivo sobre la base de llama.cpp
- Si hay una GPU Nvidia, el instalador instala automáticamente el plugin GGML optimizado para CUDA
- El plugin GGML para Mac OS usa la API Metal para ejecutar cargas de inferencia en el motor de procesamiento neuronal integrado de M1/M2/M3
- La compilación para CPU en Linux usa OpenBLAS para detectar y aprovechar automáticamente capacidades modernas de cálculo en CPU como AVX y SIMD
Este enfoque busca lograr portabilidad entre hardware y plataformas de IA heterogéneos sin sacrificar rendimiento

Trabajo futuro

Las herramientas GGML de WasmEdge ya están disponibles y las usan clientes cloud native, pero todavía están en una etapa temprana
Las áreas donde se puede contribuir son las siguientes
- Agregar plugins GGML para más plataformas de hardware y sistemas operativos
  - Hay interés en TPU, NPU ARM y chips de IA especializados en Linux y Windows
- Soportar más configuraciones de llama.cpp
  - Actualmente, solo algunas opciones de configuración pueden pasarse desde Wasm al plugin GGML
  - El objetivo es soportar todas las opciones que ofrece GGML
- Soportar la API WASI NN en otros lenguajes compatibles con Wasm
  - Hay interés en Go, Zig, Kotlin, JavaScript, C y C++

Soporte para modelos más allá de LLM

WasmEdge y WASI NN también permiten crear apps de inferencia basadas en modelos populares de IA más allá de LLM, como una alternativa a Python ligera, rápida, portable y segura
mediapipe-rs ofrece una API Rust+Wasm para la familia de modelos Tensorflow de Google mediapipe
WasmEdge YOLO es un proyecto de API Rust+Wasm para manejar el modelo PyTorch YOLOv8
WasmEdge ADAS demo es un ejemplo que realiza segmentación de caminos para autos autónomos con modelos Intel OpenVINO
WasmEdge Document AI ofrecerá una API Rust+Wasm para OCR y una familia de modelos de procesamiento de documentos
Las discusiones y contribuciones relacionadas con WasmEdge pueden hacerse en WasmEdge Discord

1 comentarios

GN⁺ 2023-11-14

Comentarios de Hacker News

Rust y WASM están bien, pero si ves el código, son apenas 150 líneas de un script básico de línea de comandos en Rust
El trabajo pesado lo resuelve una sola línea que entrega el modelo al backend de WASI-NN, que en este caso lo proporciona el runtime de WasmEdge
Pero WasmEdge no es Rust sino C++, y en este caso casi no hay ventajas aportadas por Rust; además, el backend también puede invocarse desde otros lenguajes como Python
- La ventaja de Rust aquí parece ser el empaquetado y la distribución
  Empaquetar Python y PyTorch en una forma que el usuario final pueda ejecutar con doble clic sigue siendo casi un desastre, y aunque el código real de alto rendimiento es C++ en ambos casos, si en vez de distribuir más de 2 GB y decenas de miles de archivos solo para mandarle unos comandos a ese C++, puedes terminar con un único ejecutable de 2 MB, entonces sí tiene sentido
Es un gran trabajo. Pero quien quiera usarlo igual tiene que descargar los archivos de pesos, que siguen siendo grandes
En esencia, hicieron una versión de llama.cpp de 2 MB, completamente portable y sin dependencias
Para un desarrollador de apps, quizá sea la forma más sencilla de empaquetar el motor de inferencia dentro de un archivo distribuible. Los pesos ya son portables y pueden descargarse cuando hagan falta, así que la parte que de verdad quieres fijar es el motor de inferencia
- Sería más útil poner en el título wasm de 2 MB
  Como dijiste, los pesos superan por mucho ese tamaño
- En mi máquina, el archivo main que compila llama.cpp pesa 1.2 MB
  El tamaño de 2 MB como tal no impresiona demasiado; el punto es que al apuntar a wasm se vuelve más portable, no que esté especialmente más comprimido
¿Esto no es al final un envoltorio sobre llama.cpp? La verdad, ya estoy bastante cansado de los proyectos que envuelven x.cpp
Llevo los últimos 6 meses desarrollando un framework de machine learning en Rust + WebGPU, y rápidamente me di cuenta de lo impresionante que es el trabajo de GG
Aún está en una etapa temprana, pero puede verse aquí:
https://www.ratchet.sh/
https://github.com/FL33TW00D/whisper-turbo
- ¿Podrías explicar más qué fue lo que te impresionó? No conozco nada de esta área, así que me cuesta entender bien su valor
- ¿Acabas de menospreciar el trabajo de otros y en el mismo comentario promocionar el tuyo? De verdad hace falta reflexionar sobre la ética de eso
- ¿Quién es GG?
wasm-nn, del que esto depende, es decir https://github.com/WebAssembly/wasi-nn, es una propuesta basada en enviar chunks arbitrarios a una implementación del proveedor. En la práctica, la API es poco más que configurar entrada, calcular y configurar salida
Así que no es portable en absoluto
Esto funciona porque se apoya en la abstracción ya implementada en llama.cpp. Aquí se ve que WasmEdge simplemente tomó ese código que mapea modelos gguf a distintos objetivos de hardware: https://github.com/WasmEdge/WasmEdge/tree/master/plugins/was...
Por lo tanto, una explicación del tipo “los desarrolladores pueden usar bindings para escribir apps de machine learning en un lenguaje de alto nivel, compilarlas a WebAssembly y luego ejecutarlas en runtimes con soporte de wasi-nn como WasmEdge” es completamente incorrecta. En la práctica, no se puede hacer así
Esto no es portable, tampoco es un sandbox ni una capa de abstracción de hardware
Aunque tengas el binario wasm, solo podrá ejecutarse si la versión del runtime en uso casualmente implementa el backend específico de ggml que necesitas, y como no hay ninguna obligación de que lo haga, lo más probable es que no sea así
E incluso si funciona, al final está llamando al código ggml de llama.cpp, así que la seguridad depende de esa librería y nada más
El artículo enfatiza demasiado la “portabilidad” y “Rust”, pero en la práctica no se ve muy clara la ventaja de ninguna de las dos
Aunque asumas que existe un runtime WASI en hardware nuevo, ¿puedes ejecutar el modelo? ¿Hay soporte para GPU? La respuesta termina siendo: “revisa si llama.cpp compila con soporte de GPU en esa plataforma, si el runtime que usas tiene el plugin de ggml y si dentro trae embebida la versión correcta de ggml. Si no, entonces no”
Si ese es el caso, ¿para qué usar WASI entonces?
Es cierto que el soporte GPU multiplataforma es difícil, pero esto se ve bastante absurdo
Imagina que WebGPU funcionara como “lánzale un chunk binario a la GPU y quizá dibuje algo si resulta ser el adecuado para el hardware actual”; pues esta estructura es básicamente eso
- Gracias por la explicación. Tenía curiosidad de dónde sale el soporte de GPU en WASM
- ¿Podrías explicar también con más detalle las implicaciones en seguridad?
¿Se puede correr esto offline en un iPhone? Así podría usarse como una especie de búsqueda básica en internet sin depender de la señal, lo que sería útil para ir de camping
- Puede ejecutarse en varios dispositivos Linux, Mac y Windows, y eso incluye Raspberry Pi y la mayoría de laptops y servidores
  Pero para cargar el modelo en sí, todavía hacen falta varios GB de memoria
- Probé este proyecto en un Pixel. Parece que también funciona en algunos iPhone/iPad
  [0] https://github.com/mlc-ai/mlc-llm
- Para ese caso de uso, probablemente convenga más descargar una versión de Wikipedia. También están temas como la entropía
- Hice una app comercial para iOS relativamente exitosa para ese caso de uso
  Al principio la hice con ggml, y después encontré mlc-llm y la porté sobre esa base
  [0]: https://apps.apple.com/us/app/private-llm/id6448106860
Seguramente seguirán apareciendo más formas de ejecutar la arquitectura Transformer de manera más eficiente y rápida en el edge, pero como los requisitos de VRAM no se pueden resolver con Rust, parece que ya se está llegando a un límite
Ese es justamente el principal cuello de botella cuando se intenta cargar un modelo lo suficientemente grande
Se puede decir “los modelos pequeños están mejorando, mira Mistral y llama 2”, pero incluso los modelos pequeños se están acercando a su límite de capacidad. Hay un límite a la cantidad de información que se puede meter en 7 mil millones de parámetros
No parece que este enfoque de IA vaya a conducir a AGI. Se siente demasiado ineficiente
- Incluso en modelos pequeños todavía parece haber bastante margen, con cosas como sistemas MoE o cargar LoRA dinámicamente
La explicación de que “el plugin GGML para Mac OS usa la API de Metal para ejecutar tareas de inferencia en el motor neuronal integrado de M1/M2/M3” no parece correcta
Que GGML use la API de Metal significa que se ejecuta en la GPU de los M1/2/3, no que se ejecute en el Neural Engine
Claro que eso ya de por sí está bien, pero si nos ponemos estrictos, esa es la precisión del asunto
- No es una observación menor en absoluto. En https://github.com/ggerganov/llama.cpp/discussions/336 se discute, de forma algo dispersa, si realmente vale la pena usar directamente el Neural Engine más allá de la GPU
No me gusta el marketing clickbait que hace parecer que, solo porque el wrapper que rodea la biblioteca principal usa otro lenguaje, el proyecto reduce el tamaño a 1/100 frente a otras soluciones o aumenta la velocidad entre 100 y 35000 veces
Además, ignora por completo las herramientas y la experiencia comunitaria acumuladas alrededor de otras soluciones
Para empezar, este proyecto está basado en llama.cpp[1], y el trabajo pesado de cargar archivos de modelo de varios GB en la GPU/CPU y ejecutarlos lo hace llama.cpp
La velocidad de inferencia no está limitada por la elección del wrapper; también hay wrappers para Go, Python, Node y Rust, o simplemente se puede usar llama.cpp directamente
El tamaño del binario tampoco importa tanto. Los archivos de modelos cuantizados comunes están en el rango de 5GB a 40GB, y por eso se necesita una máquina con una GPU potente o con 16 a 64GB de RAM
[1] https://github.com/ggerganov/llama.cpp
Si una gran parte del tamaño al final son los pesos del modelo entrenado, ¿cómo se puede reducir el tamaño en varios órdenes de magnitud sin perder precisión?
- Creo que reducir el tamaño sin pérdida de precisión es difícil. Aun así, GGUF cuantizado es excelente
  Parece que esos 2MB se refieren al tamaño del programa sin incluir el modelo. Da la impresión de que ejecuta llama.cpp con wasm y un servidor Rust que lo corre
  Me gusta el pequeño llama.cpp/examples/server, así que lo tengo embebido en FreeChat, pero siempre es bueno tener más opciones de herramientas
  Revisé y el ejecutable arm64/x86 que yo embebo actualmente mide 4.2MB. FreeChat pesa 12.1MB, pero como el modelo base ronda los 3GB, no me preocupo mucho por una diferencia de 2MB
  [0]: https://github.com/ggerganov/llama.cpp/tree/master/examples/...
- Si te refieres a reducir en varios órdenes de magnitud, sin pérdida de precisión, el tamaño del modelo en sí, es decir, de los pesos entrenados, entonces ese es otro problema difícil aparte
  Este artículo trata de reducir 100 veces el tamaño de la app de inferencia
No veo por qué usar esto en vez de usar llama.cpp directamente
- Pista: la moneda de la economía de “Rewrite-it-in-Rust” no es realmente ejecutar algo
- llama.cpp normalmente tiene que compilarse por separado según el sistema operativo y la arquitectura, para Windows, macOS, Linux, etc., así que tiene menos portabilidad
  El artículo también explica que aprovecha la aceleración por hardware en dispositivos con aceleradores de hardware heterogéneos
  Eso significa que un programa compilado a wasm puede usar de forma eficiente distintos recursos de hardware, como GPU o chips especiales de IA, en varios dispositivos
  Si se implementara directamente en C++, podría requerir optimizaciones o versiones separadas para cada tipo de hardware para lograr un rendimiento similar

Inferencia rápida y portable de Llama2 en edge heterogéneo

Ejecutar inferencia de Llama2 con Rust+Wasm

Basado en llama.cpp y GGUF

Procedimiento de ejecución

Configuración del modelo y ejemplos de rendimiento

Servidor API compatible con OpenAI y apps LLM

Límites del stack de inferencia en Python

Ventajas de Rust+Wasm

Ultraligero

Ejecución rápida

Portabilidad

Facilidad de configuración, desarrollo y despliegue

Seguridad y preparación para la nube

Estructura del programa de inferencia en Rust

Despliegue en la nube y en edge

Trabajo futuro

Soporte para modelos más allá de LLM

Lecturas relacionadas

1 comentarios

Comentarios de Hacker News