AXLearn：异构基础设施上的模块化大型模型训练

我们设计并实现了AXLearn，一个生产级深度学习系统，旨在促进大型深度学习模型的可扩展和高性能训练。与其他最先进的深度学习系统相比，AXLearn独特地侧重于模块化以及对异构硬件基础设施的支持。AXLearn内部软件组件之间的接口遵循严格的封装，允许不同组件进行组装，从而促进在异构计算基础设施上的快速模型开发和实验。我们引入了一种通过代码行数（LoC）复杂性来量化模块化的新颖方法，该方法表明我们的系统在扩展组件时能保持恒定复杂性，而其他系统则呈现线性或二次复杂性。这使得将诸如旋转位置嵌入（RoPE）之类的功能集成到AXLearn中时，只需10行代码即可跨数百个模块实现，而其他系统则需要数百行。同时，AXLearn与最先进的训练系统相比，保持了同等的性能。最后，我们分享了AXLearn的开发和运行经验。