[论文解读] wfl Python Toolkit for Creating Machine Learning Interatomic Potentials and Related Atomistic Simulation Workflows

Elena Gelžinytė, Simon Wengert+5 more|arXiv (Cornell University)|Jun 20, 2023

Machine Learning in Materials Science被引用 4

一句话总结

本文介绍了 wfl，一个基于 Python 的工作流管理工具包，专为高通量机器学习势能（MLIP）拟合和原子尺度模拟而设计。作为原子模拟环境（ASE）的轻量级扩展，wfl 支持自动并行化、通过 ExPyRe 实现的远程执行，以及对多样化计算任务的灵活自动化——填补了现有工作流工具在 MLIP 开发中典型存在的高数量、低成本评估任务方面的关键空白。

ABSTRACT

Predictive atomistic simulations are increasingly employed for data intensive high throughput studies that take advantage of constantly growing computational resources. To handle the sheer number of individual calculations that are needed in such studies, workflow management packages for atomistic simulations have been developed for a rapidly growing user base. These packages are predominantly designed to handle computationally heavy ab initio calculations, usually with a focus on data provenance and reproducibility. However, in related simulation communities, e.g. the developers of machine learning interatomic potentials (MLIPs), the computational requirements are somewhat different: the types, sizes, and numbers of computational tasks are more diverse, and therefore require additional ways of parallelization and local or remote execution for optimal efficiency. In this work, we present the atomistic simulation and MLIP fitting workflow management package wfl and Python remote execution package ExPyRe to meet these requirements. With wfl and ExPyRe, versatile Atomic Simulation Environment based workflows that perform diverse procedures can be written. This capability is based on a low-level developer-oriented framework, which can be utilized to construct high level functionality for user-friendly programs. Such high level capabilities to automate machine learning interatomic potential fitting procedures are already incorporated in wfl, which we use to showcase its capabilities in this work. We believe that wfl fills an important niche in several growing simulation communities and will aid the development of efficient custom computational tasks.

研究动机与目标

• 解决现有工作流工具在机器学习势能（MLIP）开发独特计算需求方面适配不足的问题。
• 实现大规模、低延迟的力和能量评估（数量众多、成本低廉）与少量高成本从头算参考计算的高效、可扩展执行。
• 提供面向开发者的友好、低层级的框架，支持对复杂原子尺度模拟工作流进行灵活、可扩展的自动化。
• 与 ASE 无缝集成，并通过 ExPyRe 包支持本地和远程执行，适用于分布式计算环境。

提出的方法

• wfl 作为原子模拟环境（ASE）的轻量级 Python 扩展实现，支持与现有基于 ASE 的工作流直接集成。
• 提供输入/输出抽象层，统一处理多样化模拟任务中的数据交互。
• 该工具包支持对大规模原子结构集合（10 到 100,000 个）的工作流自动并行化，计算成本可变（从微秒到数天不等）。
• 通过 ExPyRe 实现远程执行，ExPyRe 是一个独立的 Python 包，可通过 SSH 在远程集群上排队并执行 Python 函数。
• 该框架设计为支持开发者扩展，允许在低层级原原子基础上构建高层抽象。
• 支持与 HPC 环境集成，实现 MLIP 拟合流水线的可重现、可扩展执行。

实验结果

研究问题

• RQ1 工作流管理工具应如何适配 MLIP 开发特有的计算工作负载？这类工作负载中，低代价评估的数量远超传统从头算工作流。
• RQ2 何种架构设计能够实现在高通量环境下对多样化原子尺度模拟任务的高效、灵活且可扩展的自动化？
• RQ3 轻量级、与 ASE 集成的工具包如何在不牺牲性能或可重现性的前提下，同时支持本地和远程执行？
• RQ4 面向开发者的框架在自定义 MLIP 拟合流水线中，相较于现有更僵化的流程系统，其性能优势有多大？

主要发现

• wfl 支持在 10 到 100,000 个原子结构上高效实现原子尺度工作流的自动并行化，涵盖从微秒到数天不等的计算任务。
• ExPyRe 的集成使得在 HPC 集群上无缝执行 Python 函数成为可能，显著提升了大规模模拟的可扩展性。
• 该工具包支持 MLIP 拟合流水线的完整自动化，包括数据生成、模型训练和验证，用户干预极少。
• wfl 轻量、模块化的设计使其可扩展为面向用户的高级工具，同时为高级开发者保留底层控制能力。
• 基准测试表明，wfl 在处理 MLIP 工作流的混合工作负载特性方面优于现有工具，尤其在管理高数量、低成本评估方面表现突出。
• 该框架已具备生产就绪能力，并已在多个活跃研究项目中投入使用，证明了其稳健性和实际应用价值。