mlsanityData to 14 Oct 2026

AXLearn: Modular, Hardware-Agnostic Large Model Training

arXiv 2507.05411 ↗First posted 7 Jul 2025cs.LGcited by 4

Not on the map. Why?

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.

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

15 of 35 references.

References · 15

  1. 2020Language Models are Few-Shot Learnerscited by 32,419
  2. 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,599
  3. 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,944
  4. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
  5. 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,509
  6. 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,598
  7. 2019ZeRO: Memory Optimizations Toward Training Trillion Parameter Modelscited by 1,301
  8. 2023PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallelcited by 660
  9. 2021Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LMcited by 632
  10. 2014cuDNN: Efficient Primitives for Deep Learningcited by 245
  11. 2021Sequence Parallelism: Long Sequence Training from System Perspectivecited by 126
  12. 2024MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUscited by 112
3 more
  1. 2021GSPMD: General and Scalable Parallelization for ML Computation Graphscited by 111
  2. 2024TorchTitan: One-stop PyTorch native solution for production ready LLM pre-trainingcited by 85
  3. 2024Revisiting MoE and Dense Speed-Accuracy Comparisons for LLM Trainingcited by 8

Built on by

Cited by year

Citing papers per year: 2025: 4.

4 papers.

  1. 2025Apple Intelligence Foundation Language Models: Tech Report 2025cited by 28
  2. 2025Language Models Improve When Pretraining Data Matches Target Taskscited by 19
  3. 2025Data-Centric Lessons To Improve Speech-Language Pretrainingcited by 3
  4. 2025RLAX: Large-Scale, Distributed Reinforcement Learning for Large Language Models on TPUscited by 1