Allex Kan, Ji Zhang, Hao Yin, Hao Li, Owen Fu, Xunxing Chen, Lin Huang, Zion Wang 等 13 位作者
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\,μ$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.
Dense retrieval models are typically obtained by fine-tuning pretrained language models on supervised retrieval data, where high-quality annotations or interaction logs are often limited and reused across many epochs. This data-constrained regime differs from the abundant-data setting assumed by standard scaling laws, making it unclear how the best attainable retrieval performance changes with model size and unique fine-tuning data. We study this question through the minimum achievable evaluation loss (MAEL), defined as the lowest evaluation loss reached during training under fixed model-size and data-budget constraints. We first show that learning-rate choices strongly affect optimization trajectories but have little effect on the final attainable MAEL: after rescaling the optimization-step axis, learning curves collapse across learning rates, and CKA analysis shows that checkpoints reaching MAEL learn highly similar representations. We then fine-tune 35 BERT-family dense retrievers on MS MARCO subsets and find that MAEL is accurately described by a simple multiplicative power law over non-embedding parameters and unique samples, outperforming commonly used additive scaling formulations. We further derive a scaling law for the training step at which MAEL is reached and show that optimal weight decay follows an approximately invariant EMA timescale, yielding lower MAEL and stronger retrieval performance on both in-domain and out-of-domain benchmarks. Finally, combining loss and step scaling suggests that, in our observed regime, increasing pretrained model size is approximately 1.46 times as fine-tuning-compute-efficient as increasing unique in-domain data for reducing MAEL.