1. 项目概述:当分子动力学遇上非冯·诺依曼架构

在计算化学和材料科学的实验室里,我们这些做模拟的人,常年被一个“鱼与熊掌”的难题困扰着: 精度与效率,你选哪个? 如果你想获得接近量子力学第一性原理(DFT)的精度,那好,准备好海量的计算资源和漫长的等待时间,一个稍大点的体系跑上几周是家常便饭。如果你想快,那就得用经验力场(EFF),牺牲精度,模拟结果可能和物理现实相去甚远。这个困境,就像让你用算盘去解一个微分方程,要么慢得离谱,要么结果不准。

直到机器学习(ML)的介入,事情才有了转机。通过用DFT数据训练神经网络(NN),我们得到了一个“黑箱”模型,它既能逼近DFT的精度,又能像经验力场一样快速评估原子间作用力。这就是机器学习分子动力学(MLMD),它一度让我们看到了曙光。然而,当我们兴冲冲地把这些模型扔到传统的CPU或GPU(也就是冯·诺依曼架构的计算机)上运行时,一个新的瓶颈出现了—— “内存墙” 。超过90%的计算时间和功耗,竟然都花在了数据在计算单元和存储单元之间的来回搬运上,真正用于计算的能量和时间少得可怜。这就像你有一个顶级厨师(计算单元),但他的厨房(存储单元)在城市的另一头,他做一道菜的大部分时间都花在了取食材和送成品的路上。

于是,我们的思路转向了 非冯·诺依曼(NvN)架构 ,或者说“存算一体”。这个想法很直接:为什么不让厨师直接在仓库里做饭?把计算单元和存储单元紧密集成,甚至融合在一起,让数据原地被处理,彻底消灭数据搬运的开销。这听起来像是为MLMD这类计算密集、数据复用率高的任务量身定做的方案。

我们团队最近完成的一项工作,就是沿着这个思路的一次大胆实践。我们设计并实现了一个 异构并行非冯·诺依曼架构系统 ,专门用于高精度、高效率的MLMD计算。这个系统的核心思想是“分而治之,异构协同”:我们将MLMD计算流程中最耗时的神经网络力评估部分,用一个乘法-累加(MAC)操作被替换为移位-累加操作的、极度精简的神经网络模型来实现,并将其固化到一块专用集成电路(ASIC)芯片中;而特征提取和运动方程积分这两个模块,则部署在一块现场可编程门阵列(FPGA)上。两者通过并行协作,共同完成一次完整的MD步进。

最让人兴奋的结果是:即便我们这块ASIC芯片采用的是相对落后的180纳米工艺,时钟频率也只有25MHz,但整个系统在完成水分子动力学模拟时,其计算速度比基于12纳米先进工艺GPU的顶尖MLMD方案(DeePMD)还要快1.6倍,而能效更是高出 102到103倍 。这个数字背后,是非冯架构消除内存墙瓶颈带来的巨大红利。这篇文章,我就来拆解一下我们这个系统从设计思路、算法优化、芯片实现到最终测试的完整过程,希望能给从事科学计算加速、边缘AI或专用硬件设计的朋友们一些启发。

2. 核心思路拆解:为何是异构并行与非冯架构?

在深入硬件细节之前,我们必须先想清楚几个根本问题:为什么MLMD需要专用硬件?为什么选择非冯·诺依曼架构?又为什么采用FPGA+ASIC的异构方案?理解这些“为什么”,是看懂后续所有技术选择的基础。

2.1 机器学习分子动力学的计算特征与瓶颈

传统的分子动力学模拟,每一步都需要计算系统中所有原子受到的力,然后根据牛顿第二定律更新原子的位置和速度。在MLMD中,这个“力”是由一个神经网络模型给出的。这个计算流程可以分解为三个核心模块:

  1. 特征提取 :将原子的直角坐标(不具备平移、旋转、置换对称性)转换为一组对称性保持的特征描述符。这一步计算量中等,涉及距离计算、函数变换等。
  2. 神经网络力评估 :将上一步得到的特征向量输入神经网络,前向传播后输出原子所受的力(三个分量)。 这是整个流程中计算密度最高、最耗时的部分 ,尤其是当体系原子数多、神经网络规模大时,其计算量可能占到总计算量的90%以上。
  3. 运动方程积分 :根据计算出的力,用Verlet等积分算法更新原子的速度和位置。这一步计算量相对较小,主要是向量运算。

当我们在CPU或GPU(冯·诺依曼架构)上运行这个流程时,瓶颈立刻显现。神经网络,尤其是其核心的乘累加(MAC)操作,虽然高度并行,但需要频繁地从内存中读取权重参数和输入数据,计算完中间结果再写回。对于MLMD这种需要连续进行成千上万步模拟的任务,这种反复的数据搬运造成了巨大的时间和能量开销,即所谓的“内存墙”瓶颈。计算单元再强大,也被低速、高功耗的内存访问拖慢了。

2.2 非冯·诺依曼架构的破局之道

非冯架构的核心思想是 存算一体 。对于MLMD中的神经网络部分,这是一个绝佳的应用场景。因为在一个特定的材料模拟中,一旦神经网络训练完成,它的权重参数在长达数万甚至百万步的模拟过程中是 固定不变 的。这是一个非常关键的特性。

在非冯架构设计中,我们可以将这些权重参数直接存储在靠近或集成在计算单元(如乘法器/加法器)的本地存储器中。当前向传播计算某一层的输出时,输入数据流经计算单元,直接从旁边的存储器中读取权重进行运算,产生的中间结果直接送入下一层的计算单元, 完全避免了将权重或中间激活值写回片外主存 。这样,计算的时间几乎完全花在了实际的算术逻辑运算上,数据搬运的延迟和功耗被降至最低。

2.3 异构并行设计的权衡与协同

理解了非冯架构的优势,下一个问题就是:为什么不全用ASIC,或者全用FPGA?

  • ASIC(专用集成电路) :性能与能效的王者。一旦流片,它在特定任务上的速度、能效和面积利用率是FPGA无法比拟的。但它缺乏灵活性,设计周期长,成本高。它最适合实现那些计算模式固定、且是性能瓶颈的部分——对我们来说,就是神经网络力评估模块。
  • FPGA(现场可编程门阵列) :灵活性的代表。我们可以用硬件描述语言(HDL)快速实现和修改逻辑电路,非常适合算法仍在迭代、或计算模式相对复杂的部分。在MLMD中,特征提取(涉及条件判断、特殊函数计算)和积分模块(规则但需处理系统状态)用FPGA实现更为合适。

因此, 异构并行 就成了自然的选择:将计算最密集、模式最固定的神经网络部分用ASIC实现,追求极致的性能和能效;将相对灵活、算法可能调整的部分用FPGA实现,保持系统的整体可配置性和开发效率。FPGA和ASIC通过高速接口协同工作,FPGA负责准备数据(特征)和收尾工作(积分),ASIC则作为加速引擎,专心致志地以最高效率进行神经网络推理。这种架构在性能和灵活性之间取得了很好的平衡。

注意 :这种异构设计对系统级的数据流和通信提出了挑战。需要仔细设计FPGA与ASIC之间的数据格式、传输协议和同步机制,确保ASIC这个“加速引擎”不会因为等待数据而空闲,FPGA也不会因为等待结果而阻塞。在我们的实现中,我们采用了流水线化和双缓冲技术来掩盖通信延迟。

3. 算法层面的极致优化:为硬件而生的神经网络

有了异构并行的非冯架构蓝图,接下来就要设计一个能完美映射到这个硬件上的算法模型。直接在硬件上跑一个标准的、使用浮点数和复杂激活函数的深度神经网络是不现实的,那会带来巨大的面积和功耗开销。我们的目标是在保证力预测精度的前提下,设计一个 硬件友好的、资源节约型 神经网络。

3.1 从浮点到定点:量化神经网络

第一步是 量化 。在软件中训练时,我们通常使用32位单精度浮点数(FP32)。但在数字电路(尤其是ASIC)中,浮点运算单元非常昂贵。因此,我们将网络中的权重和激活值全部转换为 定点数 。我们选择了13位有符号定点数格式(1位符号位,2位整数位,10位小数位)。这样,所有的加法和乘法(后续会被替换)都可以用更高效的整数运算单元来实现,大幅减少了硬件资源消耗。

但简单的定点化会带来精度损失。为此,我们采用了 预训练策略 :先用FP32训练一个高精度的基线模型,然后以这个模型为起点,进行定点量化微调。这相当于给量化过程提供了一个良好的初始化点,显著缓解了精度下降问题。

3.2 设计轻量级激活函数

标准的神经网络在回归任务中常使用双曲正切函数(tanh)作为激活函数。然而,在硬件中实现tanh非常复杂,通常需要查找表(LUT)或坐标旋转数字计算机(CORDIC)迭代算法,消耗大量逻辑资源和时钟周期。

我们设计了一个分段线性近似函数来替代tanh: ϕ(x) = 1, if x >= 2; ϕ(x) = x/2 + 0.5, if -2 < x < 2; ϕ(x) = 0, if x <= -2 这个函数的妙处在于,它在关键区间[-2, 2]内是线性变换,除法 /2 可以用简单的右移一位操作来实现。整个函数只需要比较器、移位器和加法器就能实现,极其硬件友好。

我们对比了使用tanh和我们的ϕ(x)的神经网络在多个分子体系(水、乙醇、硅等)上的力预测精度。如表1所示,两者的均方根误差(RMSE)差异极小(小于0.51 meV/Å),但硬件开销天差地别。我们用逻辑综合工具评估,实现tanh需要约5万个晶体管,而实现ϕ(x)仅需约4千个晶体管, 硬件开销降低了92% 。这是一个用极小的精度代价换取巨大硬件收益的典型例子。

表1:不同激活函数下的力预测误差对比

体系 tanh(x) RMSE (meV/Å) ϕ(x) RMSE (meV/Å) 差异
25.04 24.83 0.21
乙醇 29.33 29.84 -0.51
甲苯 53.15 52.70 0.45

3.3 乘法-累加(MAC)的终结:移位-累加网络

神经网络的核心运算是MAC: output += weight * input 。乘法器在数字电路中是面积和功耗的大户。我们提出了一个更激进的想法: 能否用移位操作完全替代乘法?

答案是肯定的,但需要从训练阶段就介入。我们的方法是将权重值量化为若干个2的整数次幂之和。具体来说,对于一个浮点权重 w ,我们将其量化为: w_q = s(w) * [2^{n1} + 2^{n2} + ... + 2^{nK}] 其中, s(w) 是符号(+1或-1), n1, n2, ..., nK 是整数(可正可负), K 是项数。

例如,权重 0.375 (二进制 0.011 )可以表示为 2^{-2} + 2^{-3} ,即 0.25 + 0.125 。在推理时,权重与输入的乘法 w_q * x 就变成了: w_q * x = s(w) * [ (x << n1) + (x << n2) + ... + (x << nK) ] 这里 << 表示左移( n 为正)或右移( n 为负)。这样,昂贵的乘法器就被廉价的移位器和加法器取代了。

我们探索了不同的 K 值(1到5)对精度的影响。如图4所示,当 K=1 2 时,精度损失较大;当 K>=3 时,量化网络的精度与原始浮点网络的精度非常接近(误差在6.5%到12%之间)。而在硬件开销上,当 K=3 时,我们的移位累加网络(SQNN)相比使用16位定点乘法的基准网络(FQNN), 晶体管数量减少了50%到70% K=3 成为了精度与硬件开销之间的一个甜蜜点。

实操心得 :这种“移位替代乘法”的技术,其精髓在于将计算复杂度从训练阶段转移。训练算法需要学习如何用几个2的幂次项来逼近最优的连续权重。这增加了训练难度,但换来的是推理阶段极致的硬件效率。对于部署在专用芯片上的模型来说,这种 trade-off 是非常值得的。

4. 硬件架构设计与芯片实现

算法优化完成后,下一步就是将其映射到具体的硬件上,构建我们设想中的异构并行非冯系统。

4.1 非冯架构的具体实现:脉动式数据流

我们的非冯架构设计体现在神经网络芯片(ASIC)内部。如图6所示,我们为网络的每一层设计了一个“计算-存储”融合单元。

  1. 权重常驻 :每一层的权重参数(在这里是移位参数 s, n1, n2, n3 )被存储在紧邻处理单元的分布式寄存器或SRAM中。
  2. 数据流驱动 :输入向量从上一级流入。对于每个神经元,其输入向量与对应的权重参数进行本地化的移位-累加操作。这个操作完全在权重存储的位置附近完成,生成的部分和立即与偏置相加。
  3. 流水线化 :激活函数单元(AU)直接集成在累加器之后。一旦某个神经元的累加完成,立即进行激活计算,结果直接作为下一层对应神经元的输入,送入下一层的计算单元。数据像在流水线上一样流动,中间结果无需写回任何全局缓存。

这种设计彻底消除了在计算单元和全局内存之间搬运权重和中间激活值的需求,实现了真正的“存算一体”,最大化地利用了数据局部性,将功耗和延迟主要用于实际计算。

4.2 ASIC芯片设计:从算法到硅片

我们选择水分子这个经典体系作为第一个验证目标。为其训练的神经网络规模较小:输入层3个神经元(对应特征维度),2个隐藏层,每层3个神经元,输出层2个神经元(预测氢原子的两个力分量,氧原子的力由牛顿第三定律得出)。

芯片的核心是矩阵单元(MU)和激活单元(AU),如图7所示。

  • 矩阵单元(MU) :负责实现 输入向量 * 权重矩阵 。每个MU包含多个移位单元(SU)。每个SU对应一个 权重-输入 对,根据权重量化后的 n1, n2, n3 三个移位值,对输入进行三次移位操作,然后根据符号 s 决定加或减,最后将结果累加。
  • 激活单元(AU) :实现我们设计的轻量级激活函数 ϕ(x) 。仅需比较器、移位器、加法器和多路选择器即可构成,电路非常简单。

我们使用SilTerra 180纳米工艺完成了芯片的流片。最终芯片的裸片面积仅为1.73平方毫米,在25MHz时钟频率下,单芯片功耗仅为8.7毫瓦。图8(c)展示了芯片的显微照片,其紧凑的布局体现了高度定制化设计的优势。

4.3 异构系统集成:FPGA与ASIC的协同

完整的MLMD计算系统如图8(a)所示,它由一片Xilinx XC7Z100 FPGA和两片我们设计的MLP ASIC芯片构成。为什么是两片ASIC?因为一个水分子有两个氢原子,它们的力可以并行计算。

系统工作流程如下:

  1. FPGA(特征提取) :接收原子的位置坐标,计算两个氢原子各自的对称性特征描述符。
  2. ASIC(并行力评估) :两套特征向量分别送入两个MLP芯片。两个芯片并行工作,各自输出一个氢原子所受的力(Fx, Fy, Fz)。
  3. FPGA(集成与更新) :FPGA接收两个芯片的结果,根据牛顿第三定律计算氧原子所受的力。然后,利用这三个力,通过Velocity-Verlet积分算法更新所有原子的位置和速度。
  4. 循环 :将更新后的原子位置作为下一帧的输入,重复步骤1-3,即可得到完整的分子运动轨迹。

在这个流程中,FPGA扮演了“调度员”和“前后处理器”的角色,而ASIC则是专职的“力计算加速器”。两者通过并行流水线的方式紧密协作,FPGA在计算当前帧的特征时,ASIC可以同时计算上一帧的力,从而隐藏部分通信和计算延迟。

5. 系统测试与性能分析

芯片做出来了,系统搭好了,最关键的问题是:它到底准不准?快不快?省不省电?

5.1 精度验证:与量子力学基准对比

精度是MLMD的立身之本。我们首先测试了MLP芯片本身的力预测精度。在25MHz下,将芯片预测的力与DFT计算的标准结果进行对比,如图9所示,两者吻合得非常好,力的均方根误差仅为7.56 meV/Å。这个误差水平足以保证后续动力学模拟的可靠性。

接着,我们运行了完整的分子动力学模拟,并从中提取了水分子的结构性质(O-H键长、H-O-H键角)和动力学性质(三种振动模式的频率)。我们将四种方法的结果进行对比:

  1. DFT :量子力学计算,作为黄金标准。
  2. vN-MLMD :同样的MLMD算法,但在冯·诺依曼架构的CPU上运行。
  3. NvN-MLMD :我们提出的异构非冯系统。
  4. DeePMD :当前最先进的基于GPU的MLMD方法。

结果如表2所示。我们的NvN-MLMD系统计算出的所有性质,与DFT结果的相对误差均不超过1.06%。更重要的是,NvN-MLMD与vN-MLMD的结果几乎完全一致。这强有力地证明了: 我们的硬件优化和架构改变,没有牺牲算法的固有精度 。硬件只是忠实地、更高效地执行了算法。

表2:不同方法计算的水分子性质对比

方法 O-H键长 (Å) H-O-H键角 (°) 对称伸缩频率 (cm⁻¹) 反对称伸缩频率 (cm⁻¹) 弯曲频率 (cm⁻¹)
DFT (基准) 0.969 104.88 4007 4241 1603
vN-MLMD 0.968 104.90 4040 4291 1619
NvN-MLMD (本工作) 0.968 104.85 4040 4274 1586
DeePMD 0.970 104.82 4003 4234 1599

5.2 速度与能效:跨越工艺代沟的胜利

性能对比的结果令人振奋(表3)。我们关注两个核心指标:计算每个原子每步所需的时间(S),以及能耗(η = S × P,P为功耗)。

  • 速度 :我们的系统(180nm, 25MHz)计算一个水分子一步所需时间为1.6×10⁻⁶秒。这比DFT快了近6个数量级。与基于12nm工艺NVIDIA V100 GPU的DeePMD相比,我们的速度仍然快了 1.6倍 。请注意,GPU的时钟频率是我们的40倍以上(GHz vs MHz),工艺也先进得多。我们能在速度和工艺均处劣势的情况下实现反超,完全归功于非冯架构消除了内存墙瓶颈,将计算资源真正用于“计算”本身。
  • 能效 :这是非冯架构最大的优势所在。我们整个系统的功耗仅为1.9W,其中ASIC芯片的功耗更是低至8.7mW。最终计算出的能耗指标η,比GPU方案低了 102到103倍 。这意味着完成同样的模拟任务,我们的系统所需的能量仅为GPU方案的百分之一到千分之一。

表3:不同方法的计算时间与能耗对比

方法 硬件平台 时间/步/原子 (S) 功耗 (P) 能耗 (η = S×P)
DFT CPU 1.9 s 230 W 4.4×10² J
vN-MLMD CPU 5.1×10⁻⁴ s 45 W 2.3×10⁻² J
DeePMD CPU 8.6×10⁻⁵ s 152 W 1.3×10⁻² J
DeePMD CPU+GPU 2.6×10⁻⁶ s 250 W 6.5×10⁻⁴ J
NvN-MLMD (本工作) ASIC+FPGA 1.6×10⁻⁶ s 1.9 W 3.0×10⁻⁶ J

深度分析 :这个对比揭示了通用计算与专用计算在能效上的本质差异。GPU为了通用性,包含了大量用于图形渲染、高精度浮点运算的复杂逻辑和高速缓存层次,这些在运行MLMD时很多都成了“冗余功耗”。而我们的ASIC是“极简主义”的产物,每一颗晶体管都为目标计算服务,没有一丝浪费。这种“量身定做”带来的能效提升是数量级的。

6. 展望与挑战:从专用原型到通用平台

这项工作是一个成功的概念验证,但它仍然是一个 专用原型系统 。要走向实际应用,还面临几个关键挑战,这也是我们后续工作的方向。

6.1 工艺升级带来的性能飞跃

目前我们使用的是180nm的成熟工艺。如果采用更先进的节点(如14nm或7nm),性能提升将来自两个方面:

  1. 时钟频率提升 :从25MHz提升到GHz级别,直接带来近两个数量级的速度提升(A1 ≈ 10²)。
  2. 集成度提升 :先进工艺下单位面积晶体管数量大幅增加(约两个数量级,A2 ≈ 10²),我们可以在单芯片内集成更多的并行计算单元(例如,同时计算更多原子的力,或部署更大、更深的网络)。

两者结合,理论上可以将计算速度再提升约4个数量级(A1 × A2 ≈ 10⁴)。这意味着单原子每步的计算时间有望从10⁻⁶秒降低到10⁻¹⁰秒量级,这将使模拟更大体系、更长时标成为可能。

6.2 通用性与可扩展性

当前的ASIC是针对特定水分子模型定制的,网络结构(3-3-3-2)是固定的。而实际的科研需求千变万化,需要模拟不同的材料、不同的分子,网络结构也需要相应调整。因此,开发一个 通用的NvN-MLMD硬件平台 是下一步的核心目标。

这需要在两个层面努力:

  • 算法软件层面 :需要开发或适配更通用、可扩展的MLMD算法框架,使其能够描述广泛的原子体系,并能将训练好的模型高效编译到硬件上。
  • 硬件架构层面 :设计支持可变网络规模(层数、神经元数)的通用存算一体架构。这可能涉及可重构的互连网络、可配置的存储阵列以及灵活的指令集,使得同一块芯片能够通过配置来运行不同规模的神经网络模型。

6.3 系统级优化与生态构建

一个完整的、用户友好的MLMD加速解决方案,远不止一块芯片。它需要:

  • 完整的软件栈 :包括模型训练工具、模型编译与量化工具、硬件驱动以及任务调度与管理软件。
  • 高效的异构编程模型 :让计算化学家能够像写Python脚本一样方便地调用我们的硬件加速器,而不需要关心底层的FPGA编程或ASIC通信细节。
  • 多芯片扩展 :对于大规模体系,需要研究如何将多个ASIC芯片和FPGA通过高速互连(如硅光互联、先进封装)组织起来,实现可扩展的并行计算。

7. 总结与个人体会

回顾这个项目,从算法创新到芯片流片,再到系统集成与测试,是一次完整的“从软件到硬件”的垂直优化实践。它让我深刻体会到,对于机器学习科学计算这类特定领域问题, 软硬件协同设计 专用架构 的威力是巨大的。

非冯·诺依曼架构不是要取代CPU/GPU,而是在一个特定的计算范式(如MLMD中固定权重的神经网络推理)下,提供一种更极致的解决方案。它用面积和灵活性的代价,换取了能效和速度的巨幅提升。这对于那些计算需求固定、且对功耗和延迟极其敏感的场景(如边缘计算、嵌入式科学仪器)具有不可替代的价值。

踩过最大的“坑”其实是在算法与硬件的接口上。最初我们试图在硬件上直接实现一个标准的MLP,结果发现资源根本不够用。这才逼着我们回过头去,从最底层的算子(乘法、激活函数)开始重新思考,如何为硬件“定制”一个神经网络。这个过程是痛苦的,但也是最有收获的。它让我明白,真正的硬件加速,不是简单地把软件算法“移植”过去,而是要根据硬件的特点,从算法层面进行重构和妥协。

最后,给也想尝试类似路径的朋友一点建议: 尽早开始软硬件联调 。不要等算法完全定型、硬件设计完毕才开始集成。用一个最小可运行的原型(比如用FPGA模拟ASIC的行为),尽早开始端到端的验证。你会提前发现很多在单独仿真中看不到的问题,比如数据精度溢出、时序不同步、通信协议错误等。这些问题越早暴露,解决的成本就越低。

这个异构并行非冯系统只是一个起点。它证明了这条技术路线的可行性,也让我们看到了专用计算在突破科学计算瓶颈方面的巨大潜力。随着算法和硬件的共同演进,我相信未来我们能够在手机大小的设备上,完成今天需要超算才能胜任的分子模拟任务,那将会彻底改变材料发现和药物设计的游戏规则。

Logo

码道开发者社区,聚焦华为云码道 CodeArts 代码智能体,沉淀 Agent、Skill、鸿蒙开发实战内容,供开发者查阅资料、交流技术、分享工程实践

更多推荐