AI 资讯

让等变Transformer走向三千万原子规模，至知发布分布式计算系统Tessera

2026-11-06 14:23 更新于 2026-11-06 14:49 阅读约 8 分钟

一颗诺如病毒的蛋白质衣壳，连同周围的水和离子，构成了一个包含约 1006 万个原子的复杂体系。要在计算机中模拟它的动态变化，需要在每一个时间步计算原子受力，并不断更新所有原子的位置。

这一模拟背后，是大规模分子计算对模型表达能力与计算效率的双重要求。随着机器学习力场向复杂分子、生物大分子及材料体系拓展，如何在保持模型表达能力的同时扩大计算规模，成为分子动力学研究的重要课题。

近日，至知创新研究院（IQuest Research）在论文《Ab Initio Accuracy Beyond the Cutoff: Distributed Message-Passing Equivariant Attention at Tens-of-Millions-of-Atom Scale》中，提出分布式计算系统 Tessera ，在 512 张高性能 GPU 上完成了约 3070 万原子的完整能量与受力计算，单步耗时 1.684 秒，扩展效率达到 89.7 % ，并进一步在千万原子级体系中完成连续分子动力学轨迹验证。

这项研究基于 UBio-MolFM 模型体系，采用 E2Former V2 架构，围绕原子分区、动态邻居搜索、跨设备通信与 GPU 计算执行展开协同优化，使具有多层消息传递能力的机器学习力场能够扩展到数千万原子的计算规模。

图：Tessera 整体技术架构

多层消息传递，为什么难以扩展到数百张 GPU？

分子动力学（Molecular Dynamics，MD）通过反复计算原子受力、更新原子位置，模拟物质体系的动态演化。机器学习力场（Machine Learning Force Fields，MLFFs）利用第一性原理计算生成的数据训练模型，以较低的计算成本预测原子能量与受力。随着等变神经网络与消息传递机制的发展，模型能够通过多层特征传播，获取更丰富的原子环境信息。

此次研究使用的 E2Former V2 包含约 2350 万参数，由三个截断半径为 5 Å 的短程等变层和一个截断半径为 8 Å 的长程等变层组成。经过四层消息传递，模型的最大理论感受野达到 23 Å ，使单个原子能够通过邻居之间的信息传播，利用更大范围内的结构信息。

这一能力也增加了分布式计算的难度。当相邻原子被分配到不同 GPU 上时，模型每传播一层，都需要交换相应的原子特征。同时，分子动力学所需的受力由势能对原子坐标的梯度计算得到，跨设备通信还需要参与反向求导过程。原子持续运动又会不断改变邻居关系，使 GPU 计算图难以长期复用。

因此，Tessera 需要同时解决三个问题：让计算任务合理分布、减少动态邻居图的重复构建，以及降低多层特征交换带来的通信开销。

怎么切、怎么复用、怎么隐藏通信？

Tessera 的核心设计，是将原子运动过程中变化频率不同的计算环节分开处理，在持续更新原子坐标的同时，尽可能复用相对稳定的分区结构、邻居拓扑和通信计划。

首先，通过空间填充曲线组织原子分区。大规模分布式模拟需要将原子分配到不同 GPU。按空间等体积划分容易受到原子密度差异影响，单纯平均分配原子又可能增加跨设备通信。Tessera 采用 Hilbert 空间填充曲线，将三维原子坐标映射为一维序列，再切分为原子数量相等的连续区间，使各 GPU 保持固定的原子数量，同时尽可能保留空间邻近关系。

这种设计让节点张量能够维持稳定形状，为 CUDA Graph 的重复执行提供基础。在每 GPU 2.5 万原子的实验中，当设备数量由 64 张增加至 512 张时，需要访问候选数据来源设备最多的 GPU，其来源设备数由 6 增至 19（含自身），呈现次线性增长趋势，体现了空间局部性对通信范围的约束。

图：Hilbert 空间分区示意

其次，将邻居拓扑与实时几何信息分离。分子动力学中的原子每一步都在移动，但较短时间内的邻居关系往往相对稳定。Tessera 据此复用已有的邻居连接关系与通信计划，同时根据最新坐标更新原子间距离。在论文评估配置中，系统每 10 个时间步重建一次邻居图；当原子位移超过 0.25 Å 时，则通过全局同步提前触发重建，原子分区每 100 步重新调整一次。

为了进一步降低邻居搜索成本，系统采用三级搜索机制：先利用设备级空间包围盒筛选可能存在相互作用的区域，再通过局部查询分块缩小候选范围，最后使用 FAISS 完成距离筛选。结合邻居缓存与 CUDA Graph 复用，Tessera 减少了反复搜索、重新构图和计算调度的开销。在 512 GPU 的完整计算流程测试中，相较每步重新构建邻居图，每 10 步重建的配置获得了约 1.87 倍加速。

图：三级搜索机制

最后，让跨设备通信与本地计算同步进行。多层等变 Transformer 在前向传播和反向求导中都需要交换远程原子特征。Tessera 通过复用通信计划、提前投影部分注意力特征及预注册通信缓冲区，减少重复传输开销。其中，发送端特征投影使单原子的通信载荷从 4864 个标量减少至 3584 个，降低约 26.3%。

在此基础上，系统使用独立 CUDA Stream，使 GPU 在等待远程数据时仍可执行本地计算，并通过配对的自动微分通信算子，让这一机制覆盖前向和反向两个计算过程。在仅切换通信重叠机制的 A/B 测试中，512 GPU、每 GPU 2.5 万原子的扩展效率由 72.9% 提升至 84.7% 。

研究还利用 Triton 定制算子、torch.compile 和双向稀疏索引优化等变注意力与梯度计算。在单 GPU 的 1 万原子测试中，相关优化使计算耗时从 363.4 毫秒降至 141.1 毫秒，峰值显存由约 66.3 GiB 降至约 15.0 GiB。

约 3070 万原子，验证大规模计算效率

在整体性能实验中，研究通过固定每张 GPU 负责的原子数量，同时增加设备数量与体系规模，评估 Tessera 的弱扩展能力。当 GPU 数量从 32 张扩展至 512 张时，每 GPU 1 万、2.5 万和 6 万原子的三组配置，分别实现了 77.3%、84.7% 和 89.7% 的扩展效率。

其中，每 GPU 6 万原子的配置达到约 3070 万原子的最大规模，完整能量与受力计算单步耗时 1.684 秒，对应每秒约 1820 万原子的计算吞吐。随着 GPU 数量与体系规模同步扩大 16 倍，该配置的单步耗时仅增至原来的 1.11 倍；三档配置对应的耗时增长范围为 1.11—1.29 倍。

需要注意，6 万原子/GPU 的配置额外启用了激活重计算，以适应显存容量，因此不同负载之间的效率差异还包含计算配置的影响。

图：缓存复用、通信重叠与弱扩展性能

这些实验验证了 Tessera 在数千万原子规模上的计算扩展能力。与此同时，大规模分子动力学还需要让原子持续运动，并检验长期执行过程中能量与结构统计量的表现。

千万原子轨迹：从水溶液到诺如病毒衣壳

为了验证持续模拟能力，研究进一步构建去中心化 GPU 分子动力学引擎，使各 GPU 分别维护本地原子状态并执行运动积分，每步仅同步固定大小的全局统计信息，每 100 步进行一次全局原子收集和重新分区。

研究团队首先在一套 256 张 GPU 的集群上，完成了约 1024 万原子 KCl 水溶液的连续模拟。实验包括 10 ps 恒温平衡和 10 ps 微正则系综模拟。在关闭 TF32 的条件下，后者测得的能量漂移约为 0.42 μeV/atom/ps 。同时，分布式模拟的 K–O、Cl–O 和 O–O 径向分布函数，与相同温度和浓度下约 1.4 万原子的单 GPU 串行参考结果基本吻合，主要峰位差异处于 0.08 Å 的统计分箱宽度内，峰高相差约 2%。

图：千万原子分布式轨迹的结构统计验证

随后，研究将 Tessera 应用于开篇展示的诺如病毒衣壳及溶剂体系，完成了约 1006 万原子、100 ps 的连续分子动力学模拟。这一体系同时包含密集的蛋白质结构及大量溶剂原子，对计算分区和动态邻居搜索提出了更加复杂的要求。

这一实验还揭示了等原子数分区的实际边界：各 GPU 分配的原子数量相近，并不意味着计算工作量完全相同。蛋白质密集区域的原子往往拥有更多邻居，需要处理更多相互作用。与原子规模相近的 KCl 水溶液相比，衣壳体系的单步计算耗时增加约 13%—33% ，为后续优化复杂生物分子体系中的计算负载均衡提供了依据。

目前，千万原子级模拟的实际推进速度仍处于每天数皮秒的水平。进一步提高长时间尺度模拟能力，将是推动这套系统应用于复杂科学问题的重要方向。

因此，接下来的问题不是能不能，而是能多快。

面向更大规模的科学计算

Tessera 将空间分区、动态图缓存、GPU 算子优化、通信计算重叠与去中心化分子动力学引擎整合到统一的分布式计算体系中，实现了多层等变 Transformer 在约三千万原子规模上的完整能量与受力计算，并通过千万原子级连续轨迹检验了系统的运行稳定性。

未来，随着复杂体系负载均衡、强扩展性能及长时间尺度模拟能力的持续优化，这套算法与系统协同设计方案有望进一步拓展机器学习力场的计算规模，为复杂分子体系研究提供计算支撑。

相关阅读