AXLearn: Entrenamiento Modular de Modelos Grandes en Infraestructura Heterogénea

AXLearn: Modular Large Model Training on Heterogeneous Infrastructure

July 7, 2025

Autores: Mark Lee, Tom Gunter, Chang Lan, John Peebles, Hanzhi Zhou, Kelvin Zou, Sneha Bangalore, Chung-Cheng Chiu, Nan Du, Xianzhi Du, Philipp Dufter, Ruixuan Hou, Haoshuo Huang, Dongseong Hwang, Xiang Kong, Jinhao Lei, Tao Lei, Meng Li, Li Li, Jiarui Lu, Zhiyun Lu, Yiping Ma, David Qiu, Vivek Rathod, Senyu Tong, Zhucheng Tu, Jianyu Wang, Yongqiang Wang, Zirui Wang, Floris Weers, Sam Wiseman, Guoli Yin, Bowen Zhang, Xiyou Zhou, Danyang Zhuo, Cheng Leong, Ruoming Pang

cs.AI

Resumen

Diseñamos e implementamos AXLearn, un sistema de aprendizaje profundo de producción que facilita el entrenamiento escalable y de alto rendimiento de modelos grandes de aprendizaje profundo. En comparación con otros sistemas de aprendizaje profundo de vanguardia, AXLearn tiene un enfoque único en la modularidad y el soporte para infraestructura de hardware heterogénea. Las interfaces internas de AXLearn entre los componentes de software siguen una encapsulación estricta, lo que permite ensamblar diferentes componentes para facilitar el desarrollo rápido de modelos y la experimentación en infraestructura de computación heterogénea. Introducimos un método novedoso para cuantificar la modularidad mediante la complejidad de Líneas-de-Código (LoC), que demuestra cómo nuestro sistema mantiene una complejidad constante a medida que escalamos los componentes del sistema, en comparación con la complejidad lineal o cuadrática en otros sistemas. Esto permite integrar características como las Incrustaciones de Posición Rotacional (RoPE) en AXLearn a través de cientos de módulos con solo 10 líneas de código, en comparación con las cientos requeridas en otros sistemas. Al mismo tiempo, AXLearn mantiene un rendimiento equivalente en comparación con los sistemas de entrenamiento más avanzados. Finalmente, compartimos nuestra experiencia en el desarrollo y operación de AXLearn.

English

We design and implement AXLearn, a production deep learning system that facilitates scalable and high-performance training of large deep learning models. Compared to other state-of-the-art deep learning systems, AXLearn has a unique focus on modularity and support for heterogeneous hardware infrastructure. AXLearn's internal interfaces between software components follow strict encapsulation, allowing different components to be assembled to facilitate rapid model development and experimentation on heterogeneous compute infrastructure. We introduce a novel method of quantifying modularity via Lines-of-Code (LoC)-complexity, which demonstrates how our system maintains constant complexity as we scale the components in the system, compared to linear or quadratic complexity in other systems. This allows integrating features such as Rotary Position Embeddings (RoPE) into AXLearn across hundred of modules with just 10 lines of code, compared to hundreds as required in other systems. At the same time, AXLearn maintains equivalent performance compared to state-of-the-art training systems. Finally, we share our experience in the development and operation of AXLearn.