Poster 26

AXLearn: Modular, Hardware-Agnostic Large Model Training

Mark Lee ⋅ Chang Lan ⋅ Tom Gunter ⋅ John Peebles ⋅ Hanzhi Zhou ⋅ Xuan Zou ⋅ Sneha Bangalore ⋅ Chung-Cheng Chiu ⋅ Nan Du ⋅ Xianzhi Du ⋅ Philipp Dufter ⋅ Liang He ⋅ 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 ⋅ Chong Wang ⋅ Jianyu Wang ⋅ Yongqiang Wang ⋅ Zirui Wang ⋅ Floris Weers ⋅ Sam Wiseman ⋅ Guoli Yin ⋅ Bowen Zhang ⋅ Xiyou Zhou ⋅ Danyang Zhuo ⋅ Cheng Leong ⋅ Ruoming Pang

Project Page

Abstract

AXLearn is a production system which 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 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 hundreds 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. We also share our experience in the development and operation of AXLearn at Apple.