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·· 3 小时前AI 评分48
Tessera：突破截断半径限制的分布式消息传递等变注意力，支持数千万原子尺度从头算精度
Ab Initio Accuracy Beyond the Cutoff: Distributed Message-Passing Equivariant Attention at Tens-of-Millions-of-Atom Scale
AI 导读
分布式运行时 Tessera 在等变消息传递 Transformer 上实现数千万原子的从头算分子动力学模拟：512 GPU 上维持 30.7M 原子、弱扩展效率 89.7%，图复用使整步时间降低 1.87 倍。
正文 · 原文
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