Lm.rs: técnica mínima de inferencia de LLM en CPU basada en Rust, sin dependencias

 (github.com/samuel-vitorino)
1 puntos por GN⁺ 2024-10-12 | 1 comentarios | Compartir por WhatsApp
  • lm.rs es un proyecto para ejecutar inferencia de modelos de lenguaje locales en CPU con Rust, y busca ser una implementación de código mínima que realiza toda la inferencia sin bibliotecas de ML
  • Está inspirado en llama2.c y llm.c de Karpathy; al principio solo soportaba Google Gemma 2, pero se amplió para soportar Llama 3.2 y entrada de imágenes con PHI-3.5
  • Con los cambios más recientes se implementó el procesamiento por lotes, lo que aceleró la codificación de imágenes hasta aproximadamente 3 veces; Llama 3.2 1B corre a 50 tok/s en la máquina de 16 núcleos del autor
  • Los modelos preparados se pueden descargar desde Hugging Face; el README recomienda usar Q8_0 e indica que la cuantización Q4_0 todavía está en mejora
  • Los usuarios pueden descargar modelos y tokenizadores en formato LMRS para compilar de inmediato, o convertir los archivos de modelos originales de Hugging Face con export.py y tokenizer.py para ejecutarlos

Qué busca lm.rs

  • lm.rs es una implementación de inferencia de modelos de lenguaje locales en CPU escrita en Rust
  • Su objetivo es ser una implementación de código mínima que realice la inferencia completa de modelos de lenguaje en CPU sin bibliotecas de ML
  • Está inspirado en llama2.c y llm.c de Karpathy
  • El README indica que el código actual “no es tan mínimo”, y que parte del código tiene margen para optimización y mejoras
  • El proyecto también fue una oportunidad para que el autor probara Rust por primera vez

Modelos soportados y expansión multimodal

  • Al principio solo soportaba el modelo Gemma 2 de Google, pero luego se agregó soporte para modelos Llama 3.2
  • Recientemente se añadió la opción de usar imágenes mediante PHI-3.5
  • Elementos de soporte destacados actualmente
    • Soporte multimodal mediante el modelo PHI-3.5-vision
    • Soporte para el modelo solo de texto PHI-3.5-mini
  • Recursos relacionados

Rendimiento y modelos preparados

Flujo de conversión de modelos

  • Si descargas modelos cuantizados y tokenizadores preparados desde Hugging Face, puedes omitir el proceso de conversión
  • Para convertir directamente los modelos publicados por Google o Meta en Hugging Face, se deben instalar dependencias adicionales de Python
pip install -r requirements.txt
  • Se descargan y usan los archivos .safetensors y config.json desde la página del modelo original
  • Para modelos multimodales como PHI3.5 Vision también se necesita el archivo config de CLIP
  • export.py convierte pesos bfloat16 al formato LMRS
python export.py --files [ordered .safetensor files] --config [model config.json] --save-path [name and path to save] --type [model type (GEMMA/LLAMA/PHI)]
  • Para exportar una versión cuantizada, se usan las banderas --quantize y --quantize-type
  • El tamaño de un modelo cuantizado int8 puede reducirse de unos 9.8G a unos 2.5G, según el tamaño del grupo
  • Los modelos multimodales deben incluir el argumento --vision-config
  • tokenizer.py convierte el modelo del tokenizador al formato de tokenizador LMRS
python tokenizer.py --model-id [huggingface model_id] --tokenizer-type [type of the tokenizer (GEMMA/LLAMA/PHI)]

Compilación y ejecución

  • El código Rust se compila con cargo, y el README especifica pasar la bandera target-cpu
RUSTFLAGS="-C target-cpu=native" cargo build --release --bin chat
  • Para activar la función multimodal, se agrega el argumento --features multimodal
  • La ejecución básica se realiza especificando el archivo de pesos del modelo
./target/release/chat --model [model weights file]
  • Se pueden usar argumentos adicionales como tokenizer, temperature, top-p y show-metrics
  • Los argumentos disponibles se consultan con --help
  • En modelos multimodales se especifica la ruta de la imagen con el argumento --image
  • Al usar PHI3.5-vision, el README recomienda temperature 0

Ejecución del backend de WebUI

  • Para ejecutar el backend de WebUI, se compila con la función backend
RUSTFLAGS="-C target-cpu=native" cargo build --release --features backend --bin backend
  • El backend multimodal activa la función backend-multimodal
  • El backend se ejecuta especificando el archivo de pesos del modelo
./target/release/backend --model [model weights file]
  • Se pueden cambiar la IP y el puerto con --ip y --port
  • También se pueden usar banderas adicionales como temperature
  • Para compatibilidad multimodal se usa la bandera --multimodal
  • Después de ejecutarlo, se puede conectar a la interfaz web

Estado del TODO y licencia

  • Elementos completados
    • Agregar otros métodos de muestreo
    • De los puntos de prueba para modelos 9B y 27B, se completó la prueba de 9B; 27B aparece marcado como probablemente demasiado lento
    • Paralelización del bucle de atención multi-head
    • Agregar métricas de rendimiento
    • Soporte para cuantización int8 e int4
  • Elementos pendientes
    • Función para proporcionar system prompts
  • La licencia es MIT

Lecturas relacionadas

1 comentarios

 
GN⁺ 2024-10-12
Opiniones de Hacker News

  • Al probar el llama3.2-1b-it-q80.lmrs de 1.2 GB en una MacBook M2 de 64 GB, se sintió bastante rápido y, según Activity Monitor, usó 1000% de CPU en 13 hilos.
    Cloné lm.rs en /tmp, lo compilé con RUSTFLAGS="-C target-cpu=native" cargo build --release --bin chat, luego bajé tokenizer.bin y llama3.2-1b-it-q80.lmrs desde Hugging Face y lo ejecuté con ./target/release/chat --model llama3.2-1b-it-q80.lmrs.

    • Preguntan si se puede ejecutarlo con ./target/release/chat --model llama3.2-1b-it-q80.lmrs --show-metrics para ver cuántos tokens por segundo produce.
    • Al ingresar un prompt sencillo de conversación en francés, al principio pareció responder, pero pronto colapsó en una salida interminable de basura incomprensible.
      Por formato solo se dejó una parte, pero era una larga secuencia continua de palabras aleatorias.
    • Da curiosidad cómo habría que entender qué tan inteligente es en comparación con el ChatGPT más reciente.

  • El artículo está muy bien escrito, y podría usar parte del código fuente en clase para explicar cómo funcionan realmente los transformers.
    El código es más concreto y detallado que los diagramas de cabezales de atención. Sin embargo, si la librería imprime directamente en stdout, puede arruinar la salida de aplicaciones, por ejemplo cuando un editor de texto ofrece revisión de estilo; sería mejor escribir en un búfer de cadenas de una instancia de logging asociada al objeto lm.rs.
    También se ve una parte donde se usa unsafe en el lector de modelos para forzar la alineación de datos, y da curiosidad si sería posible hacerlo sin forzar tanto, sin unsafe.

    • Es mejor llamar a un callback de usuario que usar un búfer de cadenas.
      Así se podría manejar, por ejemplo, mostrando los logs en una GUI.

  • Ya hay varias herramientas en Rust para tareas con LLM y carga de modelos.
    Incluyen funciones como seleccionar automáticamente el modelo cuantizado más grande según la memoria disponible, extraer el tokenizador desde gguf o insertar prompts. Con esto quizá se puedan eliminar algunas dependencias de Python.
    Actualmente está orientado a soportar llama.cpp, pero esto también resulta bastante interesante. También preguntan si hay planes de soporte para restricciones gramaticales (grammar).
    https://github.com/ShelbyJenkins/llm_client

  • La expresión no dependency del título no es muy clara.
    Al verla por primera vez, pensé que podría ser no_std, pero en realidad no es no_std y parece tener algunas dependencias. Quizá quiera decir que todas son dependencias de Rust.

    • Poner títulos es difícil. Lo que se quería decir era que no hay dependencias de deep learning como PyTorch, CUDA u ONNX, y que toda la lógica está autocontenida.
      Para ser transparentes, hay 5 dependencias básicas de Rust, y entre ellas chrono y clap deberían gestionarse con feature flags para la función de chat. Las otras 3 son crates utilitarios para sacar un poco más de rendimiento del hardware: rayon para facilitar la paralelización, wide para ayudar con SIMD y memmap2 para mapear archivos de modelo en memoria.
    • Al ver el README, requirements.txt parece requerir PyTorch y varias dependencias de Python, y ese es el único lugar de la página donde aparece la palabra “dependency”, así que la expresión del título resulta bastante confusa.
      El proyecto en sí parece usar simplemente el subtítulo “Minimal LLM inference in Rust”. Al mirar el historial de Git, la cuenta que publicó esto parece ser contribuidora, pero no la autora principal; ayudaría que explicaran qué significa exactamente zero dependencies.
    • Quizá habría tenido sentido si el título original hubiera sido algo como “sin dependencias de hardware” o “sin dependencias de GPU”.
      Por desgracia, HN a veces elimina palabras de los títulos sin mucha razón ni transparencia.
    • Me pregunto si cargo de Rust ya se volvió casi como npm.
      No entiendo cómo puede llamarse sin dependencias si tiene 16 dependencias.

  • Hace tiempo hice algo parecido, pero el rendimiento era decepcionante comparado con código C/C++ corriendo en CPU.
    Eso también significa que no sabía bien cómo hacer que Rust fuera rápido. Sería bueno tener benchmarks de varias implementaciones en Rust.
    La implementación de inferencia de LLM podría convertirse en el nuevo “Hello, world!” para programadores serios.
    https://github.com/gip/yllama.rs

    • Yo también hice algo antes como una experiencia similar de “Hello, world”.
      https://github.com/crabml/crabml
      Usé directamente algunas instrucciones SIMD, y parecía que el rendimiento podía igualar al de llama.cpp. La clave parece estar en usar SIMD para la multiplicación de matrices cuantizadas y, al dividir el trabajo entre hilos, usar un bucle de espera activa en vez de variables de condición.
      Sin embargo, no he podido actualizarlo por un tiempo porque no tuve tiempo para seguir trabajando en la inferencia de modelos cuantizados en GPU con Vulkan.

  • Es interesante que ya use Dioxus, y me pregunto si WASM también podría entrar en la hoja de ruta.
    Si fuera posible correr un LLM ligero como RWKV en el navegador, el navegador podría abrir una nueva categoría de funciones sin llamar a una API SaaS.

    • He probado un poco con esto.
      https://github.com/maedoc/rwkv.js
      Usé Rwkv.cpp compilado con Emscripten, pero todavía no resolví bien la parte del tokenizador. Aun así, 1.6B RWKV6 parece lo suficientemente usable para un navegador exclusivamente offline.
      Para chat general le falta capacidad, pero para usos como RAG podría ser bastante suficiente.
    • La biblioteca en sí parece poder compilarse a WASM con muy pocos cambios.
      Las dependencias obligatorias rayon y wide ya soportan WASM, y si se cambia el tipo Mmap de transformer.rs por &[u8], también se puede eliminar memmap2.
      Sin embargo, RWKV tiene una arquitectura totalmente distinta, así que habría que implementarlo todo de nuevo, y parece muy poco probable que entre en la hoja de ruta.

  • Me pregunto si todas estas implementaciones están limitadas a CPU.
    Preguntan si, en caso de tener una buena GPU, lo correcto sería buscar otra alternativa.

    • Sí. Este proyecto funciona en CPU, por lo que no usa la GPU para el cálculo.
      Si quieres probar un framework de Rust con soporte para GPU, vale la pena mirar Candle https://github.com/huggingface/candle/tree/main.
    • Todo está implementado en CPU y, al menos por ahora, no hay ninguna aceleración por GPU.
      Si el objetivo es ejecutarlo en la práctica, aun usando solo CPU probablemente convenga más usar una alternativa, es decir, llama.cpp. Este proyecto se parece más a material educativo que muestra cómo funciona por dentro cuando se eliminan las capas complejas del ecosistema.
      Los LLM parecen magia por sus resultados, pero desde el punto de vista del código son bastante simples.
    • Según la GPU, puede haber una diferencia de 10 a 20 veces.
      En el lado de Rust hay wrappers de llama.cpp como mi llm_client, y proyectos basados en Candle como mistral.rs y Kalosm.
      Mi proyecto también intenta ofrecer una implementación con mistral.rs, pero todavía no he migrado por completo desde llama.cpp. Una implementación completa en Rust tiene grandes ventajas, como acelerar los tiempos de instalación. En este momento mi crate requiere clonar y compilar, así que aunque está automatizado en macOS, Windows y Linux, agrega alrededor de 1 minuto al tiempo de compilación.
    • Es CPU, sí, pero lo más importante es el ancho de banda de memoria.
      Por ejemplo, una RTX 3090 tiene casi 1 TB/s de ancho de banda de memoria. Para igualar eso habría que usar al menos 12 canales de la DDR5 de prueba de concepto más rápida del planeta.
      Si tienes una GPU dedicada, usar una implementación que la aproveche es un mundo completamente distinto. La razón por la que las cifras de inferencia de LLM en Apple Silicon son impresionantes también es su arquitectura de memoria unificada CPU-GPU de alto ancho de banda; si no recuerdo mal, era de unos 400 GB/s.
    • Depende del caso. Los buenos modelos son grandes y también tienen altos requisitos de memoria.
      Incluso una 4090 no tiene tanta memoria según los estándares de los LLM. La GPU será más rápida, pero es posible que no puedas cargar modelos grandes.

  • Me pregunto qué valor tiene esto en comparación con llama.cpp.

    • Puede ser más fácil integrarlo con otros proyectos en Rust.
    • Al tener menos funciones, la base de código es más limpia.

  • Es genial, y felicitaciones por crear tu primera biblioteca en Rust, pero para un uso local serio es indispensable el soporte de Metal/CUDA.

    • Usar CUDA va en contra del propósito de este proyecto, así que no encaja desde el inicio.
      Dicho eso, aunque no soy el autor principal sino colaborador, estoy experimentando con obtener cierta aceleración por GPU mediante wgpu. El autor principal quiere mantener la complejidad bajo control, así que no sé hasta dónde llegará realmente.

  • La pasión de la comunidad de Rust por reescribir casi todo me parece interesante y se agradece.