OpenXLA, un proyecto open source para acelerar y simplificar el aprendizaje automático

OpenXLA es un ecosistema de compilador de ML de código abierto desarrollado conjuntamente

Hace poco, las mayores empresas dedicadas al desarrollo en el campo del aprendizaje automático presentaron el proyecto OpenXLA, destinado al desarrollo conjunto de herramientas para compilar y optimizar modelos para sistemas de aprendizaje automático.

El proyecto se ha hecho cargo del desarrollo de herramientas que permitan unificar la compilación de modelos preparados en los frameworks TensorFlow, PyTorch y JAX para un entrenamiento y ejecución eficiente en diferentes GPUs y aceleradores especializados. Empresas como Google, NVIDIA, AMD, Intel, Meta, Apple, Arm, Alibaba y Amazon se sumaron al trabajo conjunto del proyecto.

El Proyecto OpenXLA proporciona un compilador de ML de última generación que puede escalar en medio de la complejidad de la infraestructura de ML. Sus pilares fundamentales son el rendimiento, la escalabilidad, la portabilidad, la flexibilidad y la extensibilidad para los usuarios. Con OpenXLA, aspiramos a aprovechar el potencial real de la IA acelerando su desarrollo y entrega.

OpenXLA permite a los desarrolladores compilar y optimizar modelos de todos los marcos de ML líderes para una capacitación y servicio eficientes en una amplia variedad de hardware. Los desarrolladores que utilicen OpenXLA verán mejoras significativas en el tiempo de capacitación, el rendimiento, la latencia de servicio y, en última instancia, el tiempo de comercialización y los costos de cómputo.

Se espera que al unir los esfuerzos de los principales equipos de investigación y representantes de la comunidad, será posible estimular el desarrollo de sistemas de aprendizaje automático y resolver el problema de la fragmentación de la infraestructura para varios marcos y equipos.

OpenXLA permite implementar soporte efectivo para varios hardware, independientemente del marco en el que se base el modelo de aprendizaje automático. Se espera que OpenXLA reduzca el tiempo de capacitación del modelo, mejore el rendimiento, reduzca la latencia, reduzca la sobrecarga informática y reduzca el tiempo de comercialización.

OpenXLA está formado por tres componentes principales, cuyo código se distribuye bajo la licencia Apache 2.0:

1. XLA (álgebra lineal acelerada) es un compilador que le permite optimizar modelos de aprendizaje automático para una ejecución de alto rendimiento en diferentes plataformas de hardware, incluidas GPU, CPU y aceleradores especializados de varios fabricantes.
2. StableHLO es una especificación e implementación básica de un conjunto de operaciones de alto nivel (HLO, High-Level Operations) para usar en modelos de sistemas de aprendizaje automático. Actúa como una capa entre los marcos de trabajo de aprendizaje automático y los compiladores que transforman el modelo para que se ejecute en un hardware específico. Se preparan capas para generar modelos en formato StableHLO para los frameworks PyTorch, TensorFlow y JAX. La suite MHLO se utiliza como base para StableHLO, que se amplía con soporte para serialización y control de versiones.
3. IREE (Entorno de ejecución de representación intermedia) es un compilador y tiempo de ejecución que convierte los modelos de aprendizaje automático en una representación intermedia universal basada en el formato MLIR (Representación intermedia de varios niveles) del proyecto LLVM. De las características, se destacan la posibilidad de precompilación (antes de tiempo), soporte para control de flujo, la capacidad de usar elementos dinámicos en modelos, optimización para diferentes CPU y GPU, y baja sobrecarga.

Por la parte de las principales ventajas de OpenXLA se menciona que se ha logrado un rendimiento óptimo sin tener que profundizar en la escritura de código específico del dispositivo, ademas de proporcionar optimizaciones listas para usar, incluida la simplificación de expresiones algebraicas, asignación eficiente de memoria, programación de ejecución, teniendo en cuenta la reducción del consumo máximo de memoria y las sobrecargas.

Otra de las ventajas, es la simplificación del escalado y paralelización de los cálculos. Es suficiente que un desarrollador agregue anotaciones para un subconjunto de tensores críticos, sobre la base de los cuales el compilador puede generar automáticamente código para computación paralela.

Tambien se destaca que se proporciona la portabilidad con soporte para múltiples plataformas de hardware, como GPU AMD y NVIDIA, CPU x86 y ARM, aceleradores de ML de TPU de Google, IPU de AWS Trainium Inferentia, Graphcore y Cerebras Wafer-Scale Engine.

Soporte para conectar extensiones con la implementación de funciones adicionales, como soporte para escribir primitivas de aprendizaje automático profundo usando CUDA, HIP, SYCL, Triton y otros lenguajes para computación paralela, asi como tambien la posibilidad de ajuste manual de cuellos de botella en modelos.

Finalmente si estás interesado en poder conocer más al respecto, puedes consultar los detalles en el siguiente enlace.