MhcPre

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	x
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×

功能说明

• 算子功能：基于一系列计算得到MHC架构中hidden层的${\mathbf{H}}_{\text{res}}^{\prime }$和${\mathbf{H}}_{\text{post}}$投影矩阵以及Attention或MLP层的输入矩阵${\mathbf{h}}_{\text{in}}$。对${\mathbf{H}}_{\text{res}}^{\prime }$矩阵执行Sinkhorn迭代归一化变换，最终得到双随机矩阵${\mathbf{H}}_{\text{res}}$；支持输出中间计算结果，用于反向梯度计算。包括sigmoid计算之后的${\mathbf{H}}_{\mathbf{l}}^{\mathbf{p}\mathbf{r}\mathbf{e}}$矩阵、$\overrightarrow{x_l^{{}^{\prime }}}$与$\phi$矩阵乘的结果，输入x的RmsNorm结果$\overrightarrow{{\mathbf{x}}_{\mathbf{l}}^{{}^{\prime }}}$、迭代过程中的中间归一化结果和$\mathbf{n}\mathbf{o}\mathbf{r}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}$和求和结果$\mathbf{s}\mathbf{u}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}$。

• 计算公式：

  $$\begin{array}{rl}\overrightarrow{x_l^{{}^{\prime }}}& =\frac{1}{\sqrt{\frac{1}{d}{\sum }_{\dim =-2,\text{keepdim}=\text{True}}{x}_i^2+ϵ}}\\ H_l^{pre}& ={\alpha }_l^{pre}\cdot (\overrightarrow{x_l^{{}^{\prime }}}{\phi }_l^{pre})+b_l^{pre}\\ H_l^{post}& ={\alpha }_l^{post}\cdot (\overrightarrow{x_l^{{}^{\prime }}}{\phi }_l^{post})+b_l^{post}\\ H_l^{res}& ={\alpha }_l^{res}\cdot (\overrightarrow{x_l^{{}^{\prime }}}{\phi }_l^{res})+b_l^{res}\\ H_l^{pre}& =\sigma (H_l^{pre})\\ H_l^{post}& =2\sigma (H_l^{post})\\ h_{in}& =\overrightarrow{x_l}{H}_l^{pre}\end{array}$$

  • 将${\mathbf{H}}_{\mathbf{l}}^{\mathbf{r}\mathbf{e}\mathbf{s}}$作为输入，Sinkhorn变换共执行$\mathbf{n}\mathbf{u}\mathbf{m}\mathbf{I}\mathbf{t}\mathbf{e}\mathbf{r}\mathbf{s}$次迭代，迭代过程中生成中间归一化结果$\mathbf{n}\mathbf{o}\mathbf{r}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[k]$和求和结果$\mathbf{s}\mathbf{u}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[k]$，最终输出最后一次迭代的$\mathbf{n}\mathbf{o}\mathbf{r}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}$作为变换结果。

    第一次迭代（初始化）：

    $$\begin{array}{rl}\mathbf{n}\mathbf{o}\mathbf{r}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[0]& =\text{softmax}({\mathbf{H}}_{\mathbf{l}}^{\mathbf{r}\mathbf{e}\mathbf{s}},\dim =-1)+ϵ,\\ \mathbf{s}\mathbf{u}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[1]& =\sum _{\dim =-2,\text{keepdim}=\text{True}}\mathbf{n}\mathbf{o}\mathbf{r}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[0]+ϵ,\\ \mathbf{n}\mathbf{o}\mathbf{r}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[1]& =\frac{\mathbf{n}\mathbf{o}\mathbf{r}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[0]}{\mathbf{s}\mathbf{u}\mathbf{m}\_\mathbf{o}\mathbf{u}\mathbf{t}[1]},\end{array}$$

    第$i$次迭代（$i=1,2,\dots ,(num\_iters-1)$）：

    $$\begin{array}{rl}\mathbf{s}\mathbf{u}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[2i]& =\sum _{\dim =-1,\text{keepdim}=\text{True}}\mathbf{n}\mathbf{o}\mathbf{r}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[2i-1]+ϵ,\\ \mathbf{n}\mathbf{o}\mathbf{r}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[2i]& =\frac{\mathbf{n}\mathbf{o}\mathbf{r}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[2i-1]}{\mathbf{s}\mathbf{u}\mathbf{m}\_\mathbf{o}\mathbf{u}\mathbf{t}[2i]},\\ \mathbf{s}\mathbf{u}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[2i+1]& =\sum _{\dim =-2,\text{keepdim}=\text{True}}\mathbf{n}\mathbf{o}\mathbf{r}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[2i]+ϵ,\\ \mathbf{n}\mathbf{o}\mathbf{r}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[2i+1]& =\frac{\mathbf{n}\mathbf{o}\mathbf{r}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[2i]}{\mathbf{s}\mathbf{u}\mathbf{m}\_\mathbf{o}\mathbf{u}\mathbf{t}[2i+1]},\end{array}$$

  • 最终输出

  $$\mathbf{n}\mathbf{o}\mathbf{r}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[2\times \mathbf{n}\mathbf{u}\mathbf{m}\_\mathbf{i}\mathbf{t}\mathbf{e}\mathbf{r}\mathbf{s}-1]$$

  $$\mathbf{s}\mathbf{u}\mathbf{m}\mathbf{O}\mathbf{u}\mathbf{t}[2\times \mathbf{n}\mathbf{u}\mathbf{m}\_\mathbf{i}\mathbf{t}\mathbf{e}\mathbf{r}\mathbf{s}-1]$$

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
x	输入	待计算数据，表示网络中mHC层的输入数据。	BFLOAT16, FLOAT16	ND
phi	输入	mHC的参数矩阵。	FLOAT32	ND
alpha	输入	mHC的缩放参数。	FLOAT32	ND
bias	输入	mHC的bias参数。	FLOAT32	ND
hcMult	可选输入	残差流数量，HC维度大小，当前仅支持4。	INT32	-
numIters	可选输入	表示sinkhorn算法迭代次数，当前仅支持20。	INT32	-
hcEps	可选输入	h_pre的sigmoid后的eps参数。	DOUBLE	-
normEps	可选输入	RmsNorm的防除零参数。	DOUBLE	-
needGrad	可选输入	是否需要输出额外属性。	BOOL	-
hIn	输出	输出的h_in作为Attention/MLP层的输入。	BFLOAT16, FLOAT16	ND
hPost	输出	输出的mHC的h_post变换矩阵。	FLOAT32	ND
hRes	输出	输出的mHC的h_res变换矩阵。	FLOAT32	ND
hPre	可选输出	需要反向时输出，做完sigmoid计算之后的hPre矩阵。	FLOAT32	ND
hcBeforeNorm	可选输出	需要反向时输出，x与phi矩阵乘的结果。	FLOAT32	ND
invRms	可选输出	需要反向时输出，RmsNorm计算得到的1/r。	FLOAT32	ND
sumOut	可选输出	需要反向时输出，每一次迭代的colSum/rowSum结果。	FLOAT32	ND
normOut	可选输出	需要反向时输出，每一次colSum/rowSum迭代后的comb结果。	FLOAT32	ND

约束说明

• n目前支持4。
• 输入x的最后一维需要满足128对齐。

调用说明

调用方式	调用样例	说明
aclnn调用	test_aclnn_mhc_pre_sinkhorn	通过aclnnMhcPreSinkHorn接口方式调用MhcPreSinkHorn算子。