- La pila Rust+WASM puede convertirse en una alternativa poderosa a Python para la inferencia de IA
- En comparación con Python, una app en Rust+Wasm ocupa 1/100 del tamaño, es 100 veces más rápida y puede ejecutarse de forma segura en cualquier lugar con aceleración total de hardware sin cambiar el código binario
- Rust es el lenguaje de la AGI
- Si escribimos un programa simple en Rust para ejecutar inferencia de Llama2 a velocidad nativa
- Después de compilarlo a Wasm, el binario pesa apenas 2 MB, pero sigue siendo completamente portable entre aceleradores de hardware heterogéneos
- Además, el runtime de Wasm (
WasmEdge) ofrece un entorno de ejecución seguro para la nube. Funciona sin problemas con herramientas de contenedores para orquestar y ejecutar aplicaciones portables en distintos dispositivos
¿Por qué no Python?
- Los LLM como llama2 normalmente se entrenan con Python (por ejemplo: PyTorch, Tensorflow, JAX)
- Pero usar Python para aplicaciones de inferencia, que representan alrededor del 95% del cómputo de IA, es un gran error
- Los paquetes de Python tienen dependencias complejas. Son difíciles de configurar y usar
- Las dependencias de Python son enormes. Las imágenes Docker para Python o PyTorch normalmente pesan varios GB o incluso decenas de GB, lo que es especialmente problemático al hacer inferencia de IA en servidores o dispositivos edge
- Python es un lenguaje muy lento. Puede ser hasta 35,000 veces más lento que lenguajes compilados como C, C++ y Rust
- Como Python es lento, la mayoría de las cargas reales deben delegarse a bibliotecas compartidas nativas bajo wrappers de Python
- Por eso, las apps de inferencia en Python son excelentes para demos, pero muy difíciles de modificar internamente para ajustarlas a requerimientos específicos del negocio
- Debido a su fuerte dependencia de bibliotecas nativas y a la complejidad de la gestión de dependencias, es muy difícil portar programas de IA en Python a múltiples dispositivos aprovechando al mismo tiempo las capacidades únicas de hardware de cada uno
- Los paquetes de Python usados habitualmente en el toolchain de LLM suelen entrar en conflicto entre sí
- Chris Lattner, conocido por LLVM, TensorFlow y el lenguaje Swift, dio esta semana una gran entrevista en un pódcast de startups, donde explicó por qué Python es excelente para entrenar modelos, pero una mala elección para aplicaciones de inferencia
Ventajas de Rust+Wasm
- La pila Rust+Wasm ofrece una infraestructura unificada de cómputo en la nube que va desde dispositivos hasta edge cloud, servidores on-premise y nube pública
- Es una alternativa sólida a la pila de Python para aplicaciones de inferencia de IA, y no sorprende que Elon Musk haya dicho que Rust es el lenguaje de la AGI
- Ultraligera: la aplicación de inferencia pesa apenas 2 MB incluyendo todas las dependencias. Eso no llega ni al 1% del tamaño típico de un contenedor de PyTorch
- Muy rápida: ofrece velocidad nativa de C/Rust en todas las partes de la aplicación de inferencia, incluyendo preprocesamiento, operaciones con tensores y posprocesamiento
- Portable: puede ejecutar la misma aplicación en bytecode Wasm en todas las principales plataformas de cómputo con soporte para aceleración de hardware heterogénea
- Fácil de configurar, desarrollar y desplegar: ya no hay dependencias complejas. Se puede compilar un solo archivo Wasm con herramientas estándar en una laptop y desplegarlo en cualquier lugar
- Segura y compatible con la nube: el runtime de Wasm está diseñado para aislar código de usuario no confiable. Puede gestionarse con herramientas de contenedores y desplegarse fácilmente en plataformas cloud-native
Demo de inferencia en Rust
- Está escrita en 40 líneas de código Rust y compilada a Wasm
- Una vez generado el bytecode Wasm, puede desplegarse en cualquier dispositivo compatible con el runtime
WasmEdge
- El plugin GGML de WasmEdge, basado en llama.cpp, aprovecha automáticamente la aceleración de hardware del dispositivo para ejecutar el modelo llama2
Siguientes pasos
- El tooling de WasmEdge GGML ya está disponible y de hecho lo usan clientes reales de cloud-native, pero todavía está en una etapa temprana
- Si te interesa contribuir al proyecto open source y ayudar a definir la dirección futura de la infraestructura de inferencia para LLM, aquí hay algunas áreas donde potencialmente puedes aportar
- Agregar plugins GGML para más hardware y plataformas de SO: como TPU, ARM NPU y otros chips especializados de IA en Linux y Windows
- Dar soporte a más configuraciones de llama.cpp: actualmente ya se pueden pasar algunas opciones de configuración desde Wasm al plugin GGML, pero se busca soportar todas las opciones que ofrece GGML
- Soportar la API WASI NN en otros lenguajes compatibles con Wasm: especialmente Go, Zig, Kotlin, JavaScript, C y C++
5 comentarios
El enfoque más realista que Apple, que actualmente va rezagada en IA, podría tomar parece ser una situación en la que un motor de inferencia de software optimizado como el del artículo principal sea acelerado por la NPU de los chips A del iPhone y se cargue el archivo completo del modelo fundacional directamente en el iPhone.
Últimamente, tanto en la academia como en la industria, el consenso es que en el rendimiento de la IA importan más los datos y el fine-tuning que el tamaño, y hay muchas opiniones de que el rendimiento del fine-tuning en Llama2 es bastante bueno, así que también parece totalmente posible seguir distribuyendo a los dispositivos edge modelos bien entrenados. En ese sentido, una estrategia de meterlo todo en el iPhone sería bastante realista.
¿WASM se está volviendo así de popular? Pero, ¿no habría que comparar Python con Mojo?
Estoy de acuerdo. Mojo todavía debe considerarse en etapa experimental.
Mojo todavía no ha sido lanzado oficialmente, pero Rust ya es un lenguaje bastante maduro y está reemplazando varias partes en el trabajo real.
GN⁺: Inferencia de Llama2 rápida y portable en edge heterogéneo
Consulta también el resumen de GN⁺ y los comentarios de Hacker News.