PIM-SHERPA: Software Method for On-device LLM Inference by Resolving PIM Memory Attribute and Layout Inconsistencies

El artículo presenta PIM-SHERPA, un método puramente de software que resuelve las inconsistencias de atributos y diseño de memoria en sistemas con procesamiento en memoria (PIM) para habilitar una inferencia eficiente de modelos de lenguaje grandes (LLM) en dispositivos, logrando ahorros significativos de capacidad sin sacrificar el rendimiento.

Lo esencial

¡Claro que sí! Imagina que este artículo es como una historia sobre cómo hacer que un cerebro artificial gigante (una Inteligencia Artificial o LLM) funcione rápido y sin gastar mucha batería en tu teléfono móvil.

Aquí tienes la explicación de PIM-SHERPA usando analogías sencillas:

🧠 El Problema: El Cerebro con "Dos Personalidades"

Imagina que tu teléfono tiene un nuevo tipo de memoria llamada PIM (Memoria con Procesador). Es como tener una biblioteca donde los libros se leen y se procesan al mismo tiempo dentro del estante, sin tener que llevarlos a una mesa de lectura. Esto es increíblemente rápido para ciertas tareas.

Pero, para que este sistema funcione, la Inteligencia Artificial tiene dos modos de trabajo que son como dos personalidades opuestas:

1. La fase de "Lectura Rápida" (Prefill): Cuando escribes una frase larga, el sistema necesita leer mucho texto de golpe. Aquí, le conviene tener los libros (los datos) en una mesa de lectura cercana (memoria caché) para poder hojearlos rápido y repetidamente.
2. La fase de "Respuesta Lenta" (Decode): Cuando el sistema genera palabra por palabra, necesita usar la magia de la biblioteca (PIM). Pero para que la magia funcione, los libros NO deben estar en la mesa de lectura; deben estar directamente en los estantes (memoria no caché) para que la biblioteca pueda actuar sobre ellos.

El conflicto:

• Si pones los libros en la mesa (caché), la magia de la biblioteca no funciona porque los libros ya están "copiados" y no llegan al estante.
• Si los dejas en el estante (no caché), la lectura rápida se vuelve lenta porque tienes que ir y volver constantemente.

La solución antigua (y cara):
Los ingenieros decían: "¡Tenemos dos copias de los libros! Una en la mesa y otra en el estante".

• El problema: ¡Esto duplica el espacio que ocupan los libros! En un teléfono con poca memoria, esto significa que no podrías tener modelos de IA grandes. Sería como tener que llevar dos maletas idénticas en lugar de una.

---

🦘 La Solución: PIM-SHERPA (El Mago del Truco)

Los autores de este paper crearon PIM-SHERPA, un método de software (sin cambiar el hardware) que resuelve este problema sin duplicar los libros. Imagina a un ayudante mágico que hace dos trucos diferentes dependiendo de la situación:

Truco 1: El "Doble Tapón" (DRAM Double Buffering - DDB)

Imagina que tienes dos cubos de agua pequeños (memoria caché) y un río grande (memoria principal).

• Mientras el sistema está usando el Cubo A para trabajar (leer libros), el ayudante mágico está corriendo y llenando el Cubo B con los libros del río para la siguiente tarea.
• En cuanto termina con el Cubo A, cambia instantáneamente al Cubo B (que ya está lleno) y empieza a llenar el Cubo A para la siguiente ronda.
• Resultado: Nunca hay tiempo muerto. El trabajo y la preparación ocurren al mismo tiempo, como un mago que hace dos trucos a la vez.

Truco 2: El "Cambio de Ropa al Instante" (Online Rearrangement - OWR)

A veces, el trabajo es tan largo (como escribir un cuento muy extenso) que tienes tiempo de sobra.

• Aquí, el ayudante no necesita dos cubos. Solo toma los libros del estante, los "cambia de ropa" (los reorganiza para que sean fáciles de leer) y los pone en el cubo justo antes de empezar a trabajar.
• Como el trabajo es tan largo, el tiempo que tardas en cambiar la ropa es insignificante comparado con el tiempo total de la tarea.

---

🎁 ¿Qué logran con esto?

1. Ahorro de Espacio (La Maleta Ligera):
   En lugar de llevar dos maletas pesadas (duplicar los datos), ahora solo llevan una maleta y un pequeño bolso de mano (los cubos pequeños).

   • El resultado: Ahorran casi un 50% de espacio en la memoria del teléfono. Esto significa que puedes tener modelos de IA más inteligentes en tu móvil sin que se llene la memoria.

2. Velocidad (El Tren de Alta Velocidad):
   Logran que la velocidad sea casi tan buena como la del sistema teórico perfecto (el que duplicaba los datos), pero sin el costo extra de memoria.

3. Universalidad:
   Funciona en teléfonos actuales (como el Samsung Galaxy S24+) sin necesidad de cambiar el chip de memoria. Es una solución de software que "hackea" el sistema para que funcione mejor.

En resumen

PIM-SHERPA es como un director de orquesta inteligente que organiza a los músicos (los datos) para que, cuando necesitan leer rápido, estén cerca, y cuando necesitan usar la magia del procesador de memoria, estén en su lugar correcto, todo sin tener que contratar a un segundo grupo de músicos (duplicar la memoria).

Gracias a esto, tu teléfono podrá tener cerebros artificiales más grandes y potentes, consumiendo menos batería y espacio. ¡Una gran victoria para la inteligencia artificial en nuestros bolsillos! 📱✨

Resumen técnico

Resumen Técnico: PIM-SHERPA

1. El Problema: Inconsistencias en Sistemas PIM para LLMs

La implementación de Grandes Modelos de Lenguaje (LLMs) en dispositivos móviles y de borde enfrenta un desafío crítico al intentar utilizar arquitecturas de Procesamiento en Memoria (PIM), específicamente LPDDR-PIM, para acelerar la fase de decodificación (decode).

El artículo identifica dos inconsistencias fundamentales entre la fase de prefill (relleno) y la fase de decodificación que los sistemas anteriores no han resuelto eficazmente:

1. Inconsistencia de Atributos de Memoria:

   • Prefill (GEMM): Es intensivo en cómputo. Para maximizar la reutilización de datos, las pesas (weights) deben residir en una región cachable (cacheable region).
   • Decodificación (GEMV): Es intensivo en ancho de banda. Para que PIM se active, la memoria debe ser no cachable (non-cacheable). Si las pesas están en caché, las solicitudes de la CPU no llegan al controlador de memoria, bloqueando la ejecución de PIM.
   • Conflicto: Un mismo modelo requiere atributos de memoria opuestos para sus dos fases principales.

2. Inconsistencia de Diseño de Pesas (Layout):

   • Host-friendly: El procesador host prefiere un diseño contiguo (intercalado entre canales y bancos) para maximizar el ancho de banda de DRAM.
   • PIM-aware: PIM requiere que una fila completa de la matriz de pesos resida dentro de un solo banco de DRAM para maximizar la utilización SIMD interna.
   • Conflicto: Los diseños óptimos para el host y para PIM son mutuamente excluyentes.

Soluciones anteriores y sus limitaciones:

• Duplicación de Pesas (ej. HBM-PIM): Mantener dos copias (una cachable/host y otra no cachable/PIM). Esto duplica el uso de memoria DRAM, haciendo inviable la ejecución de modelos grandes en dispositivos móviles con memoria limitada (8-12 GB).
• Reordenamiento en Línea con Hardware (ej. FACIL): Requiere modificaciones al controlador de memoria (hardware) para mapear direcciones dinámicamente, lo que limita su aplicabilidad en sistemas comerciales donde el hardware es fijo.

---

2. Metodología: PIM-SHERPA

PIM-SHERPA es un método exclusivamente de software diseñado para resolver ambas inconsistencias sin modificar el hardware. Utiliza búferes pequeños en la región cachable y realiza copias de memoria "swizzled" (desordenadas/reorganizadas) en tiempo de ejecución.

Propone dos enfoques principales:

A. Double Buffering de DRAM (DDB):

• Concepto: Reserva un pequeño búfer de DRAM en la región cachable (tamaño equivalente a las pesas de una o dos capas).
• Funcionamiento: Mientras se ejecuta el GEMM de la capa actual en el búfer 0, el sistema pre-carga (prefetch) y reorganiza las pesas de la siguiente capa desde la región no cachable al búfer 1.
• Ventaja: Oculta la latencia de la reorganización de datos detrás del tiempo de cómputo, logrando una ejecución paralela.
• Sincronización: Alterna entre dos búferes en capas sucesivas del decodificador.

B. Reordenamiento en Línea con Copia de Memoria Swizzled (OWR):

• Concepto: Diseñado para secuencias de entrada largas (tendencia en LLMs interactivos y RAG).
• Funcionamiento: Realiza la copia y reorganización de las pesas necesarias (de formato PIM a formato host) inmediatamente antes de ejecutar el GEMM de cada capa, utilizando un solo búfer cachable.
• Ventaja: Elimina la necesidad de sincronización compleja o balanceo de carga entre hilos de copia y cómputo, reduciendo la complejidad de implementación.
• Eficiencia: A medida que aumenta la longitud de la secuencia de entrada, el tiempo de GEMM domina, haciendo que el costo fijo de la reorganización sea insignificante.

Técnica Clave: Swizzled Memory Copy (SMC)
Es el mecanismo que transforma las pesas del formato "PIM-aware" (filas completas por banco) al formato "host-friendly" (intercalado) durante la copia de memoria. Esto permite que el kernel estándar de GEMM se ejecute sin modificaciones, manteniendo la portabilidad.

---

3. Contribuciones Clave

1. Identificación del Problema: Es el primer trabajo que identifica formalmente la inconsistencia de atributos de memoria y diseño de pesas entre las fases de prefill y decode en sistemas PIM habilitados para dispositivos.
2. Soluciones de Software Puro: Propone DDB y OWR, que resuelven estos conflictos sin duplicar pesas ni modificar hardware, utilizando solo pequeños búferes adicionales.
3. Análisis de Sobrecarga: Cuantifica cómo la sobrecarga del reordenamiento en línea se amortiza a medida que aumenta la longitud de la secuencia de entrada (SL), demostrando su viabilidad en escenarios reales de LLMs interactivos.
4. Validación en Plataforma Comercial: Implementación y evaluación en un dispositivo real (Samsung Galaxy S24+), demostrando la viabilidad práctica.

---

4. Resultados Experimentales

Los experimentos se realizaron utilizando modelos Llama 3.2 (1B y 3B) en un Samsung Galaxy S24+ con arquitectura Exynos 2400 y emulación de LPDDR5X-PIM.

• Ahorro de Capacidad de Memoria (DRAM):

  • Comparado con la duplicación de pesas (WD), PIM-SHERPA logra un ahorro de memoria del 47.8% al 49.7%.
  • Esto permite ejecutar modelos que de otro modo no cabrían en la memoria del dispositivo (ej. Llama 3.2 3B requiere ~6GB; con duplicación serían 12GB, excediendo la capacidad típica de móviles).
  • El uso de búferes adicionales es mínimo (aprox. 32 MB), comparable a las soluciones de hardware como FACIL-O.

• Rendimiento (TTFT - Time to First Token):

  • DDB: Logra un rendimiento casi idéntico a la solución teórica óptima (FACIL-O) para longitudes de secuencia de entrada (SL) superiores a 128 tokens, ocultando eficazmente la latencia de copia.
  • OWR: En secuencias largas (SL > 192), su rendimiento se acerca al de DDB y FACIL-O, alcanzando ~16.7 TPS en el modelo 3B.
  • Ambos métodos mantienen la aceleración de PIM, superando significativamente a los sistemas sin PIM (hasta 3.3x de velocidad en algunos casos).

• Validación de Funcionalidad:

  • Se validó la corrección de los resultados comparando las salidas de PIM-SHERPA con ExecuTorch estándar y simuladores PIM, confirmando que la reorganización de datos no introduce errores.

---

5. Significado e Impacto

• Viabilidad Comercial: PIM-SHERPA elimina la barrera principal para la adopción de PIM en dispositivos móviles: el alto costo de memoria requerido para mantener múltiples copias de los modelos.
• Independencia de Hardware: Al ser una solución puramente de software, es aplicable a dispositivos comerciales existentes sin necesidad de esperar a nuevos controladores de memoria o chips PIM personalizados.
• Escalabilidad: La metodología es universal y puede adaptarse a GPUs y NPUs que poseen jerarquías de memoria con caché, sugiriendo un camino claro para la optimización de inferencia de LLMs en el futuro.
• Adaptación a Tendencias: Se alinea perfectamente con la tendencia hacia secuencias de entrada más largas en aplicaciones de IA generativa interactiva y RAG (Generación Aumentada por Recuperación), donde el costo de reorganización se vuelve marginal.

En conclusión, PIM-SHERPA demuestra que es posible desbloquear el potencial de la memoria PIM en dispositivos de borde para LLMs grandes, resolviendo conflictos de sistema críticos mediante inteligencia de software en lugar de duplicación de hardware.