arXiv cs.AI· Allex Kan, Ji Zhang, Hao Yin, Hao Li, Owen Fu, Xunxing Chen, Lin Huang, Zion Wang, Lucas Li, Jiajun Cheng, Arthur Jiang, Junping Zhao, Jia Zhang·· 11 小时前AI 评分42
Tessera 分布式运行时实现千万级原子尺度下的从头算精度分子动力学模拟
Ab Initio Accuracy Beyond the Cutoff: Distributed Message-Passing Equivariant Attention at Tens-of-Millions-of-Atom Scale
AI 导读
Tessera 分布式运行时在 512 块 GPU 上以 89.7% 的弱扩展效率维持 3070 万原子的等变消息传递 Transformer 分子动力学模拟，单步耗时仅随机器与问题规模扩大 16 倍而增长 1.11–1.29 倍。
正文 · AI 翻译
作者：Allex Kan、Ji Zhang、Hao Yin、Hao Li、Owen Fu、Xunxing Chen、Lin Huang、Zion Wang、Lucas Li、Jiajun Cheng、Arthur Jiang、Junping Zhao、Jia Zhang
摘要：分子动力学通过原子受力和轨迹积分来预测分子运动。密度泛函理论（DFT）提供第一性原理力，但其 $O(N^3)$ 的开销将模拟规模限制在数千个原子；经验力场可扩展到数百万原子，但无法处理键断裂和电荷转移。等变消息传递 Transformer 以线性开销学习从头算势能。堆叠层可将信息传播到单个截断半径之外，四层即可获得高达 23Å 的感受野。在这一前沿，规模受限于通信：各层跨设备交换边界特征，力的反向传播会重复这一交换，而原子移动会改变邻居图和张量形状。我们提出 Tessera，一种在重建之间固定工作负载形状的分布式运行时。它沿空间填充曲线将原子排序为等数量分片，确保每个 rank 的原子数和张量形状完全一致。对等集合的规模按 $K^{0.5}$ 而非 $K$ 扩展。保守的邻居皮肤在更新之间保持拓扑不变，使每个重建区间可作为单个 CUDA Graph 重放。融合内核、三级邻居搜索和重叠的 All-to-All 共同构成完整方案。在 512 块 GPU 上，Tessera 以 89.7% 的弱扩展效率维持 3070 万原子（三种每 GPU 负载下为 77--90%）；当机器和问题规模扩大 16 倍时，每步时间仅增加 1.11--1.29 倍。图复用将完整步时间缩短 $1.87\times$。在设备数量增加 $8\times$ 的情况下，邻居搜索开销增长 $1.7\times$，而 Morton 序为 $3.0\times$。在一个实时 $10.24$M 原子的盐溶液中，NVE 能量漂移在 $10$ ps 内为 $0.42\,\mu$eV/原子/ps，结构也与串行参考一致。一个 1006 万原子的诺如病毒衣壳在 256 块 A100 GPU 上以约 $8$ ps/天维持 100 ps 的 NVT 模拟。在超过 200,000 步中，没有出现断键或异常近距原子对，且仅使用三个编译变体，没有图中断或即时回退。
| 学科： | 分布式、并行与集群计算（cs.DC）；人工智能（cs.AI）；生物物理（physics.bio-ph） |
| 引用为： | arXiv:2611.06122 [cs.DC] |
| （或本版本的 arXiv:2611.06122v1 [cs.DC]） | |
| https://doi.org/10.48550/arXiv.2611.06122 由 DataCite 颁发的 arXiv DOI（待注册） |
提交历史
来自：Jia Zhang [查看邮箱]
[v1]
2026年11月4日星期三 08:07:22 UTC（102 KB）
来源：arXiv cs.AI · arxiv.org