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·· 12 小时前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 倍。
正文 · 原文
Authors: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
Abstract:Molecular dynamics predicts molecular motion from atomic forces and trajectory integration. Density functional theory (DFT) provides first-principles forces, but its $O(N^3)$ cost limits simulations to thousands of atoms; empirical force fields scale to millions but miss bond breaking and charge transfer. Equivariant message-passing Transformers learn the ab initio potential at linear cost. Stacked layers propagate information beyond a single cutoff, giving four layers a receptive field up to 23Å. At this frontier, scale is limited by communication: layers exchange boundary features across devices, force backpropagation repeats the exchange, and moving atoms change neighbor graphs and tensor shapes. We present Tessera, a distributed runtime that fixes workload shapes between rebuilds. It sorts atoms along a space-filling curve into equal-count shards, ensuring identical atom counts and tensor shapes per rank. Peer sets scale as $K^{0.5}$, not $K$. A conservative neighbor skin preserves topology across updates, enabling each rebuild interval to replay as one CUDA Graph. Fused kernels, three-level neighbor search, and overlapped All-to-All complete it. On 512 GPUs, Tessera sustains 30.7M atoms at 89.7% weak-scaling efficiency (77--90% across three per-GPU loads); per-step time rises 1.11--1.29x as machine and problem grow 16x. Graph reuse cuts full-step time $1.87\times$. Across an $8\times$ device increase, neighbor-search cost grows $1.7\times$, versus $3.0\times$ for Morton order. In a live $10.24$M-atom salt solution, NVE energy drift is $0.42\,\mu$eV/atom/ps over $10$ ps, and structure matches a serial reference. A 10.06M-atom Norovirus capsid sustains 100 ps of NVT at $\approx 8$ ps/day on 256 A100 GPUs. Over 200,000 steps, it has no broken bonds or anomalously close pairs and uses three compiled variants, with no graph breaks or eager fallbacks.
| Subjects: | Distributed, Parallel, and Cluster Computing (cs.DC); Artificial Intelligence (cs.AI); Biological Physics (physics.bio-ph) |
| Cite as: | arXiv:2611.06122 [cs.DC] |
| (or arXiv:2611.06122v1 [cs.DC] for this version) | |
| https://doi.org/10.48550/arXiv.2611.06122 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Jia Zhang [view email]
[v1]
Wed, 4 Nov 2026 08:07:22 UTC (102 KB)
来源：arXiv cs.AI · arxiv.org