美国能源部关注的核物理逆问题求解:SAGIPS实现效率 80 倍提升
仅用于站内搜索,没有排版格式,具体信息请跳转上方微信公众号内链接
将ScienceAI设为星标
第一时间掌握
新鲜的AIforScience资讯
编辑丨%
破解生产研究中的难题通常需要一点侦探思维,也就是由果溯因。这样的思考方式被称为逆向思考,而在核物理中,通过分析粒子相互作用的后果了解原子结果的过程也就被称为逆过程。
核物理问题中,分析所用数据集大小几乎比单个GPU内存大4~6个数量级,传统直方图法更会弄丢部分关键信息,但这却是核物理中逆过程需要解决的困难。
美国杰斐逊国家加速器设施(ThomasJeffersonNationalAcceleratorFacility)主导了一项研究,开发了一种人工智能(AI)技术,该技术可以在大规模超级计算机上可靠地解决这类难题。
论文链接:https ://iopscience. iop.org/article/10. 1088/2632-2153/adc8fb
SAGIPS系统
此工作流的主要目的是解决逆问题。但是,它的设计使其也可以解决通用优化和控制问题。
图1:SAGIPS工作流程及其所有模块和依赖项的示意图。
SAGIPS使用生成对抗网络(GAN),这是对抗式的神经网络,它们相互作用以产生有意义的数据。一个不断地试图欺骗另一个,另一个试图发现假货。
生成器与判别器的相互作用产生了绝佳的训练效果,这种对抗性过程促使两个网络进行改进,从而生成模拟真实数据分布的高质量合成数据。
而SAGIPS所使用的模块由PyTorch编写,其要求将所有张量显式加载到GPU或CPU内存中。故而,在Polaris超级计算机集群上,SAGIPS使用了400个GPU处理相应问题。
「这项技术与可用的计算资源呈线性扩展,这意味着我们可以在更大的集群上处理更大的问题,」JeffersonLab的数据科学主管、该论文的合著者MalachiSchram说。「这就是它的核心。」
分布训练
在工作流程的早期,实现试图跨多个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工作流程的合适选项。
小结
论文合著者MalachiSchram表示:「将其中一些算法用于不同的项目非常容易,这最大限度地减少了重新设计和重建软件堆栈的间接成本。」
这套生成式逆向问题解决算法已经得到了美国能源部的关注,在未来的改进方向中,其将会尝试在更复杂和资源密集型的逆问题里发挥自己的潜能。
人工智能×[生物神经科学数学物理化学材料]
「ScienceAI」关注人工智能与其他前沿技术及基础科学的交叉研究与融合发展。
欢迎关注标星,并点击右下角点赞和在看。
点击阅读原文,加入专业从业者社区,以获得更多交流合作机会及服务。