文件最后提交记录最后更新时间
8 天前
30 天前
3 天前
2 天前
7 天前
5 个月前
4 天前
README

MhcPre

产品支持情况

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

功能说明

  • 算子功能:基于一系列计算得到MHC架构中hidden层的HresH^{res}HpostH^{post}投影矩阵以及Attention或MLP层的输入矩阵hinh^{in}

  • 计算公式

    其中,xFlat表示将x的最后两维n和D视作长度为nD的向量,gammaFlat表示将gamma视作长度为nD的向量,@表示矩阵乘法,⊙\odot表示逐元素乘法。

    invRms=(mean(xFlat2)+normEps)−12xGamma={xFlat⊙gammaFlat,gamma≠nullxFlat,gamma=nullhMix=xGamma@phiTw=hMix⊙invRms(pPre,pPost,pRes)={split(w,(n,n,n2)),alpha.shape=(3)(split(w,(n,n)),0),alpha.shape=(2)hPre=σ(pPre⊙alpha0+bias0)+hcEpshPost={2σ(pPost⊙alpha1+bias1),alpha.shape=(3)2σ(pPost⊙alpha1+bias1)+hcEps,alpha.shape=(2)hRes={pRes⊙alpha2+bias2,alpha.shape=(3)0,alpha.shape=(2)hInd=∑i=0n−1hPreixi,d\begin{aligned} invRms &= \left(mean(xFlat^{2}) + normEps\right)^{-\frac{1}{2}}\\ xGamma &= \begin{cases} xFlat \odot gammaFlat, & gamma \ne null \\ xFlat, & gamma = null \end{cases}\\ hMix &= xGamma @ phi^{T}\\ w &= hMix \odot invRms\\ (pPre, pPost, pRes) &= \begin{cases} split(w, (n, n, n^{2})), & alpha.shape=(3) \\ (split(w, (n, n)), 0), & alpha.shape=(2) \end{cases}\\ hPre &= \sigma(pPre \odot alpha0 + bias0) + hcEps\\ hPost &= \begin{cases} 2\sigma(pPost \odot alpha1 + bias1), & alpha.shape=(3) \\ 2\sigma(pPost \odot alpha1 + bias1) + hcEps, & alpha.shape=(2) \end{cases}\\ hRes &= \begin{cases} pRes \odot alpha2 + bias2, & alpha.shape=(3) \\ 0, & alpha.shape=(2) \end{cases}\\ hIn_{d} &= \sum_{i=0}^{n-1} hPre_{i} x_{i,d} \end{aligned}

参数说明

参数名 输入/输出/属性 描述 数据类型 数据格式
x 输入 待计算数据,表示网络中mHC层的输入数据,对应公式中的x。 BFLOAT16, FLOAT16 ND
phi 输入 mHC的参数矩阵,对应公式中的phi。 FLOAT32 ND
alpha 输入 mHC的缩放参数,对应公式中的alpha。 FLOAT32 -
bias 输入 mHC的bias参数,对应公式中的bias。 FLOAT32 -
gamma 可选输入 表示进行RmsNorm计算的缩放因子,对应公式中的gamma。 FLOAT32 ND
out_flag 可选输入 表示是否输出公式中的invRms、hMix和hPre,默认为0表示不输出,为1表示全部输出。 INT64 -
norm_eps 可选输入 RmsNorm的防除零参数,对应公式中的normEps。 FLOAT -
hc_eps 可选输入 h_pre的sigmoid后的eps参数,以及当alpha.shape=(2)时h_post的sigmoid后的eps参数,对应公式中的hcEps。 FLOAT -
op_impl_mode 可选输入 指定MhcPre算子的计算模式,0表示Cube使用FP32模式计算,1表示Cube使用HF32模式计算,默认值为0。 INT64 -
h_in 输出 输出的h_in作为Attention/MLP层的输入,对应公式中的hIn。 BFLOAT16, FLOAT16 ND
h_post 输出 输出的mHC的h_post变换矩阵,对应公式中的hPost。 FLOAT32 ND
h_res 输出 输出的mHC的h_res变换矩阵(未做sinkhorn变换),对应公式中的hRes。 FLOAT32 ND
inv_rms 可选输出 RmsNorm计算得到的1/r,对应公式中的invRms。 FLOAT32 ND
h_mix 可选输出 xGamma与phi矩阵乘的结果,对应公式中的hMix。 FLOAT32 ND
h_pre 可选输出 做完sigmoid计算之后的h_pre矩阵,对应公式中的hPre。 FLOAT32 ND

约束说明

  • Ascend 950PR/Ascend 950DT:
    • n目前支持4、6、8。
    • D支持1~16384范围以内,需满足D为16对齐。
  • Atlas A3 训练系列产品/Atlas A3 推理系列产品、Atlas A2 训练系列产品/Atlas A2 推理系列产品:
    • 参数op_impl_mode仅支持配置为0。
    • n目前支持4。
    • D支持100000范围以内,需满足D为128对齐。

调用说明

调用方式 调用样例 说明
aclnn调用 test_aclnn_mhc_pre 通过aclnnMhcPre接口方式调用MhcPre算子。
aclnn调用 test_aclnn_mhc_pre_v2 通过aclnnMhcPreV2接口方式调用MhcPre算子,并通过opImplMode选择Cube的FP32或HF32计算模式。