有限元求解器搬进可微图程序:同一离散下,单块H200每步比32核CPU快181倍
核心概要
共轭传热与浮力流的有限元求解器被改写为自动微分框架中的固定权重图程序,装配算子与代数多重格点都按稀疏图执行,不做任何训练。在离散、网格、时间步和容差完全相同的条件下,同一问题在单块H200上每步最多比32核CPU快181倍,压力求解单独快18.7倍。反向模式梯度与中心差分一致,隐式微分让一次梯度的代价稳定在约4.9次前向求解。
深度剖析
共轭传热这类把流体浮力流动与相邻固体导热耦合起来的计算,现在可以在GPU上以端到端可微的方式跑完整个时间步,而不必另写伴随程序。 此前的做法要么停留在CPU上的稀疏线性求解,要么把算子写成笛卡尔网格上的卷积,难以处理贴体网格和弯曲的固液界面。 在离散、网格、时间步、步数和压力容差都固定、只更换线性代数后端与执行设备的对照中,二维共轭腔的每步CPU/GPU时间比从1.25升到23.9,三维共轭热沉在1.2×107耦合自由度上达到181倍。
求解器把装配好的有限元算子和多重格点循环表示为固定权重的图运算,节点是自由度,边是非零矩阵耦合,权重由离散格式给定。 这些权重不是训练出来的,图神经网络在这里只描述执行结构;压力求解用代数多重格点构造的图U-Net作为预条件,外层仍是共轭梯度。 制造解在三维中恢复出理论收敛阶,速度与温度约2.95、压力2.03,与理论{3,3,2}相差不超过0.06;图V循环与参考实现吻合到fp64舍入误差。
整条计算链可反向求导,热壁努塞尔数对固液导热比的敏感度直接由反向模式给出,无需手工推导伴随。 对截断迭代直接展开求导只能缓慢逼近收敛解的敏感度,而且内存随迭代次数增长。 在规定的每步40次前向迭代下,隐式微分在导热比100处的相对误差约5×10^-4,展开求导为25%;隐式路径峰值内存恒为2.7 GiB,展开路径在320次迭代时达到41.6 GiB。
启示与展望
演示范围是二维与三维的Boussinesq自然对流、共轭腔、湍流瑞利–贝努瓦胞、分支热沉和单缸发动机几何,全部在双精度下运行,分布式结果限于单节点八块H200。作者把同离散对照(Tier B)与完整求解器栈对照(Tier C)分开标注,只有前者支持关于执行设备的因果陈述。
作者说明两级粗修正在每个进程上复制,单节点可行,多节点规模会重新引入去膨胀所避免的内存开销;展开求导的内存约每翻倍网格增长2.6倍,因此共轭热研究停在9.4×104自由度。分布式演示因原地halo交换打断计算图而限于单GPU可微,强扩展在八GPU上仅约15%并行效率。另有一个基准单元中CPU有限元参考比其余三者低0.35%,作者将其作为未解决的离群点报告。
