AXLearn: Modulaire training van grote modellen op heterogene infrastructuur

AXLearn: Modular Large Model Training on Heterogeneous Infrastructure

July 7, 2025

Auteurs: Mark Lee, Tom Gunter, Chang Lan, John Peebles, Hanzhi Zhou, Kelvin Zou, Sneha Bangalore, Chung-Cheng Chiu, Nan Du, Xianzhi Du, Philipp Dufter, Ruixuan Hou, Haoshuo Huang, Dongseong Hwang, Xiang Kong, Jinhao Lei, Tao Lei, Meng Li, Li Li, Jiarui Lu, Zhiyun Lu, Yiping Ma, David Qiu, Vivek Rathod, Senyu Tong, Zhucheng Tu, Jianyu Wang, Yongqiang Wang, Zirui Wang, Floris Weers, Sam Wiseman, Guoli Yin, Bowen Zhang, Xiyou Zhou, Danyang Zhuo, Cheng Leong, Ruoming Pang

cs.AI

Samenvatting

We ontwerpen en implementeren AXLearn, een productiesysteem voor deep learning dat schaalbare en hoogwaardige training van grote deep learning-modellen mogelijk maakt. In vergelijking met andere state-of-the-art deep learning-systemen heeft AXLearn een unieke focus op modulariteit en ondersteuning voor heterogene hardware-infrastructuur. De interne interfaces tussen softwarecomponenten in AXLearn volgen strikte encapsulatie, waardoor verschillende componenten kunnen worden samengesteld om snelle modelontwikkeling en experimenten op heterogene rekeninfrastructuur te faciliteren. We introduceren een nieuwe methode om modulariteit te kwantificeren via Lines-of-Code (LoC)-complexiteit, wat aantoont hoe ons systeem een constante complexiteit behoudt terwijl we de componenten in het systeem opschalen, in tegenstelling tot lineaire of kwadratische complexiteit in andere systemen. Hierdoor kunnen functies zoals Rotary Position Embeddings (RoPE) in AXLearn worden geïntegreerd over honderden modules met slechts 10 regels code, vergeleken met honderden die in andere systemen nodig zijn. Tegelijkertijd behoudt AXLearn een vergelijkbare prestatieniveau als state-of-the-art trainingssystemen. Tot slot delen we onze ervaringen met de ontwikkeling en het operationele beheer van AXLearn.

English

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.