TorchBraid：利用MPI和GPU加速实现深度神经网络的高性能层并行训练

《ACM Transactions on Mathematical Software》：TorchBraid: High-Performance Layer-Parallel Training of Deep Neural Networks with MPI and GPU Acceleration

TorchBraid通过整合PyTorch和XBraid实现层并行训练，支持MPI和GPU加速，解决了传统DNN训练的资源利用率瓶颈，在图像分类和循环神经网络中验证了高效性，并首次提出层并行与数据并行的混合算法，同时开发适应层并行的批归一化等标准技术，3D图像分类场景下速度提升达10倍。

　　

摘要

TorchBraid是一种高性能的深度神经网络（DNN）层并行训练实现方式，支持基于MPI的并行处理和GPU加速。层并行训练旨在克服DNN在前向传播和反向传播过程中固有的序列化问题，这种序列化限制了计算资源在高扩展性场景下的利用率。为此，TorchBraid将PyTorch神经网络框架与先进的XBraid时间并行库相结合。本文介绍了TorchBraid的用途和性能，并探讨了在自动微分与层并行训练结合过程中遇到的算法挑战及其解决方案。实验结果涵盖了使用GPU加速与否的情况，涉及Tiny ImageNet、MNIST图像分类数据集以及循环神经网络。总体而言，TorchBraid能够在高扩展性和低扩展性场景下实现DNN的快速训练。

除了TorchBraid软件本身，文章还详细介绍了几种层并行算法的新进展。首次提出了层并行算法与数据并行算法的集成方案，展示了这种组合在计算上的优势。针对层并行训练，还开发了批量归一化等标准深度学习技术。最后，提出了一种将层并行技术与空间粗化相结合的新方法，用于加速3D图像分类任务的训练，其速度比串行执行快约10倍。