Page 484 - 《软件学报》2025年第8期

王昊天 等: MTTorch: 面向 MT-3000 芯片和 Transformer 模型的 PyTorch 算子库实现与优化 3907

 
 K ∑  K ∑
 ( )
  log


dx i = y i × dy i − dy j ×y j  ,dx = dy i −exp(y i )× dy j (5)

 i
 
j=1 j=1
计算得到. 对于前向传播需要进行规约操作求出数组的最大值和总和, 再按照公式进行向量化计算. 对于反向
传播则只需要规约求和直接计算即可, 因为 y i 为 softmax 之后的计算结果, 取值范围为 [0–1], 不会造成溢出.

3.6.2.2 Layer_norm
PyTorch 中的 Layer_Normalization (层归一化) 是一种用于深度神经网络的归一化技术, 其在每个样本的特征
维度上进行归一化. Layer_norm 适用于处理变长序列数据, 如文本数据, 可以提高模型的稳定性和收敛速度, 是大
型语言模型中必不可少的一环. 计算过程涉及将输入张量沿指定维度计算均值和标准差, 并使用它们进行归一化.
其计算公式为:

x− E (x)
y = √ ×γ +β (6)
Var(x)−ϵ
其中, E(x) 代表数据的均值, Var(x) 代表数据的标准差, ϵ 是一个用于保持数值稳定的常数, γ 和 β 是该层的可学习
参数.
在传统的计算均值和标准差的方法中, 通常需要将所有数据相加, 然后除以数据个数来得到均值, 再遍历每一
个样本数据求得标准差. 然而, 当数据量非常大时, 将所有数据相加可能导致精度溢出的问题, 同时进行两次遍历
也会造成较大的时间开销. 为了解决这个问题, 并且实现并行化计算, 本文在 DSP 上实现了适用于多核并行的
Welford 迭代计算算法. 程序只需要对数据进行一次遍历, 并利用迭代公式:

 x k − Mean(x k−1 )
 Mean(x k ) = Mean(x k−1 )+



 k




 S (x k ) = S (x k−1 )+(x k − Mean(x k−1 ))×(x k − Mean(x k )) (7)





 S (x k )

 Var(x k ) =


k
即可计算 k 个数据的均值和标准差, 避免了精度溢出的问题. 其中, S(x) 是迭代计算标准差的中间结果. 此外,
Welford 算法的合并公式:


 delta = Mean(x 2 )− Mean(x 1 )





 n 2
 2

Var(x) = Var(x 1 )+Var(x 2 )+delta ×n 1 × (8)

 n 1 +n 2



 n 2

 Mean(x) = Mean(x 1 )+delta×

n 1 +n 2
也使得可以将数据进行分段处理, 不同的数据段间互不干涉, 从而可以通过一次规约来求得结果, 其中 n 1 和
n 2 代表要合并的两段子数据的长度.
具体来说, 在核内分段计算阶段利用迭代公式将每个核内的数据在不同向量间进行迭代, 最终得到长为 32 的迭
代结果. 接着, 在 VPU 内规约和 DSP 簇核间规约时采用合并公式, 将不同数据段迭代得到的结果最终规约为一个值.

3.6.3 访存优化
在 MT-3000 DSP 上, 每个核都有若干个直接存储器存取 (DMA). 使用 DMA 在各级存储之间搬移数据无需占
用计算核心, 可以使得计算和访存最大化的重叠, 提高性能.
乒乓算法是一种经典的访存优化技术, 用于实现异步情况下计算和数据传输之间的重叠. 本文在访存时采用
乒乓算法, 通过手动编程实现通信和计算最大限度的互相隐藏, 从而取得良好的加速效果, 乒乓算法具体流程如算
法 4 所示. DSP 上实现乒乓操作的核心是在 VPU 的 AM 存储器上申请数据所需大小 2 倍的空间, 以存放 2 份相
同大小但不同时序的数据互为对方的缓冲. 具体来说: 除去第 1 轮次的数据读取和最后一轮次的数据写回, DSP 核
心在进行本轮的数据计算同时将下一轮数据通过异步 DMA 加载到 AM 中, dma_p2p & dma_wait; VPU 计算完成
后, 再通过异步 DMA 方式将本线程负责的数据块从 AM 拷贝回 DDR. 异步 DMA 方式将计算与主存访问相重叠.