AXLearn: Modular Large Model Training on Heterogeneous Infrastructure

2025-07-07

AI Summary

To address the low modularity and high scalability cost of large-model training systems on heterogeneous hardware, this paper proposes a modular training system architecture characterized by high cohesion and low coupling. Our approach strictly encapsulates component interfaces, introduces a lines-of-code-based metric for quantifying module complexity, and integrates a heterogeneous-resource-aware scheduling mechanism—enabling constant-complexity system growth (e.g., integrating RoPE requires only ~10 lines of code across hundreds of modules). Experiments demonstrate that the system achieves training performance comparable to mainstream frameworks while significantly improving modularity, accelerating feature integration by an order of magnitude, and substantially reducing maintenance overhead. This work establishes a novel, evolution-friendly engineering paradigm for large-scale deep learning systems.

Technology Category

• Machine Learning: Scalability of ML Systems
• Computer Vision: Large Vision Models
• Natural Language Processing: (Large) Language Models

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

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.

Problem

• Enables scalable training of large deep learning models
• Focuses on modularity and heterogeneous hardware support
• Reduces code complexity for integrating new features

Innovation

• Modular design for scalable model training
• Heterogeneous hardware infrastructure support
• Low-code integration across hundreds modules