AXLearn: Modular, Hardware-Agnostic Large Model Training

2025/07/07 by Mark Lee, Lee, Mark, Chang Lan +67 · 4 citations
Computer Science · #Advanced Neural Network Applications #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning and Data Classification #Parallel Computing and Optimization Techniques

paper · pdf · doi:10.48550/arxiv.2507.05411

openalex publication_date 2025/07/07 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/30

Abstract

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.

Citations

• TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training
• Revisiting MoE and Dense Speed-Accuracy Comparisons for LLM Training
• MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs
• Efficient Memory Management for Large Language Model Serving with PagedAttention
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• Sequence Parallelism: Long Sequence Training from System Perspective
• GSPMD: General and Scalable Parallelization for ML Computation Graphs
• Efficient Large-Scale Language Model Training on GPU Clusters Using\n Megatron-LM
• Language Models are Few-Shot Learners
• Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
• cuDNN: Efficient Primitives for Deep Learning

Cited by

• RLAX: Large-Scale, Distributed Reinforcement Learning for Large Language Models on TPUs
• Data-Centric Lessons To Improve Speech-Language Pretraining
• Apple Intelligence Foundation Language Models: Tech Report 2025
• Language Models Improve When Pretraining Data Matches Target Tasks

Related

• Systems for Parallel and Distributed Large-Model Deep Learning Training
• Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training
• Varuna: Scalable, Low-cost Training of Massive Deep Learning Models
• Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning
• Whale: Efficient Giant Model Training over Heterogeneous GPUs