El problema fundamental que Janas-LLM aborda es la democratización del acceso a la inferencia de modelos de lenguaje grandes (LLM), especialmente aquellos con arquitecturas Mixture-of-Experts (MoE), en hardware de consumo. Tradicionalmente, la ejecución de LLMs de decenas o cientos de miles de millones de parámetros ha requerido infraestructuras de GPU de alta gama y gran capacidad de memoria. Janas-LLM desafía esta limitación al demostrar que, con una gestión de memoria y E/S de disco optimizada, es posible realizar inferencia de modelos de 80B en una laptop con 32 GB de RAM y una NVMe SSD.
Esto es relevante ahora debido a la creciente complejidad y tamaño de los LLMs, que a menudo superan la capacidad de memoria de las GPUs de consumo. La arquitectura MoE, aunque prometedora para la eficiencia de entrenamiento y la capacidad de los modelos, presenta desafíos únicos en inferencia debido a la necesidad de cargar dinámicamente subconjuntos de pesos (expertos). Janas-LLM resuelve esto mediante un enfoque de streaming y caching inteligente, conectando la actualidad técnica de los LLMs con principios fundamentales de gestión de memoria y rendimiento de E/S.
Arquitectura del Sistema
Janas-LLM está implementado en C, sin dependencias externas más allá de la biblioteca estándar de C. Su arquitectura se centra en la gestión eficiente de los pesos del modelo, que residen en el disco (NVMe SSD) y se cargan bajo demanda. Para modelos MoE, solo los 'expertos' necesarios para cada token se transmiten desde el SSD. Los expertos que se utilizan con frecuencia se mantienen en una caché en memoria RAM, que Janas-LLM gestiona de forma adaptativa, llenando la memoria libre disponible. Utiliza O_DIRECT para lecturas de disco, minimizando la sobrecarga del caché del sistema operativo y maximizando el rendimiento del NVMe.
La inferencia aprovecha tanto la CPU (con instrucciones como AVX2, FMA, F16C, AVX-VNNI) como la GPU integrada (a través de Vulkan). El motor mide dinámicamente las capacidades del hardware subyacente y adapta su estrategia de ejecución (por ejemplo, el número de hilos de CPU, el uso de la GPU para el prompt inicial) para optimizar el rendimiento. Los modelos se almacenan en un formato propietario .jns (versión 2), que está diseñado para facilitar este streaming eficiente. La arquitectura también incluye un bloque de predicción de borradores (--mtp) para acelerar la generación de texto, donde un modelo más pequeño o una parte del modelo principal predice los siguientes tokens, que luego son verificados por el modelo completo.
Flujo de Inferencia de Token con Janas-LLM (MoE)
- 1 Solicitud de Token El modelo necesita generar el siguiente token.
- 2 Selección de Expertos El router del modelo MoE identifica los expertos necesarios para el token act...
- 3 Carga de Expertos (NVMe) Los pesos de los expertos requeridos se leen desde el NVMe SSD usando O_DIRECT.
- 4 Cache de Expertos (RAM) Los expertos cargados se almacenan en una caché en RAM para futuras reutiliza...
- 5 Cálculo (CPU/GPU) Los cálculos del token se realizan utilizando los expertos cargados, aprovech...
- 6 Generación de Token El token resultante se produce y se añade a la secuencia de salida.
| Capa | Tecnología | Justificación |
|---|---|---|
| compute | Intel Core Ultra 9 185H (CPU) | Procesamiento principal de la inferencia, incluyendo instrucciones AVX2, FMA, F16C, AVX-VNNI. |
| compute | Intel Arc GPU (integrada) | Aceleración de la inferencia, especialmente para el procesamiento de prompts, utilizando el backend Vulkan. |
| storage | NVMe SSD (1 TB) | Almacenamiento de los pesos del modelo, con streaming de expertos bajo demanda para modelos MoE. Lecturas O_DIRECT. 5.7 GB/s de lectura |
| cache | DDR5 RAM (32 GB) | Caché de expertos en memoria para reducir la latencia de acceso a disco, gestionada dinámicamente. |
| orchestration | Janas-LLM Engine (C) | Motor de inferencia que gestiona la carga de modelos, el scheduling de tareas en CPU/GPU y la gestión de memoria/disco. vs llama.cpp |
Trade-offs
Ganancias
- ▲ Capacidad de ejecutar modelos muy grandes (80B MoE) en hardware de consumo
- ▲▲ Rendimiento de inferencia superior en la máquina de desarrollo comparado con alternativas
- ▲ Eficiencia de recursos (memoria, disco) mediante streaming y caching adaptativo
Costes
- △ Posible inestabilidad del sistema si no hay suficiente headroom de memoria/disco
- △ Dependencia del rendimiento del NVMe SSD para modelos que no caben en RAM
- △ Complejidad de la implementación en C sin dependencias para optimización de bajo nivel
Fundamentos Teóricos
El desafío de ejecutar modelos grandes con recursos limitados se conecta con principios fundamentales de la jerarquía de memoria y la gestión de la paginación, estudiados desde los inicios de la computación. Conceptos como el 'working set' de un proceso, popularizado por Peter J. Denning en los años 60 y 70, son directamente aplicables: Janas-LLM intenta mantener el 'working set' de expertos (los más utilizados) en la memoria RAM, mientras que el resto permanece en el almacenamiento secundario. La optimización del acceso a disco mediante O_DIRECT y el streaming de datos se relaciona con técnicas de 'out-of-core computation', donde los algoritmos están diseñados para operar en conjuntos de datos que no caben completamente en la memoria principal.
La gestión de la caché de expertos es una aplicación práctica de algoritmos de reemplazo de caché (como LRU o LFU, aunque el artículo no especifica el algoritmo exacto), fundamentales en sistemas operativos y bases de datos. La idea de usar un modelo más pequeño o una parte del modelo para generar borradores que luego son verificados por un modelo más grande se alinea con el concepto de 'speculative execution' o 'drafting', una técnica que ha ganado tracción en la investigación de inferencia de LLMs para reducir la latencia, similar a cómo las CPUs usan la ejecución especulativa para mejorar el rendimiento.