Firefox Acelera su IA Local: Hasta 10 Veces Más Rápido con ONNX Runtime C++

¿Te has dado cuenta de que Firefox está cada vez más inteligente? Detrás de funciones como la descripción automática de imágenes en PDFs y el agrupamiento inteligente de pestañas se encuentra el Firefox AI Runtime. Sin embargo, Mozilla no se conformaba con el rendimiento inicial y se puso manos a la obra para optimizarlo. ¿El resultado? ¡Aceleración de hasta 10 veces en la inferencia gracias a la integración de ONNX Runtime C++!

El Punto de Partida: Transformers.js y ONNX Runtime Web

Inicialmente, Firefox AI Runtime se basaba en Transformers.js, la contraparte en JavaScript de la popular librería de Python Hugging Face. Este framework utilizaba onnxruntime-web, una versión de ONNX Runtime compilada para WebAssembly (WASM). El ciclo de inferencia típico era el siguiente:

1. Preprocesamiento en JavaScript (tokenización, modelado de tensores).
2. Ejecución del modelo en WASM.
3. Postprocesamiento de vuelta en JavaScript.

A pesar de las cachés, este proceso implicaba cruzar múltiples capas, y el cuello de botella se encontraba en las multiplicaciones de matrices, implementadas con SIMD genérico al ejecutarse en la CPU.

¿Por Qué WASM No Era Suficiente?

Aunque WASM SIMD es una buena solución, no puede competir con las instrucciones específicas del hardware, como NEON en Apple Silicon o AVX-512 en los chips Intel modernos. Firefox Translations ya demostró que el uso de código nativo acelera las cosas mediante el uso de WASM built-ins, pequeños hooks que permiten a WASM llamar a C++ compilado con esas instrucciones. Sin embargo, aplicar este truco a ONNX era inviable debido al gran número de operadores.

El Salto a ONNX Runtime C++

La solución fue reemplazar onnxruntime-web con su contraparte nativa en C++. Transformers.js se comunica con ONNX Runtime a través de una interfaz simple: crea una sesión, envía un Tensor y recibe un resultado. Esto facilitó el cambio de backend sin modificar el código de las funcionalidades. Los pasos fueron:

1. Integrar ONNX Runtime C++ directamente en el árbol de código de Firefox.
2. Exponerlo a JavaScript a través de una fina capa WebIDL.
3. Conectar Transformers.js al nuevo backend.

Desde la perspectiva de una función como la de texto alternativo en PDF, nada cambió; sigue llamando a await pipeline(...). La diferencia es que ahora los tensores van directamente al código nativo.

Integración en el Sistema de Compilación

Debido a que ONNX Runtime no soporta todas las configuraciones de compilación de Firefox y es una gran cantidad de código, se optó por no agregarlo directamente al árbol. En su lugar, se utiliza un flag de configuración para proporcionar una versión compilada de ONNX Runtime, que se descarga automáticamente desde Taskcluster (donde se compila para una selección de configuraciones soportadas) o es proporcionada por desarrolladores externos. Esto ofrece flexibilidad y reduce los tiempos de compilación.

La Recompensa: Mayor Velocidad y Eficiencia

Al ser el backend nativo un reemplazo directo, se puede habilitar función por función y recopilar datos reales. Las primeras pruebas muestran una inferencia de 2 a 10 veces más rápida, sin el overhead de la inicialización de WASM. Por ejemplo, la sugerencia de temas para el Agrupamiento Inteligente de Pestañas, que podía ser lenta al principio, ahora es mucho más rápida. La función de texto alternativo para PDF.js también se beneficia enormemente, con una latencia que se redujo de 3.5 segundos a 350 milisegundos en el mismo hardware.

Próximos Pasos: Más Optimización y Soporte para GPU

Mozilla planea implementar gradualmente este nuevo backend en más funciones a lo largo del verano. Además, con la API de C++ a mano, tienen previsto abordar algunos problemas pendientes y habilitar el soporte para GPU. También se están explorando optimizaciones como:

1. Multihilo en DequantizeLinear: Distribuir el trabajo entre los núcleos para acelerar esta operación.
2. Transposición de matrices multihilo: Utilizar un esquema de transposición en mosaico con caché y SIMD para acelerar esta operación.
3. Caché del grafo compilado: Almacenar en caché el grafo compilado para evitar la recompilación en cada lanzamiento.

Conclusión

La migración a ONNX Runtime C++ ha demostrado ser un éxito, permitiendo mejorar significativamente el rendimiento de las funcionalidades de IA en Firefox sin necesidad de modificar el código de las mismas. Los resultados son visibles desde el punto de vista de la experiencia de usuario, y Mozilla confía en que aún hay mucho margen de mejora en el futuro, haciendo que las funciones basadas en Machine Learning sean más eficientes y accesibles.

Si tienes ideas, preguntas o informes de errores, puedes contactar a Mozilla en el canal #firefox-ai de Discord o a través de Bugzilla.

Fuente: Mozilla Blog