# AXLearn: Modular, Hardware-Agnostic Large Model Training

arXiv: https://arxiv.org/abs/2507.05411 · HTML: https://mlsanity.com/paper/2507.05411/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 35 references that are in the map.
- First posted on arXiv 7 Jul 2025.

## Builds on (most cited first, 15 of 15)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [cuDNN: Efficient Primitives for Deep Learning](https://mlsanity.com/paper/1410.0759/) (2014)
- [Sequence Parallelism: Long Sequence Training from System Perspective](https://mlsanity.com/paper/2105.13120/) (2021)
- [MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs](https://mlsanity.com/paper/2402.15627/) (2024)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training](https://mlsanity.com/paper/2410.06511/) (2024)
- [Revisiting MoE and Dense Speed-Accuracy Comparisons for LLM Training](https://mlsanity.com/paper/2405.15052/) (2024)

## Built on by (most cited first, 4 of 4)

- [Apple Intelligence Foundation Language Models: Tech Report 2025](https://mlsanity.com/paper/2507.13575/) (2025)
- [Language Models Improve When Pretraining Data Matches Target Tasks](https://mlsanity.com/paper/2507.12466/) (2025)
- [Data-Centric Lessons To Improve Speech-Language Pretraining](https://mlsanity.com/paper/2510.20860/) (2025)
- [RLAX: Large-Scale, Distributed Reinforcement Learning for Large Language Models on TPUs](https://mlsanity.com/paper/2512.06392/) (2025)

## Abstract (arXiv)

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.