Merged
fix(mhc): 新增 TextGatedResidualGrad 和Readme文档 #15537
azyuliAdd on 6 days ago
fix(mhc): 新增 TextGatedResidualGrad 和Readme文档 #15537
Merged
Total: 3 files changed+671-0
| @@ -0,0 +1,145 @@ | |||
| 1 | +# TextGatedResidualGrad | ||
| 2 | + | ||
| 3 | +## 产品支持情况 | ||
| 4 | + | ||
| 5 | +| 产品 | 是否支持 | | ||
| 6 | +| :--- | :---: | | ||
| 7 | +| <term>Ascend 950PR&950DT系列产品</term> | √ | | ||
| 8 | +| <term>Atlas A3系列产品</term> | √ | | ||
| 9 | +| <term>Atlas A2系列产品</term> | √ | | ||
| 10 | +| <term>Atlas 200I/500 A2推理产品</term> | × | | ||
| 11 | +| <term>Atlas推理系列产品</term> | × | | ||
| 12 | +| <term>Atlas训练系列产品</term> | × | | ||
| 13 | + | ||
| 14 | +## 功能说明 | ||
| 15 | + | ||
| 16 | +`TextGatedResidualGrad` 是 `TextGatedResidual` 的反向算子。根据混合输出和注入输出的上游梯度，以及前向保存的张量，计算输入 `x` 和权重 `wDown`、`wUp`、`wGate`、`gamma` 的梯度。 | ||
| 17 | + | ||
| 18 | +记 $B$ 为批大小，$S$ 为序列长度，$C$ 为通道数，$H$ 为每通道的特征数，$R$ 为下投影维度，$D=C\times H$。$\odot$ 表示逐元素乘法，$@$ 表示矩阵乘法，$\sigma$ 表示 Sigmoid 函数。 | ||
| 19 | + | ||
| 20 | +对应的前向计算为： | ||
| 21 | + | ||
| 22 | +$$ | ||
| 23 | +\begin{aligned} | ||
| 24 | +\mathrm{invRms}_{b,s,c} &= \left(\frac{1}{H}\sum_{j=0}^{H-1}x_{b,s,c,j}^{2}+\epsilon\right)^{-1/2} \\ | ||
| 25 | +\mathrm{xNorm} &= x\odot\mathrm{invRms}\odot(1+\gamma) \\ | ||
| 26 | +\mathrm{downPreAct} &= \mathrm{xNorm}@\mathrm{wDown}^{T} \\ | ||
| 27 | +\mathrm{act} &= \mathrm{SiLU}(\mathrm{downPreAct}/C) \\ | ||
| 28 | +\mathrm{up} &= \mathrm{act}@\mathrm{wUp}^{T} \\ | ||
| 29 | +\mathrm{mixed}_{b,s,j} &= \frac{1}{C}\sum_{c=0}^{C-1}\mathrm{xNorm}_{b,s,c,j}\,\sigma(\mathrm{up}_{b,s,c,j}) \\ | ||
| 30 | +\mathrm{inject} &= 2\sigma\!\left((\mathrm{xNorm}@\mathrm{wGate}^{T})/C\right) | ||
| 31 | +\end{aligned} | ||
| 32 | +$$ | ||
| 33 | + | ||
| 34 | +其中，归一化及通道求和时将最后一维 $D$ 视为 $[C,H]$，矩阵乘法时仍使用最后一维 $D$。`invRms` 在每个通道的 $H$ 个特征上广播，`gamma` 在 $B$、$S$ 维上广播。$\epsilon$ 是前向归一化使用的参数，本接口通过 `invRms` 接收对应的前向结果。 | ||
| 35 | + | ||
| 36 | +给定 `gradMixed` 和 `gradInject`，定义 | ||
| 37 | + | ||
| 38 | +$$ | ||
| 39 | +L=\sum_{b,s,j}\mathrm{gradMixed}_{b,s,j}\,\mathrm{mixed}_{b,s,j} | ||
| 40 | + +\sum_{b,s,c}\mathrm{gradInject}_{b,s,c}\,\mathrm{inject}_{b,s,c}. | ||
| 41 | +$$ | ||
| 42 | + | ||
| 43 | +本算子的输出为： | ||
| 44 | + | ||
| 45 | +$$ | ||
| 46 | +dx=\frac{\partial L}{\partial x},\quad | ||
| 47 | +dwDown=\frac{\partial L}{\partial wDown},\quad | ||
| 48 | +dwUp=\frac{\partial L}{\partial wUp},\quad | ||
| 49 | +dwGate=\frac{\partial L}{\partial wGate},\quad | ||
| 50 | +dgamma=\frac{\partial L}{\partial\gamma}. | ||
| 51 | +$$ | ||
| 52 | + | ||
| 53 | +`wGate` 缺省时不包含注入分支及其梯度；`gamma` 缺省时归一化缩放系数为 1，不输出 `dgamma`。`invRms`、`inject`、`downPreAct`、`up` 必须来自同一次前向计算，并与本次传入的输入、权重对应。 | ||
| 54 | + | ||
| 55 | +## 参数说明 | ||
| 56 | + | ||
| 57 | +| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 | shape | | ||
| 58 | +| :--- | :--- | :--- | :--- | :--- | :--- | | ||
| 59 | +| gradMixed | 输入 | 混合输出 mixed 的上游梯度。 | FLOAT16、BFLOAT16 | ND | [B,S,H] | | ||
| 60 | +| x | 输入 | 前向输入。 | FLOAT16、BFLOAT16 | ND | [B,S,D] | | ||
| 61 | +| wDown | 输入 | 下投影权重。 | FLOAT16、BFLOAT16 | ND | [R,D] | | ||
| 62 | +| wUp | 输入 | 上投影权重。 | FLOAT16、BFLOAT16 | ND | [D,R] | | ||
| 63 | +| invRms | 输入 | 前向保存的各通道均方根倒数。 | FLOAT32 | ND | [B,S,C] | | ||
| 64 | +| downPreAct | 输入 | 前向保存的下投影结果 xNorm@wDownᵀ。 | FLOAT16、BFLOAT16 | ND | [B,S,R] | | ||
| 65 | +| up | 输入 | 前向保存的上投影结果 act@wUpᵀ。 | FLOAT16、BFLOAT16 | ND | [B,S,D] | | ||
| 66 | +| gradInject | 可选输入 | 注入输出 inject 的上游梯度。 | FLOAT16、BFLOAT16 | ND | [B,S,C] | | ||
| 67 | +| wGate | 可选输入 | 注入分支权重。 | FLOAT16、BFLOAT16 | ND | [C,D] | | ||
| 68 | +| gamma | 可选输入 | 归一化缩放权重，实际缩放系数为 1+gamma。 | FLOAT32 | ND | [D] | | ||
| 69 | +| inject | 可选输入 | 前向保存的注入输出 2*sigmoid((xNorm@wGateᵀ)/C)。 | FLOAT16、BFLOAT16 | ND | [B,S,C] | | ||
| 70 | +| channels | 属性 | 通道数 C。 | INT64 | - | 标量 | | ||
| 71 | +| dx | 输出 | x 的梯度。 | 同 x | ND | [B,S,D] | | ||
| 72 | +| dwDown | 输出 | wDown 的梯度。 | FLOAT16、BFLOAT16 | ND | [R,D] | | ||
| 73 | +| dwUp | 输出 | wUp 的梯度。 | FLOAT16、BFLOAT16 | ND | [D,R] | | ||
| 74 | +| dwGate | 可选输出 | wGate 的梯度。 | FLOAT16、BFLOAT16 | ND | [C,D] | | ||
| 75 | +| dgamma | 可选输出 | gamma 的梯度。 | FLOAT32 | ND | [D] | | ||
| 76 | + | ||
| 77 | +## 约束说明 | ||
| 78 | + | ||
| 79 | +- `C`、`H`、`R` 为正整数；`H`、`R` 均须为 16 的整数倍。 | ||
| 80 | + | ||
| 81 | +通道数须满足 `C ≤ 20`。 | ||
| 82 | + | ||
| 83 | +`H` 的取值范围按产品确定。令 | ||
| 84 | + | ||
| 85 | +<!-- npu="950" id7 --> | ||
| 86 | +- <term>Ascend 950PR&950DT系列产品</term>： | ||
| 87 | + | ||
| 88 | + $$ | ||
| 89 | + K = 62448-16\left\lceil C/8\right\rceil | ||
| 90 | + $$ | ||
| 91 | + | ||
| 92 | +<!-- end id7 --> | ||
| 93 | +<!-- npu="A3,910b" id8 --> | ||
| 94 | +- <term>Atlas A2系列产品</term>、<term>Atlas A3系列产品</term>： | ||
| 95 | + | ||
| 96 | + $$ | ||
| 97 | + K = 39664-16\left\lceil C/8\right\rceil | ||
| 98 | + $$ | ||
| 99 | + | ||
| 100 | +<!-- end id8 --> | ||
| 101 | + | ||
| 102 | +`H`、`R` 须同时满足： | ||
| 103 | + | ||
| 104 | +$$ | ||
| 105 | +\boxed{ | ||
| 106 | +\begin{aligned} | ||
| 107 | +13H &\le K,\\ | ||
| 108 | +3H+8R &\le K,\\ | ||
| 109 | +7H+6R &\le K. | ||
| 110 | +\end{aligned} | ||
| 111 | +} | ||
| 112 | +$$ | ||
| 113 | + | ||
| 114 | +因此，给定 `C` 后，`H` 的上限为 | ||
| 115 | + | ||
| 116 | +$$ | ||
| 117 | +H_{\max}=16\left\lfloor\frac{K}{208}\right\rfloor. | ||
| 118 | +$$ | ||
| 119 | + | ||
| 120 | +`R` 须同时满足上述关系式和 `R ≤ 640`。 | ||
| 121 | + | ||
| 122 | +其中，$\lceil\cdot\rceil$ 表示向上取整，$\lfloor\cdot\rfloor$ 表示向下取整。上述关系适用于 FLOAT16、BFLOAT16，以及所有合法的可选参数组合。 | ||
| 123 | + | ||
| 124 | +例如，`C=4、H=2560、R=320` 适用于上述三类产品。其他 `C/H/R` 组合须满足上述约束。调用时仍需通过第一段接口 `aclnnTextGatedResidualGradGetWorkspaceSize` 完成参数校验，并按返回的 `workspaceSize` 申请内存。 | ||
| 125 | + | ||
| 126 | +- 所有 FLOAT16/BFLOAT16 输入以及 `dx`、`dwDown`、`dwUp`、`dwGate` 的数据类型必须一致；各权重梯度与对应权重的数据类型相同。`invRms`、`gamma` 和 `dgamma` 为 FLOAT32。 | ||
| 127 | +- 支持非连续输入 Tensor。输入输出的 shape 必须符合参数表，同名维度的大小必须一致。 | ||
| 128 | +- 支持 `B=0` 或 `S=0`：`dx` 为空，已提供的非空权重梯度和 `dgamma` 全部置零，其余参数仍须满足上述约束。 | ||
| 129 | +- 输出梯度覆盖写入，调用前无需初始化，也不会累加到输出 Tensor 的原有数据中。 | ||
| 130 | +- 默认执行确定性计算。 | ||
| 131 | + | ||
| 132 | +ACLNN 调用时，可选参数须满足以下要求： | ||
| 133 | + | ||
| 134 | +| 功能 | 提供时 | 缺省时 | | ||
| 135 | +| :--- | :--- | :--- | | ||
| 136 | +| 注入分支 | `wGate`、`gradInject`、`inject`、`dwGate` 四者必须同时提供 | 四者必须全部传入 `nullptr` | | ||
| 137 | +| 归一化缩放 | `gamma`、`dgamma` 必须同时提供 | 两者必须全部传入 `nullptr`，缩放系数取 1 | | ||
| 138 | + | ||
| 139 | +两组可选参数相互独立。`nullptr` 表示省略该功能，不能用 shape 为 `[0]` 的占位 Tensor 代替。 | ||
| 140 | + | ||
| 141 | +## 调用说明 | ||
| 142 | + | ||
| 143 | +| 调用方式 | 调用样例 | 说明 | | ||
| 144 | +| :--- | :--- | :--- | | ||
| 145 | +| aclnn 调用 | [接口调用示例](./docs/aclnnTextGatedResidualGrad.md#调用示例) | 创建输入输出 Tensor，调用两段式接口获取并执行计算。 | | ||
| @@ -0,0 +1,145 @@ | |||
| 1 | +# TextGatedResidualGrad | ||
| 2 | + | ||
| 3 | +## Supported Products | ||
| 4 | + | ||
| 5 | +| Product | Supported | | ||
| 6 | +| :--- | :---: | | ||
| 7 | +| <term>Ascend 950PR & 950DT Series</term> | √ | | ||
| 8 | +| <term>Atlas A3 Series</term> | √ | | ||
| 9 | +| <term>Atlas A2 Series</term> | √ | | ||
| 10 | +| <term>Atlas 200I/500 A2 Inference Products</term> | × | | ||
| 11 | +| <term>Atlas Inference Series</term> | × | | ||
| 12 | +| <term>Atlas Training Series</term> | × | | ||
| 13 | + | ||
| 14 | +## Function | ||
| 15 | + | ||
| 16 | +`TextGatedResidualGrad` is the backward operator of `TextGatedResidual`. It computes gradients of input `x` and weights `wDown`, `wUp`, `wGate`, and `gamma` from the upstream gradients of the mixed and injection outputs and tensors saved during the forward pass. | ||
| 17 | + | ||
| 18 | +Let $B$ denote the batch size, $S$ the sequence length, $C$ the number of channels, $H$ the number of features per channel, $R$ the down-projection dimension, and $D=C\times H$. $\odot$ denotes element-wise multiplication, $@$ denotes matrix multiplication, and $\sigma$ denotes the sigmoid function. | ||
| 19 | + | ||
| 20 | +The corresponding forward computation is: | ||
| 21 | + | ||
| 22 | +$$ | ||
| 23 | +\begin{aligned} | ||
| 24 | +\mathrm{invRms}_{b,s,c} &= \left(\frac{1}{H}\sum_{j=0}^{H-1}x_{b,s,c,j}^{2}+\epsilon\right)^{-1/2} \\ | ||
| 25 | +\mathrm{xNorm} &= x\odot\mathrm{invRms}\odot(1+\gamma) \\ | ||
| 26 | +\mathrm{downPreAct} &= \mathrm{xNorm}@\mathrm{wDown}^{T} \\ | ||
| 27 | +\mathrm{act} &= \mathrm{SiLU}(\mathrm{downPreAct}/C) \\ | ||
| 28 | +\mathrm{up} &= \mathrm{act}@\mathrm{wUp}^{T} \\ | ||
| 29 | +\mathrm{mixed}_{b,s,j} &= \frac{1}{C}\sum_{c=0}^{C-1}\mathrm{xNorm}_{b,s,c,j}\,\sigma(\mathrm{up}_{b,s,c,j}) \\ | ||
| 30 | +\mathrm{inject} &= 2\sigma\!\left((\mathrm{xNorm}@\mathrm{wGate}^{T})/C\right) | ||
| 31 | +\end{aligned} | ||
| 32 | +$$ | ||
| 33 | + | ||
| 34 | +The last dimension $D$ is interpreted as $[C,H]$ for normalization and summation across channels, and remains $D$ for matrix multiplication. `invRms` is broadcast across the $H$ features of each channel; `gamma` is broadcast across the $B$ and $S$ dimensions. $\epsilon$ is the parameter used for forward normalization. This interface receives the corresponding forward result through `invRms`. | ||
| 35 | + | ||
| 36 | +Given `gradMixed` and `gradInject`, define | ||
| 37 | + | ||
| 38 | +$$ | ||
| 39 | +L=\sum_{b,s,j}\mathrm{gradMixed}_{b,s,j}\,\mathrm{mixed}_{b,s,j} | ||
| 40 | + +\sum_{b,s,c}\mathrm{gradInject}_{b,s,c}\,\mathrm{inject}_{b,s,c}. | ||
| 41 | +$$ | ||
| 42 | + | ||
| 43 | +The outputs of this operator are: | ||
| 44 | + | ||
| 45 | +$$ | ||
| 46 | +dx=\frac{\partial L}{\partial x},\quad | ||
| 47 | +dwDown=\frac{\partial L}{\partial wDown},\quad | ||
| 48 | +dwUp=\frac{\partial L}{\partial wUp},\quad | ||
| 49 | +dwGate=\frac{\partial L}{\partial wGate},\quad | ||
| 50 | +dgamma=\frac{\partial L}{\partial\gamma}. | ||
| 51 | +$$ | ||
| 52 | + | ||
| 53 | +When `wGate` is omitted, the injection branch and its gradients are excluded. When `gamma` is omitted, the normalization scale is 1 and `dgamma` is not produced. `invRms`, `inject`, `downPreAct`, and `up` must come from the same forward pass and correspond to the supplied input and weights. | ||
| 54 | + | ||
| 55 | +## Parameters | ||
| 56 | + | ||
| 57 | +| Name | Input/Output/Attribute | Description | Data Type | Format | Shape | | ||
| 58 | +| :--- | :--- | :--- | :--- | :--- | :--- | | ||
| 59 | +| gradMixed | Input | Upstream gradient of the mixed output mixed. | FLOAT16, BFLOAT16 | ND | [B,S,H] | | ||
| 60 | +| x | Input | Forward input. | FLOAT16, BFLOAT16 | ND | [B,S,D] | | ||
| 61 | +| wDown | Input | Down-projection weight. | FLOAT16, BFLOAT16 | ND | [R,D] | | ||
| 62 | +| wUp | Input | Up-projection weight. | FLOAT16, BFLOAT16 | ND | [D,R] | | ||
| 63 | +| invRms | Input | Per-channel reciprocal root mean square saved during the forward pass. | FLOAT32 | ND | [B,S,C] | | ||
| 64 | +| downPreAct | Input | Down-projection result xNorm@wDownᵀ saved during the forward pass. | FLOAT16, BFLOAT16 | ND | [B,S,R] | | ||
| 65 | +| up | Input | Up-projection result act@wUpᵀ saved during the forward pass. | FLOAT16, BFLOAT16 | ND | [B,S,D] | | ||
| 66 | +| gradInject | Optional input | Upstream gradient of the injection output inject. | FLOAT16, BFLOAT16 | ND | [B,S,C] | | ||
| 67 | +| wGate | Optional input | Injection branch weight. | FLOAT16, BFLOAT16 | ND | [C,D] | | ||
| 68 | +| gamma | Optional input | Normalization scale weight; the actual scale is 1+gamma. | FLOAT32 | ND | [D] | | ||
| 69 | +| inject | Optional input | Injection output 2*sigmoid((xNorm@wGateᵀ)/C) saved during the forward pass. | FLOAT16, BFLOAT16 | ND | [B,S,C] | | ||
| 70 | +| channels | Attribute | Number of channels C. | INT64 | - | Scalar | | ||
| 71 | +| dx | Output | Gradient of x. | Same as x | ND | [B,S,D] | | ||
| 72 | +| dwDown | Output | Gradient of wDown. | FLOAT16, BFLOAT16 | ND | [R,D] | | ||
| 73 | +| dwUp | Output | Gradient of wUp. | FLOAT16, BFLOAT16 | ND | [D,R] | | ||
| 74 | +| dwGate | Optional output | Gradient of wGate. | FLOAT16, BFLOAT16 | ND | [C,D] | | ||
| 75 | +| dgamma | Optional output | Gradient of gamma. | FLOAT32 | ND | [D] | | ||
| 76 | + | ||
| 77 | +## Constraints | ||
| 78 | + | ||
| 79 | +- `C`, `H`, and `R` must be positive integers; both `H` and `R` must be multiples of 16. | ||
| 80 | + | ||
| 81 | +The number of channels must satisfy `C ≤ 20`. | ||
| 82 | + | ||
| 83 | +The range of `H` depends on the product. Define | ||
| 84 | + | ||
| 85 | +<!-- npu="950" id7 --> | ||
| 86 | +- <term>Ascend 950PR & 950DT Series</term>: | ||
| 87 | + | ||
| 88 | + $$ | ||
| 89 | + K = 62448-16\left\lceil C/8\right\rceil | ||
| 90 | + $$ | ||
| 91 | + | ||
| 92 | +<!-- end id7 --> | ||
| 93 | +<!-- npu="A3,910b" id8 --> | ||
| 94 | +- <term>Atlas A2 Series</term>, <term>Atlas A3 Series</term>: | ||
| 95 | + | ||
| 96 | + $$ | ||
| 97 | + K = 39664-16\left\lceil C/8\right\rceil | ||
| 98 | + $$ | ||
| 99 | + | ||
| 100 | +<!-- end id8 --> | ||
| 101 | + | ||
| 102 | +`H` and `R` must satisfy all of the following: | ||
| 103 | + | ||
| 104 | +$$ | ||
| 105 | +\boxed{ | ||
| 106 | +\begin{aligned} | ||
| 107 | +13H &\le K,\\ | ||
| 108 | +3H+8R &\le K,\\ | ||
| 109 | +7H+6R &\le K. | ||
| 110 | +\end{aligned} | ||
| 111 | +} | ||
| 112 | +$$ | ||
| 113 | + | ||
| 114 | +Therefore, for a given `C`, the upper bound of `H` is | ||
| 115 | + | ||
| 116 | +$$ | ||
| 117 | +H_{\max}=16\left\lfloor\frac{K}{208}\right\rfloor. | ||
| 118 | +$$ | ||
| 119 | + | ||
| 120 | +`R` must satisfy both the inequalities above and `R ≤ 640`. | ||
| 121 | + | ||
| 122 | +Here, $\lceil\cdot\rceil$ denotes rounding up and $\lfloor\cdot\rfloor$ denotes rounding down. These relationships apply to FLOAT16, BFLOAT16, and all valid combinations of optional parameters. | ||
| 123 | + | ||
| 124 | +For example, `C=4, H=2560, R=320` is valid for all three supported product series. Other `C/H/R` combinations must satisfy the constraints above. Call the first-stage interface `aclnnTextGatedResidualGradGetWorkspaceSize` to validate parameters and allocate memory according to the returned `workspaceSize`. | ||
| 125 | + | ||
| 126 | +- All FLOAT16/BFLOAT16 inputs and `dx`, `dwDown`, `dwUp`, and `dwGate` must have the same data type. Each weight gradient must have the same data type as its corresponding weight. `invRms`, `gamma`, and `dgamma` are FLOAT32. | ||
| 127 | +- Non-contiguous input tensors are supported. Input and output shapes must match the parameter table, and dimensions with the same name must have the same size. | ||
| 128 | +- `B=0` or `S=0` is supported: `dx` is empty, and all supplied non-empty weight gradients and `dgamma` are set to zero. Other parameters must still satisfy the constraints above. | ||
| 129 | +- Output gradients are overwritten. They do not require initialization and are not accumulated into existing output tensor values. | ||
| 130 | +- Computation is deterministic by default. | ||
| 131 | + | ||
| 132 | +For ACLNN calls, optional parameters must satisfy the following requirements: | ||
| 133 | + | ||
| 134 | +| Feature | When enabled | When omitted | | ||
| 135 | +| :--- | :--- | :--- | | ||
| 136 | +| Injection branch | `wGate`, `gradInject`, `inject`, and `dwGate` must all be supplied | All four must be `nullptr` | | ||
| 137 | +| Normalization scaling | `gamma` and `dgamma` must both be supplied | Both must be `nullptr`; the scale is 1 | | ||
| 138 | + | ||
| 139 | +The two groups of optional parameters are independent. `nullptr` omits the corresponding feature; a placeholder tensor with shape `[0]` cannot be used instead. | ||
| 140 | + | ||
| 141 | +## Calling the Operator | ||
| 142 | + | ||
| 143 | +| Method | Example | Description | | ||
| 144 | +| :--- | :--- | :--- | | ||
| 145 | +| aclnn call | [Interface example](./docs/aclnnTextGatedResidualGrad.md#调用示例) | Create input and output tensors, then call the two-stage interface to prepare and execute the computation. | | ||
| @@ -0,0 +1,381 @@ | |||
| 1 | +# aclnnTextGatedResidualGrad | ||
| 2 | + | ||
| 3 | +## 产品支持情况 | ||
| 4 | + | ||
| 5 | +<!-- npu="950" id1 --> | ||
| 6 | +- <term>Ascend 950PR&950DT系列产品</term>：支持 | ||
| 7 | +<!-- end id1 --> | ||
| 8 | +<!-- npu="A3" id2 --> | ||
| 9 | +- <term>Atlas A3系列产品</term>：支持 | ||
| 10 | +<!-- end id2 --> | ||
| 11 | +<!-- npu="910b" id3 --> | ||
| 12 | +- <term>Atlas A2系列产品</term>：支持 | ||
| 13 | +<!-- end id3 --> | ||
| 14 | +<!-- npu="310b" id4 --> | ||
| 15 | +- <term>Atlas 200I/500 A2推理产品</term>：不支持 | ||
| 16 | +<!-- end id4 --> | ||
| 17 | +<!-- npu="310p" id5 --> | ||
| 18 | +- <term>Atlas推理系列产品</term>：不支持 | ||
| 19 | +<!-- end id5 --> | ||
| 20 | +<!-- npu="910" id6 --> | ||
| 21 | +- <term>Atlas训练系列产品</term>：不支持 | ||
| 22 | +<!-- end id6 --> | ||
| 23 | + | ||
| 24 | +## 功能说明 | ||
| 25 | + | ||
| 26 | +`TextGatedResidualGrad` 是 `TextGatedResidual` 的反向算子。根据混合输出和注入输出的上游梯度，以及前向保存的张量，计算输入 `x` 和权重 `wDown`、`wUp`、`wGate`、`gamma` 的梯度。 | ||
| 27 | + | ||
| 28 | +记 $B$ 为批大小，$S$ 为序列长度，$C$ 为通道数，$H$ 为每通道的特征数，$R$ 为下投影维度，$D=C\times H$。$\odot$ 表示逐元素乘法，$@$ 表示矩阵乘法，$\sigma$ 表示 Sigmoid 函数。 | ||
| 29 | + | ||
| 30 | +对应的前向计算为： | ||
| 31 | + | ||
| 32 | +$$ | ||
| 33 | +\begin{aligned} | ||
| 34 | +\mathrm{invRms}_{b,s,c} &= \left(\frac{1}{H}\sum_{j=0}^{H-1}x_{b,s,c,j}^{2}+\epsilon\right)^{-1/2} \\ | ||
| 35 | +\mathrm{xNorm} &= x\odot\mathrm{invRms}\odot(1+\gamma) \\ | ||
| 36 | +\mathrm{downPreAct} &= \mathrm{xNorm}@\mathrm{wDown}^{T} \\ | ||
| 37 | +\mathrm{act} &= \mathrm{SiLU}(\mathrm{downPreAct}/C) \\ | ||
| 38 | +\mathrm{up} &= \mathrm{act}@\mathrm{wUp}^{T} \\ | ||
| 39 | +\mathrm{mixed}_{b,s,j} &= \frac{1}{C}\sum_{c=0}^{C-1}\mathrm{xNorm}_{b,s,c,j}\,\sigma(\mathrm{up}_{b,s,c,j}) \\ | ||
| 40 | +\mathrm{inject} &= 2\sigma\!\left((\mathrm{xNorm}@\mathrm{wGate}^{T})/C\right) | ||
| 41 | +\end{aligned} | ||
| 42 | +$$ | ||
| 43 | + | ||
| 44 | +其中，归一化及通道求和时将最后一维 $D$ 视为 $[C,H]$，矩阵乘法时仍使用最后一维 $D$。`invRms` 在每个通道的 $H$ 个特征上广播，`gamma` 在 $B$、$S$ 维上广播。$\epsilon$ 是前向归一化使用的参数，本接口通过 `invRms` 接收对应的前向结果。 | ||
| 45 | + | ||
| 46 | +给定 `gradMixed` 和 `gradInject`，定义 | ||
| 47 | + | ||
| 48 | +$$ | ||
| 49 | +L=\sum_{b,s,j}\mathrm{gradMixed}_{b,s,j}\,\mathrm{mixed}_{b,s,j} | ||
| 50 | + +\sum_{b,s,c}\mathrm{gradInject}_{b,s,c}\,\mathrm{inject}_{b,s,c}. | ||
| 51 | +$$ | ||
| 52 | + | ||
| 53 | +本算子的输出为： | ||
| 54 | + | ||
| 55 | +$$ | ||
| 56 | +dx=\frac{\partial L}{\partial x},\quad | ||
| 57 | +dwDown=\frac{\partial L}{\partial wDown},\quad | ||
| 58 | +dwUp=\frac{\partial L}{\partial wUp},\quad | ||
| 59 | +dwGate=\frac{\partial L}{\partial wGate},\quad | ||
| 60 | +dgamma=\frac{\partial L}{\partial\gamma}. | ||
| 61 | +$$ | ||
| 62 | + | ||
| 63 | +`wGate` 缺省时不包含注入分支及其梯度；`gamma` 缺省时归一化缩放系数为 1，不输出 `dgamma`。`invRms`、`inject`、`downPreAct`、`up` 必须来自同一次前向计算，并与本次传入的输入、权重对应。 | ||
| 64 | + | ||
| 65 | +## 函数原型 | ||
| 66 | + | ||
| 67 | +本接口采用[两段式调用](../../../docs/zh/context/two_phase_api.md)：先调用 `aclnnTextGatedResidualGradGetWorkspaceSize` 获取工作空间大小和执行器，再调用 `aclnnTextGatedResidualGrad` 执行计算。 | ||
| 68 | + | ||
| 69 | +```cpp | ||
| 70 | +aclnnStatus aclnnTextGatedResidualGradGetWorkspaceSize( | ||
| 71 | + const aclTensor* gradMixed, const aclTensor* x, const aclTensor* wDown, const aclTensor* wUp, | ||
| 72 | + const aclTensor* invRms, const aclTensor* downPreAct, const aclTensor* up, const aclTensor* gradInject, | ||
| 73 | + const aclTensor* wGate, const aclTensor* gamma, const aclTensor* inject, int64_t channels, aclTensor* dx, | ||
| 74 | + aclTensor* dwDown, aclTensor* dwUp, aclTensor* dwGate, aclTensor* dgamma, uint64_t* workspaceSize, | ||
| 75 | + aclOpExecutor** executor); | ||
| 76 | + | ||
| 77 | +aclnnStatus aclnnTextGatedResidualGrad(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, | ||
| 78 | + aclrtStream stream); | ||
| 79 | +``` | ||
| 80 | + | ||
| 81 | +## aclnnTextGatedResidualGradGetWorkspaceSize | ||
| 82 | + | ||
| 83 | +- **参数说明** | ||
| 84 | + | ||
| 85 | +| 参数名 | 输入/输出 | 描述 | 使用说明 | 数据类型 | 数据格式 | shape | 非连续 Tensor | | ||
| 86 | +| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :---: | | ||
| 87 | +| gradMixed（const aclTensor*） | 输入 | 混合输出 mixed 的上游梯度。 | 与 x 数据类型一致。 | FLOAT16、BFLOAT16 | ND | [B,S,H] | √ | | ||
| 88 | +| x（const aclTensor*） | 输入 | 前向输入。 | D=C*H。 | FLOAT16、BFLOAT16 | ND | [B,S,D] | √ | | ||
| 89 | +| wDown（const aclTensor*） | 输入 | 下投影权重。 | 与 x 数据类型一致。 | FLOAT16、BFLOAT16 | ND | [R,D] | √ | | ||
| 90 | +| wUp（const aclTensor*） | 输入 | 上投影权重。 | 与 x 数据类型一致。 | FLOAT16、BFLOAT16 | ND | [D,R] | √ | | ||
| 91 | +| invRms（const aclTensor*） | 输入 | 前向保存的各通道均方根倒数。 | 对应前向的归一化结果，包含前向 epsilon 的作用。 | FLOAT32 | ND | [B,S,C] | √ | | ||
| 92 | +| downPreAct（const aclTensor*） | 输入 | 前向保存的下投影结果 xNorm@wDownᵀ。 | 保存未除以 C、未经过 SiLU 的结果；与 x 数据类型一致。 | FLOAT16、BFLOAT16 | ND | [B,S,R] | √ | | ||
| 93 | +| up（const aclTensor*） | 输入 | 前向保存的上投影结果 act@wUpᵀ。 | 保存 Sigmoid 之前的值；与 x 数据类型一致。 | FLOAT16、BFLOAT16 | ND | [B,S,D] | √ | | ||
| 94 | +| gradInject（const aclTensor*） | 可选输入 | 注入输出 inject 的上游梯度。 | 与 x 数据类型一致；与 wGate、inject、dwGate 同时提供或同时缺省。 | FLOAT16、BFLOAT16 | ND | [B,S,C] | √ | | ||
| 95 | +| wGate（const aclTensor*） | 可选输入 | 注入分支权重。 | 与 x 数据类型一致；与 gradInject、inject、dwGate 同时提供或同时缺省。 | FLOAT16、BFLOAT16 | ND | [C,D] | √ | | ||
| 96 | +| gamma（const aclTensor*） | 可选输入 | 归一化缩放权重，实际缩放系数为 1+gamma。 | 传入原始 gamma；与 dgamma 同时提供或同时缺省。 | FLOAT32 | ND | [D] | √ | | ||
| 97 | +| inject（const aclTensor*） | 可选输入 | 前向保存的注入输出 2*sigmoid((xNorm@wGateᵀ)/C)。 | 与 x 数据类型一致；与 wGate、gradInject、dwGate 同时提供或同时缺省。 | FLOAT16、BFLOAT16 | ND | [B,S,C] | √ | | ||
| 98 | +| channels（int64_t） | 输入 | 通道数 C。 | C 为正整数，C ≤ 20。 | INT64 | - | 标量 | - | | ||
| 99 | +| dx（aclTensor*） | 输出 | x 的梯度。 | shape、数据类型与 x 一致。 | 同 x | ND | [B,S,D] | √ | | ||
| 100 | +| dwDown（aclTensor*） | 输出 | wDown 的梯度。 | shape 和数据类型与 wDown 一致。 | FLOAT16、BFLOAT16 | ND | [R,D] | √ | | ||
| 101 | +| dwUp（aclTensor*） | 输出 | wUp 的梯度。 | shape 和数据类型与 wUp 一致。 | FLOAT16、BFLOAT16 | ND | [D,R] | √ | | ||
| 102 | +| dwGate（aclTensor*） | 可选输出 | wGate 的梯度。 | shape 和数据类型与 wGate 一致；与 wGate、gradInject、inject 同时提供或同时缺省。 | FLOAT16、BFLOAT16 | ND | [C,D] | √ | | ||
| 103 | +| dgamma（aclTensor*） | 可选输出 | gamma 的梯度。 | shape 与 gamma 一致；与 gamma 同时提供或同时缺省。 | FLOAT32 | ND | [D] | √ | | ||
| 104 | +| workspaceSize（uint64_t*） | 输出 | 所需工作空间的字节数。 | 不能为空指针；按返回值申请 Device 内存。 | - | - | - | - | | ||
| 105 | +| executor（aclOpExecutor**） | 输出 | 本次调用的执行器。 | 不能为空指针；用于第二段接口。 | - | - | - | - | | ||
| 106 | + | ||
| 107 | +- **返回值** | ||
| 108 | + | ||
| 109 | +返回 `aclnnStatus`，参见 [aclnn 返回码](../../../docs/zh/context/aclnn_return_code.md)。 | ||
| 110 | + | ||
| 111 | +| 返回值 | 错误码 | 描述 | | ||
| 112 | +| :--- | ---: | :--- | | ||
| 113 | +| ACLNN_SUCCESS | 0 | 参数校验及执行准备成功。 | | ||
| 114 | +| ACLNN_ERR_PARAM_NULLPTR | 161001 | 必选 Tensor、workspaceSize 或 executor 为空指针，或可选参数未按要求成组提供。 | | ||
| 115 | +| ACLNN_ERR_PARAM_INVALID | 161002 | Tensor 的数据类型、维度或 shape 不符合参数表；channels 不在 [1,20] 范围内或不能整除 D；R 不在 [1,640] 范围内；B、S 为负数或 B*S 超出 INT64 表示范围。 | | ||
| 116 | +| ACLNN_ERR_INNER_CREATE_EXECUTOR | 561101 | 执行器创建失败。 | | ||
| 117 | +| ACLNN_ERR_INNER_NULLPTR | 561103 | 请求的形状组合不满足规格约束，或计算所需 Tensor 创建失败。 | | ||
| 118 | + | ||
| 119 | +## aclnnTextGatedResidualGrad | ||
| 120 | + | ||
| 121 | +- **参数说明** | ||
| 122 | + | ||
| 123 | +| 参数名 | 输入/输出 | 描述 | | ||
| 124 | +| :--- | :--- | :--- | | ||
| 125 | +| workspace（void*） | 输入 | Device 侧工作空间地址；workspaceSize 为 0 时可传入 nullptr。 | | ||
| 126 | +| workspaceSize（uint64_t） | 输入 | 第一段接口返回的工作空间字节数。 | | ||
| 127 | +| executor（aclOpExecutor*） | 输入 | 第一段接口返回的执行器。 | | ||
| 128 | +| stream（aclrtStream） | 输入 | 执行计算的 Stream。 | | ||
| 129 | + | ||
| 130 | +- **返回值** | ||
| 131 | + | ||
| 132 | +返回 `aclnnStatus`，参见 [aclnn 返回码](../../../docs/zh/context/aclnn_return_code.md)。 | ||
| 133 | + | ||
| 134 | +| 返回值 | 错误码 | 描述 | | ||
| 135 | +| :--- | ---: | :--- | | ||
| 136 | +| ACLNN_SUCCESS | 0 | 计算任务提交成功。 | | ||
| 137 | +| ACLNN_ERR_INNER | 561000 | 计算任务提交失败。 | | ||
| 138 | + | ||
| 139 | +第二段接口异步提交计算。读取输出前，应调用 `aclrtSynchronizeStream` 等待任务完成并检查其返回值。 | ||
| 140 | + | ||
| 141 | +## 约束说明 | ||
| 142 | + | ||
| 143 | +- `C`、`H`、`R` 为正整数；`H`、`R` 均须为 16 的整数倍。 | ||
| 144 | + | ||
| 145 | +通道数须满足 `C ≤ 20`。 | ||
| 146 | + | ||
| 147 | +`H` 的取值范围按产品确定。令 | ||
| 148 | + | ||
| 149 | +<!-- npu="950" id7 --> | ||
| 150 | +- <term>Ascend 950PR&950DT系列产品</term>： | ||
| 151 | + | ||
| 152 | + $$ | ||
| 153 | + K = 62448-16\left\lceil C/8\right\rceil | ||
| 154 | + $$ | ||
| 155 | + | ||
| 156 | +<!-- end id7 --> | ||
| 157 | +<!-- npu="A3,910b" id8 --> | ||
| 158 | +- <term>Atlas A2系列产品</term>、<term>Atlas A3系列产品</term>： | ||
| 159 | + | ||
| 160 | + $$ | ||
| 161 | + K = 39664-16\left\lceil C/8\right\rceil | ||
| 162 | + $$ | ||
| 163 | + | ||
| 164 | +<!-- end id8 --> | ||
| 165 | + | ||
| 166 | +`H`、`R` 须同时满足： | ||
| 167 | + | ||
| 168 | +$$ | ||
| 169 | +\boxed{ | ||
| 170 | +\begin{aligned} | ||
| 171 | +13H &\le K,\\ | ||
| 172 | +3H+8R &\le K,\\ | ||
| 173 | +7H+6R &\le K. | ||
| 174 | +\end{aligned} | ||
| 175 | +} | ||
| 176 | +$$ | ||
| 177 | + | ||
| 178 | +因此，给定 `C` 后，`H` 的上限为 | ||
| 179 | + | ||
| 180 | +$$ | ||
| 181 | +H_{\max}=16\left\lfloor\frac{K}{208}\right\rfloor. | ||
| 182 | +$$ | ||
| 183 | + | ||
| 184 | +`R` 须同时满足上述关系式和 `R ≤ 640`。 | ||
| 185 | + | ||
| 186 | +其中，$\lceil\cdot\rceil$ 表示向上取整，$\lfloor\cdot\rfloor$ 表示向下取整。上述关系适用于 FLOAT16、BFLOAT16，以及所有合法的可选参数组合。 | ||
| 187 | + | ||
| 188 | +例如，`C=4、H=2560、R=320` 适用于上述三类产品。其他 `C/H/R` 组合须满足上述约束。调用时仍需通过第一段接口 `aclnnTextGatedResidualGradGetWorkspaceSize` 完成参数校验，并按返回的 `workspaceSize` 申请内存。 | ||
| 189 | + | ||
| 190 | +- 所有 FLOAT16/BFLOAT16 输入以及 `dx`、`dwDown`、`dwUp`、`dwGate` 的数据类型必须一致；各权重梯度与对应权重的数据类型相同。`invRms`、`gamma` 和 `dgamma` 为 FLOAT32。 | ||
| 191 | +- 支持非连续输入 Tensor。输入输出的 shape 必须符合参数表，同名维度的大小必须一致。 | ||
| 192 | +- 支持 `B=0` 或 `S=0`：`dx` 为空，已提供的非空权重梯度和 `dgamma` 全部置零，其余参数仍须满足上述约束。 | ||
| 193 | +- 输出梯度覆盖写入，调用前无需初始化，也不会累加到输出 Tensor 的原有数据中。 | ||
| 194 | +- 默认执行确定性计算。 | ||
| 195 | + | ||
| 196 | +可选参数的使用方式如下： | ||
| 197 | + | ||
| 198 | +| 功能 | 提供时 | 缺省时 | | ||
| 199 | +| :--- | :--- | :--- | | ||
| 200 | +| 注入分支 | `wGate`、`gradInject`、`inject`、`dwGate` 四者必须同时提供 | 四者必须全部传入 `nullptr` | | ||
| 201 | +| 归一化缩放 | `gamma`、`dgamma` 必须同时提供 | 两者必须全部传入 `nullptr`，缩放系数取 1 | | ||
| 202 | + | ||
| 203 | +两组可选参数相互独立。`nullptr` 表示省略该功能，不能用 shape 为 `[0]` 的占位 Tensor 代替。 | ||
| 204 | + | ||
| 205 | +## 调用示例 | ||
| 206 | + | ||
| 207 | +以下示例使用 BFLOAT16，`B=1、S=2、C=4、H=16、R=16`，同时计算注入分支和 gamma 梯度。 | ||
| 208 | +输入 `x`、全部权重及 `gamma` 为 0，前向取 `epsilon=1e-6`，因此 `invRms=1000`、`downPreAct=0`、`up=0`、`inject=1`；两个上游梯度均为 1。 | ||
| 209 | +预期 `dx` 的所有元素为 125，`dwDown`、`dwUp`、`dwGate`、`dgamma` 的所有元素为 0。 | ||
| 210 | + | ||
| 211 | +编译与运行方式参见[编译与运行样例](../../../docs/zh/context/compile_and_run_sample.md)。请根据安装方式设置本算子的头文件和库搜索路径。 | ||
| 212 | + | ||
| 213 | +```cpp | ||
| 214 | +#include <cstdint> | ||
| 215 | +#include <cstdio> | ||
| 216 | +#include <cstring> | ||
| 217 | +#include <vector> | ||
| 218 | +#include "acl/acl.h" | ||
| 219 | +#include "aclnnop/aclnn_text_gated_residual_grad.h" | ||
| 220 | + | ||
| 221 | +namespace { | ||
| 222 | +constexpr int32_t DEVICE_ID = 0; | ||
| 223 | +constexpr int64_t B = 1; | ||
| 224 | +constexpr int64_t S = 2; | ||
| 225 | +constexpr int64_t C = 4; | ||
| 226 | +constexpr int64_t H = 16; | ||
| 227 | +constexpr int64_t R = 16; | ||
| 228 | +constexpr int64_t D = C * H; | ||
| 229 | +constexpr uint16_t BF16_ZERO = 0x0000; | ||
| 230 | +constexpr uint16_t BF16_ONE = 0x3F80; | ||
| 231 | +constexpr uint32_t BF16_SHIFT = 16; | ||
| 232 | +constexpr float SAVED_INV_RMS = 1000.0F; | ||
| 233 | + | ||
| 234 | +bool Check(int status, const char* action) | ||
| 235 | +{ | ||
| 236 | + if (status != ACL_SUCCESS) { | ||
| 237 | + std::printf("%s failed: %d\n", action, status); | ||
| 238 | + return false; | ||
| 239 | + } | ||
| 240 | + return true; | ||
| 241 | +} | ||
| 242 | + | ||
| 243 | +struct Resources { | ||
| 244 | + bool initialized = false; | ||
| 245 | + bool deviceSet = false; | ||
| 246 | + aclrtStream stream = nullptr; | ||
| 247 | + std::vector<aclTensor*> tensors; | ||
| 248 | + std::vector<void*> buffers; | ||
| 249 | + | ||
| 250 | + ~Resources() | ||
| 251 | + { | ||
| 252 | + if (stream != nullptr) { | ||
| 253 | + aclrtSynchronizeStream(stream); | ||
| 254 | + } | ||
| 255 | + for (auto* tensor : tensors) { | ||
| 256 | + aclDestroyTensor(tensor); | ||
| 257 | + } | ||
| 258 | + for (auto* buffer : buffers) { | ||
| 259 | + aclrtFree(buffer); | ||
| 260 | + } | ||
| 261 | + if (stream != nullptr) { | ||
| 262 | + aclrtDestroyStream(stream); | ||
| 263 | + } | ||
| 264 | + if (deviceSet) { | ||
| 265 | + aclrtResetDevice(DEVICE_ID); | ||
| 266 | + } | ||
| 267 | + if (initialized) { | ||
| 268 | + aclFinalize(); | ||
| 269 | + } | ||
| 270 | + } | ||
| 271 | +}; | ||
| 272 | + | ||
| 273 | +template <typename T> | ||
| 274 | +aclTensor* CreateTensor(Resources& resources, const std::vector<int64_t>& shape, aclDataType dtype, T value, | ||
| 275 | + bool isInput, void** outputAddress = nullptr) | ||
| 276 | +{ | ||
| 277 | + int64_t count = 1; | ||
| 278 | + std::vector<int64_t> strides(shape.size()); | ||
| 279 | + for (size_t axis = shape.size(); axis > 0; --axis) { | ||
| 280 | + strides[axis - 1] = count; | ||
| 281 | + count *= shape[axis - 1]; | ||
| 282 | + } | ||
| 283 | + const size_t bytes = static_cast<size_t>(count) * sizeof(T); | ||
| 284 | + void* address = nullptr; | ||
| 285 | + if (!Check(aclrtMalloc(&address, bytes, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc")) { | ||
| 286 | + return nullptr; | ||
| 287 | + } | ||
| 288 | + resources.buffers.push_back(address); | ||
| 289 | + if (isInput) { | ||
| 290 | + std::vector<T> host(static_cast<size_t>(count), value); | ||
| 291 | + if (!Check(aclrtMemcpy(address, bytes, host.data(), bytes, ACL_MEMCPY_HOST_TO_DEVICE), "aclrtMemcpy")) { | ||
| 292 | + return nullptr; | ||
| 293 | + } | ||
| 294 | + } | ||
| 295 | + auto* tensor = aclCreateTensor(shape.data(), shape.size(), dtype, strides.data(), 0, ACL_FORMAT_ND, shape.data(), | ||
| 296 | + shape.size(), address); | ||
| 297 | + if (tensor == nullptr) { | ||
| 298 | + std::printf("aclCreateTensor failed\n"); | ||
| 299 | + return nullptr; | ||
| 300 | + } | ||
| 301 | + resources.tensors.push_back(tensor); | ||
| 302 | + if (outputAddress != nullptr) { | ||
| 303 | + *outputAddress = address; | ||
| 304 | + } | ||
| 305 | + return tensor; | ||
| 306 | +} | ||
| 307 | +} // namespace | ||
| 308 | + | ||
| 309 | +int main() | ||
| 310 | +{ | ||
| 311 | + Resources resources; | ||
| 312 | + if (!Check(aclInit(nullptr), "aclInit")) { | ||
| 313 | + return 1; | ||
| 314 | + } | ||
| 315 | + resources.initialized = true; | ||
| 316 | + if (!Check(aclrtSetDevice(DEVICE_ID), "aclrtSetDevice")) { | ||
| 317 | + return 1; | ||
| 318 | + } | ||
| 319 | + resources.deviceSet = true; | ||
| 320 | + if (!Check(aclrtCreateStream(&resources.stream), "aclrtCreateStream")) { | ||
| 321 | + return 1; | ||
| 322 | + } | ||
| 323 | + | ||
| 324 | + auto* gradMixed = CreateTensor(resources, {B, S, H}, ACL_BF16, BF16_ONE, true); | ||
| 325 | + auto* gradInject = CreateTensor(resources, {B, S, C}, ACL_BF16, BF16_ONE, true); | ||
| 326 | + auto* x = CreateTensor(resources, {B, S, D}, ACL_BF16, BF16_ZERO, true); | ||
| 327 | + auto* wDown = CreateTensor(resources, {R, D}, ACL_BF16, BF16_ZERO, true); | ||
| 328 | + auto* wUp = CreateTensor(resources, {D, R}, ACL_BF16, BF16_ZERO, true); | ||
| 329 | + auto* wGate = CreateTensor(resources, {C, D}, ACL_BF16, BF16_ZERO, true); | ||
| 330 | + auto* gamma = CreateTensor(resources, {D}, ACL_FLOAT, 0.0F, true); | ||
| 331 | + auto* invRms = CreateTensor(resources, {B, S, C}, ACL_FLOAT, SAVED_INV_RMS, true); | ||
| 332 | + auto* inject = CreateTensor(resources, {B, S, C}, ACL_BF16, BF16_ONE, true); | ||
| 333 | + auto* downPreAct = CreateTensor(resources, {B, S, R}, ACL_BF16, BF16_ZERO, true); | ||
| 334 | + auto* up = CreateTensor(resources, {B, S, D}, ACL_BF16, BF16_ZERO, true); | ||
| 335 | + | ||
| 336 | + void* dxAddress = nullptr; | ||
| 337 | + auto* dx = CreateTensor(resources, {B, S, D}, ACL_BF16, BF16_ZERO, false, &dxAddress); | ||
| 338 | + auto* dwDown = CreateTensor(resources, {R, D}, ACL_BF16, BF16_ZERO, false); | ||
| 339 | + auto* dwUp = CreateTensor(resources, {D, R}, ACL_BF16, BF16_ZERO, false); | ||
| 340 | + auto* dwGate = CreateTensor(resources, {C, D}, ACL_BF16, BF16_ZERO, false); | ||
| 341 | + auto* dgamma = CreateTensor(resources, {D}, ACL_FLOAT, 0.0F, false); | ||
| 342 | + for (auto* tensor : {gradMixed, gradInject, x, wDown, wUp, wGate, gamma, invRms, inject, downPreAct, up, dx, dwDown, | ||
| 343 | + dwUp, dwGate, dgamma}) { | ||
| 344 | + if (tensor == nullptr) { | ||
| 345 | + return 1; | ||
| 346 | + } | ||
| 347 | + } | ||
| 348 | + | ||
| 349 | + uint64_t workspaceSize = 0; | ||
| 350 | + aclOpExecutor* executor = nullptr; | ||
| 351 | + auto status = aclnnTextGatedResidualGradGetWorkspaceSize(gradMixed, x, wDown, wUp, invRms, downPreAct, up, | ||
| 352 | + gradInject, wGate, gamma, inject, C, dx, dwDown, dwUp, | ||
| 353 | + dwGate, dgamma, &workspaceSize, &executor); | ||
| 354 | + if (!Check(status, "aclnnTextGatedResidualGradGetWorkspaceSize")) { | ||
| 355 | + return 1; | ||
| 356 | + } | ||
| 357 | + void* workspace = nullptr; | ||
| 358 | + if (workspaceSize > 0) { | ||
| 359 | + if (!Check(aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc workspace")) { | ||
| 360 | + return 1; | ||
| 361 | + } | ||
| 362 | + resources.buffers.push_back(workspace); | ||
| 363 | + } | ||
| 364 | + if (!Check(aclnnTextGatedResidualGrad(workspace, workspaceSize, executor, resources.stream), | ||
| 365 | + "aclnnTextGatedResidualGrad") || | ||
| 366 | + !Check(aclrtSynchronizeStream(resources.stream), "aclrtSynchronizeStream")) { | ||
| 367 | + return 1; | ||
| 368 | + } | ||
| 369 | + | ||
| 370 | + uint16_t first = 0; | ||
| 371 | + if (!Check(aclrtMemcpy(&first, sizeof(first), dxAddress, sizeof(first), ACL_MEMCPY_DEVICE_TO_HOST), | ||
| 372 | + "copy dx to host")) { | ||
| 373 | + return 1; | ||
| 374 | + } | ||
| 375 | + const uint32_t bits = static_cast<uint32_t>(first) << BF16_SHIFT; | ||
| 376 | + float value = 0.0F; | ||
| 377 | + std::memcpy(&value, &bits, sizeof(value)); | ||
| 378 | + std::printf("dx[0] = %.1f (expected 125.0)\n", value); | ||
| 379 | + return 0; | ||
| 380 | +} | ||
| 381 | +``` | ||