AXLearn : Entraînement de grands modèles modulaires sur une infrastructure hétérogène

AXLearn: Modular Large Model Training on Heterogeneous Infrastructure

July 7, 2025

papers.authors: 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

papers.abstract

Nous concevons et mettons en œuvre AXLearn, un système de deep learning de production qui facilite l'entraînement scalable et performant de modèles de deep learning de grande taille. Par rapport à d'autres systèmes de deep learning de pointe, AXLearn se distingue par son accent unique sur la modularité et son support pour une infrastructure matérielle hétérogène. Les interfaces internes d'AXLearn entre les composants logiciels suivent une encapsulation stricte, permettant à différents composants d'être assemblés pour faciliter le développement rapide de modèles et l'expérimentation sur des infrastructures de calcul hétérogènes. Nous introduisons une méthode novatrice de quantification de la modularité via la complexité en lignes de code (LoC), qui démontre comment notre système maintient une complexité constante lors de la mise à l'échelle des composants, contrairement à une complexité linéaire ou quadratique dans d'autres systèmes. Cela permet d'intégrer des fonctionnalités telles que les Rotary Position Embeddings (RoPE) dans AXLearn à travers des centaines de modules avec seulement 10 lignes de code, contre des centaines nécessaires dans d'autres systèmes. Parallèlement, AXLearn maintient des performances équivalentes par rapport aux systèmes d'entraînement de pointe. Enfin, nous partageons notre expérience dans le développement et l'exploitation d'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.