_version_ 1866911456106119168
author Lee, Mark
Lan, Chang
Gunter, Tom
Peebles, John
Zhou, Hanzhi
Zou, Kelvin
Bangalore, Sneha
Chiu, Chung-Cheng
Du, Nan
Du, Xianzhi
Dufter, Philipp
Hou, Ruixuan
Huang, Haoshuo
Hwang, Dongseong
Kong, Xiang
Lei, Jinhao
Lei, Tao
Li, Meng
Li, Li
Lu, Jiarui
Lu, Zhiyun
Ma, Yiping
Qiu, David
Rathod, Vivek
Tong, Senyu
Tu, Zhucheng
Wang, Jianyu
Wang, Yongqiang
Wang, Zirui
Weers, Floris
Wiseman, Sam
Yin, Guoli
Zhang, Bowen
Zhou, Xiyou
Zhuo, Danyang
Leong, Cheng
Pang, Ruoming
author_facet Lee, Mark
Lan, Chang
Gunter, Tom
Peebles, John
Zhou, Hanzhi
Zou, Kelvin
Bangalore, Sneha
Chiu, Chung-Cheng
Du, Nan
Du, Xianzhi
Dufter, Philipp
Hou, Ruixuan
Huang, Haoshuo
Hwang, Dongseong
Kong, Xiang
Lei, Jinhao
Lei, Tao
Li, Meng
Li, Li
Lu, Jiarui
Lu, Zhiyun
Ma, Yiping
Qiu, David
Rathod, Vivek
Tong, Senyu
Tu, Zhucheng
Wang, Jianyu
Wang, Yongqiang
Wang, Zirui
Weers, Floris
Wiseman, Sam
Yin, Guoli
Zhang, Bowen
Zhou, Xiyou
Zhuo, Danyang
Leong, Cheng
Pang, Ruoming
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2507_05411
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle AXLearn: Modular, Hardware-Agnostic Large Model Training
Lee, Mark
Lan, Chang
Gunter, Tom
Peebles, John
Zhou, Hanzhi
Zou, Kelvin
Bangalore, Sneha
Chiu, Chung-Cheng
Du, Nan
Du, Xianzhi
Dufter, Philipp
Hou, Ruixuan
Huang, Haoshuo
Hwang, Dongseong
Kong, Xiang
Lei, Jinhao
Lei, Tao
Li, Meng
Li, Li
Lu, Jiarui
Lu, Zhiyun
Ma, Yiping
Qiu, David
Rathod, Vivek
Tong, Senyu
Tu, Zhucheng
Wang, Jianyu
Wang, Yongqiang
Wang, Zirui
Weers, Floris
Wiseman, Sam
Yin, Guoli
Zhang, Bowen
Zhou, Xiyou
Zhuo, Danyang
Leong, Cheng
Pang, Ruoming
Machine Learning
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.
title AXLearn: Modular, Hardware-Agnostic Large Model Training
topic Machine Learning
url https://arxiv.org/abs/2507.05411