AXLearn: Модульное обучение больших моделей на гетерогенной инфраструктуре
AXLearn: Modular Large Model Training on Heterogeneous Infrastructure
July 7, 2025
Авторы: 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
Аннотация
Мы разработали и реализовали AXLearn — производственную систему глубокого обучения, которая обеспечивает масштабируемое и высокопроизводительное обучение крупных моделей глубокого обучения. В отличие от других современных систем глубокого обучения, AXLearn уделяет особое внимание модульности и поддержке гетерогенной аппаратной инфраструктуры. Внутренние интерфейсы между программными компонентами AXLearn строго инкапсулированы, что позволяет гибко комбинировать различные компоненты для ускоренной разработки моделей и экспериментов на гетерогенных вычислительных инфраструктурах. Мы представляем новый метод количественной оценки модульности с использованием сложности, измеряемой в строках кода (LoC), который демонстрирует, как наша система сохраняет постоянную сложность при масштабировании компонентов, в отличие от линейной или квадратичной сложности в других системах. Это позволяет интегрировать такие функции, как Rotary Position Embeddings (RoPE), в AXLearn через сотни модулей всего с 10 строками кода, тогда как в других системах для этого требуются сотни строк. При этом AXLearn сохраняет производительность на уровне современных систем обучения. Наконец, мы делимся опытом разработки и эксплуатации 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.