AXLearn: Modular, Hardware-Agnostic Large Model Training
Abstract
AXLearn is a production system which facilitates scalable and high-performance training of large deep learning models. Compared to other state-of-art deep learning systems, AXLearn has a unique focus on modularity and support for hardware-agnostic training. AXLearn's internal interfaces between software components follow strict encapsulation, allowing different components to be assembled to facilitate rapid model development and experimentation on different hardware infrastructure. AXLearn maintains constant complexity as we scale the components in the system, compared to linear or quadratic complexity in state-of-the-art training 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 at Apple.
Keywords
Cite
@article{arxiv.2507.05411,
title = {AXLearn: Modular, Hardware-Agnostic Large Model Training},
author = {Mark Lee and Chang Lan and Tom Gunter and John Peebles and Hanzhi Zhou and Kelvin Zou and Sneha Bangalore and Chung-Cheng Chiu and Nan Du and Xianzhi Du and Philipp Dufter and Ruixuan Hou and Haoshuo Huang and Dongseong Hwang and Xiang Kong and Jinhao Lei and Tao Lei and Meng Li and Li Li and Jiarui Lu and Zhiyun Lu and Yiping Ma and David Qiu and Vivek Rathod and Senyu Tong and Zhucheng Tu and Jianyu Wang and Yongqiang Wang and Zirui Wang and Floris Weers and Sam Wiseman and Guoli Yin and Bowen Zhang and Xiyou Zhou and Danyang Zhuo and Cheng Leong and Ruoming Pang},
journal= {arXiv preprint arXiv:2507.05411},
year = {2026}
}