AXLearn: Modular Large Model Training on Heterogeneous Infrastructure
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.
686e2a81a5f0f70d9de40c59
Mark Lee
686e2a81a5f0f70d9de40c5a
Tom Gunter
686e2a81a5f0f70d9de40c5b
Chang Lan
686e2a81a5f0f70d9de40c5c
John Peebles
686e2a81a5f0f70d9de40c5d
Hanzhi Zhou
686e2a81a5f0f70d9de40c5e
Kelvin Zou
686e2a81a5f0f70d9de40c5f
Sneha Bangalore
686e2a81a5f0f70d9de40c60
Chung-Cheng Chiu
686e2a81a5f0f70d9de40c61
Nan Du
686e2a81a5f0f70d9de40c62
Xianzhi Du
686e2a81a5f0f70d9de40c63
Philipp Dufter
686e2a81a5f0f70d9de40c64
Ruixuan Hou
686e2a81a5f0f70d9de40c65
Haoshuo Huang
686e2a81a5f0f70d9de40c66
Dongseong Hwang
686e2a81a5f0f70d9de40c67
Xiang Kong
686e2a81a5f0f70d9de40c68
Jinhao Lei
686e2a81a5f0f70d9de40c69
Tao Lei
686e2a81a5f0f70d9de40c6a
Meng Li
686e2a81a5f0f70d9de40c6b
Li Li
686e2a81a5f0f70d9de40c6c
Jiarui Lu
686e2a81a5f0f70d9de40c6d
Zhiyun Lu
686e2a81a5f0f70d9de40c6e
Yiping Ma
686e2a81a5f0f70d9de40c6f
David Qiu
686e2a81a5f0f70d9de40c70
Vivek Rathod
686e2a81a5f0f70d9de40c71
Senyu Tong
686e2a81a5f0f70d9de40c72
Zhucheng Tu
686e2a81a5f0f70d9de40c73
Jianyu Wang
686e2a81a5f0f70d9de40c74
Yongqiang Wang
686e2a81a5f0f70d9de40c75
Zirui Wang
686e2a81a5f0f70d9de40c76
Floris Weers
686e2a81a5f0f70d9de40c77
Sam Wiseman
686e2a81a5f0f70d9de40c78
Guoli Yin
686e2a81a5f0f70d9de40c79
Bowen Zhang
686e2a81a5f0f70d9de40c7a
Xiyou Zhou
686e2a81a5f0f70d9de40c7b
Danyang Zhuo
686e2a81a5f0f70d9de40c7c
Cheng Leong
686e2a81a5f0f70d9de40c7d
Ruoming Pang
AXLearn: Modular Large Model Training on Heterogeneous Infrastructure
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.
AXLearn is a modular deep learning system designed for scalable and high-performance training on heterogeneous hardware, maintaining constant complexity and equivalent performance to state-of-the-art systems.
AXLearn modularity heterogeneous hardware infrastructure encapsulation Lines-of-Code (LoC)-complexity Rotary Position Embeddings (RoPE)