[论文解读] The Case for Strong Scaling in Deep Learning: Training Large 3D CNNs with Hybrid Parallelism

Yosuke Oyama, Naoya Maruyama+7 more|arXiv (Cornell University)|Jul 25, 2020

Advanced Neural Network Applications参考文献 54被引用 9

一句话总结

本文提出了一种端到端的混合并行训练框架，通过结合数据并行与空间并行，将单个样本分布到多个GPU上，从而实现对大型3D卷积神经网络（3D CNNs）的强可扩展性。通过将并行化扩展至小批量大小之外，并利用空间分区优化I/O，该方法实现了高性能，支持在全分辨率512³宇宙学数据上进行训练，与低分辨率训练相比，CosmoFlow的预测精度提升了10倍。

ABSTRACT

We present scalable hybrid-parallel algorithms for training large-scale 3D convolutional neural networks. Deep learning-based emerging scientific workflows often require model training with large, high-dimensional samples, which can make training much more costly and even infeasible due to excessive memory usage. We solve these challenges by extensively applying hybrid parallelism throughout the end-to-end training pipeline, including both computations and I/O. Our hybrid-parallel algorithm extends the standard data parallelism with spatial parallelism, which partitions a single sample in the spatial domain, realizing strong scaling beyond the mini-batch dimension with a larger aggregated memory capacity. We evaluate our proposed training algorithms with two challenging 3D CNNs, CosmoFlow and 3D U-Net. Our comprehensive performance studies show that good weak and strong scaling can be achieved for both networks using up 2K GPUs. More importantly, we enable training of CosmoFlow with much larger samples than previously possible, realizing an order-of-magnitude improvement in prediction accuracy.

研究动机与目标

• 解决在高维科学数据（如宇宙学模拟和3D医学影像）上训练大型3D CNN时面临的内存与计算瓶颈。
• 克服传统数据并行的局限性，后者因大样本带来的单GPU内存压力而难以有效扩展。
• 通过引入空间并行，将单个样本在多个GPU间分区，实现强可扩展性——在不增加小批量大小的前提下减少训练时间。
• 通过在数据加载中应用相同的混合并行技术，实现端到端的I/O流水线扩展，最大限度减少大样本的I/O瓶颈。
• 证明全分辨率训练能显著提升模型精度，尤其在需要长距离特征学习的科学任务中。

提出的方法

• 提出一种混合并行算法，结合数据并行与空间并行，将单个3D样本分布到多个GPU上，降低单GPU的内存使用量。
• 通过将特征图的空间维度（如高度、宽度、深度）在处理单元间分区，实现3D卷积的空间分区，从而在小批量限制之外实现强可扩展性。
• 利用空间分区优化I/O流水线，实现可扩展的集体与异步I/O，使其随GPU数量增长而扩展。
• 在LBANN深度学习框架中扩展自定义内核与通信模式，以支持混合并行训练与I/O，确保通信开销最低。
• 使用边界层交换（halo exchanges）在空间分区之间传递梯度，相比通道/特征并行中使用的完整张量通信，显著减少数据移动。
• 采用基于模型的性能分析，识别并缓解3D CNN在计算与I/O方面的可扩展瓶颈。

实验结果

研究问题

• RQ1 通过将并行化从数据并行扩展到单个样本的空间分区，是否可以在3D CNN训练中实现强可扩展性？
• RQ2 结合空间与数据并行的混合并行训练是否能在高达2,048个GPU上实现高效扩展，同时保持高吞吐量与低通信开销？
• RQ3 通过在计算中使用的相同混合并行技术，是否能有效扩展I/O流水线，以应对大尺寸3D数据样本？
• RQ4 尽管面临更高的内存与计算需求，使用全分辨率3D数据（如512³）进行训练是否能显著提升模型精度，优于低分辨率数据训练？
• RQ5 在使用混合并行训练与I/O时，CosmoFlow与3D U-Net等3D CNN的性能特征与可扩展效率如何？

主要发现

• 在512³全分辨率样本上训练CosmoFlow，其均方误差（MSE）达到0.0018，相比在128³样本上训练（MSE = 0.0289）实现了10倍的预测精度提升。
• CosmoFlow模型在64小批量大小下，从512个GPU扩展到2048个GPU时实现了1.77倍加速，证明了在不增加小批量大小情况下的强可扩展性。
• 3D U-Net模型在从256个GPU扩展到512个GPU时实现了1.42倍加速，证实了在不同架构与工作负载下均存在强可扩展行为。
• CosmoFlow与3D U-Net在2,048个GPU内均表现出良好的弱可扩展性与强可扩展性，性能分析显示计算与I/O组件均具有高可扩展效率。
• 混合并行I/O流水线成功随GPU数量扩展，即使在常见于科学3D数据的千兆字节级样本下，也有效最小化了I/O瓶颈。