小内存占用机器学习原子间势解锁多组分块体材料分子动力学

摘要

块体材料需在大空间尺度（约10^9原子或更多）上进行分子动力学（MD）模拟。现有机器学习原子间势（MLIP）虽将MD扩展到该尺度，但即使单组分体系也需数万块GPU；多组分模拟几乎不可行，因为现有MLIP的HBM占用在单组分体系中已很大，在多组分场景下会爆炸式增长。本文提出一种HBM占用不到现有MLIP 3%的MLIP：引入物理和化学知识降低特征向量维度，并将所有内核融合为单个mega-kernel，从而消除中间张量。评估中，该MLIP用144块NVIDIA A100 GPU，对含1.14×10^9个原子的6组分块体体系完成MD模拟；此前该尺度仅限单组分体系，且通常需要配备数万块GPU的高端超算。