AXLearn: Modular Large Model Training on Heterogeneous Infrastructure
Authors: name: Mark Lee; name: Tom Gunter; name: Chang Lan; name: John Peebles; name: Hanzhi Zhou; name: Kelvin Zou; name: Sneha Bangalore; name: Chung-Cheng Chiu; name: Nan Du; name: Xianzhi Du; name: Philipp Dufter; name: Ruixuan Hou; name: Haoshuo Huang; name: Dongseong Hwang; name: Xiang Kong; name: Jinhao Lei; name: Tao Lei; name: Meng Li; name: Li Li; name: Jiarui Lu; name: Zhiyun Lu; name: Yiping Ma; name: David Qiu; name: Vivek Rathod; name: Senyu Tong; name: Zhucheng Tu; name: Jianyu Wang; name: Yongqiang Wang; name: Zirui Wang; name: Floris Weers; name: Sam Wiseman; name: Guoli Yin; name: Bowen Zhang; name: Xiyou Zhou; name: Danyang Zhuo; name: Cheng Leong; name: Ruoming Pang
Published: 2025-07-07 18:50:58

Abstract: 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.