AMD, Arm, AWS, Google, NVIDIA, Intel, Tesla, SambaNova y otros gigantes tecnológicos se unen para descifrar el código de las gigantescas cargas de trabajo en IA. A medida que los modelos de IA se vuelven más complejos y requieren mayores capacidades de cómputo, la necesidad de una infraestructura eficiente, escalable y compatible con diferentes tipos de hardware nunca ha sido mayor. OpenXLA es un marco de compilación de aprendizaje profundo diseñado para acelerar y escalar masivamente modelos de IA en una amplia gama de hardware, desde GPUs y CPUs hasta chips especializados como los TPUs de Google y Trainium de AWS. Es compatible con populares marcos de modelado como JAX, PyTorch y TensorFlow, y ofrece un rendimiento líder en su clase. OpenXLA es la infraestructura de aceleración preferida para productos de IA de escala global como Amazon.com Search, Google Gemini, los vehículos autónomos de Waymo, y Grok de x.AI.
El pasado 25 de abril, el OpenXLA Dev Lab recibió a más de 100 expertos en ML de 10 países, representando a líderes de la industria como AMD, Arm, AWS, ByteDance, Cerebras, Cruise, Google, NVIDIA, Intel, Tesla y SambaNova, entre otros. Este evento de un día, dirigido a proveedores de hardware de IA e ingenieros de infraestructura, rompió con el molde de las cumbres anteriores de OpenXLA, enfocándose exclusivamente en «Sesiones de laboratorio» y tutoriales prácticos. La energía del evento fue palpable, con desarrolladores trabajando codo a codo, aprendiendo y colaborando en desafíos prácticos y posibilidades emocionantes para la infraestructura de IA.
El Dev Lab se centró en tres aspectos clave: educar y empoderar a los desarrolladores sobre los flujos de trabajo y funciones avanzadas de OpenXLA a través de tutoriales prácticos; ofrecer orientación experta mediante horarios de oficina personalizados dirigidos por expertos de OpenXLA para ayudar a los desarrolladores a refinar sus ideas y contribuciones; y fomentar la comunidad alentando la colaboración, el intercambio de conocimientos y la creación de conexiones duraderas entre las mentes brillantes de la comunidad de OpenXLA.
Los tutoriales ofrecidos incluyeron temas como la integración de un compilador y entorno de ejecución de IA en PJRT, la extracción de gráficos StableHLO y una introducción al cuantizador StableHLO, la optimización del auto-sharding de PyTorch/XLA y la programación de cómputo y comunicación con XLA.
Durante las sesiones de laboratorio, se ofrecieron horarios de oficina específicos para AMD, Arm, AWS, ByteDance, Intel, NVIDIA, SambaNova, Tesla y otros. Los ingenieros de OpenXLA estuvieron disponibles para proporcionar soporte dedicado y resolver puntos conflictivos y diseños específicos. Además, se llevaron a cabo mesas redondas informativas sobre temas más amplios como la optimización del rendimiento de ML en GPU, JAX y PyTorch-XLA GPU.
El evento comenzó con una conferencia magistral a cargo de Robert Hundt, ingeniero distinguido de Google, quien presentó los ambiciosos planes de OpenXLA para 2024, centrados en tres áreas principales: el entrenamiento a gran escala, el rendimiento de cómputo en GPU y PyTorch, y la modularidad y extensibilidad.
OpenXLA está introduciendo poderosas características para permitir el entrenamiento de modelos a escalas récord. Una de las adiciones más destacadas es Shardonnay, una herramienta que optimiza y divide las cargas de trabajo de IA grandes entre múltiples unidades de procesamiento. Además, los desarrolladores pueden esperar una suite de características diseñadas para optimizar la superposición de cómputo y comunicación, incluyendo la estimación automática de latencia guiada por perfiles, el pipelining colectivo y combinadores de colectivas basados en heurísticas.
OpenXLA también ha logrado avances significativos en el rendimiento, particularmente en GPUs con modelos generativos basados en PyTorch. PyTorch-XLA GPU está ahora a la par con TorchInductor para modelos de gráficos completos de TorchBench y tiene una tasa de aprobación de TorchBench dentro del 5% de TorchInductor.
Además de las mejoras en el rendimiento, OpenXLA está comprometido con hacer que todo el stack sea más modular y extensible. Varias iniciativas planeadas para 2024 incluyen fortalecer los contratos de interfaz de los módulos, mejorar el intercambio de código entre plataformas y habilitar un flujo de compilación de alto nivel compartido a través de la configuración de tiempo de ejecución y registros de componentes.
El éxito de Alibaba con PyTorch XLA FSDP dentro de su marco TorchAcc es un ejemplo claro de los beneficios de la modularidad y extensibilidad de OpenXLA. Al aprovechar estas características, Alibaba logró un rendimiento de vanguardia para el modelo LLaMa 2 13B, superando el benchmark anterior establecido por Megatron.
Para aquellos que se perdieron el Dev Lab, aún pueden acceder a los tutoriales de StableHLO en openxla.org, así como al Gist de GitHub para la sesión de PJRT. Además, las grabaciones de la keynote y los tutoriales están disponibles en el canal de YouTube de OpenXLA. Únete a nuestra comunidad global, ya seas un experto en sistemas de IA, desarrollador de modelos, estudiante o principiante, hay un lugar para ti en nuestro ecosistema innovador.
vía: Google Blog Open Source