AXLearn: Modular Large Model Training on Heterogeneous Infrastructure

我们设计并实现了 AXLearn，这是一个用于大规模深度学习模型可扩展、高性能训练的生产级深度学习系统。与其他最先进的深度学习系统相比，AXLearn 独特地专注于模块化设计以及对异构硬件基础设施的支持。AXLearn 内部软件组件之间的接口遵循严格的封装原则，使得不同组件可以灵活组合，从而加快在异构计算基础设施上的模型开发与实验迭代速度。我们提出了一种通过代码行数（LoC）复杂度来量化模块化程度的新方法，表明随着系统中组件规模的扩大，我们的系统能够保持恒定的复杂度，而其他系统则通常呈现线性甚至二次增长的复杂度。这种特性使得像旋转位置嵌入（RoPE）这样的功能可以在 AXLearn 中仅通过 10 行代码就集成到数百个模块中，而在其他系统中则通常需要数百行代码。同时，AXLearn 的训练性能与当前最先进的训练系统相当。最后，我们分享了在 AXLearn 开发与运维过程中的实践经验。

论文旨在解决当前深度学习系统在扩展大规模模型训练时面临的可扩展性、高性能以及模块化支持不足的问题。同时，验证通过严格封装的内部接口和模块化设计，可以实现快速模型开发和实验，并在异构硬件基础设施上保持系统复杂度的可控性。这是一个在当前深度学习系统研究中较为新颖的问题，特别是在模块化与异构硬件支持的结合方面。

论文提出了一种名为AXLearn的生产级深度学习系统，其核心思想是通过严格的模块封装和接口设计实现系统的高度模块化，从而支持在异构硬件上进行高效的大规模模型训练。关键创新点在于引入了一种基于代码行复杂度（LoC-complexity）的模块化量化方法，使得系统复杂度在组件扩展时保持常数级增长，而非传统的线性或二次增长。

{AXLearn能够在不牺牲性能的前提下显著提升系统的模块化程度。,实验显示，在集成如RoPE等复杂特性时，AXLearn仅需少量代码即可完成跨多个模块的集成。,论文通过实际模型开发和部署经验分享了系统在工业级应用中的可行性。,未来值得深入研究的方向包括进一步优化异构硬件调度机制以及扩展LoC复杂度模型在更大规模系统中的应用。}

{"DeepSpeed: System Optimization for Large-Scale Model Training","TensorFlow: A System for Large-Scale Machine Learning","PyTorch: An Imperative Style, High-Performance Deep Learning Platform","JAX: Composable Transformations of Python+NumPy Programs","Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism"}