AXLearn: Modulares Training großer Modelle auf heterogener Infrastruktur

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

Wir entwerfen und implementieren AXLearn, ein produktionsreifes Deep-Learning-System, das skalierbares und leistungsstarkes Training großer Deep-Learning-Modelle ermöglicht. Im Vergleich zu anderen modernsten Deep-Learning-Systemen zeichnet sich AXLearn durch einen besonderen Fokus auf Modularität und die Unterstützung heterogener Hardware-Infrastrukturen aus. Die internen Schnittstellen zwischen den Softwarekomponenten von AXLearn folgen einer strengen Kapselung, wodurch verschiedene Komponenten zusammengesetzt werden können, um eine schnelle Modellentwicklung und Experimentation auf heterogenen Recheninfrastrukturen zu ermöglichen. Wir führen eine neuartige Methode zur Quantifizierung von Modularität über die Komplexität der Codezeilen (Lines-of-Code, LoC) ein, die zeigt, wie unser System eine konstante Komplexität beibehält, während wir die Komponenten im System skalieren, im Gegensatz zu linearer oder quadratischer Komplexität in anderen Systemen. Dies ermöglicht die Integration von Funktionen wie Rotary Position Embeddings (RoPE) in AXLearn über Hunderte von Modulen mit nur 10 Codezeilen, verglichen mit Hunderten, die in anderen Systemen erforderlich sind. Gleichzeitig bietet AXLearn eine gleichwertige Leistung im Vergleich zu modernsten Trainingssystemen. Abschließend teilen wir unsere Erfahrungen in der Entwicklung und dem Betrieb von 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.