通过优化数据读写,让分子动力学的图神经网络快10倍且省内存。
FlashSchNet: Fast and Accurate Coarse-Grained Neural Network Molecular Dynamics
- 用四类优化技术融合计算与存储,减少GPU内存访问次数。
- 单卡运行64个并行模拟,达1000纳秒/天,速度超基线6.5倍。
- 适合需要高精度且高效模拟蛋白质的科研人员使用。
图神经网络势函数(如SchNet)通过学习多体相互作用提升了分子动力学(MD)模拟的精度和可迁移性,但因内核碎片化和内存密集型流水线导致速度慢于经典力场。本文指出关键问题是缺乏对输入输出的感知设计,未充分考虑GPU高带宽内存(HBM)与片上SRAM间的数据读写。为此提出FlashSchNet,基于四项技术:(1)闪速径向基函数,将距离计算、高斯基展开与余弦包络融合为单次分块处理,每距离仅算一次并复用;(2)闪速消息传递,融合截断、邻居收集、滤波乘法与归约,避免在HBM中生成边张量;(3)闪速聚合,改用CSR段归约重构散加操作,将原子写入降低至特征维度倍数,实现前向与反向无竞争积累;(4)通道级16位量化,利用SchNet MLP权重的低通道动态范围,进一步提升吞吐而几乎不损失精度。在单张NVIDIA RTX PRO 6000上,FlashSchNet在包含269个粗粒度粒子的蛋白系统上,对64个并行副本实现1000纳秒/天的总模拟吞吐,较基线CGSchNet快6.5倍,峰值内存降低80%,超越经典力场(如MARTINI),同时保持SchNet级别的精度与可迁移性。
原文摘要 · Abstract (English)
Graph neural network (GNN) potentials such as SchNet improve the accuracy and transferability of molecular dynamics (MD) simulation by learning many-body interactions, but remain slower than classical force fields due to fragmented kernels and memory-bound pipelines that underutilize GPUs. We show that a missing principle is making GNN-MD IO-aware, carefully accounting for reads and writes between GPU high-bandwidth memory (HBM) and on-chip SRAM. We present FlashSchNet, an efficient and accurate IO-aware SchNet-style GNN-MD framework built on four techniques: (1) flash radial basis, which fuses pairwise distance computation, Gaussian basis expansion, and cosine envelope into a single tiled pass, computing each distance once and reusing it across all basis functions; (2) flash message passing, which fuses cutoff, neighbor gather, filter multiplication, and reduction to avoid materializing edge tensors in HBM; (3) flash aggregation, which reformulates scatter-add via CSR segment reduce, reducing atomic writes by a factor of feature dimension and enabling contention-free accumulation in both forward and backward passes; (4) channel-wise 16-bit quantization that exploits the low per-channel dynamic range in SchNet MLP weights to further improve throughput with negligible accuracy loss. On a single NVIDIA RTX PRO 6000, FlashSchNet achieves 1000 ns/day aggregate simulation throughput over 64 parallel replicas on coarse-grained (CG) protein containing 269 beads (6.5x faster than CGSchNet baseline with 80% reduction of peak memory), surpassing classical force fields (e.g. MARTINI) while retaining SchNet-level accuracy and transferability.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。