Modelos de lenguaje grande en el dispositivo con MediaPipe y TensorFlow Lite

MAR 07, 2024
Mark Sherwood Senior Product Manager
Juhyun Lee Staff Software Engineer

TensorFlow Lite ha sido una herramienta poderosa para el aprendizaje automático en el dispositivo desde su lanzamiento en 2017 y MediaPipe amplió aún más ese poder en 2019 ya que comenzó a admitir canalizaciones de AA completas. Si bien estas herramientas se centraron inicialmente en modelos más pequeños en el dispositivo, en la actualidad TensorFlow Lite marca un cambio dramático con la API de Inferencia experimental MediaPipe LLM.

Esta nueva versión permite que los modelos de lenguaje grande (LLM) se ejecuten de forma completa en el dispositivo en todas las plataformas. Esta capacidad nueva es particularmente transformadora si tenemos en cuenta las demandas de memoria y computación de los modelos de lenguaje grande, que son cien veces más grandes que los modelos tradicionales en el dispositivo. Las optimizaciones en toda la pila en el dispositivo hacen que estos cambios sean posibles, incluidas las operaciones nuevas, la cuantificación, el almacenamiento en caché y la compartición de peso.

La API de Inferencia de MediaPipe LLM experimental multiplataforma, diseñada para optimizar la integración de modelo de lenguaje grande en el dispositivo para desarrolladores web, admite Web, iOS y Android con soporte inicial para cuatro modelos de lenguaje grande disponibles abiertamente: Gemma, Phi 2, Falcon y Stable LM. Esta API ofrece a los investigadores y desarrolladores la flexibilidad de prototipar y probar modelos de lenguaje grande populares disponibles de manera abierta en el dispositivo.

En Android, la API de Inferencia de MediaPipe LLM está diseñada solo para uso experimental y de investigación. Las aplicaciones de producción con modelos de lenguaje grande pueden usar la API de Gemini o Gemini Nano en el dispositivo por medio de Android AICore. AICore es la nueva capacidad a nivel de sistema introducida en Android 14 para brindar soluciones impulsadas por Gemini para dispositivos de gama alta, incluidas integraciones con los últimos aceleradores de AA, adaptadores LoRA optimizados para casos de uso y filtros de seguridad. Para comenzar a usar Gemini Nano en el dispositivo con tu app, aplica a la vista previa de acceso anticipado.

API de Inferencia de modelo de lenguaje grande

A partir de hoy, puedes probar la API de Inferencia de MediaPipe LLM en nuestra Web de demostración o al compilar nuestras apps de demostración. Puedes probarla e integrarla en tus proyectos por medio de la Web, Android o iOS SDKs.

El uso de la API de Inferencia de modelo de lenguaje grande permite incorporar modelos de lenguaje grande en el dispositivo en solo unos pocos pasos. Estos pasos se aplican en la Web, iOS y Android, aunque el SDK y la API nativa serán específicos de la plataforma. Los siguientes códigos de ejemplo muestran el SDK web.

  1. Elige pesos de modelo compatibles con una de nuestras arquitecturas de modelo compatibles

2. Convierte los pesos del modelo en un Flatbuffer de TensorFlow Lite utilizando el paquete MediaPipe Python

desde mediapipe.tasks.python.genai import converter 
 
config = converter.ConversionConfig(...)
converter.convert_checkpoint(config)

3. Incluye el SDK de Inferencia de modelo de lenguaje grande en tu aplicación

importar { FilesetResolver, LlmInference } desde "https://cdn.jsdelivr.net/npm/@mediapipe/tasks-genai"

4. Aloja el Flatbuffer de TensorFlow Lite junto con tu aplicación

5. Utiliza la API de Inferencia de modelo de lenguaje grande para tomar un mensaje de texto y obtener una respuesta de texto de tu modelo

const fileset  = await FilesetResolver.forGenAiTasks("https://cdn.jsdelivr.net/npm/@mediapipe/tasks-genai/wasm");
const llmInference = await LlmInference.createFromModelPath(fileset, model.bin);
const responseText = await llmInference.generateResponse ("Hola, encantado de conocerte");
document.getElementById('output').textContent = responseText;

Consulta la documentación y los ejemplos de códigos para obtener una descripción detallada de cada uno de estos pasos.

Estos son los gifs en tiempo real de Gemma 2B que se ejecutan mediante la API de Inferencia MediaPipe LLM.

Moving image of Gemma 2B running on-device in browser via the MediaPipe LLM Inference API
Gemma 2B se ejecuta en el dispositivo en el navegador por medio de la API de Inferencia MediaPipe LLM
Moving image of Gemma 2B running on-device on iOS (left) and Android (right) via the MediaPipe LLM Inference API
Gemma 2B se ejecuta en el dispositivo en iOS (izquierda) y Android (derecha) con la API de Inferencia MediaPipe LLM

Modelos

Nuestra versión inicial admite las siguientes cuatro arquitecturas de modelo. Cualquier peso de modelo compatible con estas arquitecturas funcionará con la API de Inferencia de modelo de lenguaje grande. Utiliza los pesos del modelo base, una versión ajustada de los pesos de la comunidad o ajusta los pesos utilizando tus propios datos.

Table showing model and parameter size across the four model architectures - Falcon 1B, Gemma 2B, Phi 2 and Stable LM 3B

Rendimiento del modelo

Mediante optimizaciones significativas, algunas de las cuales se detallan a continuación, la API de Inferencia MediaPipe LLM puede ofrecer una latencia de vanguardia en el dispositivo, centrándose en la CPU y la GPU para admitir varias plataformas. Para un rendimiento sostenido en un entorno de producción en teléfonos premium seleccionados, Android AICore puede aprovechar los aceleradores neuronales específicos del hardware.

Al medir la latencia para un modelo de lenguaje grande, se deben tener en cuenta algunos términos y mediciones. El tiempo hasta el primer token y la velocidad de decodificación serán los dos más significativos, ya que miden la rapidez con la que obtienes el inicio de tu respuesta y la rapidez con la que se genera la respuesta una vez que comienza.

Table showing latency measurements for model performance

La velocidad de prellenado y la velocidad de decodificación dependen del modelo, el hardware y los tokens máximos. También pueden cambiar según la carga actual del dispositivo.

Las siguientes velocidades se tomaron en dispositivos de gama alta utilizando un máximo de 1280 tokens, un indicador de entrada de 1024 tokens y cuantificación de peso int8. La excepción es Gemma 2B (int4), que se encuentra aquí en Kaggle, que utiliza una cuantificación de peso mixta de 4/8 bits.


Puntos de referencia

Graph of prefill performance in tokens per second
Graph of decode performance in tokens per second
En la GPU, Falcon 1B y Phi 2 usan activaciones de fp32, mientras que Gemma y StableLM 3B usan activaciones de fp16, ya que estos últimos modelos mostraron una mayor robustez a la pérdida de precisión de acuerdo con nuestros estudios de evaluación de calidad. Se eligió el tipo de datos de activación de bits más bajo que mantuvo la calidad del modelo para cada uno. Ten en cuenta que Gemma 2B (int4) era el único modelo que podíamos ejecutar en iOS debido a sus limitaciones de memoria, y que también estamos trabajando para habilitar otros modelos en iOS.

Optimizaciones de rendimiento

Para lograr los números de rendimiento anteriores, se realizaron innumerables optimizaciones en MediaPipe, TensorFlow Lite, XNNPack (nuestra biblioteca de operadores de red neuronal de CPU) y en nuestro tiempo de ejecución acelerado por GPU. Las siguientes son algunas optimizaciones seleccionadas que resultaron en mejoras significativas en el rendimiento.

Compartición de pesos: el proceso de inferencia de modelo de lenguaje grande comprende 2 fases: una fase de prellenado y una fase de decodificación. Tradicionalmente, esta configuración requeriría 2 contextos de inferencia separados, cada uno administrando recursos de forma independiente para su modelo de AA correspondiente. Dadas las demandas de memoria de los modelos de lenguaje grande, agregamos una función que permite compartir los pesos y la caché KV en contextos de inferencia. Aunque compartir ponderaciones puede parecer sencillo, tiene importantes implicaciones de rendimiento cuando se comparte entre operaciones vinculadas a la computación y a la memoria. En los escenarios típicos de inferencia de AA, donde los pesos no se comparten con otros operadores, se configuran de forma meticulosa para cada operador completamente conectado por separado a fin de garantizar un rendimiento óptimo. Compartir pesos con otro operador implica una pérdida de optimización por operador y esto exige la creación de implementaciones nuevas del kernel que puedan ejecutarse de manera eficiente incluso con pesos subóptimos.

Operaciones optimizadas totalmente conectadas: la operación FULLY_CONNECTED de XNNPack tuvo dos optimizaciones significativas para la inferencia de modelo de lenguaje grande. En primer lugar, la cuantificación de rango dinámico combina a la perfección los beneficios computacionales y de memoria de la cuantificación de enteros completos con las ventajas de precisión de la inferencia de punto flotante. La utilización de pesos int8/int4 no solo mejora la capacidad de procesamiento de la memoria, sino que también logra un rendimiento notable, en especial con la decodificación eficiente y en registro de pesos de 4 bits que requieren solo una instrucción adicional. En segundo lugar, aprovechamos activamente las instrucciones I8MM en las CPU ARM v9 que habiliten la multiplicación de una matriz int8 de 2x8 por una matriz int8 de 8x2 en una sola instrucción, lo que resulta en el doble de velocidad que la implementación basada en productos de punto de NEON.

Equilibrio entre cómputo y memoria: cuando perfilamos la inferencia de modelo de lenguaje grande, identificamos distintas limitaciones para ambas fases: la fase de prellenado enfrenta restricciones impuestas por la capacidad de cómputo, mientras que la fase de decodificación está limitada por el ancho de banda de la memoria. En consecuencia, cada fase emplea diferentes estrategias para la descuantificación de los pesos int8/int4 compartidos. En la fase de prellenado, cada operador de convolución primero descuantifica los pesos en valores de punto flotante antes del cálculo primario, lo que garantiza un rendimiento óptimo para convoluciones computacionalmente intensivas. Por el contrario, la fase de decodificación minimiza el ancho de banda de la memoria ya que agrega el cálculo de descuantificación a las principales operaciones de convolución matemática.

Flowchart showing compute-intensive prefill phase and memory-intensive decode phase, highlighting difference in performance bottlenecks

Operadores personalizados: para la inferencia de modelo de lenguaje grande acelerada por GPU en el dispositivo, confiamos ampliamente en las operaciones personalizadas para mitigar la ineficiencia causada por numerosos sombreadores pequeños. Estas operaciones personalizadas permiten que las fusiones especiales del operador y varios parámetros de modelo de lenguaje grande, como el ID de token, el tamaño del parche de secuencia y los parámetros de muestreo, se empaqueten en un tensor personalizado y especializado que se utiliza dentro de estas operaciones especializadas en particular.

Pseudodinamismo: en el bloque de atención, nos encontramos con operaciones dinámicas que aumentan con el tiempo a medida que crece el contexto. Dado que nuestro tiempo de ejecución de GPU carece de asistencia para operaciones/tensores dinámicos, optamos por operaciones fijas con un tamaño máximo de caché predefinido. Para reducir la complejidad computacional, se introduce un parámetro que permite omitir ciertos cálculos de valor o el procesamiento de datos reducidos.

Diseño de caché KV optimizado: dado que las entradas en la caché KV sirven en última instancia como ponderaciones para convoluciones, empleadas en lugar de multiplicaciones de matrices, las almacenamos en un diseño especializado adaptado para ponderaciones de convolución. Este ajuste estratégico elimina la necesidad de conversiones adicionales o la dependencia de diseños no optimizados y, por lo tanto, contribuye a un proceso más eficiente y ágil.



Lo que viene

Estamos felices con las optimizaciones y el rendimiento de la versión experimental actual de la API de inferencia MediaPipe LLM. Esto es solo el comienzo, ya que a lo largo de 2024 nos expandiremos a más plataformas y modelos, ofreceremos herramientas de conversión más amplias, componentes gratuitos en el dispositivo, tareas de alto nivel y mucho más.

Descubre la muestra oficial en GitHub para probar todo lo que acabas de aprender y lee nuestra documentación oficial para obtener aún más detalles. Echa un vistazo al canal de YouTube de Google para desarrolladores a fin de obtener actualizaciones e instructivos.



Agradecimientos

Nos gustaría agradecer a todos los miembros del equipo que contribuyeron a este trabajo: T.J. Alumbaugh, Alek Andreev, Frank Ban, Jeanine Banks, Frank Barchard, Pulkit Bhuwalka, Buck Bourdon, Maxime Brénon, Chuo-Ling Chang, Lin Chen, Linkun Chen, Yu-hui Chen, Nikolai Chinaev, Clark Duvall, Rosário Fernandes, Mig Gerard, Matthias Grundmann, Ayush Gupta, Mohammadreza Heydary, Ekaterina Ignasheva, Ram Iyengar, Grant Jensen, Alex Kanaukou, Prianka Liz Kariat, Alan Kelly, Kathleen Kenealy, Ho Ko, Sachin Kotwani, Andrei Kulik, Yi-Chun Kuo, Khanh LeViet, Yang Lu, Lalit Singh Manral, Tyler Mullen, Karthik Raveendran, Raman Sarokin, Sebastian Schmidt, Kris Tonthat, Lu Wang, Zoe Wang, Tris Warkentin, Geng Yan, Tenghui Zhu, y el equipo de Gemma.