AXLearn: Addestramento di Modelli di Grandi Dimensioni su Infrastruttura Eterogenea con Approccio Modulare
AXLearn: Modular Large Model Training on Heterogeneous Infrastructure
July 7, 2025
Autori: 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
Abstract
Progettiamo e implementiamo AXLearn, un sistema di deep learning di produzione che facilita l'addestramento scalabile e ad alte prestazioni di modelli di deep learning di grandi dimensioni. Rispetto ad altri sistemi di deep learning all'avanguardia, AXLearn si distingue per un focus unico sulla modularità e sul supporto per infrastrutture hardware eterogenee. Le interfacce interne di AXLearn tra i componenti software seguono un rigoroso incapsulamento, consentendo a diversi componenti di essere assemblati per facilitare lo sviluppo rapido di modelli e la sperimentazione su infrastrutture di calcolo eterogenee. Introduciamo un metodo innovativo per quantificare la modularità tramite la complessità delle Linee di Codice (LoC), che dimostra come il nostro sistema mantenga una complessità costante mentre scaliamo i componenti del sistema, rispetto a una complessità lineare o quadratica in altri sistemi. Ciò consente di integrare funzionalità come gli Embedding di Posizione Rotazionali (RoPE) in AXLearn attraverso centinaia di moduli con sole 10 righe di codice, rispetto alle centinaia richieste in altri sistemi. Allo stesso tempo, AXLearn mantiene prestazioni equivalenti rispetto ai sistemi di addestramento all'avanguardia. Infine, condividiamo la nostra esperienza nello sviluppo e nel funzionamento di 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.