AXLearn: Treinamento de Modelos Grandes de Forma Modular em Infraestrutura 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

Resumo

Projetamos e implementamos o AXLearn, um sistema de aprendizado profundo de produção que facilita o treinamento escalável e de alto desempenho de grandes modelos de aprendizado profundo. Em comparação com outros sistemas de aprendizado profundo de ponta, o AXLearn tem um foco único em modularidade e suporte para infraestrutura de hardware heterogênea. As interfaces internas do AXLearn entre componentes de software seguem uma encapsulação rigorosa, permitindo que diferentes componentes sejam montados para facilitar o desenvolvimento rápido de modelos e experimentação em infraestruturas de computação heterogêneas. Introduzimos um método inovador de quantificar a modularidade por meio da complexidade de Linhas de Código (LoC), que demonstra como nosso sistema mantém uma complexidade constante à medida que escalamos os componentes do sistema, em comparação com a complexidade linear ou quadrática em outros sistemas. Isso permite integrar recursos como Embeddings de Posição Rotacional (RoPE) no AXLearn em centenas de módulos com apenas 10 linhas de código, em contraste com as centenas necessárias em outros sistemas. Ao mesmo tempo, o AXLearn mantém um desempenho equivalente em comparação com os sistemas de treinamento mais avançados. Por fim, compartilhamos nossa experiência no desenvolvimento e operação do 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.