Apple Open-Sources AXLearn, the System Training Siri Across 10,000 Experiments

Apple has published a research paper and fully open-sourced AXLearn, the internal deep learning training system it has used to train models powering features for over a billion users. The library sits on top of JAX and XLA, and its central pitch is simple but ambitious: write your model once, run it efficiently on whatever hardware is cheapest or most available, without touching a single line of model code.

The problem every ML team hits at scale

Most training frameworks are quietly coupled to a specific hardware vendor. Megatron-LM has a vested interest in optimizing for Nvidia GPUs, while Haiku, Flax, Pax, and MaxText are mostly optimized for Google TPUs. For a company like Apple that trains thousands of models across multiple cloud providers simultaneously, that kind of lock-in is a liability , both operationally and financially.

Apple took the alternate approach of developing a hardware-agnostic system that can adapt seamlessly to multiple platforms. For example, while AWS Trainium2 did not exist when development began, AXLearn is one of the first deep learning systems that supports Trainium2 at scale.

Modularity as a first-class engineering goal

The headline technical contribution of the paper is not a new architecture or training trick , it is a rigorous argument that modularity should be measured, not assumed. The team introduces a metric called LoC-complexity: how many lines of code must change across a codebase when you add a new feature, as the number of modules scales asymptotically.

They introduce a novel method of quantifying modularity via Lines-of-Code (LoC)-complexity, which demonstrates how their system maintains constant complexity as they 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 hundreds of modules with just 10 lines of code, compared to hundreds as required in other systems.

To make this concrete: adding Mixture-of-Experts (MoE) , routing tokens to specialized sub-networks instead of a single feed-forward layer , to a production codebase with 20 model variants would cost an estimated 4,000 lines of code in DeepSpeed. In AXLearn, the same change costs zero lines in existing modules, because the MoE layer is a drop-in swap at the config level.