美国能源部关注的核物理逆问题求解：SAGIPS实现效率 80 倍提升

将 ScienceAI 设为星标

第一时间掌握

新鲜的 AI for Science 资讯

破解生产研究中的难题通常需要一点侦探思维，也就是由果溯因。这样的思考方式被称为逆向思考，而在核物理中，通过分析粒子相互作用的后果了解原子结果的过程也就被称为逆过程。

核物理问题中，分析所用数据集大小几乎比单个 GPU 内存大 4~6 个数量级，传统直方图法更会弄丢部分关键信息，但这却是核物理中逆过程需要解决的困难。

美国杰斐逊国家加速器设施（Thomas Jefferson National Accelerator Facility）主导了一项研究，开发了一种人工智能 （AI） 技术，该技术可以在大规模超级计算机上可靠地解决这类难题。

该研究以「SAGIPS: a physics-inspired scalable asynchronous generative inverse-problem solver」为题，发布在《Machine Learning》。

论文链接：https://iopscience.iop.org/article/10.1088/2632-2153/adc8fb

SAGIPS 系统

该系统称为 SAGIPS（Scalable Asynchronous Generative Inverse Problem Solver）。它依赖于高性能计算和生成式 AI 模型，这些模型可以根据算法训练的数据生成新的文本、图像或视频。

此工作流的主要目的是解决逆问题。但是，它的设计使其也可以解决通用优化和控制问题。

图 1：SAGIPS 工作流程及其所有模块和依赖项的示意图。

SAGIPS 使用生成对抗网络（GAN），这是对抗式的神经网络，它们相互作用以产生有意义的数据。一个不断地试图欺骗另一个，另一个试图发现假货。

生成器与判别器的相互作用产生了绝佳的训练效果，这种对抗性过程促使两个网络进行改进，从而生成模拟真实数据分布的高质量合成数据。

而 SAGIPS 所使用的模块由 PyTorch 编写，其要求将所有张量显式加载到 GPU 或 CPU 内存中。故而，在 Polaris 超级计算机集群上，SAGIPS 使用了 400 个 GPU 处理相应问题。

「这项技术与可用的计算资源呈线性扩展，这意味着我们可以在更大的集群上处理更大的问题，」Jefferson Lab 的数据科学主管、该论文的合著者 Malachi Schram 说。「这就是它的核心。」

分布训练

在工作流程的早期，实现试图跨多个 GPU 训练生成器和判别器，但观察到的扩展行为并不乐观。于是团队将目光转移到了并行训练策略上。

在 HPC 系统上训练 GAN 工作流程通常有两种选择：集成分析与异步数据并行训练。前者不使用通信（即 GAN 在单个 GPU 上彼此独立地训练），而后者在工作流之间传输生成器梯度。

在分布式计算中，为并行运行程序而生成的单个进程通常称为 rank。SAGIPS 工作流在多个 GPU 上并行运行，下文中将交替使用术语 GPU 和 rank。

SAGIPS 采用环形拓扑，每个 GPU 仅与相邻两个节点通信。以 400GPU 为例，通信次数从 16 万次降至 800 次，理论通信复杂度从 O (N²) 优化至 O (N)，这是开销降低的核心机制。

图 2：12 个等级之间的 ring-all-reduce 通信的示意图。

SAGIPS 还具备远程内存访问（RMA），允许 GPU 直接访问彼此的数据并共享自己的数据。这种操作大大减少了通信瓶颈并加快了处理速度。

通过内部组（即将可用的 rank 分为若干组），每个内组使用自己的环全减少机制进行通信，该系统成功减少了通讯开销，又引入了外部组解决不同节点之间的传输问题。根据古斯塔夫森定律，随着计算机资源的可用增加，系统的运行速度将会更快。

实验验证

SAGIPS 得到了 SciDAC 通过 QuantOm 项目（量子色动力学核断层扫描）的支持，并在 Polaris 集群上得到了验证。

扩展实验基于一个循环闭合测试，类似于典型的核物理分析，目的是推断那些不能直接测量的量。

图 3：用于运行缩放实验的环闭合测试的示意图。

进行分布训练时，可以立即注意到，传统异步环形全减少（ARAR）的训练时间几乎呈线性增长，而 ARAR 和 RMA-ARAR 分析（两者都包含分组）则几乎与参与的计算节点数量无关，呈线性增长。

由图 4 可知，从 4 块 GPU 增加到 400 块 GPU 时，传统 ARAR 的分析速率增益约为 40。本次验证中使用的分组机制使这一增益翻倍。

图 4：总训练时间与用于在 Polaris 上训练分布式 GAN 的等级数的函数。

在团队展示的示例中，即便是一个简单的通道，团队也能节省将近 48 分钟时间。但他们强调，单个 GPU 集成分析也不是在真实物理数据上运行 GAN 工作流程的合适选项。

小结

论文合著者 Malachi Schram 表示：「将其中一些算法用于不同的项目非常容易，这最大限度地减少了重新设计和重建软件堆栈的间接成本。」

这套生成式逆向问题解决算法已经得到了美国能源部的关注，在未来的改进方向中，其将会尝试在更复杂和资源密集型的逆问题里发挥自己的潜能。