ADCME MPI：分布式机器学习计算工程-并行化数值PDE求解器和DNN的框架优势与应用

--NN你好OBSN∇ADCME MPI：分布式机器学习计算工程徐开来，1埃里克·达夫1，2kailaix，stanford.edu1计算与数学工程研究所，斯坦福大学，斯坦福，加利福尼亚州94305，美国;2机械工程，斯坦福大学，斯坦福，CA 94305，美国摘要我们提出了一个用于在并行计算环境中训练与偏微分方程（PDE）耦合的深度神经网络（DNN）的框架。与DNN的大多数分布式计算框架不同，我们的重点是在正向和伴随计算中并行化数值求解器和DNN。我们的并行计算模型将数据通信视为数值模拟计算图中的一个节点。我们的模型的优点是，数据通信和计算是干净的分离，这使得更好的灵活性，模块化和软件的可测试性。我们在一个大规模问题上展示了我们的方法，并表明我们可以通过使用并行数值PDE求解器来实现大幅加速，同时训练与PDE耦合的DNN。介绍深度神经网络（DNN）已被证明对于解决计算工程中的逆问题非常有效（Raissi，Perdikaris和Karniadakis2019; Meng等人2020; Pakravan等人2020）。在我们之前 的 工 作 中 （Xu ， Huang ， and Darve 2020; Xu andDarve 2019; Huang et al. 2020; Fan et al. 2020），我们成功地将数值求解器和深度神经网络结合起来，用于数据驱动的逆建模。在数学上，我们考虑隐式模型F（f，u）= 0，其中f是未知函数，我们使用DNN对其进行近似。我们表示DNNθ，其中θ是神经网络的权重和偏置;u是一个函数，它依赖于f通过F（f，u）= 0。我们给出函数u的一些（部分）观测值uobs，用于优化θ和最小化差异在f和θ之间。将优化问题转化为一个偏微分方程约束优化问题最小L（u）（u是θ的间接函数）在本文中，我们使用平方损失函数L（u）=22这种制剂的一个优点是已知的物理学，例如由偏微分方程描述的物理定律，在最大程度上被保留，并且用发展良好且有效的数值求解器来求解。同时，我们可以利用DNN的近似能力。为了解决这个优化问题，可以首先数值求解方程（1）中的物理约束F（θ，u）= 0，然后将解u（θ）代入L（u）。这导致了一个无约束优化问题minL（θ）=L（u（θ））θ我 们 开 发 了 Julia （ Bezanson 等 人 。 2017 ） 库 ， AD-CME（Xu和Darve 2020 a），使用TensorFlow（Abadi等人2016）自动区分后端来解决这种类型的问题。ADCME使用计算图表达数值求解器和DNN在那里-前，梯度可自动计算θLθ（θ）通过反向传播梯度通过数值求解器1和DNN。在本文中，我们可以互换地使用然而，这种方法的一个挑战是，对于大规模的问题，数值求解器的内存和计算成本是过高的。在现代分布式内存高性能计算（HPC）架构上解决此类大规模问题的实际标准是消息传递接口（MPI）（Gabriel等人2004; Gropp，Thakur，and Lusk 1999）。ADCME使用的TensorFlow后端最初是为深度学习/机器学习而设计的。尽管在扩展TensorFlow以进行机器学习模型的分布式训练方面有很多工作，但一些关键功能，如分布式θ使得F（Nθ，u）= 0（一）线性代数和区域分解的方法，用于解决科学计算问题，目前还很缺乏。这页-L（u）是一个损失函数，它衡量假设观测值和实际观测值之间的差异。为了考试-版权所有© 2021本文由其作者。知识共享许可署名4.0国际（CC BY 4.0）per是关于将MPI功能整合到ADCME中，以实现分布式内存的可扩展性和灵活性1有关梯度反向传播如何在ADCME中用于数值求解器的详细信息，我们请读者参阅（Xu and Darve 2020 b）。NN参数NN参数其主要思想是通过将网格或矩阵拆分到不同的MPI秩（或处理器）上来并行化数值求解器然后，将数据通信节点插入到计算图中。 因为对于我们的计算工程应用，DNN通常很小，所以它们在每个处理器上都是重复的。每个计算图包括一组“通信”节点（图1）。1-top），这在单处理器计算图中是不存在的。 这些操作符调用MPI调用，并负责不同计算机节点之间的数据通信。在梯度反向传播过程中，我们需要反转数据流方向和数据通信算子的操作（Wang和Poten 2015; Utke等人，2015年）。2009）（Fig.1-底部）。和小批量优化算法，如随机梯度下降（Bottou 2010）。每个处理器计算预测和梯度，这些预测和梯度在一个或多个处理器上聚合。为了在有限的带宽环境中进一步扩展，实现了参数服务器（Li等人，2013，2014），其中每个服务器存储一部分参数还有大量关于模型并行化的工作，通过将DNN分成多个部分来并行计算（Chen，Yang和Cheng 2018; Hewett和Grady II 2020）。计算工程中的并行计算模型与深度学习有很大不同计算工程应用程序的特点是在网格中的相邻点之间进行数据通信（区域分解），NN参数最后时间步计算系数数据传输更新状态变量处理器0的计算图下一时间步骤矩阵的不同部分从计算图的角度来看，这种模式表明，除了在最后减少梯度之外，还有更多的通信这激励我们设计新的分布式计算模型，用于计算工程逆建模应用。MPI通信计算系数更新处理器1的计算图方法最后时间步最后时间步最后时间步数据传输计算系数数据传输手动MPI伴随计算系数数据传输状态变量更新状态变量更新状态变量下一时间步骤处理器0上的计算图下一时间步骤处理器1上的计算图下一时间步骤ADCME MPI旨在提供一个模块化，高效，灵活实现分布式计算的反求建模。从概念上讲，我们可以将数据通信操作视为计算图中的节点：它们类似于计算节点（例如，线性求解器），除了它们的责任是调用MPI调用并反向传播反向模式自动差分中的梯度（Baydin等人，2017年）。 该方案对ADCME库作了很大的改进，因为要将单处理机程序转换为多处理机程序，用户只需根据需要插入数据通信节点，而原程序的大部分内容不作改动。在本节中，我们简要描述了我们在ADCME MPI中的贡献，以扩展其分布式计算能力，以耦合DNN和PDE求解器。图1：用于计算工程的分布式机器学习的范例。数据通信操作被视为计算图中的节点它们与其他计算操作员是分开的这使得它很容易重用现有的串行实现的数值偏微分方程求解器。分布式计算模型有 许 多 现 有 的 工 作 和 软 件 用 于 使 用 深 度 神 经 网 络（ Griebel 和 Zum-busch1999; Notay 1995; Douglas ，Haase 和 Langer 2003 ） 和 数 值 PDE （ Bekkerman ，Bilenko和Langford 2011; Jordan和Mitchell 2015）的分布式计算由于目标应用程序的不同特征，这两个领域具有ADCME MPI包含一种混合模型，该模型适用于计算工程中的逆建模（1）的方法需要上述两种模型的组合在深度学习中，一个主要的挑战是数据集太大，无法放入内存。因此，数据集和计算负载都分布在不同的机器MPI APIADCME MPI 提 供 了 一 组 常 用 的 MPI 工 具 ， 如 mpibcast、mpi gather、mpi send等。这些操作符是标准MPI API的包装器然而，这些运算符也是“可微的”，因为它们可以处理梯度反向传播。梯度反向传播功能利用了前向和后向MPI调用之间存在一一对应的事实。例如，前向光环交换为了实现网格中相邻块之间的通信，ADCME为晕圈交换模式提供了数据通信算子的有效实现（图2）。Halo交换模式在科学计算中非常常见，因为数值求解器通常涉及相邻点之间的通信（Bianco 2014）。我们使用非阻塞发送/接收策略。在梯度反向传播阶段，这种发送/接收顺序被颠倒。基于自动微分的梯度反向传播算法或物理限制学习发送->接收Bcast->减少NN参数正演计算标志∇埃什基 .Σ.Σ≡Σ原始矩阵MPI_Isend/MPI_Irecv转置每个块数据交换Deadlock无死锁图2：左：区域分解和晕圈交换模式。两个相邻的补丁交换边界信息。右：使用非阻塞发送/接收来避免死锁。矩阵转置稀疏矩阵是计算工程中非常重要的工具。ADCME MPI将大型稀疏矩阵存储为CSR矩阵，并且每个MPI处理器存储具有连续行索引的行的一部分对于反向模式自动微分，矩阵转置是梯度反向传播中常见的算子。例如，假设前向计算是（x是输入，y是输出，A是矩阵）y=Ax（2）图3：分布式稀疏矩阵的转置 每个块与对应的转置块交换非零条目。原始矩阵和结果矩阵都以CSR格式存储。尽管已有许多分布式优化算法，在这项工作中，我们采用了一种简单的方法：聚集梯度θfi（θ）和更新θ根处理器上。图4展示了如何将现有的优化器转换为支持MPI的优化器。 基本思想是让根处理器通知工作处理器是否计算损失函数或梯度。然后，根处理器和工人将合作执行相同的例程，从而确保集体MPI调用的正确性。师傅给定损失函数L（y），梯度反向传播计算对于k = 1，2，3，flag =同步_OBJmpi_sync！（国旗）while truempi_sync！（国旗）if（flag==0）L（y（x）= L（y）（y）=Af =compute_objective_function（x）…compute_objective_function（x）elseif（flag==BLOG_GRAD）∂x∂y∂x∂y这里，n（y）是行向量，因此标志=BLOG_GRADmpi_sync！（国旗）dx =compute_gradient（x）…x = xcompute_gradient（x）else断裂末端端L（y（x）T阿克斯=AT（y）T埃什基标志=OPTIMIZATION_STOPmpi_sync！（国旗）需要矩阵向量乘法，其中矩阵为霸王分布式稀疏矩阵的转置分三步实现（图1）。（3）：1. 每个MPI处理器拥有的子矩阵被分成子块。Meta信息（例如，每个块中的非零数目）被收集。2. 每个块B与应该放置BT的目标块交换Meta信息图4：在我们的框架中将串行优化器重构为基于MPI的优化器。数字基准作为演示，我们提出了一个基准测试结果与ADCMEMPI。算例表明，ADCME引入的开销与实际计算相比非常小控制方程由泊松方程给出3. 每个子块被转置并且数据被传送到目标块。n·（κ（x）<$u（x））=f（x）x∈nu（x）= 0x∈N（三）分布式优化一般来说，我们问题的目标函数可以写成局部目标函数NminL（θ）=fi（θ）θi=1其中fi是局部目标函数，N是处理器的数量这里，f（x）1和κ（x）由深度神经网络近似。κ（x）=Nθ（x）其中θ是神经网络的权重和偏差。在均匀网格上用有限差分法对方程进行离散，离散后得到线性方程组A（θ）uf（4）u是解向量，f是源向量。注意A是一个稀疏矩阵，它的元素依赖于θ。发送发送RecvRecv伊森德伊森德伊雷奇夫伊雷奇夫等等时间（秒）×××稀疏矩阵是使用可微晕交换算子构造的，并存储为分布式CSR 矩阵 。方程 （4 ）使 用 Hypre （ Falgout 和 Yang2002）中的代数多重网格方法求解。 在梯度反向传播过程中，我们需要求解一个具有系数矩阵AT的线性系统。该矩阵是使用上一节中描述的技术获得的在强标度实验中，我们考虑一个固定的403020100149162536496481100 400 900 1600 2500 3600数量的处理器问题大小1，800 1，800（网格大小，这意味着矩阵大小约为3200万3200万）。 在弱缩放实验中，每个MPI处理器拥有一个300 300块。例如，一个有3，600个处理器的问题，其问题大小为90，000×3，600×30.3亿。我们首先考虑弱标度的情况。我们考虑两种情况：每个MPI等级有1个核心或4个核心。在后一种情况下，TensorFlow后端享有以下好处：10向前向后86420916 25 36 49 64 81 100 400 900数量的处理器互并行性，其中计算图中的独立运算符可以同时执行然而，4个核心并不能保证4倍的加速;性能取决于独立任务的可用性、调度冲突、资源争用等。图5显示了前向计算以及梯度反向传播的运行时间。有两个重要的观察结果：1. 通过在每个处理器上使用更多的内核，运行时间显著减少。例如，通过将每个处理器从1个核心切换到4个核心，向后的运行时间从30秒减少到大约10秒。2. 向后传递的运行时间通常小于向前计算的两倍。虽然反向传递需要求解两个线性系统（其中一个在正向计算中），但反向传播中的AMG（代数多重网格）线性求解器可能收敛得更快，因此成本可能低于正向传递。图5：前向计算和梯度反向传播的弱缩放运行时顶部：每个MPI处理器有1个内核;底部：每个MPI处理器有4个内核。我们看到在64核附近有一个跳跃，因为这是网络通信开始发生的地方（回想一下，我们的每个CPU都有32个核，每个节点有两个CPU）。 图显示了弱缩放的结果，每个MPI处理器解决了相同大小的局部问题。0.70.60.50.40.30.20.10.0149162536496481100 400 900 1600 2500 3600数量的处理器此外，我们在图6中显示了开销，它被定义为总运行时间和Hypre线性求解器时间之间的差异，用于向前和向后计算。我们看到，与总时间相比，开销相当小，特别是当问题大小很大时。这0。150。15· 10−20916 25 36 49 64 81 100 400 900数量的处理器表明ADCME MPI实现是非常有效的。图7、我们考虑强定标。 在这种情况下，我们固定了整个问题的大小，并将网格分割到不同的MPI处理器上。图7显示了forward的运行时图6：在Hypre中求解矩阵的ADCME开销顶部和底部的条形图对应于图1中的顶部和底部条形图。五、正演计算计算和梯度反向传播。利用100多个MPI，可以将代价昂贵的梯度反向传播算 法 的 运 行 时 间 缩 短 2 0 倍 以上101101处理器. 图8显示了加速和效率。我们可以看到，4个核心的运行时间比1个核心小二、10010−1100 101 102数量的处理器100100 101 102数量的处理器2在多核集群上为Hypre AMG求解器的混合编程模型（MPI和 OpenMP ） 找 到 一 个 扩 展 的 最 佳 点 是 具 有 挑 战 性 的（Baker，Schulz和Yang 2010）。TensorFlow后端的内部和内部并行性也增加了寻找扩展策略的难度。图7：前向计算和梯度反向传播的迭代。图显示了强缩放的结果，整个问题大小是固定的，我们增加了MPI处理器的数量（每个处理器有1个或4个核心）。正倒向正倒向正倒向梯度反向传播总时间（单核）总时间（四核）超线性求解器（1核）超线性求解器（4核）时间（秒）时间（秒）时间（秒）时间（秒）时间（秒）30201000 20 40 6080数量的处理器1000 20 40 60 80数量的处理器博图湖2010年。随机梯度下降的大规模机器学习。在COMPSTAT'2010的会议记录斯普林格。陈春C.的;杨角L.的;和Cheng，H.- Y. 2018.在多gpu平台上通过模型并行实现arXiv预印本arXiv：1809.02839。图8：泊松方程并行计算的加速比和效率结论介绍了ADCME MPI的功能。我们的基准测试结果表明，ADCME引入的分布式计算程序的开销是非常小的计算时间相比。 ADCME MPI分布式计算解决方案非常灵活，允许用户自行使用自定义并行算法或库。 随着能够收集大量数据的实验技术的出现，基于深度神经网络的数据驱动建模将成为科学发现的重要工具。 不断增长的数据集和问题规模增加了另一个层面的挑战。因此，在可预见的未来，正在进行的关于计算工程机器学习中的分布式计算的工作仍然是一个重要且有前途的方向。确认这项工作得到了能源部（DOE）高级科学计算研究办公室（ASCR）内的应用数学计划的支持，通过多尺度和多物理问题研究中心（DE-SC 0019453）的数学和物理信息学习机合作实验室。引用Abadi，M.;Barham，P.;陈杰;陈志;Davis，A.;迪恩J.道：Devin，M.; Ghemawat，S.; Irving，G.; Isard，M.;等，2016年。Tensorflow：一个大规模机器学习系统。第12届{USENIX}操作系统研讨会Design and Implementation（OSDI），265贝克，A. H.的; Schulz，M.; Yang，U. M. 2010. 代数多重网格求解器在多核集群上的性能。在计算科学的高性能计算国际会议上，102斯普林格。Baydin，A.G.地; 珀尔马特湾一、Radul，A.一、还有西斯金德J. M. 2017年。机器学习中的自动微分：综述。TheJournal of Machine Learning Research18（1）：5595-5637.Bekkerman，R.;Bilenko，M.;和Langford，J.2011年。扩展机器学习：并行和分布式方法。北京：清华大学出版社.Bezanson，J.; Edelman，A.; Karpinski，S.;和Shah，V.B.2017.朱莉娅：一种新的数值计算方法。SIAM review59（1）：65Bianco，M. 2014. 晕圈交换模式的界面。www. 白濑里。eu/IMG/pdf/wp86。PDF.程湾，澳-地N. 2006年。前向和伴随MPI通信例程之间的对偶性.道格拉斯角C.的; Haase，G.;和Langer，U.2003年。椭圆偏微分方程解算器及其并行化教程。暹罗。法尔古河D.的; Yang，U. M. 2002. hypre：一个高性能预处理器库。在国际计算科学会议上，632-641。斯普林格。Fan，T.; Xu，K.; Pathak，J.; 和Darve，E. 2020. 使用深度神经网络求解稳态Navier-Stokes方程的反问题arXiv预印本arXiv：2008.13074。Gabriel，E.;Fagg，G.E.的; Bosilca，G.;Angskun，T.;当加拉，J. J.; Squyres ， J. M.; Sahay ， V.; Kambadur ， P.;Barrett，B.; Lumsdaine，A.;等人2004年。开放MPI：下一代MPI实现的目标、概念在欧洲并行虚拟机/消息传递接口用户斯普林格。Griebel，M.;和Zumbusch，G. 1999.基于散列和空间填充曲线的自适应偏微分方程求解器中的并行多重网格并行计算25（7）：827Gropp，W.; Thakur，R.;和Lusk，E. 1999. 使用MPI-2：消息传递接口的高级功能。MIT Press.休伊特河J.道：和Grady II，T. J. 2020年。深度学习中的线性代数方法。arXiv预印本arXiv：2006.03108。Huang ， 黄 氏 拟 谷 盗 D.Z 的 ; Xu ， K.;Farhat ， C.; 和Darve，E.2020年。使用深度神经网络从间接观测中学习本构关系。计算物理学报109491.Jordan，M.一、和Mitchell，T. M. 2015.机器学习：趋势，观点和前景。 Science 349（6245）：255-260.李，M.; Andersen，D. G.地;帕克，J.W.; Smola，A. J.道：艾哈迈德，一、Josifovski，V.;朗，J.; Shekita，E. J.道：和Su，B.- Y.2014. 使用参数服务器扩展分布式机器学习。 在第11届{USENIX}操作研讨会上系统设计和实现（OSDI），583李，M.;周，L.;杨志; Li，A.; Xia，F.; Andersen，D. G.地;和Smola，A.2013年。分布式机器学习的参数服务器在大学习NIPS研讨会，第6卷，2。孟 ， X.; Li ， Z.; Zhang ， D.; Karniadakis ， G. E. 2020.PPINN：用于时间依赖PDE的Parareal physics-informedneural network 。应用力学与工程中的计算机方法370：113250。不，Y。1995. 一个高效的并行离散偏微分方程求解器。并行计算21（11）：1725正向（1个核心）向后（1个核心）正向（4核）向后（4核）加速比效率Pakravan，S.; Mistani，P. A.;阿拉贡-卡尔沃湾一、和Gi-bou，F.2020年。用物理感知神经网络解决逆偏微分方程问题arXiv预印本arXiv：2001.03608。Raissi ， M.; Perdikaris ， P.; Karniadakis ， G. E. 2019.Physics-informed neural networks：一个深度学习框架，用于解决涉及非线性偏微分方程的正问题和逆问题。计算物理学杂志378：686-707。Towara，M.; Schanen，M.;和Naumann，U. 2015. MPI-并行离散伴随OpenFOAM。InICCS，19Utke ， J.; Hascoet ， L.; Heimbach ， P.; Hill ， C.;Hovland，P.; 和Naumann ，U. 2009. 向可 邻接 MPI发展。2009年IEEE International Symposium on ParallelDistributed Processing，1-8。美国电气与电子工程师协会。王 ，M.; 和Poten ，A.2015 年 。用MPI 实 现高 阶自 动Dep3：v0.Xu，K.;和Darve，E.2019年。微分方程反问题的神经网络方法。arXiv预印本arXiv：1901.07758。Xu，K.;和Darve，E. 2020年a。ADCME：使用深度神经网络学习空间变化的物理场Xu，K.;和Darve，E.2020年b。从稀疏观测数据驱动逆建模的物理约束学习。arXiv预印本arXiv：2002.10521。Xu，K.; Huang，黄氏拟谷盗D. Z的;和Darve，E. 2020.使用对称正定神经网络学习构造关系。arXiv预印本arXiv：2004.00265。