AXLearn: Modular Large Model Training on Heterogeneous Infrastructure

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

论文试图解决深度学习系统在支持大规模、高性能训练大模型的同时，如何实现系统的模块化与对异构硬件基础设施的支持这一问题。该问题在当前的深度学习系统中未能很好地解决，尤其是在快速开发和实验迭代方面存在不足。

提出了一种名为AXLearn的新型深度学习系统，其核心思路是通过严格的软件组件接口封装来实现系统的高度模块化，并引入一种基于代码行复杂度（LoC-complexity）量化模块化的方法。相比现有系统，AXLearn能够在扩展系统组件时保持恒定的复杂度，而其他系统通常呈现线性或二次增长的复杂度。

{"创新地提出了用Lines-of-Code (LoC)-complexity来衡量系统模块化程度的新方法。",实现了高效的模块集成，例如在数百个模块中集成RoPE仅需10行代码，而其他系统可能需要数百行。,系统性能与当前最先进的训练系统相当，同时兼具高可扩展性和灵活性。,论文分享了在实际开发和操作AXLearn中的经验，为后续研究和工业应用提供了宝贵的参考。,实验设计充分验证了系统的高效性和灵活性，但未明确提及是否开源代码及具体使用的数据集。}

{Google的JAX及其生态系统，强调高性能计算与模块化编程。,"PyTorch Lightning: A Highly Modular Framework for Rapid Experimentation in Deep Learning.","TensorFlow Extended (TFX): A TensorFlow-Based Production-Scale Machine Learning Platform.","DeepSpeed: A System for Training Large-Scale Deep Learning Models Efficiently.","FairScale: A PyTorch Extension Library for Model Parallelism and Optimization."}