Google AI nos trae un avance tecnológico fascinante en su blog. Su equipo de ingenieros, Juhyun Lee y Raman Sarokin, han desarrollado un método para acelerar la implementación de grandes modelos de difusión (LDMs) en dispositivos móviles, logrando que modelos de tamaño considerable se ejecuten en smartphones modernos de forma eficiente.
El Reto de los LDMs en Dispositivos Móviles
Los LDMs son modelos para la generación de imágenes de gran tamaño. Implementarlos en un entorno móvil supone un reto considerable, debido a sus elevados requerimientos de memoria y computación. Para superar estos obstáculos, el equipo de Google AI se centró en optimizar la eficiencia de la memoria y reducir la latencia total de la inferencia de ML.
Técnicas de Optimización: Eficiencia de la Memoria
Para lograr la ejecución optimizada de los LDMs, Google AI ha implementado una serie de técnicas. Entre ellas, el mejoramiento del módulo de atención para mayor eficiencia de la memoria. Este módulo es crucial en el modelo de denoiser del LDM, permitiendo a la inteligencia artificial focalizarse en partes específicas de la entrada. Dos de las técnicas de optimización empleadas son el softmax parcialmente fusionado y la FlashAttention.
Softmax Parcialmente Fusionado
El softmax parcialmente fusionado es una optimización que ahorra la necesidad de extensas lecturas y escrituras de memoria entre el softmax y la multiplicación de matrices en el módulo de atención. Esta técnica consigue reducir de manera significativa la cantidad de memoria requerida.
FlashAttention
Por otro lado, la FlashAttention es un algoritmo de atención exacto y consciente del I/O, que disminuye el número de accesos a la memoria de alta velocidad de la GPU. Este método resultó ser especialmente eficiente en ciertos tamaños de SRAM y con un número considerable de registros.
Técnicas de Optimización: Convoluciones Rápidas Winograd
Las convoluciones rápidas Winograd son otra técnica que Google AI ha empleado para optimizar los LDMs. A pesar de un incremento en el consumo de memoria y los errores numéricos, estas convoluciones demostraron ser eficaces para acelerar los procesos de los LDMs, logrando un equilibrio entre la eficiencia computacional y el uso de la memoria.
Técnicas de Optimización: Fusión de Operadores Especializados
Otro hallazgo importante fue la necesidad de fusiones más amplias para capas comunes en LDMs que las proporcionadas por los motores de inferencia de ML para GPU. Como solución, se desarrollaron implementaciones especializadas para una mayor gama de operadores neuronales, como la Unidad Lineal de Error Gaussiano (GELU) y la capa de normalización de grupo.
Resultados Notables
Aplicando estas optimizaciones, Google AI ha logrado ejecutar modelos de difusión en alta resolución y de tamaño considerable en dispositivos móviles modernos en menos de 12 segundos. Estos resultados marcan un hito importante en la implementación de ML en dispositivos móviles.
El logro de Google AI demuestra que los desafíos de implementar grandes modelos de ML en dispositivos móviles no son insuperables. Enfocándose en la utilización óptima de la memoria y el equilibrio entre la eficiencia de la ALU y la memoria, el equipo ha logrado un tiempo de latencia para la inferencia de ML sin precedentes. Este avance representa un paso significativo hacia una tecnología móvil más potente y eficiente. Para más detalles sobre este trabajo, puedes visitar el artículo original en Google AI blog.