[论文解读] Distributed Machine Learning for Computational Engineering using MPI

Kailai Xu, Weiqiang Zhu+1 more|arXiv (Cornell University)|Nov 2, 2020

Reservoir Engineering and Simulation Methods参考文献 42被引用 5

一句话总结

该论文提出了一种分布式机器学习框架，通过MPI将深度神经网络（DNNs）与偏微分方程（PDE）求解器集成，将数据通信视为计算图中的第一类节点。通过在分布式内存系统上实现并行的前向与伴随计算，该方法实现了强大的弱缩放性能，并在使用100个处理器时实现了超过20倍的加速，展示了大规模PDE约束DNNs的可扩展反演建模能力。

ABSTRACT

We propose a framework for training neural networks that are coupled with partial differential equations (PDEs) in a parallel computing environment. Unlike most distributed computing frameworks for deep neural networks, our focus is to parallelize both numerical solvers and deep neural networks in forward and adjoint computations. Our parallel computing model views data communication as a node in the computational graph for numerical simulations. The advantage of our model is that data communication and computing are cleanly separated and thus provide better flexibility, modularity, and testability. We demonstrate using various large-scale problems that we can achieve substantial acceleration by using parallel solvers for PDEs in training deep neural networks that are coupled with PDEs.

研究动机与目标

• 解决使用反向模式自动微分训练与PDE耦合的DNNs时面临的内存与计算瓶颈。
• 为计算工程中的大规模反演问题提供可扩展的分布式优化能力。
• 在统一的计算图抽象下，整合并行PDE求解器与分布式DNN训练，并将通信视为第一类节点。
• 克服通过MPI数据通信和混合并行（MPI + 多线程）进行梯度反向传播的挑战。

提出的方法

• 将数据通信操作（如MPI广播、边界交换）作为计算图中的显式节点，支持端到端的自动微分。
• 将DNNs与PDE求解器（如有限元/有限差分法）整合到统一的计算图中，通过反向模式AD进行梯度计算。
• 采用粗粒度自动微分方法，操作于高层抽象（如矩阵求解器、PDE算子），而非基本算术运算。
• 在ADCME库中实现自定义MPI原语，以支持通过通信模式的梯度反向传播。
• 将计算与通信逻辑解耦，以提升在混合MPI+多线程环境中的模块化、可测试性与灵活性。
• 应用域分解技术，将计算域划分为MPI进程，通过边界交换支持基于模板的PDE求解器。

实验结果

研究问题

• RQ1 能否将分布式PDE求解器中的数据通信建模为计算图中的第一类节点，以支持自动微分？
• RQ2 如何将反向模式自动微分扩展至处理混合MPI+多线程环境中的MPI通信模式梯度？
• RQ3 使用该框架训练与PDE耦合的DNNs时，可实现何种程度的强缩放与弱缩放性能？
• RQ4 该框架在跨数千个核心实现大规模并行的同时，能否保持数值精度与收敛性？
• RQ5 与传统数据并行或模型并行的DNN训练相比，该方法在物理信息学习中的性能表现如何？

主要发现

• 对于声波方程，该框架在100个处理器（3200个核心）上实现了超过20倍的加速，弱缩放性能几乎无下降。
• 弱缩放结果显示，声波方程从1个处理器扩展到100个处理器时，运行时间仅增加2倍，表明具有极佳的可扩展性。
• 对于弹性波方程，尽管内存与状态变量复杂度更高，该方法在100个处理器上仍保持强可扩展性，实现超过20倍的加速。
• 计算图抽象实现了通信与计算的清晰分离，提升了模块化与可测试性。
• 该框架成功支持对DNN参数与MPI通信操作的梯度反向传播，实现了端到端的梯度计算。
• 该实现展示了在涉及PDE与DNN的大规模反演问题（如波传播中的参数识别）中的实际可行性。