AI & ML
Ab Initio Accuracy Beyond the Cutoff: Distributed Message-Passing Equivariant Attention at Tens-of-Millions-of-Atom Scale
arXiv cs.DC0 views
arXiv:2611.06122v1 Announce Type: new 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{\AA}. 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.
Source:https://arxiv.org/abs/2611.06122PreviousvExpert: Virtualizing Expert Storage for Adaptive Load Balancing in Distributed MoE InferenceNextCARE-FL: Federated Fine-Tuning of Multimodal LLMs in Resource-Constrained Edge Networks
Related posts
Scalable Source Selection for Machine Learning Tasks
0arXiv cs.DB
AvalancheBench: Evaluating Enterprise Agent-Driven Semantic Query Processing Through Latent World Recovery
0arXiv cs.DB
ConStruM: A Structure-Guided LLM Framework for Context-Aware Schema Matching
0arXiv cs.DB
From Certified Pairs to a Certified Hub: Reasoner-Checked Crosswalks for the Information Exchange Standard
0arXiv cs.DB
A Comparative Study of Importance Attribution to Columns in Deep Learning over Relational Databases
0arXiv cs.DB
Community Search under Member-Specific Tolerances
0arXiv cs.DB