AXLearn: Modular Large Model Training on Heterogeneous Infrastructure
This addresses the problem of efficient and flexible model development for researchers and engineers working with large-scale AI, though it appears incremental as it builds on existing training systems with a focus on modularity.
The authors tackled the challenge of scalable and high-performance training for large deep learning models by developing AXLearn, a modular system that supports heterogeneous hardware infrastructure, achieving equivalent performance to state-of-the-art systems while enabling features like Rotary Position Embeddings to be integrated with just 10 lines of code compared to hundreds in other systems.
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.